[mlir][Rewrite] Add support for using an operation with no results as location

[zero9178]

Summary

Prior to this patch, using an operation without any results as the location would result in the generation of invalid C++ code. It'd try to format using the result values, which would would end up being an empty string for an operation without any.
This patch fixes that issue by instead using getValueAndRangeUse which handles both ranges as well as the case for an op without any results.

[jpienaar]

Nice, thanks!

mlir/lib/TableGen/Pattern.cpp

if (index >= 0) {
       LLVM_DEBUG(llvm::dbgs() << repl << " (SingleResult)\n");
       return std::string(repl);
     }
 
     // If this op has no result at all but still we bind a symbol to it, it
     // means we want to capture the op itself.
     if (op->getNumResults() == 0) {
       LLVM_DEBUG(llvm::dbgs() << name << " (Op)\n");
-      return std::string(name);
+      return formatv(fmt, name);
     }
 
     // We are referencing all results of the multi-result op. A specific result
     // can either be a value or a range. Then join them with `separator`.
     SmallVector<std::string, 4> values;
     values.reserve(op->getNumResults());
 
     for (int i = 0, e = op->getNumResults(); i < e; ++i) {

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

def TestLocationSrcOp : TEST_Op<"loc_src"> {
   let results = (outs I32:$output);
 }
 
 def TestLocationDstOp : TEST_Op<"loc_dst", [SameOperandsAndResultType]> {
   let arguments = (ins I32:$input);
   let results = (outs I32:$output);
 }
 
+def TestLocationSrcNoResOp : TEST_Op<"loc_src_no_res"> {
+  let arguments = (ins I32:$input);
+  let results = (outs);
+}
+
+def TestLocationDstNoResOp : TEST_Op<"loc_dst_no_res"> {
+  let arguments = (ins I32:$input);
+  let results = (outs);
+}
+
 //===----------------------------------------------------------------------===//
 // Test Patterns
 //===----------------------------------------------------------------------===//
 
 def OpA : TEST_Op<"op_a"> {
   let arguments = (ins I32, I32Attr:$attr);
   let results = (outs I32);
 }
def : Pat<(TestLocationSrcOp:$res1
            (TestLocationSrcOp:$res2
             (TestLocationSrcOp:$res3 $input))),
           (TestLocationDstOp
             (TestLocationDstOp
               (TestLocationDstOp $input, (location $res1)),
               (location "named")),
             (location "fused", $res2, $res3))>;
 
+// Test that we can use the location of an op without results
+def : Pat<(TestLocationSrcNoResOp:$loc
+            (TestLocationSrcOp (TestLocationSrcOp $input))),
+          (TestLocationDstNoResOp $input, (location $loc))>;
+
 //===----------------------------------------------------------------------===//
 // Test Patterns (Type Builders)
 
 def SourceOp : TEST_Op<"source_op"> {
   let arguments = (ins AnyInteger:$arg, AnyI32Attr:$tag);
   let results = (outs AnyInteger);
 }
 

mlir/test/mlir-tblgen/pattern.mlir

func @verifyFusedLocs(%arg0 : i32) -> i32 {
   return %result : i32
 }
 
 // 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")
+  "test.loc_src_no_res"(%2) : (i32) -> () loc("loc4")
 
   // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc("loc1")
   // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc("named")
   // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc(fused<"fused">["loc2", "loc3"])
+  // CHECK: "test.loc_dst_no_res"({{.*}}) : (i32) -> () loc("loc4")
   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

StringRef PatternEmitter::handleReplaceWithValue(DagNode tree) {
   }
 
   return tree.getArgName(0);
 }
 
 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);
+    const auto *const lookupFmt = "{0}.getLoc()";
+    return symbolInfoMap.getValueAndRangeUse(tree.getArgName(idx), lookupFmt);
   };
 
   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");