Don't count attributes when addressing operands.

[ijsung]

Summary

Fixes out-of-bound access in generated nested DAG rewriter matching code.

[jpienaar]

Could you add a test where this would have triggered OOB access?

[tpopp]

mlir/tools/mlir-tblgen/RewriterGen.cpp

372

It might be slightly easier to read this by having a nextOperand++ instead of this numPrevAttrs computation because then it's more like an iterator than arithmetic to understand and the value is only used in one place.

[ijsung]

This comment was removed by ijsung.

[ijsung]

This comment was removed by ijsung.

[ijsung]

Addressed comments and added a test.

PTAL

[ijsung]

Small test case indenting fix.

[tpopp]

mlir/test/mlir-tblgen/rewriter-indexing.td

2

This is an internal detail. Can you run this test through presubmit in the work codebase? We sometimes have to work around lit implementation details and it can take time, so if you find the test failing, I can start to work on the fix in advance.

40

Does this one gain anything? Ideally we avoid these ifdefs because they make it more confusing.

mlir/tools/mlir-tblgen/RewriterGen.cpp

383

Arg. I didn't realize you would have to increment nextOperand in two places anyway. Sorry for making you turn a 2 line change into a different 2 line change.

[ijsung]

Addressed more comments on the test.

[ijsung]

Thanks & PTAL

mlir/test/mlir-tblgen/rewriter-indexing.td

2

Ran through

blaze test //third_party/llvm/llvm-project/mlir/test/mlir-tblgen/...

and verified test working as expected.

40

Replaced with just one FileCheck run.

mlir/tools/mlir-tblgen/RewriterGen.cpp

383

Ack. It looks slightly better than old version (as there's no confusing subtraction) anyways.

[tpopp]

mlir/test/mlir-tblgen/rewriter-indexing.td

12–13

Maybe move these 2 lines to below the COp definition to make it clear that all the actual tests are only below that point.

46

Delete the newline to be consistent with spacing above.

[ijsung]

Addressed comments on the test.

[ijsung]

Thanks for reviewing!

[ijsung]

For the Author property please use the following:

Name: Ray (I-Jui) Sung
Email: [e-mail]

Thanks!

[jpienaar]

Thanks!

mlir/test/mlir-tblgen/rewriter-indexing.td

+// RUN: mlir-tblgen -gen-rewriters -I %S/../../include %s | FileCheck %s
+
+include "mlir/IR/OpBase.td"
+
+def Test_Dialect : Dialect {
+  let name = "test";
+}
+class NS_Op<string mnemonic, list<OpTrait> traits> :
+    Op<Test_Dialect, mnemonic, traits>;
+
+def AOp : NS_Op<"a_op", []> {
+  let arguments = (ins
+    AnyInteger:$any_integer
+  );
+
+  let results = (outs AnyInteger);
+}
+
+def BOp : NS_Op<"b_op", []> {
+  let arguments = (ins
+    AnyAttr: $any_attr,
+    AnyInteger
+  );
+}
+
+// Tests dag operand indexing for ops with mixed attr and operand.
+// ---
+
+def COp : NS_Op<"c_op", []> {
+  let arguments = (ins
+    AnyAttr: $any_attr1,
+    AnyInteger,
+    AnyAttr: $any_attr2,
+    AnyInteger
+  );
+}
+
+// Only operand 0 should be addressed during matching.
+// CHECK: struct test1 : public ::mlir::RewritePattern {
+// CHECK: castedOp0.getODSOperands(0).begin()).getDefiningOp()
+def test1 : Pat<(BOp $attr, (AOp $input)),
+                (BOp $attr, $input)>;
+
+// Only operand 0 and 1 should be addressed during matching.
+// CHECK: struct test2 : public ::mlir::RewritePattern {
+// CHECK: castedOp0.getODSOperands(0);
+// CHECK: castedOp0.getODSOperands(1).begin()).getDefiningOp()
+def test2 : Pat<(COp $attr1, $op1, $attr2, (AOp $op2)),
+                (BOp $attr1, $op2)>;

mlir/tools/mlir-tblgen/RewriterGen.cpp

PrintFatalError(loc, formatv("op '{0}' argument number mismatch: {1} in "
                                  op.getNumArgs()));
   }
 
   // If the operand's name is set, set to that variable.
   auto name = tree.getSymbol();
   if (!name.empty())
     os << formatv("{0} = {1};\n", name, castedName);
 
-  for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
+  for (int i = 0, e = tree.getNumArgs(), nextOperand = 0; i != e; ++i) {
     auto opArg = op.getArg(i);
     std::string argName = formatv("op{0}", depth + 1);
 
     // Handle nested DAG construct first
     if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       if (auto *operand = opArg.dyn_cast<NamedTypeConstraint *>()) {
         if (operand->isVariableLength()) {
           auto error = formatv("use nested DAG construct to match op {0}'s "
                                "variadic operand #{1} unsupported now",
                                op.getOperationName(), i);
           PrintFatalError(loc, error);
         }
       }
       os << "{\n";
 
+      // Attributes don't count for getODSOperands.
       os.indent() << formatv(
           "auto *{0} = "
           "(*{1}.getODSOperands({2}).begin()).getDefiningOp();\n",
-          argName, castedName, i);
+          argName, castedName, nextOperand++);
       emitMatch(argTree, argName, depth + 1);
       os << formatv("tblgen_ops[{0}] = {1};\n", ++opCounter, argName);
       os.unindent() << "}\n";
       continue;
     }
 
     // Next handle DAG leaf: operand or attribute
     if (opArg.is<NamedTypeConstraint *>()) {
+      // emitOperandMatch's argument indexing counts attributes.
       emitOperandMatch(tree, castedName, i, depth);
+      ++nextOperand;
     } else if (opArg.is<NamedAttribute *>()) {
       emitAttributeMatch(tree, opName, i, depth);
     } else {
       PrintFatalError(loc, "unhandled case when matching op");
     }
   }
   LLVM_DEBUG(llvm::dbgs() << "done emitting match for op '"
                           << op.getOperationName() << "' at depth " << depth