[DRR] Allow capturing and referencing no-result ops

Previously when we bind a symbol to an op in DRR, it means to capture
the op's result(s) and later references will be expanded to result(s).
This means for ops without result, we are replacing the symbol with
nothing. This CL treats non-result op capturing and referencing as a
special case to mean the op itself.

PiperOrigin-RevId: 275269702

diff --git a/mlir/lib/TableGen/Pattern.cpp b/mlir/lib/TableGen/Pattern.cpp
index ec9af4c..9db085a 100644 (file)
--- a/mlir/lib/TableGen/Pattern.cpp
+++ b/mlir/lib/TableGen/Pattern.cpp
@@ -269,6 +269,13 @@ std::string tblgen::SymbolInfoMap::SymbolInfo::getValueAndRangeUse(
       return 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 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;

diff --git a/mlir/test/lib/TestDialect/TestOps.td b/mlir/test/lib/TestDialect/TestOps.td
index 32d7755..d8a0e57 100644 (file)
--- a/mlir/test/lib/TestDialect/TestOps.td
+++ b/mlir/test/lib/TestDialect/TestOps.td
@@ -556,6 +556,14 @@ def : Pattern<
     // Use bound symbols in additional constraints
     [(HasOneUse $res_a)]>;
 
+def OpSymbolBindingNoResult : TEST_Op<"symbol_binding_no_result", []> {
+  let arguments = (ins I32:$operand);
+}
+
+// Test that we can bind to an op without results and reference it later.
+def : Pat<(OpSymbolBindingNoResult:$op $operand),
+          (NativeCodeCall<"handleNoResultOp($_builder, $0)"> $op)>;
+
 //===----------------------------------------------------------------------===//
 // Test Patterns (Attributes)
 

diff --git a/mlir/test/lib/TestDialect/TestPatterns.cpp b/mlir/test/lib/TestDialect/TestPatterns.cpp
index 5a88ce8..052d7b9 100644 (file)
--- a/mlir/test/lib/TestDialect/TestPatterns.cpp
+++ b/mlir/test/lib/TestDialect/TestPatterns.cpp
@@ -30,6 +30,12 @@ static void createOpI(PatternRewriter &rewriter, Value *input) {
   rewriter.create<OpI>(rewriter.getUnknownLoc(), input);
 }
 
+void handleNoResultOp(PatternRewriter &rewriter, OpSymbolBindingNoResult op) {
+  // Turn the no result op to a one-result op.
+  rewriter.create<OpSymbolBindingB>(op.getLoc(), op.operand()->getType(),
+                                    op.operand());
+}
+
 namespace {
 #include "TestPatterns.inc"
 } // end anonymous namespace

diff --git a/mlir/test/mlir-tblgen/pattern.mlir b/mlir/test/mlir-tblgen/pattern.mlir
index 004c5a0..7f0e1d1 100644 (file)
--- a/mlir/test/mlir-tblgen/pattern.mlir
+++ b/mlir/test/mlir-tblgen/pattern.mlir
@@ -83,6 +83,13 @@ func @symbolBinding(%arg0: i32) -> i32 {
   return %0: i32
 }
 
+// CHECK-LABEL: symbolBindingNoResult
+func @symbolBindingNoResult(%arg0: i32) {
+  // CHECK: test.symbol_binding_b
+  "test.symbol_binding_no_result"(%arg0) : (i32) -> ()
+  return
+}
+
 //===----------------------------------------------------------------------===//
 // Test Attributes
 //===----------------------------------------------------------------------===//