[mlir] Add a new `ConstantLike` trait to better identify operations that represent a "constant".

[rriddle]

Summary

The current mechanism for identifying is a bit hacky and extremely adhoc, i.e. we explicit check 1-result, 0-operand, no side-effect, and always foldable and then assume that this is a constant. Adding a trait adds structure to this, and makes checking for a constant much more efficient as we can guarantee that all of these things have already been verified.

Depends On D76019

[jpienaar]

Seems reasonable :) Thanks

mlir/include/mlir/IR/Matchers.h

51

Update comment?

mlir/include/mlir/IR/OpDefinition.h

916

What would happen if this wasn't true for the op? What error/assert/failure would be triggered?

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1551

I think we also are missing documentation on the ordering effect of traits and even discussion pending on if order should matter ...

[rriddle]

Resolve comments

[rriddle]

mlir/include/mlir/IR/OpDefinition.h

916

An assert in m_Constant(i.e. all of the current users) would fire when the folding fails.

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1551

Yes, I find the behavior surprising sometimes when my mental model conflicts with what the implementation is. I'll try to draft something in a followup for Traits.md.

mlir/examples/toy/Ch7/include/toy/Ops.td

 // Toy Operations
 //===----------------------------------------------------------------------===//
 
 // We define a toy operation by inheriting from our base 'Toy_Op' class above.
 // Here we provide the mnemonic and a list of traits for the operation. The
 // constant operation is marked as 'NoSideEffect' as it is a pure operation
 // and may be removed if dead.
 def ConstantOp : Toy_Op<"constant",
-    [NoSideEffect, DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
+    [ConstantLike, NoSideEffect,
+     DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
   // Provide a summary and description for this operation. This can be used to
   // auto-generate documentation of the operations within our dialect.
   let summary = "constant";
   let description = [{
     Constant operation turns a literal into an SSA value. The data is attached
     to the operation as an attribute. For example:
 
     ```mlir
def StructAccessOp : Toy_Op<"struct_access", [NoSideEffect]> {
   ];
 
   let verifier = [{ return ::verify(*this); }];
 
   // Set the folder bit so that we can fold constant accesses.
   let hasFolder = 1;
 }
 
-def StructConstantOp : Toy_Op<"struct_constant", [NoSideEffect]> {
+def StructConstantOp : Toy_Op<"struct_constant", [ConstantLike, NoSideEffect]> {
   let summary = "struct constant";
   let description = [{
     Constant operation turns a literal struct value into an SSA value. The data
     is attached to the operation as an attribute. The struct constant is encoded
     as an array of other constant values. For example:
 
     ```mlir
       %0 = toy.struct_constant [

mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td

def SPV_AddressOfOp : SPV_Op<"_address_of", [InFunctionScope, NoSideEffect]> {
   let builders = [OpBuilder<[{Builder *builder, OperationState &state,
                               spirv::GlobalVariableOp var}]>];
 
   let assemblyFormat = "$variable attr-dict `:` type($pointer)";
 }
 
 // -----
 
-def SPV_ConstantOp : SPV_Op<"constant", [NoSideEffect]> {
+def SPV_ConstantOp : SPV_Op<"constant", [ConstantLike, NoSideEffect]> {
   let summary = "The op that declares a SPIR-V normal constant";
 
   let description = [{
     This op declares a SPIR-V normal constant. SPIR-V has multiple constant
     instructions covering different constant types:
 
     * `OpConstantTrue` and `OpConstantFalse` for boolean constants
     * `OpConstant` for scalar constants

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

def CondBranchOp : Std_Op<"cond_br",
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
 def ConstantOp : Std_Op<"constant",
-    [NoSideEffect, DeclareOpInterfaceMethods<OpAsmOpInterface>]> {
+    [ConstantLike, NoSideEffect, DeclareOpInterfaceMethods<OpAsmOpInterface>]> {
   let summary = "constant";
 
   let arguments = (ins AnyAttr:$value);
   let results = (outs AnyType);
 
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, Attribute value",
     [{ build(builder, result, value.getType(), value); }]>];

mlir/include/mlir/IR/Matchers.h

bool match(const Attribute &attr) {
     if (auto intAttr = attr.dyn_cast<AttrClass>()) {
       *bind_value = intAttr.getValue();
       return true;
     }
     return false;
   }
 };
 
-/// The matcher that matches a constant foldable operation that has no side
-/// effect, no operands and produces a single result.
+/// The matcher that matches operations that have the `ConstantLike` trait.
+struct constant_op_matcher {
+  bool match(Operation *op) { return op->hasTrait<OpTrait::ConstantLike>(); }
+};
+
+/// The matcher that matches operations that have the `ConstantLike` trait, and
+/// binds the folded attribute value.
 template <typename AttrT> struct constant_op_binder {
   AttrT *bind_value;
 
   /// Creates a matcher instance that binds the constant attribute value to
   /// bind_value if match succeeds.
   constant_op_binder(AttrT *bind_value) : bind_value(bind_value) {}
   /// Creates a matcher instance that doesn't bind if match succeeds.
   constant_op_binder() : bind_value(nullptr) {}
 
   bool match(Operation *op) {
-    if (op->getNumOperands() > 0 || op->getNumResults() != 1)
-      return false;
-    if (!op->hasNoSideEffect())
+    if (!op->hasTrait<OpTrait::ConstantLike>())
       return false;
 
+    // Fold the constant to an attribute.
     SmallVector<OpFoldResult, 1> foldedOp;
-    if (succeeded(op->fold(/*operands=*/llvm::None, foldedOp))) {
-      if (auto attr = foldedOp.front().dyn_cast<Attribute>()) {
-        if (auto attrT = attr.dyn_cast<AttrT>()) {
-          if (bind_value)
-            *bind_value = attrT;
-          return true;
-        }
-      }
-    }
+    LogicalResult result = op->fold(/*operands=*/llvm::None, foldedOp);
+    (void)result;
+    assert(succeeded(result) && "expected constant to be foldable");
+
+    if (auto attr = foldedOp.front().get<Attribute>().dyn_cast<AttrT>()) {
+        *bind_value = attr;
     return false;
   }
 };
 
 /// The matcher that matches a constant scalar / vector splat / tensor splat
 /// integer operation and binds the constant integer value.
 struct constant_int_op_binder {
   IntegerAttr::ValueType *bind_value;
bool match(Operation *op) {
     return res;
   }
   std::tuple<OperandMatchers...> operandMatchers;
 };
 
 } // end namespace detail
 
 /// Matches a constant foldable operation.
-inline detail::constant_op_binder<Attribute> m_Constant() {
-  return detail::constant_op_binder<Attribute>();
+inline detail::constant_op_matcher m_Constant() {
+  return detail::constant_op_matcher();
 }
 
 /// Matches a value from a constant foldable operation and writes the value to
 /// bind_value.
 template <typename AttrT>
 inline detail::constant_op_binder<AttrT> m_Constant(AttrT *bind_value) {
   return detail::constant_op_binder<AttrT>(bind_value);
 }

mlir/include/mlir/IR/OpBase.td

 
 // Op supports operand broadcast behavior.
 def ResultsBroadcastableShape :
   NativeOpTrait<"ResultsBroadcastableShape">;
 // TODO: Alias of the above, remove post integrate.
 def Broadcastable : NativeOpTrait<"ResultsBroadcastableShape">;
 // X op Y == Y op X
 def Commutative  : NativeOpTrait<"IsCommutative">;
+// Op behaves like a constant.
+def ConstantLike : NativeOpTrait<"ConstantLike">;
 // Op behaves like a function.
 def FunctionLike : NativeOpTrait<"FunctionLike">;
 // Op is isolated from above.
 def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;
 // Op results are float or vectors/tensors thereof.
 def ResultsAreFloatLike : NativeOpTrait<"ResultsAreFloatLike">;
 // Op has the same operand type.
 def SameTypeOperands : NativeOpTrait<"SameTypeOperands">;

mlir/include/mlir/IR/OpDefinition.h

 template <typename ConcreteType>
 class SameTypeOperands : public TraitBase<ConcreteType, SameTypeOperands> {
 public:
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifySameTypeOperands(op);
   }
 };
 
+/// This class provides the API for a sub-set of ops that are known to be
+/// constant-like. These are non-side effecting operations with one result and
+/// zero operands that can always be folded to a specific attribute value.
+template <typename ConcreteType>
+class ConstantLike : public TraitBase<ConcreteType, ConstantLike> {
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    static_assert(ConcreteType::template hasTrait<OneResult>(),
+                  "expected operation to produce one result");
+    static_assert(ConcreteType::template hasTrait<ZeroOperands>(),
+                  "expected operation to take zero operands");
+    // TODO: We should verify that the operation can always be folded, but this
+    // requires that the attributes of the op already be verified. We should add
+    // support for verifying traits "after" the operation to enable this use
+    // case.
+    return success();
+  }
+};
+
 /// This class provides the API for ops that are known to be isolated from
 /// above.
 template <typename ConcreteType>
 class IsIsolatedFromAbove
     : public TraitBase<ConcreteType, IsIsolatedFromAbove> {
 public:
   static AbstractOperation::OperationProperties getTraitProperties() {
     return static_cast<AbstractOperation::OperationProperties>(

mlir/lib/IR/Builders.cpp

for (auto &it : llvm::enumerate(foldResults)) {
     auto *constOp = dialect->materializeConstant(
         cstBuilder, attr, op->getResult(it.index()).getType(), op->getLoc());
     if (!constOp) {
       // Erase any generated constants.
       for (Operation *cst : generatedConstants)
         cst->erase();
       return cleanupFailure();
     }
-    assert(matchPattern(constOp, m_Constant(&attr)));
+    assert(matchPattern(constOp, m_Constant()));
 
     generatedConstants.push_back(constOp);
     results.push_back(constOp->getResult(0));
   }
 
   // If we were successful, insert any generated constants.
   for (Operation *cst : generatedConstants)
     insert(cst);
 
   return success();
 }

mlir/lib/Transforms/Utils/FoldUtils.cpp

static Operation *materializeConstant(Dialect *dialect, OpBuilder &builder,
                                       Attribute value, Type type,
                                       Location loc) {
   auto insertPt = builder.getInsertionPoint();
   (void)insertPt;
 
   // Ask the dialect to materialize a constant operation for this value.
   if (auto *constOp = dialect->materializeConstant(builder, value, type, loc)) {
     assert(insertPt == builder.getInsertionPoint());
-    assert(matchPattern(constOp, m_Constant(&value)));
+    assert(matchPattern(constOp, m_Constant()));
     return constOp;
   }
 
   // If the dialect is unable to materialize a constant, check to see if the
   // standard constant can be used.
   if (ConstantOp::isBuildableWith(value, type))
     return builder.create<ConstantOp>(loc, type, value);
   return nullptr;

mlir/test/Dialect/Linalg/invalid.mlir

%0 = linalg.generic {
     ^bb(%i: f32):
       linalg.yield %i: f32
   }: memref<?xf32, affine_map<(i)[off]->(off + i)>> -> f32
 }
 
 // -----
 
 func @generic_fun_result_0_element_type(%arg0: memref<?xf32>) {
-  // expected-error @+1 {{'linalg.dot' op expected 3 or more operands}}
+  // expected-error @+1 {{'linalg.dot' op expected 3 operands, but found 2}}
   linalg.dot(%arg0, %arg0): memref<?xf32>, memref<?xf32>
 }
 
 // -----
 
 // expected-error @+1 {{unknown Linalg type}}
 !invalid_type = type !linalg.unknown
 

mlir/test/IR/traits.mlir

 func @failedSameOperandElementType(%t1f: tensor<1xf32>, %t1i: tensor<1xi32>) {
   // expected-error@+1 {{requires the same element type for all operands}}
   "test.same_operand_element_type"(%t1f, %t1i) : (tensor<1xf32>, tensor<1xi32>) -> tensor<1xf32>
 }
 
 // -----
 
 func @failedSameOperandAndResultElementType_no_operands() {
-  // expected-error@+1 {{expected 1 or more operands}}
+  // expected-error@+1 {{expected 2 operands, but found 0}}
   "test.same_operand_element_type"() : () -> tensor<1xf32>
 }
 
 // -----
 
 func @failedSameOperandElementType_scalar_type_mismatch(%si: i32, %sf: f32) {
   // expected-error@+1 {{requires the same element type for all operands}}
   "test.same_operand_element_type"(%sf, %si) : (f32, i32) -> tensor<1xf32>

mlir/test/mlir-tblgen/op-decl.td

 // CHECK:   AOpOperandAdaptor(ArrayRef<Value> values);
 // CHECK:   ArrayRef<Value> getODSOperands(unsigned index);
 // CHECK:   Value a();
 // CHECK:   ArrayRef<Value> b();
 // CHECK: private:
 // CHECK:   ArrayRef<Value> tblgen_operands;
 // CHECK: };
 
-// CHECK: class AOp : public Op<AOp, OpTrait::AtLeastNResults<1>::Impl, OpTrait::ZeroSuccessor, OpTrait::HasNoSideEffect, OpTrait::AtLeastNOperands<1>::Impl
+// CHECK: class AOp : public Op<AOp, OpTrait::AtLeastNResults<1>::Impl, OpTrait::ZeroSuccessor, OpTrait::AtLeastNOperands<1>::Impl, OpTrait::HasNoSideEffect
 // CHECK: public:
 // CHECK:   using Op::Op;
 // CHECK:   using OperandAdaptor = AOpOperandAdaptor;
 // CHECK:   static StringRef getOperationName();
 // CHECK:   Operation::operand_range getODSOperands(unsigned index);
 // CHECK:   Value a();
 // CHECK:   Operation::operand_range b();
 // CHECK:   Operation::result_range getODSResults(unsigned index);

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

void OpEmitter::genTraits() {
   // Add return size trait.
   addSizeCountTrait(opClass, "Result", numResults, numVariadicResults);
 
   // Add successor size trait.
   unsigned numSuccessors = op.getNumSuccessors();
   unsigned numVariadicSuccessors = op.getNumVariadicSuccessors();
   addSizeCountTrait(opClass, "Successor", numSuccessors, numVariadicSuccessors);
 
-  // Add the native and interface traits.
-  for (const auto &trait : op.getTraits()) {
-    if (auto opTrait = dyn_cast<tblgen::NativeOpTrait>(&trait))
-      opClass.addTrait(opTrait->getTrait());
-    else if (auto opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait))
-      opClass.addTrait(opTrait->getTrait());
-  }
-
   // Add variadic size trait and normal op traits.
   int numOperands = op.getNumOperands();
   int numVariadicOperands = op.getNumVariadicOperands();
 
   // Add operand size trait.
   if (numVariadicOperands != 0) {
     if (numOperands == numVariadicOperands)
       opClass.addTrait("OpTrait::VariadicOperands");
     else
       opClass.addTrait("OpTrait::AtLeastNOperands<" +
                        Twine(numOperands - numVariadicOperands) + ">::Impl");
   } else {
     switch (numOperands) {
     case 0:
       opClass.addTrait("OpTrait::ZeroOperands");
       break;
     case 1:
       opClass.addTrait("OpTrait::OneOperand");
       break;
     default:
       opClass.addTrait("OpTrait::NOperands<" + Twine(numOperands) + ">::Impl");
       break;
     }
   }
+
+  // Add the native and interface traits.
+  for (const auto &trait : op.getTraits()) {
+    if (auto opTrait = dyn_cast<tblgen::NativeOpTrait>(&trait))
+      opClass.addTrait(opTrait->getTrait());
+    else if (auto opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait))
+      opClass.addTrait(opTrait->getTrait());
+  }
 }
 
 void OpEmitter::genOpNameGetter() {
   auto &method = opClass.newMethod("StringRef", "getOperationName",
                                    /*params=*/"", OpMethod::MP_Static);
   method.body() << "  return \"" << op.getOperationName() << "\";\n";
 }