[mlir][drr] Allow specifying string in location

[jpienaar]

Summary

The string in the location is used to provide metadata for the fused location
or create a NamedLoc. This allows tagging individual locations to convey
additional rewrite information.

[antiagainst]

Cool! :)

mlir/docs/DeclarativeRewrites.md

681

Update example to reflect this change?

mlir/tools/mlir-tblgen/RewriterGen.cpp

688

Debugging leftover?

[jpienaar]

Adding example & removing leftover debugging

mlir/docs/DeclarativeRewrites.md

 directive to provide finer control.
 
 `location` is of the following syntax:
 
 ```tablgen
 (location $symbol0, $symbol1, ...)
 ```
 
-where all `$symbol` should be bound previously in the pattern.
+where all `$symbol` should be bound previously in the pattern and one optional
+string may be specified as an attribute. The following locations are creted:
+
+*   If only 1 symbol is specified then that symbol's location is used,
+*   If multiple are specified then a fused location is created;
+*   If no symbol is specified then string must be specified and a NamedLoc is
+    created instead;
 
 `location` must be used as the last argument to an op creation. For example,
 
 ```tablegen
 def : Pat<(LocSrc1Op:$src1 (LocSrc2Op:$src2 ...),
-          (LocDst1Op (LocDst2Op ..., (location $src2)))>;
+          (LocDst1Op (LocDst2Op ..., (location $src2)), (location "outer"))>;
 ```
 
 In the above pattern, the generated `LocDst2Op` will use the matched location
-of `LocSrc2Op` while the root `LocDst1Op` node will still se the fused location
-of all source Ops.
+of `LocSrc2Op` while the root `LocDst1Op` node will used the named location
+`outer`.
 
 ### `replaceWithValue`
 
 The `replaceWithValue` directive is used to eliminate a matched op by replacing
 all of it uses with a captured value. It is of the following syntax:
 
 ```tablegen
 (replaceWithValue $symbol)

mlir/include/mlir/TableGen/Pattern.h

public:
   bool isNativeCodeCall() const;
 
   // Returns true if this DAG leaf is specifying a constant attribute.
   bool isConstantAttr() const;
 
   // Returns true if this DAG leaf is specifying an enum attribute case.
   bool isEnumAttrCase() const;
 
+  // Returns true if this DAG leaf is specifying a string attribute.
+  bool isStringAttr() const;
+
   // Returns this DAG leaf as a constraint. Asserts if fails.
   Constraint getAsConstraint() const;
 
   // Returns this DAG leaf as an constant attribute. Asserts if fails.
   ConstantAttr getAsConstantAttr() const;
 
   // Returns this DAG leaf as an enum attribute case.
   // Precondition: isEnumAttrCase()
   EnumAttrCase getAsEnumAttrCase() const;
 
   // Returns the matching condition template inside this DAG leaf. Assumes the
   // leaf is an operand/attribute matcher and asserts otherwise.
   std::string getConditionTemplate() const;
 
   // Returns the native code call template inside this DAG leaf.
   // Precondition: isNativeCodeCall()
   StringRef getNativeCodeTemplate() const;
 
+  // Returns the string associated with the leaf.
+  // Precondition: isStringAttr()
+  std::string getStringAttr() const;
+
   void print(raw_ostream &os) const;
 
 private:
   // Returns true if the TableGen Init `def` in this DagLeaf is a DefInit and
   // also a subclass of the given `superclass`.
   bool isSubClassOf(StringRef superclass) const;
 
   const llvm::Init *def;

mlir/lib/TableGen/Pattern.cpp

 bool tblgen::DagLeaf::isConstantAttr() const {
   return isSubClassOf("ConstantAttr");
 }
 
 bool tblgen::DagLeaf::isEnumAttrCase() const {
   return isSubClassOf("EnumAttrCaseInfo");
 }
 
+bool tblgen::DagLeaf::isStringAttr() const {
+  return isa<llvm::StringInit>(def) || isa<llvm::CodeInit>(def);
+}
+
 tblgen::Constraint tblgen::DagLeaf::getAsConstraint() const {
   assert((isOperandMatcher() || isAttrMatcher()) &&
          "the DAG leaf must be operand or attribute");
   return Constraint(cast<llvm::DefInit>(def)->getDef());
 }
 
 tblgen::ConstantAttr tblgen::DagLeaf::getAsConstantAttr() const {
   assert(isConstantAttr() && "the DAG leaf must be constant attribute");
std::string tblgen::DagLeaf::getConditionTemplate() const {
   return getAsConstraint().getConditionTemplate();
 }
 
 llvm::StringRef tblgen::DagLeaf::getNativeCodeTemplate() const {
   assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
   return cast<llvm::DefInit>(def)->getDef()->getValueAsString("expression");
 }
 
+std::string tblgen::DagLeaf::getStringAttr() const {
+  assert(isStringAttr() && "the DAG leaf must be string attribute");
+  return def->getAsUnquotedString();
+}
 bool tblgen::DagLeaf::isSubClassOf(StringRef superclass) const {
   if (auto *defInit = dyn_cast_or_null<llvm::DefInit>(def))
     return defInit->getDef()->isSubClassOf(superclass);
   return false;
 }
 
 void tblgen::DagLeaf::print(raw_ostream &os) const {
   if (def)

mlir/test/lib/Dialect/Test/TestOps.td

 // Test Patterns (Location)
 
 // Test that we can specify locations for generated ops.
 def : Pat<(TestLocationSrcOp:$res1
            (TestLocationSrcOp:$res2
             (TestLocationSrcOp:$res3 $input))),
           (TestLocationDstOp
             (TestLocationDstOp
-              (TestLocationDstOp $input, (location $res1))),
-      (location $res2, $res3))>;
+              (TestLocationDstOp $input, (location $res1)),
+              (location "named")),
+            (location "fused", $res2, $res3))>;
 
 //===----------------------------------------------------------------------===//
 // Test Legalization
 //===----------------------------------------------------------------------===//
 
 def Test_LegalizerEnum_Success : StrEnumAttrCase<"Success">;
 def Test_LegalizerEnum_Failure : StrEnumAttrCase<"Failure">;
 

mlir/test/mlir-tblgen/pattern.mlir

 
 // CHECK-LABEL: verifyDesignatedLoc
 func @verifyDesignatedLoc(%arg0 : i32) -> i32 {
   %0 = "test.loc_src"(%arg0) : (i32) -> i32 loc("loc3")
   %1 = "test.loc_src"(%0) : (i32) -> i32 loc("loc2")
   %2 = "test.loc_src"(%1) : (i32) -> i32 loc("loc1")
 
   // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc("loc1")
-  // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc(fused[
-  // CHECK-SAME: "loc1"
-  // CHECK-SAME: "loc3"
-  // CHECK-SAME: "loc2"
-  // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc(fused["loc2", "loc3"])
+  // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc("named")
+  // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc(fused<"fused">["loc2", "loc3"])
   return %1 : i32
 }
 
 // CHECK-LABEL: verifyZeroResult
 func @verifyZeroResult(%arg0 : i32) {
   // CHECK: "test.op_i"(%arg0) : (i32) -> ()
   "test.op_h"(%arg0) : (i32) -> ()
   return

mlir/tools/mlir-tblgen/RewriterGen.cpp

private:
 
   // Emits the C++ statement to replace the matched DAG with a value built via
   // calling native C++ code.
   std::string handleReplaceWithNativeCodeCall(DagNode resultTree);
 
   // Returns the symbol of the old value serving as the replacement.
   StringRef handleReplaceWithValue(DagNode tree);
 
-  // Returns the symbol of the value whose location to use.
-  std::string handleUseLocationOf(DagNode tree);
+  // Returns the location value to use.
+  std::string handleLocationDirective(DagNode tree);
 
   // Emits the C++ statement to build a new op out of the given DAG `tree` and
   // returns the variable name that this op is assigned to. If the root op in
   // DAG `tree` has a specified name, the created op will be assigned to a
   // variable of the given name. Otherwise, a unique name will be used as the
   // result value name.
   std::string handleOpCreation(DagNode tree, int resultIndex, int depth);
 
StringRef PatternEmitter::handleReplaceWithValue(DagNode tree) {
 
   if (!tree.getSymbol().empty()) {
     PrintFatalError(loc, "cannot bind symbol to replaceWithValue");
   }
 
   return tree.getArgName(0);
 }
 
-std::string PatternEmitter::handleUseLocationOf(DagNode tree) {
+std::string PatternEmitter::handleLocationDirective(DagNode tree) {
   assert(tree.isLocationDirective());
   auto lookUpArgLoc = [this, &tree](int idx) {
     const auto *const lookupFmt = "(*{0}.begin()).getLoc()";
     return symbolInfoMap.getAllRangeUse(tree.getArgName(idx), lookupFmt);
   };
 
-  if (tree.getNumArgs() != 1) {
-    std::string ret;
-    llvm::raw_string_ostream os(ret);
-    os << "rewriter.getFusedLoc({";
-    for (int i = 0, e = tree.getNumArgs(); i != e; ++i)
-      os << (i ? ", " : "") << lookUpArgLoc(i);
-    os << "})";
-    return os.str();
-  }
+  if (tree.getNumArgs() == 0)
+    llvm::PrintFatalError(
+        "At least one argument to location directive required");
 
   if (!tree.getSymbol().empty())
     PrintFatalError(loc, "cannot bind symbol to location");
 
+  if (tree.getNumArgs() == 1) {
+    DagLeaf leaf = tree.getArgAsLeaf(0);
+    if (leaf.isStringAttr())
+      return formatv("mlir::NameLoc::get(rewriter.getIdentifier(\"{0}\"), "
+                     "rewriter.getContext())",
+                     leaf.getStringAttr())
+          .str();
     return lookUpArgLoc(0);
   }
 
+  std::string ret;
+  llvm::raw_string_ostream os(ret);
+  std::string strAttr;
+  os << "rewriter.getFusedLoc({";
+  bool first = true;
+  for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
+    DagLeaf leaf = tree.getArgAsLeaf(i);
+    // Handle the optional string value.
+    if (leaf.isStringAttr()) {
+      if (!strAttr.empty())
+        llvm::PrintFatalError("Only one string attribute may be specified");
+      strAttr = leaf.getStringAttr();
+      continue;
+    }
+    os << (first ? "" : ", ") << lookUpArgLoc(i);
+    first = false;
+  }
+  os << "}";
+  if (!strAttr.empty()) {
+    os << ", rewriter.getStringAttr(\"" << strAttr << "\")";
+  }
+  os << ")";
+  return os.str();
+}
+
 std::string PatternEmitter::handleOpArgument(DagLeaf leaf,
                                              StringRef patArgName) {
   if (leaf.isConstantAttr()) {
     auto constAttr = leaf.getAsConstantAttr();
     return handleConstantAttr(constAttr.getAttribute(),
                               constAttr.getConstantValue());
   }
   if (leaf.isEnumAttrCase()) {
std::string PatternEmitter::handleOpCreation(DagNode tree, int resultIndex,
   auto numOpArgs = resultOp.getNumArgs();
   auto numPatArgs = tree.getNumArgs();
 
   // Get the location for this operation if explicitly provided.
   std::string locToUse;
   if (numPatArgs != 0) {
     if (auto lastArg = tree.getArgAsNestedDag(numPatArgs - 1))
       if (lastArg.isLocationDirective())
-        locToUse = handleUseLocationOf(lastArg);
+        locToUse = handleLocationDirective(lastArg);
   }
 
   auto inPattern = numPatArgs - !locToUse.empty();
   if (numOpArgs != inPattern) {
     PrintFatalError(loc,
                     formatv("resultant op '{0}' argument number mismatch: "
                             "{1} in pattern vs. {2} in definition",
                             resultOp.getOperationName(), inPattern, numOpArgs));