[Mlir-commits] [mlir] 7b196f1 - [mlir][Rewrite] Add support for using an operation with no results as location

[Markus Böck]

Author: Markus Böck
Date: 2022-02-03T15:08:09+01:00
New Revision: 7b196f1b093b49e5ab00c0483a8cda48c05e3e98

URL: https://github.com/llvm/llvm-project/commit/7b196f1b093b49e5ab00c0483a8cda48c05e3e98
DIFF: https://github.com/llvm/llvm-project/commit/7b196f1b093b49e5ab00c0483a8cda48c05e3e98.diff

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

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.

Differential Revision: https://reviews.llvm.org/D118885

Added: 
    

Modified: 
    mlir/lib/TableGen/Pattern.cpp
    mlir/test/lib/Dialect/Test/TestOps.td
    mlir/test/mlir-tblgen/pattern.mlir
    mlir/tools/mlir-tblgen/RewriterGen.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/TableGen/Pattern.cpp b/mlir/lib/TableGen/Pattern.cpp
index e7df9af87f17c..9574e7081beea 100644
--- a/mlir/lib/TableGen/Pattern.cpp
+++ b/mlir/lib/TableGen/Pattern.cpp
@@ -324,7 +324,7 @@ std::string SymbolInfoMap::SymbolInfo::getValueAndRangeUse(
     // 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

diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 6e14b62a3e73c..efaf2e1bd96f9 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -690,6 +690,16 @@ def TestLocationDstOp : TEST_Op<"loc_dst", [SameOperandsAndResultType]> {
   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
 //===----------------------------------------------------------------------===//
@@ -1375,6 +1385,11 @@ def : Pat<(TestLocationSrcOp:$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)
 

diff  --git a/mlir/test/mlir-tblgen/pattern.mlir b/mlir/test/mlir-tblgen/pattern.mlir
index 65bf11d8d14c7..9d64865014308 100644
--- a/mlir/test/mlir-tblgen/pattern.mlir
+++ b/mlir/test/mlir-tblgen/pattern.mlir
@@ -15,10 +15,12 @@ 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
 }
 

diff  --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index e9ff5431a3ab0..7c18b5ef4969e 100644
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -1148,8 +1148,8 @@ StringRef PatternEmitter::handleReplaceWithValue(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);
+    const auto *const lookupFmt = "{0}.getLoc()";
+    return symbolInfoMap.getValueAndRangeUse(tree.getArgName(idx), lookupFmt);
   };
 
   if (tree.getNumArgs() == 0)