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

373

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

1 ↗

(On Diff #307768)

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.

39 ↗

(On Diff #307768)

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

1 ↗

(On Diff #307768)

Ran through

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

and verified test working as expected.

39 ↗

(On Diff #307768)

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

11–12 ↗

(On Diff #308129)

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

45 ↗

(On Diff #308129)

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/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(), numPrevAttrs = 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";
 
+      // We need to subtract the number of attributes before this operand to get
+      // the index in the operand list.
       os.indent() << formatv(
           "auto *{0} = "
           "(*{1}.getODSOperands({2}).begin()).getDefiningOp();\n",
-          argName, castedName, i);
+          argName, castedName, i - numPrevAttrs);
       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(tree, castedName, i, depth);
     } else if (opArg.is<NamedAttribute *>()) {
       emitAttributeMatch(tree, opName, i, depth);
+      // getODSOperands do not count Attributes,
+      // so we need to track how many we have seen so far.
+      ++numPrevAttrs;
     } else {
       PrintFatalError(loc, "unhandled case when matching op");
     }
   }
   LLVM_DEBUG(llvm::dbgs() << "done emitting match for op '"
                           << op.getOperationName() << "' at depth " << depth
                           << '\n');
 }