Split `ElementwiseMappable` trait into four more precise traits.

[frgossen]

Summary

Some elementwise operations are not scalarizable, vectorizable, or tensorizable.
Split ElementwiseMappable trait into the following, more precise traits.

• Elementwise
• Scalarizable
• Vectorizable
• Tensorizable

This allows for reuse of Elementwise in dialects like HLO.

[frgossen]

https://llvm.discourse.group/t/rfc-split-trait-elementwisemappable-into-elementwise-scalarizable-vectorizable-and-tensorizable/2897

[herhut]

mlir/include/mlir/IR/OpBase.td

1794

Would it make sense to define something like ElementwiseMappable.traits as a shortcut to all of these? As is done in https://reviews.llvm.org/D97636?

[frgossen]

Address comments

[frgossen]

mlir/include/mlir/IR/OpBase.td

1794

Thanks!

[herhut]

Just nits for comments. Looks good, thanks!

mlir/include/mlir/IR/OpBase.td

1792

nit: can be applied to ...

Also, replicate the comment for both entries rather than reuse it here to make it obvious to which def it applies.

1797

nit: Tensorizable?

mlir/include/mlir/IR/OpDefinition.h

1288

Can you copy the comment to the other def and split it? So that automated tooling can provide a description for the def?

[frgossen]

Address comments

[silvas]

mlir/include/mlir/IR/OpDefinition.h

1216–1217

don't mention scalars. this trait simply constrains the ways that ops can use vectors or tensors, if they have one as an operand.

1258

s/per/of the/

1262

"based on the same op's behavior on scalars".

1285

perhaps mention that this is just Tensorizable with s/tensor/vector. (that way folks can read the larger documentation on Tensorizable to understand Vectorizable)

1288

s/arguments/elements/

1297

s/select_scalar_pred above/scalar_pred below/

1306

stray ).

[frgossen]

See https://reviews.llvm.org/D97841

mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td

 
 class Complex_Op<string mnemonic, list<OpTrait> traits = []>
     : Op<Complex_Dialect, mnemonic, traits>;
 
 // Base class for standard arithmetic operations on complex numbers with a
 // floating-point element type. These operations take two operands and return
 // one result, all of which must be complex numbers of the same type.
 class ComplexArithmeticOp<string mnemonic, list<OpTrait> traits = []> :
-    Complex_Op<mnemonic,
-       !listconcat(traits, [NoSideEffect,
-                            SameOperandsAndResultType,
-                            DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
-                            ElementwiseMappable])> {
+    Complex_Op<mnemonic, traits # [NoSideEffect, SameOperandsAndResultType,
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
+    ElementwiseMappable.traits> {
   let arguments = (ins Complex<AnyFloat>:$lhs, Complex<AnyFloat>:$rhs);
   let results = (outs Complex<AnyFloat>:$result);
   let assemblyFormat = "$lhs `,` $rhs  attr-dict `:` type($result)";
   let verifier = ?;
 }
 
 //===----------------------------------------------------------------------===//
 // AddOp

mlir/include/mlir/Dialect/Math/IR/MathOps.td

 include "mlir/Dialect/Math/IR/MathBase.td"
 include "mlir/Interfaces/VectorInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 class MathOp<string mnemonic, list<OpTrait> traits = []>
     : Op<Math_Dialect, mnemonic, traits # [NoSideEffect]>;
 
 class FloatUnaryOp<string mnemonic, list<OpTrait> traits = []> :
-    MathOp<mnemonic,
-      traits # [DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
-                ElementwiseMappable,
-                SameOperandsAndResultType]> {
+    MathOp<mnemonic, traits #
+    [DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
+    SameOperandsAndResultType] # ElementwiseMappable.traits> {
   let arguments = (ins FloatLike:$operand);
 
   let results = (outs FloatLike:$result);
 
   let assemblyFormat = "$operand attr-dict `:` type($result)";
 }
 
 class FloatBinaryOp<string mnemonic, list<OpTrait> traits = []> :
-    MathOp<mnemonic,
-      traits # [DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
-                ElementwiseMappable,
-                SameOperandsAndResultType]> {
+    MathOp<mnemonic, traits # [
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
+    SameOperandsAndResultType] # ElementwiseMappable.traits> {
   let arguments = (ins FloatLike:$lhs, FloatLike:$rhs);
   let results = (outs FloatLike:$result);
 
   let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($result)";
 }
 
 //===----------------------------------------------------------------------===//
 // AtanOp

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

class CastOp<string mnemonic, list<OpTrait> traits = []> :
   }];
 
   // Cast operations are fully verified by its traits.
   let verifier = ?;
 }
 
 // Base class for arithmetic cast operations.
 class ArithmeticCastOp<string mnemonic, list<OpTrait> traits = []> :
-    CastOp<mnemonic,
-    !listconcat(traits, [ElementwiseMappable,
-                         DeclareOpInterfaceMethods<VectorUnrollOpInterface>])> {
+    CastOp<mnemonic, traits # [
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
+    ElementwiseMappable.traits> {
 }
 
 // Base class for unary ops. Requires single operand and result. Individual
 // classes will have `operand` accessor.
 class UnaryOp<string mnemonic, list<OpTrait> traits = []> :
     Op<StandardOps_Dialect, mnemonic, !listconcat(traits, [NoSideEffect])> {
   let results = (outs AnyType);
   let printer = [{
     return printStandardUnaryOp(this->getOperation(), p);
   }];
 }
 
 class UnaryOpSameOperandAndResultType<string mnemonic,
                                       list<OpTrait> traits = []> :
     UnaryOp<mnemonic, !listconcat(traits, [SameOperandsAndResultType])> {
   let parser = [{
     return impl::parseOneResultSameOperandTypeOp(parser, result);
   }];
 }
 
 class FloatUnaryOp<string mnemonic, list<OpTrait> traits = []> :
-    UnaryOpSameOperandAndResultType<mnemonic,
-      !listconcat(traits,
-                  [DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
-                   ElementwiseMappable])>,
-    Arguments<(ins FloatLike:$operand)>;
+    UnaryOpSameOperandAndResultType<mnemonic, traits # [
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
+    ElementwiseMappable.traits>, Arguments<(ins FloatLike:$operand)>;
 
 // Base class for standard arithmetic operations.  Requires operands and
 // results to be of the same type, but does not constrain them to specific
 // types.
 class ArithmeticOp<string mnemonic, list<OpTrait> traits = []> :
-    Op<StandardOps_Dialect, mnemonic,
-       !listconcat(traits, [NoSideEffect,
-                            SameOperandsAndResultType,
-                            DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
-                            ElementwiseMappable])> {
+    Op<StandardOps_Dialect, mnemonic, traits # [NoSideEffect,
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
+    ElementwiseMappable.traits> {
 
   let results = (outs AnyType:$result);
 
   let parser = [{
     return impl::parseOneResultSameOperandTypeOp(parser, result);
   }];
 
   let printer = [{

 def CmpFPredicateAttr : I64EnumAttr<
     "CmpFPredicate", "",
     [CMPF_P_FALSE, CMPF_P_OEQ, CMPF_P_OGT, CMPF_P_OGE, CMPF_P_OLT, CMPF_P_OLE,
      CMPF_P_ONE, CMPF_P_ORD, CMPF_P_UEQ, CMPF_P_UGT, CMPF_P_UGE, CMPF_P_ULT,
      CMPF_P_ULE, CMPF_P_UNE, CMPF_P_UNO, CMPF_P_TRUE]> {
   let cppNamespace = "::mlir";
 }
 
-def CmpFOp : Std_Op<"cmpf",
-    [NoSideEffect, SameTypeOperands, ElementwiseMappable,
-     DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
-     TypesMatchWith<
-       "result type has i1 element type and same shape as operands",
-       "lhs", "result", "getI1SameShape($_self)">]> {
+def CmpFOp : Std_Op<"cmpf", [NoSideEffect, SameTypeOperands,
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>, TypesMatchWith<
+    "lhs", "result", "getI1SameShape($_self)">] # ElementwiseMappable.traits> {
   let summary = "floating-point comparison operation";
   let description = [{
     The `cmpf` operation compares its two operands according to the float
     comparison rules and the predicate specified by the respective attribute.
     The predicate defines the type of comparison: (un)orderedness, (in)equality
     and signed less/greater than (or equal to) as well as predicates that are
     always true or false.  The operands must have the same type, and this type
     must be a float type, or a vector or tensor thereof.  The result is an i1,

 
 def CmpIPredicateAttr : I64EnumAttr<
     "CmpIPredicate", "",
     [CMPI_P_EQ, CMPI_P_NE, CMPI_P_SLT, CMPI_P_SLE, CMPI_P_SGT,
      CMPI_P_SGE, CMPI_P_ULT, CMPI_P_ULE, CMPI_P_UGT, CMPI_P_UGE]> {
   let cppNamespace = "::mlir";
 }
 
-def CmpIOp : Std_Op<"cmpi",
-    [NoSideEffect, SameTypeOperands, ElementwiseMappable,
-     DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
-     TypesMatchWith<
-       "result type has i1 element type and same shape as operands",
-       "lhs", "result", "getI1SameShape($_self)">]> {
+def CmpIOp : Std_Op<"cmpi", [NoSideEffect, SameTypeOperands,
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>, TypesMatchWith<
+    "lhs", "result", "getI1SameShape($_self)">] # ElementwiseMappable.traits> {
   let summary = "integer comparison operation";
   let description = [{
     The `cmpi` operation is a generic comparison for integer-like types. Its two
     arguments can be integers, vectors or tensors thereof as long as their types
     match. The operation produces an i1 for the former case, a vector or a
     tensor of i1 with the same shape as inputs in the other cases.
 
     Its first argument is an attribute that defines which type of comparison is
def ReturnOp : Std_Op<"return", [NoSideEffect, HasParent<"FuncOp">,
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
 //===----------------------------------------------------------------------===//
 // SelectOp
 //===----------------------------------------------------------------------===//
 
 def SelectOp : Std_Op<"select", [NoSideEffect,
-     AllTypesMatch<["true_value", "false_value", "result"]>,
-     ElementwiseMappable, DeclareOpInterfaceMethods<VectorUnrollOpInterface>]> {
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
+    ElementwiseMappable.traits> {
   let summary = "select operation";
   let description = [{
     The `select` operation chooses one value based on a binary condition
     supplied as its first operand. If the value of the first operand is `1`,
     the second operand is chosen, otherwise the third operand is chosen.
     The second and the third operand must have the same type.
 
     The operation applies to vectors and tensors elementwise given the _shape_
let description = [{
     ```
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SignExtendIOp
 //===----------------------------------------------------------------------===//
 
-def SignExtendIOp : Std_Op<"sexti",
-    [NoSideEffect, ElementwiseMappable,
-    DeclareOpInterfaceMethods<VectorUnrollOpInterface>]> {
+def SignExtendIOp : Std_Op<"sexti", [NoSideEffect,
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
+    ElementwiseMappable.traits> {
   let summary = "integer sign extension operation";
   let description = [{
     The integer sign extension operation takes an integer input of
     width M and an integer destination type of width N. The destination
     bit-width must be larger than the input bit-width (N > M).
     The top-most (N - M) bits of the output are filled with copies
     of the most-significant bit of the input.
 
def TransposeOp : Std_Op<"transpose", [NoSideEffect]>,
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TruncateIOp
 //===----------------------------------------------------------------------===//
 
-def TruncateIOp : Std_Op<"trunci",
-  [NoSideEffect, ElementwiseMappable,
-   DeclareOpInterfaceMethods<VectorUnrollOpInterface>,]> {
+def TruncateIOp : Std_Op<"trunci", [NoSideEffect,
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
+    ElementwiseMappable.traits> {
   let summary = "integer truncation operation";
   let description = [{
     The integer truncation operation takes an integer input of
     width M and an integer destination type of width N. The destination
     bit-width must be smaller than the input bit-width (N < M).
     The top-most (N - M) bits of the input are discarded.
 
     Example:
def XOrOp : IntBinaryOp<"xor", [Commutative]> {
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ZeroExtendIOp
 //===----------------------------------------------------------------------===//
 
-def ZeroExtendIOp : Std_Op<"zexti",
-  [NoSideEffect, ElementwiseMappable,
-   DeclareOpInterfaceMethods<VectorUnrollOpInterface>,]> {
+def ZeroExtendIOp : Std_Op<"zexti", [NoSideEffect,
+    DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
+    ElementwiseMappable.traits> {
   let summary = "integer zero extension operation";
   let description = [{
     The integer zero extension operation takes an integer input of
     width M and an integer destination type of width N. The destination
     bit-width must be larger than the input bit-width (N > M).
     The top-most (N - M) bits of the output are filled with zeros.
 
     Example:

mlir/include/mlir/IR/OpBase.td

 // Op has the same operand and result element type (or type itself, if scalar).
 def SameOperandsAndResultElementType :
   NativeOpTrait<"SameOperandsAndResultElementType">;
 // Op is a terminator.
 def Terminator : NativeOpTrait<"IsTerminator">;
 // Op can be safely normalized in the presence of MemRefs with
 // non-identity maps.
 def MemRefsNormalizable : NativeOpTrait<"MemRefsNormalizable">;
-// Op can be systematically interconverted between scalar and vector/tensor
-// form by mapping elementwise based on the type.
-def ElementwiseMappable : NativeOpTrait<"ElementwiseMappable">;
+// Op is elementwise on tensor/vector operands and results.
+def Elementwise : NativeOpTrait<"Elementwise">;
+// Elementwise op can be applied to scalars instead tensor/vector operands.
+def Scalarizable : NativeOpTrait<"Scalarizable">;
+// Elementwise op can be applied all-tensor/vector operands.
+def Vectorizable : NativeOpTrait<"Vectorizable">;
+def Tensorizable : NativeOpTrait<"Tensorizable">;
+
+// Group together `Elementwise`, `Scalarizable`, `Vectorizable`, and
+// `Vectorizable` for convenience.
+def ElementwiseMappable {
+  list<OpTrait> traits = [
+    Elementwise,
+    Scalarizable,
+    Vectorizable,
+    Tensorizable,
+  ];
+}
 
 // Op's regions have a single block with the specified terminator.
 class SingleBlockImplicitTerminator<string op>
     : ParamNativeOpTrait<"SingleBlockImplicitTerminator", op>;
 
 // Op's parent operation is the provided one.
 class HasParent<string op>
     : ParamNativeOpTrait<"HasParent", op>;

mlir/include/mlir/IR/OpDefinition.h

 LogicalResult verifyIsTerminator(Operation *op);
 LogicalResult verifyZeroSuccessor(Operation *op);
 LogicalResult verifyOneSuccessor(Operation *op);
 LogicalResult verifyNSuccessors(Operation *op, unsigned numSuccessors);
 LogicalResult verifyAtLeastNSuccessors(Operation *op, unsigned numSuccessors);
 LogicalResult verifyOperandSizeAttr(Operation *op, StringRef sizeAttrName);
 LogicalResult verifyResultSizeAttr(Operation *op, StringRef sizeAttrName);
 LogicalResult verifyNoRegionArguments(Operation *op);
-LogicalResult verifyElementwiseMappable(Operation *op);
+LogicalResult verifyElementwise(Operation *op);
 } // namespace impl
 
 /// Helper class for implementing traits.  Clients are not expected to interact
 /// with this directly, so its members are all protected.
 template <typename ConcreteType, template <typename> class TraitType>
 class TraitBase {
 protected:
   /// Return the ultimate Operation being worked on.

 // values of `MemRef` type where those references can be 'normalized'.
 // TODO: Right now, the operands of an operation are either all normalizable,
 // or not. In the future, we may want to allow some of the operands to be
 // normalizable.
 template <typename ConcrentType>
 struct MemRefsNormalizable
     : public TraitBase<ConcrentType, MemRefsNormalizable> {};
 
-/// This trait tags scalar ops that also can be applied to vectors/tensors, with
-/// their semantics on vectors/tensors being elementwise application.
+/// This trait tags element-wise ops that operate on scalars, vectors, or
+/// tensors.
 ///
 /// NOTE: Not all ops that are "elementwise" in some abstract sense satisfy this
-/// trait. In particular, broadcasting behavior is not allowed. This trait
-/// describes a set of invariants that allow systematic
-/// vectorization/tensorization, and the reverse, scalarization. The properties
-/// needed for this also can be used to implement a number of
-/// transformations/analyses/interfaces.
-///
-/// An `ElementwiseMappable` op must satisfy the following properties:
-///
-/// 1. If any result is a vector (resp. tensor), then at least one operand must
-/// be a vector (resp. tensor).
-/// 2. If any operand is a vector (resp. tensor), then there must be at least
-/// one result, and all results must be vectors (resp. tensors).
-/// 3. The static types of all vector (resp. tensor) operands and results must
-/// have the same shape.
-/// 4. In the case of tensor operands, the dynamic shapes of all tensor operands
-/// must be the same, otherwise the op has undefined behavior.
-/// 5. ("systematic scalarization" property) If an op has vector/tensor
-/// operands/results, then the same op, with the operand/result types changed to
-/// their corresponding element type, shall be a verifier-valid op.
-/// 6. The semantics of the op on vectors (resp. tensors) shall be the same as
-/// applying the scalarized version of the op for each corresponding element of
-/// the vector (resp. tensor) operands in parallel.
-/// 7. ("systematic vectorization/tensorization" property) If an op has
-/// scalar operands/results, the op shall remain verifier-valid if all scalar
-/// operands are replaced with vectors/tensors of the same shape and
-/// corresponding element types.
-///
-/// Together, these properties provide an easy way for scalar operations to
-/// conveniently generalize their behavior to vectors/tensors, and systematize
-/// conversion between these forms.
+/// trait. In particular, broadcasting behavior is not allowed.
 ///
-/// Examples:
-/// ```
-/// %scalar = "std.addf"(%a, %b) : (f32, f32) -> f32
-/// // Applying the systematic vectorization/tensorization property, this op
-/// // must also be valid:
-/// %tensor = "std.addf"(%a_tensor, %b_tensor)
-///           : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>)
+/// An `Elementwise` op must satisfy the following properties:
 ///
-/// // These properties generalize well to the cases of non-scalar operands.
-/// %select_scalar_pred = "std.select"(%pred, %true_val, %false_val)
-///                       : (i1, tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
-/// // Applying the systematic vectorization / tensorization property, this
-/// // op must also be valid:
-/// %select_tensor_pred = "std.select"(%pred_tensor, %true_val, %false_val)
-///                       : (tensor<?xi1>, tensor<?xf32>, tensor<?xf32>)
-///                       -> tensor<?xf32>
-/// // Applying the systematic scalarization property, this op must also
-/// // be valid.
-/// %select_scalar = "std.select"(%pred, %true_val_scalar, %false_val_scalar)
-///                  : (i1, f32, f32) -> f32
+/// 1. If any result is a vector/tensor then at least one operand must also be a
+///    vector/tensor.
+/// 2. If any operand is a vector/tensor then there must be at least one result
+///    and all results must be vectors/tensors.
+/// 3. All operand and result vector/tensor types must be of the same shape. The
+///    shape may be dynamic in which case the op's behaviour is undefined for
+///    non-matching shapes.
+/// 4. The operation must be elementwise on its vector/tensor operands and
+///    results. When applied to single-element vectors/tensors, the result must
+///    be the same per elememnt.
+///
+/// TODO: Avoid hardcoding vector/tensor, and generalize this trait to a new
+/// interface `ElementwiseTypeInterface` that describes the container types for
+/// which the operation is elementwise.
+///
+/// Rationale:
+/// - 1. and 2. guarantee a well-defined iteration space and exclude the cases
+///   of 0 non-scalar operands or 0 non-scalar results, which complicate a
+///   generic definition of the iteration space.
+/// - 3. guarantees that folding can be done across scalars/vectors/tensors with
+///   the same pattern, as otherwise lots of special handling for type
+///   mismatches would be needed.
+/// - 4. guarantees that no error handling is needed. Higher-level dialects
+///   should reify any needed guards or error handling code before lowering to
+///   an `Elementwise` op.
+template <typename ConcreteType>
+struct Elementwise : public TraitBase<ConcreteType, Elementwise> {
+  static LogicalResult verifyTrait(Operation *op) {
+    return ::mlir::OpTrait::impl::verifyElementwise(op);
+  }
+};
+
+/// This trait tags `Elementwise` operatons that can be systematically
+/// scalarized. All vector/tensor operands and results are then replaced by
+/// scalars of the respective element type. Semantically, this is the operation
+/// on a single element per vector/tensor.
+///
+/// Rationale:
+/// Allow to define the vector/tensor semantics of elementwise operations based
+/// on scalars. This provides a constructive procedure for IR transformations
+/// to, e.g., create scalar loop bodies from tensor ops.
+///
+/// Example:
+/// ```
+/// %tensor_select = "std.select"(%pred_tensor, %true_val, %false_val)
+///                      : (tensor<?xi1>, tensor<?xf32>, tensor<?xf32>)
+///                      -> tensor<?xf32>
 /// ```
+/// can be scalarized to
 ///
-/// TODO: Avoid hardcoding vector/tensor, and generalize this to any type
-/// implementing a new "ElementwiseMappableTypeInterface" that describes types
-/// for which it makes sense to apply a scalar function to each element.
+/// ```
+/// %scalar_select = "std.select"(%pred, %true_val_scalar, %false_val_scalar)
+///                      : (i1, f32, f32) -> f32
+/// ```
+template <typename ConcreteType>
+struct Scalarizable : public TraitBase<ConcreteType, Scalarizable> {
+  static LogicalResult verifyTrait(Operation *op) {
+    static_assert(
+        ConcreteType::template hasTrait<Elementwise>(),
+        "`Scalarizable` trait is only applicable to `Elementwise` ops.");
+    return success();
+  }
+};
+
+/// These traits tag `Elementwise` operatons that can be systematically
+/// vectorized/tensorized. All scalar operands and results are then replaced by
+/// tensors/vectors with the respective element type. Semantically, this is the
+/// operation on multiple arguments simultaneously.
 ///
 /// Rationale:
-/// - 1. and 2. guarantee a well-defined iteration space for 6.
-///   - These also exclude the cases of 0 non-scalar operands or 0 non-scalar
-///     results, which complicate a generic definition of the iteration space.
-/// - 3. guarantees that folding can be done across scalars/vectors/tensors
-///   with the same pattern, as otherwise lots of special handling of type
-///   mismatches would be needed.
-/// - 4. guarantees that no error handling cases need to be considered.
-///   - Higher-level dialects should reify any needed guards / error handling
-///   code before lowering to an ElementwiseMappable op.
-/// - 5. and 6. allow defining the semantics on vectors/tensors via the scalar
-///   semantics and provide a constructive procedure for IR transformations
-///   to e.g. create scalar loop bodies from tensor ops.
-/// - 7. provides the reverse of 5., which when chained together allows
-///   reasoning about the relationship between the tensor and vector case.
-///   Additionally, it permits reasoning about promoting scalars to
-///   vectors/tensors via broadcasting in cases like `%select_scalar_pred`
-///   above.
+/// Provide the reverse to `Scalarizable` which, when chained together, allows
+/// reasoning about the relationship between the tensor and vector case.
+/// Additionally, it permits reasoning about promoting scalars to
+/// vectors/tensors via broadcasting in cases like `%select_scalar_pred` above.
+///
+/// Examples:
+/// ```
+/// %scalar = "std.addf"(%a, %b) : (f32, f32) -> f32
+/// ```
+/// can be tensorized to
+/// ```
+/// %tensor = "std.addf"(%a, %b) : (tensor<?xf32>, tensor<?xf32>)
+///               -> tensor<?xf32>)
+/// ```
+///
+/// ```
+/// %scalar_pred = "std.select"(%pred, %true_val, %false_val)
+///                    : (i1, tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
+/// ```
+/// can be tensorized to
+/// ```
+/// %tensor_pred = "std.select"(%pred, %true_val, %false_val)
+///                    : (tensor<?xi1>, tensor<?xf32>, tensor<?xf32>)
+///                    -> tensor<?xf32>
+/// ```
+template <typename ConcreteType>
+struct Vectorizable : public TraitBase<ConcreteType, Vectorizable> {
+  static LogicalResult verifyTrait(Operation *op) {
+    static_assert(
+        ConcreteType::template hasTrait<Elementwise>(),
+        "`Vectorizable` trait is only applicable to `Elementwise` ops.");
+    return success();
+  }
+};
 template <typename ConcreteType>
-struct ElementwiseMappable
-    : public TraitBase<ConcreteType, ElementwiseMappable> {
+struct Tensorizable : public TraitBase<ConcreteType, Tensorizable> {
   static LogicalResult verifyTrait(Operation *op) {
-    return ::mlir::OpTrait::impl::verifyElementwiseMappable(op);
+    static_assert(
+        ConcreteType::template hasTrait<Elementwise>(),
+        "`Tensorizable` trait is only applicable to `Elementwise` ops.");
+    return success();
   }
 };
 
+/// Together, `Elementwise`, `Scalarizable`, `Vectorizable`, and `Tensorizable`
+/// provide an easy way for scalar operations to conveniently generalize their
+/// behavior to vectors/tensors, and systematize conversion between these forms.
+bool hasElementwiseMappableTraits(Operation *op);
+
 } // end namespace OpTrait
 
 //===----------------------------------------------------------------------===//
 // Internal Trait Utilities
 //===----------------------------------------------------------------------===//
 
 namespace op_definition_impl {
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp

 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/StandardOps/Utils/Utils.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 static bool isElementwiseMappableOpOnRankedTensors(Operation *op) {
-  if (!op->hasTrait<OpTrait::ElementwiseMappable>())
+  if (!OpTrait::hasElementwiseMappableTraits(op))
     return false;
 
   // TODO: The conversion pattern can be made to work for `any_of` here, but
   // it's more complex as it requires tracking which operands are scalars.
   return llvm::all_of(op->getOperandTypes(),
                       [](Type type) { return type.isa<RankedTensorType>(); });
 }
 

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

vectorizeOneOp(OpBuilder &builder, Operation *op,
   }
 
   // 2. Constant ops don't get vectorized but rather broadcasted at their users.
   // Clone so that the constant is not confined to the linalgOp block .
   if (isa<ConstantOp>(op))
     return VectorizationResult{VectorizationStatus::NewOp, builder.clone(*op)};
 
   // 3. Only ElementwiseMappable are allowed in the generic vectorization.
-  if (!op->hasTrait<OpTrait::ElementwiseMappable>())
+  if (!OpTrait::hasElementwiseMappableTraits(op))
     return VectorizationResult{VectorizationStatus::Failure, nullptr};
 
   // 4. Generic vectorization path for ElementwiseMappable ops.
   //   a. first get the first max ranked shape.
   SmallVector<int64_t, 4> firstMaxRankedShape;
   for (Value operand : op->getOperands()) {
     auto vt = bvm.lookup(operand).getType().dyn_cast<VectorType>();
     if (vt && firstMaxRankedShape.size() < vt.getShape().size())

 }
 
 /// Detect whether `r` has only ConstantOp, ElementwiseMappable and YieldOp.
 static bool hasOnlyScalarElementwiseOp(Region &r) {
   if (!llvm::hasSingleElement(r))
     return false;
   for (Operation &op : r.front()) {
     if (!(isa<ConstantOp, linalg::YieldOp>(op) ||
-          op.hasTrait<OpTrait::ElementwiseMappable>()) ||
+          OpTrait::hasElementwiseMappableTraits(&op)) ||
         llvm::any_of(op.getResultTypes(),
                      [](Type type) { return !type.isIntOrIndexOrFloat(); }))
       return false;
   }
   return true;
 }
 
 // Return true if the op is an element-wise linalg op.

mlir/lib/IR/Operation.cpp

 static bool areSameShapedTypeIgnoringElementType(ShapedType a, ShapedType b) {
   if (a.getTypeID() != b.getTypeID())
     return false;
   if (!a.hasRank())
     return !b.hasRank();
   return a.getShape() == b.getShape();
 }
 
-LogicalResult OpTrait::impl::verifyElementwiseMappable(Operation *op) {
+LogicalResult OpTrait::impl::verifyElementwise(Operation *op) {
   auto isMappableType = [](Type type) {
     return type.isa<VectorType, TensorType>();
   };
   auto resultMappableTypes = llvm::to_vector<1>(
       llvm::make_filter_range(op->getResultTypes(), isMappableType));
   auto operandMappableTypes = llvm::to_vector<2>(
       llvm::make_filter_range(op->getOperandTypes(), isMappableType));
 
if (!areSameShapedTypeIgnoringElementType(type.cast<ShapedType>(),
                                   "the same shape and base type: found "
                                << type << " and " << mustMatchType;
     }
   }
 
   return success();
 }
 
+bool OpTrait::hasElementwiseMappableTraits(Operation *op) {
+  return op->hasTrait<Elementwise>() && op->hasTrait<Scalarizable>() &&
+         op->hasTrait<Vectorizable>() && op->hasTrait<Tensorizable>();
+}
+
 //===----------------------------------------------------------------------===//
 // BinaryOp implementation
 //===----------------------------------------------------------------------===//
 
 // These functions are out-of-line implementations of the methods in BinaryOp,
 // which avoids them being template instantiated/duplicated.
 
 void impl::buildBinaryOp(OpBuilder &builder, OperationState &result, Value lhs,

mlir/test/lib/Dialect/Test/TestOps.td

 //===----------------------------------------------------------------------===//
 
 def SameOperandElementTypeOp : TEST_Op<"same_operand_element_type",
     [SameOperandsElementType]> {
   let arguments = (ins AnyType, AnyType);
   let results = (outs AnyType);
 }
 
-def SameOperandAndResultElementTypeOp : TEST_Op<"same_operand_and_result_element_type",
+def SameOperandAndResultElementTypeOp :
+    TEST_Op<"same_operand_and_result_element_type",
     [SameOperandsAndResultElementType]> {
   let arguments = (ins Variadic<AnyType>);
   let results = (outs Variadic<AnyType>);
 }
 
 def SameOperandShapeOp : TEST_Op<"same_operand_shape", [SameOperandsShape]> {
   let arguments = (ins Variadic<AnyShaped>);
 }
 
 def SameOperandAndResultShapeOp : TEST_Op<"same_operand_and_result_shape",
     [SameOperandsAndResultShape]> {
   let arguments = (ins Variadic<AnyShaped>);
   let results = (outs Variadic<AnyShaped>);
 }
 
 def SameOperandAndResultTypeOp : TEST_Op<"same_operand_and_result_type",
     [SameOperandsAndResultType]> {
   let arguments = (ins Variadic<AnyType>);
   let results = (outs Variadic<AnyType>);
 }
 
 def ElementwiseMappableOp : TEST_Op<"elementwise_mappable",
-    [ElementwiseMappable]> {
+    ElementwiseMappable.traits> {
   let arguments = (ins Variadic<AnyType>);
   let results = (outs Variadic<AnyType>);
 }
 
 def ArgAndResHaveFixedElementTypesOp :
     TEST_Op<"arg_and_res_have_fixed_element_types",
       [PredOpTrait<"fixed type combination",
          And<[ElementTypeIsPred<"x", I32>,