[mlir] Add $_op hook to RewriterGen.

[lucyrfox]

Summary

The $_op hook makes it possible to access information about the current operation in a TableGen pattern. This is already implemented in OpDefinitionsGen; this change adds it to RewriterGen as well.

This is needed for a lowering in the TF dialect.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[mehdi_amini]

Can you expand on the commit description of what is this about / why are you doing this?

(also is this something that a test can be written for?)

[lucyrfox]

Added tests.

[merge_guards_bot]

Unit tests: pass. 62430 tests passed, 0 failed and 845 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[antiagainst]

mlir/test/lib/TestDialect/TestOps.td

608

I think this is misleading at least, if not problematic. The current way to support this is that we capture both the op's result and its operands and then pass them in as positional parameters so CheckResultCount can be written as

def CheckResultCount : Constraint<CPred<"$0.getOperation()->getNumResults() == $1.getInt()">>;

and used as

[(CheckResultCount $op, $result_count)]

One thing to note that we can have multiple source ops matched (as the source pattern can match a DAG of ops) so it's unclear which matchec op $_op will point to, so it's misleading at least. I think the above explicitly capturing is a better way to handle this case.

[lucyrfox]

mlir/test/lib/TestDialect/TestOps.td

608

Thanks for the reply, Lei.

$0.getOperation() doesn't compile -- $0, or $op, is an OpResult -- but CheckResultCount can be written as:

def CheckResultCount : Constraint<CPred<"$0.getOwner()->getNumResults() == $1.getInt()">>;

This works for the use case I have in the TF dialect as well. --> I'll abandon this change as the $_op hook is not necessary.

mlir/test/lib/TestDialect/TestOps.td

def OpNativeCodeCall3 : TEST_Op<"native_code_call3"> {
   let arguments = (ins I32:$input);
   let results = (outs I32);
 }
 // Test that NativeCodeCall is not ignored if it is not used to directly
 // replace the matched root op.
 def : Pattern<(OpNativeCodeCall3 $input),
               [(NativeCodeCall<"createOpI($_builder, $0)"> $input), (OpK)]>;
 
+// Test that the $_op hook works.
+def OpHook : TEST_Op<"op_hook"> {
+  let arguments = (ins
+    I32:$input,
+    I64Attr:$result_count
+  );
+  let results = (outs I32);
+}
+
+def CheckResultCount : Constraint<CPred<"$_op.getOperation()->getNumResults() == $0.getInt()">>;
+
+def : Pat<(OpHook:$op $input, $result_count),
+            (OpNativeCodeCall3 $input),
+            [(CheckResultCount $result_count)]>;
+
 // Test AllAttrConstraintsOf.
 def OpAllAttrConstraint1 : TEST_Op<"all_attr_constraint_of1"> {
   let arguments = (ins I64ArrayAttr:$attr);
   let results = (outs I32);
 }
 def OpAllAttrConstraint2 : TEST_Op<"all_attr_constraint_of2"> {
   let arguments = (ins I64ArrayAttr:$attr);
   let results = (outs I32);

mlir/test/mlir-tblgen/pattern.mlir

 // CHECK-LABEL: verifyAuxiliaryNativeCodeCall
 func @verifyAuxiliaryNativeCodeCall(%arg0: i32) -> (i32) {
   // CHECK: test.op_i
   // CHECK: test.op_k
   %0 = "test.native_code_call3"(%arg0) : (i32) -> (i32)
   return %0 : i32
 }
 
+// CHECK-LABEL: verifyOpHookValid
+func @verifyOpHookValid(%arg0: i32) -> (i32) {
+  // CHECK: test.op_i
+  // CHECK: test.op_k
+  %0 = "test.op_hook"(%arg0) {result_count = 1} : (i32) -> (i32)
+  return %0 : i32
+}
+
+// CHECK-LABEL: verifyOpHookInvalid
+func @verifyOpHookInvalid(%arg0: i32) -> (i32) {
+  // The result count is 1, not 2, so lowering constraint should fail.
+  // CHECK: test.op_hook
+  %0 = "test.op_hook"(%arg0) {result_count = 2} : (i32) -> (i32)
+  return %0 : i32
+}
+
 // CHECK-LABEL: verifyAllAttrConstraintOf
 func @verifyAllAttrConstraintOf() -> (i32, i32, i32) {
   // CHECK: "test.all_attr_constraint_of2"
   %0 = "test.all_attr_constraint_of1"() {attr = [0, 1]} : () -> (i32)
   // CHECK: "test.all_attr_constraint_of1"
   %1 = "test.all_attr_constraint_of1"() {attr = [0, 2]} : () -> (i32)
   // CHECK: "test.all_attr_constraint_of1"
   %2 = "test.all_attr_constraint_of1"() {attr = [-1, 1]} : () -> (i32)

mlir/tools/mlir-tblgen/RewriterGen.cpp

 void PatternEmitter::emitMatchLogic(DagNode tree) {
   LLVM_DEBUG(llvm::dbgs() << "--- start emitting match logic ---\n");
   emitOpMatch(tree, 0);
 
   for (auto &appliedConstraint : pattern.getConstraints()) {
     auto &constraint = appliedConstraint.constraint;
     auto &entities = appliedConstraint.entities;
 
+    fmtCtx.withOp("op");
+
     auto condition = constraint.getConditionTemplate();
     auto cmd = "if (!({0})) return matchFailure();\n";
 
     if (isa<TypeConstraint>(constraint)) {
       auto self = formatv("({0}.getType())",
                           symbolInfoMap.getValueAndRangeUse(entities.front()));
       os.indent(4) << formatv(cmd,
                               tgfmt(condition, &fmtCtx.withSelf(self.str())));