Add ElementwiseMappable trait and apply it to std elementwise ops.

[silvas]

Summary

This patch adds an ElementwiseMappable trait as discussed in the RFC
here:
https://llvm.discourse.group/t/rfc-std-elementwise-ops-on-tensors/2113/23

This trait can power a number of transformations and analyses.
A subsequent patch adds a convert-elementwise-to-linalg pass exhibits
how this trait allows writing generic transformations.
See https://reviews.llvm.org/D90354 for that patch. I have split it out for clarity.

This trait slightly changes some verifier messages, but the diagnostics
are usually about as good. I fiddled with the ordering of the trait in
the .td file trait lists to minimize the changes here.

[silvas]

split out the convert-elementwise-to-linalg pass

[silvas]

@bondhugula I'd love to get your review on the definition of the ElementwiseMappable trait.

[stellaraccident]

This is very nice, and the corresponding RFC seems to have gotten to a nice conclusion.

mlir/include/mlir/IR/OpBase.td

1753–1754

I feel like this description could stand some truing up, at least having the words that make up the trait:

Op can be systematically interconverted between scalar and vector/tensor
form [by mapping elementwise based on the type].

mlir/include/mlir/IR/OpDefinition.h

1331

s/shall be same/shall be the same/

[silvas]

address comments

[mehdi_amini]

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

3651

How does SameOperandsAndResultShape relates to ElementwiseMappable? Is it a subset by any chance?

mlir/include/mlir/IR/OpDefinition.h

1387

Concrete instead of Concrent?

[rriddle]

Please add documentation to https://mlir.llvm.org/docs/Traits/

mlir/lib/IR/Operation.cpp

1082

nit: isa is variadic: isa<TensorType, VectorType>().

[bondhugula]

mlir/include/mlir/IR/OpDefinition.h

1366

Side note: would ElementwiseMappableTypeInterface == isa<ShapedType>()? - the way it is now (dense shaped types). One could perhaps introduce a new type in the hierarchy sitting between ShapedType and Type to handle dense + sparse like the way ElementsAttribute works. What interface methods do you specifically see being added to ElementwiseMappableTypeInterface? If there are none, it could just sit in the type inheritance hierarchy?

[silvas]

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

3651

ElementwiseMappable permits select %scalar_pred, %true_tensor, %false_tensor, but SameOperandsAndResultShape does not.

SameOperandsAndResultShape actually doesn't check for "same", only "compatible" (which includes broadcast compatible). So SameOperandsAndResultShape allows "foo"(%0) : tensor<2xf32> -> tensor<?xf32>, but ElementwiseMappable does not (as described in the rationale, this makes folding not generalize well from scalars/vectors to tensors).

So neither is strictly redundant from the other one in general. I think in this case, due to there being only a single operand/result, ElementwiseMappable imposes strictly stronger invariants.

I didn't want to touch the SameOperandsAndResultShape/SameOpearndsAndResultTypes situation in this patch; there's a lot to unwind there. See e.g. this thread https://llvm.discourse.group/t/sameoperandandresulttypes-name-confusing/1313/4

mlir/include/mlir/IR/OpDefinition.h

1366

ElementwiseMappableTypeInterface is not just ShapedType. In general, it could support things like "database columns" or abstract "streams".

So far, one interface method that seems like it will be needed is bool canSpeculateTraversal(). E.g. addf %lhs, %rhs : vector<4xf32> can be speculated with standard poison-like reasoning that LLVM uses. But addf %lhs, %rhs : tensor<?xf32> cannot be speculated, because the traversal itself might have undefined behavior if %lhs and %rhs dynamically have mismatching sizes.

[silvas]

Address comments.

[silvas]

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

3651

Correction, SameOperandsAndResultShape doesn't allow broadcast compatible. But it does allow e.g. 2 vs ? like in my example.

[mehdi_amini]

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

3651

Basically what I was after was: "If an operation already defines ElementwiseMappable, what extra verification does it get from SameOperandsAndResultShape which isn't already covered" ?

[silvas]

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

3651

All I can think of is that SameOperandsAndResultShape would prohibit ops like select %scalar_pred, %true_tensor, %false_tensor.

[mehdi_amini]

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

3651

OK! Should we remove SameOperandsAndResultShape in cases like here then?

[silvas]

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

3651

Will do.

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

class CastOp<string mnemonic, list<OpTrait> traits = []> :
   let printer = [{
     return printStandardCastOp(this->getOperation(), p);
   }];
   let verifier = [{ return ::verifyCastOp(*this); }];
 
   let hasFolder = 1;
 }
 
+// Base class for arithmetic cast operations.
+class ArithmeticCastOp<string mnemonic, list<OpTrait> traits = []> :
+    CastOp<mnemonic, !listconcat(traits, [ElementwiseMappable])> {
+}
+
 // 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>])>,
+                  [DeclareOpInterfaceMethods<VectorUnrollOpInterface>,
+                   ElementwiseMappable])>,
     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.  Individual classes will have `lhs` and `rhs` accessor to operands.
 class ArithmeticOp<string mnemonic, list<OpTrait> traits = []> :
     Op<StandardOps_Dialect, mnemonic,
-       !listconcat(traits, [NoSideEffect, SameOperandsAndResultType])> {
+       !listconcat(traits, [NoSideEffect,
+                            SameOperandsAndResultType,
+                            ElementwiseMappable])> {
 
   let results = (outs AnyType);
 
   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, SameOperandsAndResultShape,
-     TypesMatchWith<
+    [NoSideEffect, SameTypeOperands,
+     SameOperandsAndResultShape, TypesMatchWith<
        "result type has i1 element type and same shape as operands",
-       "lhs", "result", "getI1SameShape($_self)">]> {
+       "lhs", "result", "getI1SameShape($_self)">, ElementwiseMappable]> {
   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, SameOperandsAndResultShape,
-     TypesMatchWith<
+    [NoSideEffect, SameTypeOperands,
+     SameOperandsAndResultShape, TypesMatchWith<
        "result type has i1 element type and same shape as operands",
-       "lhs", "result", "getI1SameShape($_self)">]> {
+       "lhs", "result", "getI1SameShape($_self)">, ElementwiseMappable]> {
   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 TensorFromElementsOp : Std_Op<"tensor_from_elements", [
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // FPExtOp
 //===----------------------------------------------------------------------===//
 
-def FPExtOp : CastOp<"fpext">, Arguments<(ins AnyType:$in)> {
+def FPExtOp : ArithmeticCastOp<"fpext">, Arguments<(ins AnyType:$in)> {
   let summary = "cast from floating-point to wider floating-point";
   let description = [{
     Cast a floating-point value to a larger floating-point-typed value.
     The destination type must to be strictly wider than the source type.
     Only scalars are currently supported.
   }];
 
   let extraClassDeclaration = [{
     /// Return true if `a` and `b` are valid operand and result pairs for
     /// the operation.
     static bool areCastCompatible(Type a, Type b);
   }];
 
   let hasFolder = 0;
 }
 
 //===----------------------------------------------------------------------===//
 // FPToSIOp
 //===----------------------------------------------------------------------===//
 
-def FPToSIOp : CastOp<"fptosi">, Arguments<(ins AnyType:$in)> {
+def FPToSIOp : ArithmeticCastOp<"fptosi">, Arguments<(ins AnyType:$in)> {
   let summary = "cast from floating-point type to integer type";
   let description = [{
     Cast from a value interpreted as floating-point to the nearest (rounding
     towards zero) signed integer value.
   }];
 
   let extraClassDeclaration = [{
     /// Return true if `a` and `b` are valid operand and result pairs for
     /// the operation.
     static bool areCastCompatible(Type a, Type b);
   }];
 
   let hasFolder = 0;
 }
 
 //===----------------------------------------------------------------------===//
 // FPToUIOp
 //===----------------------------------------------------------------------===//
 
-def FPToUIOp : CastOp<"fptoui">, Arguments<(ins AnyType:$in)> {
+def FPToUIOp : ArithmeticCastOp<"fptoui">, Arguments<(ins AnyType:$in)> {
   let summary = "cast from floating-point type to integer type";
   let description = [{
     Cast from a value interpreted as floating-point to the nearest (rounding
     towards zero) unsigned integer value.
   }];
 
   let extraClassDeclaration = [{
     /// Return true if `a` and `b` are valid operand and result pairs for
     /// the operation.
     static bool areCastCompatible(Type a, Type b);
   }];
 
   let hasFolder = 0;
 }
 
 //===----------------------------------------------------------------------===//
 // FPTruncOp
 //===----------------------------------------------------------------------===//
 
-def FPTruncOp : CastOp<"fptrunc">, Arguments<(ins AnyType:$in)> {
+def FPTruncOp : ArithmeticCastOp<"fptrunc">, Arguments<(ins AnyType:$in)> {
   let summary = "cast from floating-point to narrower floating-point";
   let description = [{
     Truncate a floating-point value to a smaller floating-point-typed value.
     The destination type must be strictly narrower than the source type.
     If the value cannot be exactly represented, it is rounded using the default
     rounding mode. Only scalars are currently supported.
   }];
 
def ImOp : Std_Op<"im",
   // `ImOp` is fully verified by its traits.
   let verifier = ?;
 }
 
 //===----------------------------------------------------------------------===//
 // IndexCastOp
 //===----------------------------------------------------------------------===//
 
-def IndexCastOp : CastOp<"index_cast">, Arguments<(ins AnyType:$in)> {
+def IndexCastOp : ArithmeticCastOp<"index_cast">, Arguments<(ins AnyType:$in)> {
   let summary = "cast between index and integer types";
   let description = [{
     Casts between integer scalars and 'index' scalars. Index is an integer of
     platform-specific bit width. If casting to a wider integer, the value is
     sign-extended. If casting to a narrower integer, the value is truncated.
   }];
 
   let extraClassDeclaration = [{
def RsqrtOp : FloatUnaryOp<"rsqrt"> {
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SelectOp
 //===----------------------------------------------------------------------===//
 
 def SelectOp : Std_Op<"select", [NoSideEffect,
-     AllTypesMatch<["true_value", "false_value", "result"]>]> {
+     AllTypesMatch<["true_value", "false_value", "result"]>,
+     ElementwiseMappable]> {
   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_
def SignedShiftRightOp : IntArithmeticOp<"shift_right_signed"> {
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SignExtendIOp
 //===----------------------------------------------------------------------===//
 
 def SignExtendIOp : Std_Op<"sexti",
-    [NoSideEffect, SameOperandsAndResultShape]> {
+    [NoSideEffect, SameOperandsAndResultShape, ElementwiseMappable]> {
   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.
 
let printer = [{
     return printStandardCastOp(this->getOperation(), p);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SIToFPOp
 //===----------------------------------------------------------------------===//
 
-def SIToFPOp : CastOp<"sitofp">, Arguments<(ins AnyType:$in)> {
+def SIToFPOp : ArithmeticCastOp<"sitofp">, Arguments<(ins AnyType:$in)> {
   let summary = "cast from integer type to floating-point";
   let description = [{
     Cast from a value interpreted as signed or vector of signed integers to the
     corresponding floating-point scalar or vector value. If the value cannot be
     exactly represented, it is rounded using the default rounding mode. Scalars
     and vector types are currently supported.
   }];
 
def TransposeOp : Std_Op<"transpose", [NoSideEffect]>,
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TruncateIOp
 //===----------------------------------------------------------------------===//
 
-def TruncateIOp : Std_Op<"trunci", [NoSideEffect, SameOperandsAndResultShape]> {
+def TruncateIOp : Std_Op<"trunci", [NoSideEffect,
+                                    SameOperandsAndResultShape,
+                                    ElementwiseMappable]> {
   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:
let printer = [{
     return printStandardCastOp(this->getOperation(), p);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // UIToFPOp
 //===----------------------------------------------------------------------===//
 
-def UIToFPOp : CastOp<"uitofp">, Arguments<(ins AnyType:$in)> {
+def UIToFPOp : ArithmeticCastOp<"uitofp">, Arguments<(ins AnyType:$in)> {
   let summary = "cast from unsigned integer type to floating-point";
   let description = [{
     Cast from a value interpreted as unsigned integer or vector of unsigned
     integers to the corresponding scalar or vector floating-point value. If the
     value cannot be exactly represented, it is rounded using the default
     rounding mode. Scalars and vector types are currently supported.
   }];
 
def XOrOp : IntArithmeticOp<"xor", [Commutative]> {
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ZeroExtendIOp
 //===----------------------------------------------------------------------===//
 
-def ZeroExtendIOp : Std_Op<"zexti", [NoSideEffect, SameOperandsAndResultShape]> {
+def ZeroExtendIOp : Std_Op<"zexti", [NoSideEffect,
+                                     SameOperandsAndResultShape,
+                                     ElementwiseMappable]> {
   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.
+def ElementwiseMappable : NativeOpTrait<"ElementwiseMappable">;
 
 // 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);
 } // 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.
+///
+/// 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 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.
+///
+/// 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>)
+///
+/// // 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
+/// ```
+///
+/// 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.
+///
+/// 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.
+template <typename ConcrentType>
+struct ElementwiseMappable
+    : public TraitBase<ConcrentType, ElementwiseMappable> {
+  static LogicalResult verifyTrait(Operation *op) {
+    return ::mlir::OpTrait::impl::verifyElementwiseMappable(op);
+  }
+};
+
 } // end namespace OpTrait
 
 //===----------------------------------------------------------------------===//
 // Operation Definition classes
 //===----------------------------------------------------------------------===//
 
 /// This provides public APIs that all operations should have.  The template
 /// argument 'ConcreteType' should be the concrete type by CRTP and the others

mlir/lib/IR/Operation.cpp

if (region.getNumArguments() != 0) {
                << region.getRegionNumber() << " should have no arguments";
       else
         return op->emitOpError("region should have no arguments");
     }
   }
   return success();
 }
 
+/// Checks if two ShapedTypes are the same, ignoring the element type.
+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) {
+  auto isMappableType = [](Type type) {
+    return type.isa<VectorType>() || type.isa<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 the op only has scalar operand/result types, then we have nothing to
+  // check.
+  if (resultMappableTypes.empty() && operandMappableTypes.empty())
+    return success();
+
+  if (!resultMappableTypes.empty() && operandMappableTypes.empty())
+    return op->emitOpError("if a result is non-scalar, then at least one "
+                           "operand must be non-scalar");
+
+  assert(!operandMappableTypes.empty());
+
+  if (resultMappableTypes.empty())
+    return op->emitOpError("if an operand is non-scalar, then there must be at "
+                           "least one non-scalar result");
+
+  if (resultMappableTypes.size() != op->getNumResults())
+    return op->emitOpError(
+        "if an operand is non-scalar, then all results must be non-scalar");
+
+  auto mustMatchType = operandMappableTypes[0].cast<ShapedType>();
+  for (auto type :
+       llvm::concat<Type>(resultMappableTypes, operandMappableTypes)) {
+    if (!areSameShapedTypeIgnoringElementType(type.cast<ShapedType>(),
+                                              mustMatchType)) {
+      return op->emitOpError() << "all non-scalar operands/results must have "
+                                  "the same shape and base type: found "
+                               << type << " and " << mustMatchType;
+    }
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // 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/Dialect/Standard/invalid.mlir

 // RUN: mlir-opt -split-input-file %s -verify-diagnostics
 
 // CHECK-LABEL: test_index_cast_shape_error
 func @test_index_cast_shape_error(%arg0 : tensor<index>) -> tensor<2xi64> {
-  // expected-error @+1 {{requires the same shape for all operands and results}}
+  // expected-error @+1 {{all non-scalar operands/results must have the same shape and base type: found 'tensor<2xi64>' and 'tensor<index>'}}
   %0 = index_cast %arg0 : tensor<index> to tensor<2xi64>
   return %0 : tensor<2xi64>
 }
 
 // -----
 
 // CHECK-LABEL: test_index_cast_tensor_error
 func @test_index_cast_tensor_error(%arg0 : tensor<index>) -> i64 {
-  // expected-error @+1 {{requires the same shape for all operands and results}}
+  // expected-error @+1 {{if an operand is non-scalar, then there must be at least one non-scalar result}}
   %0 = index_cast %arg0 : tensor<index> to i64
   return %0 : i64
 }
 
 // -----
 
 func @dynamic_tensor_from_elements(%m : index)
     -> tensor<?x3x?xf32> {

mlir/test/IR/invalid-ops.mlir

^bb0(%cond : i1, %t : i32, %f : i64):
   // expected-error@+1 {{all of {true_value, false_value, result} have same type}}
   %r = "std.select"(%cond, %t, %f) : (i1, i32, i64) -> i32
 }
 
 // -----
 
 func @func_with_ops(vector<12xi1>, vector<42xi32>, vector<42xi32>) {
 ^bb0(%cond : vector<12xi1>, %t : vector<42xi32>, %f : vector<42xi32>):
-  // expected-error@+1 {{expected condition type to have the same shape as the result type, expected 'vector<42xi1>', but got 'vector<12xi1>'}}
+  // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'vector<42xi32>' and 'vector<12xi1>'}}
   %r = "std.select"(%cond, %t, %f) : (vector<12xi1>, vector<42xi32>, vector<42xi32>) -> vector<42xi32>
 }
 
 // -----
 
 func @func_with_ops(tensor<12xi1>, tensor<42xi32>, tensor<42xi32>) {
 ^bb0(%cond : tensor<12xi1>, %t : tensor<42xi32>, %f : tensor<42xi32>):
-  // expected-error@+1 {{expected condition type to have the same shape as the result type, expected 'tensor<42xi1>', but got 'tensor<12xi1>'}}
+  // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'tensor<42xi32>' and 'tensor<12xi1>'}}
   %r = "std.select"(%cond, %t, %f) : (tensor<12xi1>, tensor<42xi32>, tensor<42xi32>) -> tensor<42xi32>
 }
 
 // -----
 
 func @invalid_cmp_shape(%idx : () -> ()) {
   // expected-error@+1 {{'lhs' must be signless-integer-like, but got '() -> ()'}}
   %cmp = cmpi "eq", %idx, %idx : () -> ()
func @fpext_f32_to_i32(%arg0 : f32) {
   // expected-error@+1 {{are cast incompatible}}
   %0 = fpext %arg0 : f32 to i32
   return
 }
 
 // -----
 
 func @fpext_vec(%arg0 : vector<2xf16>) {
-  // expected-error@+1 {{requires the same shape for all operands and results}}
+  // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'vector<3xf32>' and 'vector<2xf16>'}}
   %0 = fpext %arg0 : vector<2xf16> to vector<3xf32>
   return
 }
 
 // -----
 
 func @fpext_vec_f32_to_f16(%arg0 : vector<2xf32>) {
   // expected-error@+1 {{are cast incompatible}}
func @fptrunc_f32_to_i32(%arg0 : f32) {
   // expected-error@+1 {{are cast incompatible}}
   %0 = fptrunc %arg0 : f32 to i32
   return
 }
 
 // -----
 
 func @fptrunc_vec(%arg0 : vector<2xf16>) {
-  // expected-error@+1 {{requires the same shape for all operands and results}}
+  // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'vector<3xf32>' and 'vector<2xf16>'}}
   %0 = fptrunc %arg0 : vector<2xf16> to vector<3xf32>
   return
 }
 
 // -----
 
 func @fptrunc_vec_f16_to_f32(%arg0 : vector<2xf16>) {
   // expected-error@+1 {{are cast incompatible}}

mlir/test/IR/traits.mlir

 
 func @failedSameOperandAndResultType_operand_result_mismatch(%t10 : tensor<10xf32>, %t20 : tensor<20xf32>) {
   // expected-error@+1 {{requires the same type for all operands and results}}
   "test.same_operand_and_result_type"(%t10, %t20) : (tensor<10xf32>, tensor<20xf32>) -> tensor<10xf32>
 }
 
 // -----
 
+func @failedElementwiseMappable_different_rankedness(%arg0: tensor<?xf32>, %arg1: tensor<*xf32>) {
+  // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'tensor<*xf32>' and 'tensor<?xf32>'}}
+  %0 = "test.elementwise_mappable"(%arg0, %arg1) : (tensor<?xf32>, tensor<*xf32>) -> tensor<*xf32>
+}
+
+// -----
+
+func @failedElementwiseMappable_different_rank(%arg0: tensor<?xf32>, %arg1: tensor<?x?xf32>) {
+  // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'tensor<?x?xf32>' and 'tensor<?xf32>'}}
+  %0 = "test.elementwise_mappable"(%arg0, %arg1) : (tensor<?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+}
+
+// -----
+
+func @failedElementwiseMappable_different_shape(%arg0: tensor<?xf32>, %arg1: tensor<5xf32>) {
+  // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'tensor<5xf32>' and 'tensor<?xf32>'}}
+  %0 = "test.elementwise_mappable"(%arg0, %arg1) : (tensor<?xf32>, tensor<5xf32>) -> tensor<?xf32>
+}
+
+// -----
+
+func @failedElementwiseMappable_different_base_type(%arg0: vector<2xf32>, %arg1: tensor<2xf32>) {
+  // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'tensor<2xf32>' and 'vector<2xf32>'}}
+  %0 = "test.elementwise_mappable"(%arg0, %arg1) : (vector<2xf32>, tensor<2xf32>) -> tensor<2xf32>
+  return
+}
+
+// -----
+
+func @failedElementwiseMappable_non_scalar_output(%arg0: vector<2xf32>) {
+  // expected-error@+1 {{if an operand is non-scalar, then there must be at least one non-scalar result}}
+  %0 = "test.elementwise_mappable"(%arg0) : (vector<2xf32>) -> f32
+  return
+}
+
+// -----
+
+func @failedElementwiseMappable_non_scalar_result_all_scalar_input(%arg0: f32) {
+  // expected-error@+1 {{if a result is non-scalar, then at least one operand must be non-scalar}}
+  %0 = "test.elementwise_mappable"(%arg0) : (f32) -> tensor<f32>
+  return
+}
+
+// -----
+
+func @failedElementwiseMappable_mixed_scalar_non_scalar_results(%arg0: tensor<10xf32>) {
+  // expected-error@+1 {{if an operand is non-scalar, then all results must be non-scalar}}
+  %0, %1 = "test.elementwise_mappable"(%arg0) : (tensor<10xf32>) -> (f32, tensor<10xf32>)
+  return
+}
+
+// -----
+
+func @failedElementwiseMappable_zero_results(%arg0: tensor<10xf32>) {
+  // expected-error@+1 {{if an operand is non-scalar, then there must be at least one non-scalar result}}
+  "test.elementwise_mappable"(%arg0) : (tensor<10xf32>) -> ()
+  return
+}
+
+// -----
+
+func @failedElementwiseMappable_zero_operands() {
+  // expected-error@+1 {{if a result is non-scalar, then at least one operand must be non-scalar}}
+  "test.elementwise_mappable"() : () -> (tensor<6xf32>)
+  return
+}
+
+// -----
+
+func @succeededElementwiseMappable(%arg0: vector<2xf32>) {
+  // Check that varying element types are allowed.
+  // CHECK: test.elementwise_mappable
+  %0 = "test.elementwise_mappable"(%arg0) : (vector<2xf32>) -> vector<2xf16>
+  return
+}
+
+// -----
+
 func @failedHasParent_wrong_parent() {
   "some.op"() ({
    // expected-error@+1 {{'test.child' op expects parent op 'test.parent'}}
     "test.child"() : () -> ()
   }) : () -> ()
 }
 
 // -----

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

 }
 
 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]> {
+  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>,
               ElementTypeIsPred<"y", F32>]>>,
       ElementTypeIs<"res", I16>]> {
   let arguments = (ins
     AnyShaped:$x, AnyShaped:$y);