fix DRR's optional operand support

[yaochengji]

Summary

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 EOp : NS_Op<"e_op", []> {
  let arguments = (ins
    AnyInteger: $any_integer,
    Optional<AnyInteger>: $optioanl_integer
  );
}

def FOp : NS_Op<"f_op", []> {
  let arguments = (ins
    AnyInteger: $any_integer,
    Optional<AnyInteger>: $optioanl_integer
  );
}

// CHECK: if (auto tmpOperand = (*arg0.begin())) {
// CHECK: if (auto tmpOperand = (arg1.size() > 0? (*arg1.begin()): nullptr)) {
def test5 : Pat<(EOp $arg0, $arg1), (FOp $arg0, $arg1)>;

Without the patch, tblgen would generate cpp code like

arg1 = castedOp0.getODSOperands(1);
tblgen_values.push_back(arg1)

and cause compilation error. Here arg1 is type ValueRange and tblgen_values is type SmallVector.

[bondhugula]

1. Please add a test case.

2. Could you update the commit summary to also describe what it's fixing?

[jpienaar]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1225

I don't know if this will work: if you have two optional args next to each other, the caller would not be able to differentiate which is specified. It would seem that passing in nullptr for those and handling it in the builder would work instead and keep ordering.

[yaochengji]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1225

Yeah, it won't work for two optional args. Let me figure out a more reasonable fix.

[yaochengji]

add an unittest and update summary.

[yaochengji]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1225

After studying the code, I find that a bottom-up change is required to make things right. You could take a look at OperationState::addOperands, it could not differentiate multi optional operands, either. Neither could the automatically generated Op' build method.

I think this patch could resolve some compilation problems of the previous version of generated cpp code. Making the tblgen handle all sort of optional operands situation is beyond its scope.

[yaochengji]

change nullptr to Value()

[bondhugula]

I just had superficial comments. Please do go with @jpienaar or one of the other reviewers' comments.

mlir/lib/TableGen/Pattern.cpp

   return isSubClassOf("EnumAttrCaseInfo");
 }
 
-bool DagLeaf::isStringAttr() const {
-  return isa<llvm::StringInit>(def);
-}
+bool DagLeaf::isStringAttr() const { return isa<llvm::StringInit>(def); }
 
 Constraint DagLeaf::getAsConstraint() const {
   assert((isOperandMatcher() || isAttrMatcher()) &&

     auto *operand = op->getArg(*argIndex).get<NamedTypeConstraint *>();
     // If this operand is variadic, then return a range. Otherwise, return the
     // value itself.
-    if (operand->isVariableLength()) {
+    if (operand->isVariadic()) {
       auto repl = formatv(fmt, name);
       LLVM_DEBUG(llvm::dbgs() << repl << " (VariadicOperand)\n");
       return std::string(repl);
+    } else if (operand->isOptional()) {
+      auto repl = formatv("({0}.size() > 0? (*{0}.begin()): nullptr)",
+                          formatv(fmt, name));
+      LLVM_DEBUG(llvm::dbgs() << repl << " (OptionalOperand)\n");
+      return std::string(repl);
     }
     auto repl = formatv(fmt, formatv("(*{0}.begin())", name));
     LLVM_DEBUG(llvm::dbgs() << repl << " (SingleOperand)\n");

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

 // CHECK: nativeCall(rewriter, odsLoc, (*v1.begin()), (*v2.begin()), (*v3.begin()), (*v4.begin()), (*v5.begin()), (*v6.begin()), (*v7.begin()), (*v8.begin()), (*v9.begin()), (*v10.begin()))
 def test4 : Pat<(DOp $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10),
                 (NativeBuilder $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10)>;
+
+// Tests dag operand indexing for ops with optional operand.
+// ---
+
+def EOp : NS_Op<"e_op", []> {
+  let arguments = (ins
+    AnyInteger: $any_integer,
+    Optional<AnyInteger>: $optioanl_integer
+  );
+}
+
+def FOp : NS_Op<"f_op", []> {
+  let arguments = (ins
+    AnyInteger: $any_integer,
+    Optional<AnyInteger>: $optioanl_integer
+  );
+}
+
+// CHECK: if (auto tmpOperand = (*arg0.begin())) {
+// CHECK: if (auto tmpOperand = (arg1.size() > 0? (*arg1.begin()): nullptr)) {
+def test5 : Pat<(EOp $arg0, $arg1), (FOp $arg0, $arg1)>;
 No newline at end of file

mlir/tools/mlir-tblgen/RewriterGen.cpp

       "if (auto tmpAttr = {1}) {\n"
       "  tblgen_attrs.emplace_back(rewriter.getIdentifier(\"{0}\"), "
       "tmpAttr);\n}\n";
+  const char *addOperandCmd = "if (auto tmpOperand = {0}) {{\n"
+                              "  tblgen_values.push_back(tmpOperand);\n}\n";
   for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {
     if (resultOp.getArg(argIndex).is<NamedAttribute *>()) {
       // The argument in the op definition.

       os << formatv("for (auto v: {0}) {{\n  tblgen_values.push_back(v);\n}\n",
                     range);
     } else {
-      os << formatv("tblgen_values.push_back(");
+      std::string tmpSymbol;
       if (node.isNestedDagArg(argIndex)) {
-        os << symbolInfoMap.getValueAndRangeUse(
-            childNodeNames.lookup(argIndex));
+        tmpSymbol =
+            symbolInfoMap.getValueAndRangeUse(childNodeNames.lookup(argIndex));
       } else {
         DagLeaf leaf = node.getArgAsLeaf(argIndex);
         auto symbol =
             symbolInfoMap.getValueAndRangeUse(node.getArgName(argIndex));
         if (leaf.isNativeCodeCall()) {
-          os << std::string(
+          tmpSymbol = std::string(
               tgfmt(leaf.getNativeCodeTemplate(), &fmtCtx.withSelf(symbol)));
         } else {
-          os << symbol;
+          tmpSymbol = symbol;
         }
       }
-      os << ");\n";
+      os << formatv(addOperandCmd, tmpSymbol);
     }
   }
 }