[llvm-branch-commits] [mlir] [mlir] Add the `useOpFoldResults` ODS dialect bit (PR #228770)

[via llvm-branch-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Victor Perez (victor-eds)

<details>
<summary>Changes</summary>

The new `OpFoldResults fold(FoldAdaptor)` form needs a matching ODS declaration. This patch adds the dialect bit `useOpFoldResults`. The bit is off by default. When a dialect sets it, `genFolderDecls` declares `OpFoldResults fold(FoldAdaptor adaptor)` for each op that does not have exactly one fixed result. An op with one fixed result keeps the `OpFoldResult` form. A dialect without the bit keeps the legacy declaration.

The `test` dialect sets the bit for testing purposes. Its three folds of ops that do not have exactly one fixed result now use the new form. The behavior of these folds does not change, with one exception: the fold of `TestOpWithVariadicResultsAndFolder` with no operands now fails. Before, it reported an in-place change, but it changed nothing (breaking the API). The greedy driver then put the op on the worklist again and folded it again, so canonicalize did not stop. A new test in `test-canonicalize.mlir` covers this case. Without this patch, canonicalize does not stop on the test file, and `timeout` ends it:

```console
$ timeout 120 mlir-opt mlir/test/Transforms/test-canonicalize.mlir -pass-pipeline='builtin.module(func.func(canonicalize))'
$ echo $?
124
```

The patch also adds test ops for the driver tests of later patches: a partial fold, a partial fold with an in-place change, a fold that keeps all results, own-result forwarding in a graph region, a partial fold with a replacement that does not materialize, an in-place change without the in-place mark, and an in-place fold of an op without results. No driver uses the partial result yet.

The new ODS test `op-fold-results.td` checks the three declarations.

Assisted by: Claude Code.

---
Full diff: https://github.com/llvm/llvm-project/pull/228770.diff


9 Files Affected:

- (modified) mlir/include/mlir/IR/DialectBase.td (+5) 
- (modified) mlir/include/mlir/TableGen/Dialect.h (+4) 
- (modified) mlir/lib/TableGen/Dialect.cpp (+4) 
- (modified) mlir/test/Transforms/test-canonicalize.mlir (+9) 
- (modified) mlir/test/lib/Dialect/Test/TestDialect.td (+1) 
- (modified) mlir/test/lib/Dialect/Test/TestOpDefs.cpp (+75-12) 
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+60-2) 
- (added) mlir/test/mlir-tblgen/op-fold-results.td (+107) 
- (modified) mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp (+2) 


``````````diff
diff --git a/mlir/include/mlir/IR/DialectBase.td b/mlir/include/mlir/IR/DialectBase.td
index da669c6ada639..c3a591a173e8f 100644
--- a/mlir/include/mlir/IR/DialectBase.td
+++ b/mlir/include/mlir/IR/DialectBase.td
@@ -99,6 +99,11 @@ class Dialect {
 
   // If this dialect can be extended at runtime with new operations or types.
   bit isExtensible = 0;
+
+  // If set, operations of this dialect that do not have exactly one fixed
+  // result declare `OpFoldResults fold(FoldAdaptor)` instead of the legacy
+  // `LogicalResult fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &)`.
+  bit useOpFoldResults = 0;
 }
 
 #endif // DIALECTBASE_TD
diff --git a/mlir/include/mlir/TableGen/Dialect.h b/mlir/include/mlir/TableGen/Dialect.h
index 68230196b17cc..3e77e1bfa4b17 100644
--- a/mlir/include/mlir/TableGen/Dialect.h
+++ b/mlir/include/mlir/TableGen/Dialect.h
@@ -92,6 +92,10 @@ class Dialect {
   /// operations or types.
   bool isExtensible() const;
 
+  /// Returns true if operations of this dialect that do not have exactly one
+  /// fixed result declare `OpFoldResults fold(FoldAdaptor)`.
+  bool useOpFoldResults() const;
+
   const llvm::DagInit *getDiscardableAttributes() const;
 
   const llvm::Record *getDef() const { return def; }
diff --git a/mlir/lib/TableGen/Dialect.cpp b/mlir/lib/TableGen/Dialect.cpp
index 84530bfa4f6fd..af24dacfb0f58 100644
--- a/mlir/lib/TableGen/Dialect.cpp
+++ b/mlir/lib/TableGen/Dialect.cpp
@@ -106,6 +106,10 @@ bool Dialect::isExtensible() const {
   return def->getValueAsBit("isExtensible");
 }
 
+bool Dialect::useOpFoldResults() const {
+  return def->getValueAsBit("useOpFoldResults");
+}
+
 const llvm::DagInit *Dialect::getDiscardableAttributes() const {
   return def->getValueAsDag("discardableAttrs");
 }
diff --git a/mlir/test/Transforms/test-canonicalize.mlir b/mlir/test/Transforms/test-canonicalize.mlir
index 81538e88beae4..e039fa1917b73 100644
--- a/mlir/test/Transforms/test-canonicalize.mlir
+++ b/mlir/test/Transforms/test-canonicalize.mlir
@@ -37,6 +37,15 @@ func.func @remove_op_with_variadic_results_and_folder(%arg0 : i32, %arg1 : i32)
   return %0, %1 : i32, i32
 }
 
+// Without operands, the fold replaces no result, so it fails.
+// CHECK-LABEL: func @keep_op_with_variadic_results_and_folder_no_operands
+func.func @keep_op_with_variadic_results_and_folder_no_operands() {
+  // CHECK-NEXT: "test.op_with_variadic_results_and_folder"() : () -> ()
+  // CHECK-NEXT: return
+  "test.op_with_variadic_results_and_folder"() : () -> ()
+  return
+}
+
 // CHECK-LABEL: func @test_commutative_multi
 // CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32, %[[ARG_1:[a-z0-9]*]]: i32)
 func.func @test_commutative_multi(%arg0: i32, %arg1: i32) -> (i32, i32) {
diff --git a/mlir/test/lib/Dialect/Test/TestDialect.td b/mlir/test/lib/Dialect/Test/TestDialect.td
index 37a263f1d10b8..04e6edf945dd9 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.td
+++ b/mlir/test/lib/Dialect/Test/TestDialect.td
@@ -24,6 +24,7 @@ def Test_Dialect : Dialect {
   let useDefaultTypePrinterParser = 0;
   let useDefaultAttributePrinterParser = 1;
   let isExtensible = 1;
+  let useOpFoldResults = 1;
   let dependentDialects = [
     "::mlir::DLTIDialect",
     "::mlir::bufferization::BufferizationDialect"
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 025fed9822317..722613b20dfe7 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -228,27 +228,90 @@ OpFoldResult TestOpConstant::fold(FoldAdaptor adaptor) { return getValue(); }
 // TestOpWithVariadicResultsAndFolder
 //===----------------------------------------------------------------------===//
 
-LogicalResult TestOpWithVariadicResultsAndFolder::fold(
-    FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) {
-  for (Value input : this->getOperands()) {
-    results.push_back(input);
-  }
-  return success();
+OpFoldResults TestOpWithVariadicResultsAndFolder::fold(FoldAdaptor adaptor) {
+  return getOperands();
 }
 
 //===----------------------------------------------------------------------===//
 // TestOpFoldUnmaterializable
 //===----------------------------------------------------------------------===//
 
-LogicalResult
-TestOpFoldUnmaterializable::fold(FoldAdaptor adaptor,
-                                 SmallVectorImpl<OpFoldResult> &results) {
+OpFoldResults TestOpFoldUnmaterializable::fold(FoldAdaptor adaptor) {
   // The unmaterializable replacement comes last, so that the rollback sees
   // both a materialized constant and an existing value.
   Builder b(getContext());
-  results.push_back(b.getI32IntegerAttr(42));
-  results.push_back(getOperand());
-  results.push_back(b.getStringAttr("unmaterializable"));
+  return {b.getI32IntegerAttr(42), getOperand(),
+          b.getStringAttr("unmaterializable")};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpPartialFold
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpPartialFold::fold(FoldAdaptor adaptor) {
+  Builder b(getContext());
+  return {b.getI32IntegerAttr(42), getOperand(), OpFoldResult()};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpPartialFoldInPlace
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpPartialFoldInPlace::fold(FoldAdaptor adaptor) {
+  if (getFolded())
+    return {};
+  setFolded(true);
+  OpFoldResults results(getOperation());
+  results.replace(getForwarded(), getOperand());
+  results.setModifiedInPlace();
+  return results;
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpFoldKeepAll
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpFoldKeepAll::fold(FoldAdaptor adaptor) {
+  return {getResult(0), getResult(1)};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpFoldForwardOperands
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpFoldForwardOperands::fold(FoldAdaptor adaptor) {
+  return {getFirst(), getSecond()};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpPartialFoldUnmaterializable
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpPartialFoldUnmaterializable::fold(FoldAdaptor adaptor) {
+  Builder b(getContext());
+  return {b.getI32IntegerAttr(42), b.getStringAttr("unmaterializable"),
+          OpFoldResult()};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpPartialFoldUnmarkedInPlace
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpPartialFoldUnmarkedInPlace::fold(FoldAdaptor adaptor) {
+  // The missing setModifiedInPlace() is the defect under test.
+  if (!getFolded())
+    setFolded(true);
+  return {getOperand(), OpFoldResult()};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpZeroResultsFoldInPlace
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpZeroResultsFoldInPlace::fold(FoldAdaptor adaptor) {
+  if (getFolded())
+    return failure();
+  setFolded(true);
   return success();
 }
 
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index bbe92cf651c2f..763d48343a625 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1685,6 +1685,65 @@ def TestOpFoldUnmaterializable : TEST_Op<"op_fold_unmaterializable"> {
   let hasFolder = 1;
 }
 
+// Op for testing a partial fold. The folder replaces result 0 with a constant
+// and result 1 with the operand, and keeps result 2.
+def TestOpPartialFold : TEST_Op<"op_partial_fold", [NoMemoryEffect]> {
+  let arguments = (ins I32:$operand);
+  let results = (outs I32, I32, I32);
+  let hasFolder = 1;
+}
+
+// Op for testing a partial fold together with an in-place change. The first
+// fold sets `folded`, marks the in-place change, replaces `forwarded` with the
+// operand, and keeps `kept`. Later folds fail.
+def TestOpPartialFoldInPlace : TEST_Op<"op_partial_fold_in_place"> {
+  let arguments = (ins I32:$operand, UnitAttr:$folded);
+  let results = (outs I32:$forwarded, I32:$kept);
+  let hasFolder = 1;
+}
+
+// Op for testing that a fold fails when it keeps every result and does not
+// change the op in place. The folder returns the results of the op.
+def TestOpFoldKeepAll : TEST_Op<"op_fold_keep_all"> {
+  let results = (outs I32, I32);
+  let hasFolder = 1;
+}
+
+// Op for testing a fold that forwards its operands. In a graph region where
+// operand i is result i of the same op, the folder returns only the results of
+// the op, and the fold must fail.
+def TestOpFoldForwardOperands : TEST_Op<"op_fold_forward_operands"> {
+  let arguments = (ins I32:$first, I32:$second);
+  let results = (outs I32, I32);
+  let hasFolder = 1;
+}
+
+// Op for testing a partial fold that fails constant materialization. The
+// folder replaces result 0 with a constant, replaces result 1 with an attribute
+// that the test dialect cannot materialize, and keeps result 2.
+def TestOpPartialFoldUnmaterializable
+    : TEST_Op<"op_partial_fold_unmaterializable"> {
+  let results = (outs I32, I32, I32);
+  let hasFolder = 1;
+}
+
+// Op for testing the check for an in-place change without an in-place mark.
+// The first fold sets `folded` but does not mark the in-place change. Every
+// fold replaces `forwarded` with the operand and keeps `kept`.
+def TestOpPartialFoldUnmarkedInPlace
+    : TEST_Op<"op_partial_fold_unmarked_in_place"> {
+  let arguments = (ins I32:$operand, UnitAttr:$folded);
+  let results = (outs I32:$forwarded, I32:$kept);
+  let hasFolder = 1;
+}
+
+// Op for testing an in-place fold of an op without results. The first fold
+// sets `folded`. Later folds fail.
+def TestOpZeroResultsFoldInPlace : TEST_Op<"op_zero_results_fold_in_place"> {
+  let arguments = (ins UnitAttr:$folded);
+  let hasFolder = 1;
+}
+
 def TestAddIOp : TEST_Op<"addi"> {
   let arguments = (ins AnyTypeOf<[I32, TestI32]>:$op1,
                        AnyTypeOf<[I32, TestI32]>:$op2);
@@ -1767,8 +1826,7 @@ def TestOpInPlaceFoldSuccess : TEST_Op<"op_in_place_fold_success"> {
   let results = (outs Variadic<I1>);
   let hasFolder = 1;
   let extraClassDefinition = [{
-    ::llvm::LogicalResult $cppClass::fold(FoldAdaptor adaptor,
-        SmallVectorImpl<OpFoldResult> &results) {
+    ::mlir::OpFoldResults $cppClass::fold(FoldAdaptor adaptor) {
       return success();
     }
   }];
diff --git a/mlir/test/mlir-tblgen/op-fold-results.td b/mlir/test/mlir-tblgen/op-fold-results.td
new file mode 100644
index 0000000000000..3b5d06c4e6012
--- /dev/null
+++ b/mlir/test/mlir-tblgen/op-fold-results.td
@@ -0,0 +1,107 @@
+// RUN: mlir-tblgen -gen-op-decls -I %S/../../include %s | FileCheck %s
+
+// Check the `fold` declaration that ODS emits with and without the dialect bit
+// `useOpFoldResults`.
+
+include "mlir/IR/OpBase.td"
+
+def Legacy_Dialect : Dialect {
+  let name = "legacy";
+  let cppNamespace = "legacy";
+}
+class Legacy_Op<string mnemonic, list<Trait> traits = []> :
+    Op<Legacy_Dialect, mnemonic, traits>;
+
+def Results_Dialect : Dialect {
+  let name = "results";
+  let cppNamespace = "results";
+  let useOpFoldResults = 1;
+}
+class Results_Op<string mnemonic, list<Trait> traits = []> :
+    Op<Results_Dialect, mnemonic, traits>;
+
+//===----------------------------------------------------------------------===//
+// Dialect without the bit
+//===----------------------------------------------------------------------===//
+
+def Legacy_SingleResultOp : Legacy_Op<"single_result"> {
+  let results = (outs AnyType:$res);
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: legacy::SingleResultOp declarations
+// CHECK: ::mlir::OpFoldResult fold(FoldAdaptor adaptor);
+
+def Legacy_TwoResultsOp : Legacy_Op<"two_results"> {
+  let results = (outs AnyType:$lhs, AnyType:$rhs);
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: legacy::TwoResultsOp declarations
+// CHECK: ::llvm::LogicalResult fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult> &results);
+
+def Legacy_VariadicResultOp : Legacy_Op<"variadic_result"> {
+  let results = (outs Variadic<AnyType>:$res);
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: legacy::VariadicResultOp declarations
+// CHECK: ::llvm::LogicalResult fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult> &results);
+
+//===----------------------------------------------------------------------===//
+// Dialect with the bit
+//===----------------------------------------------------------------------===//
+
+def Results_MixedResultsOp : Results_Op<"mixed_results"> {
+  let results = (outs AnyType:$fixed, Variadic<AnyType>:$variadic);
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::MixedResultsOp declarations
+// CHECK-NOT: SmallVectorImpl<::mlir::OpFoldResult>
+// CHECK: ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
+
+def Results_OptionalResultOp : Results_Op<"optional_result"> {
+  let results = (outs Optional<AnyType>:$res);
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::OptionalResultOp declarations
+// CHECK-NOT: SmallVectorImpl<::mlir::OpFoldResult>
+// CHECK: ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
+
+// The bit does not change the declaration of an op with exactly one fixed
+// result.
+def Results_SingleResultOp : Results_Op<"single_result"> {
+  let results = (outs AnyType:$res);
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::SingleResultOp declarations
+// CHECK: ::mlir::OpFoldResult fold(FoldAdaptor adaptor);
+
+def Results_TwoResultsOp : Results_Op<"two_results"> {
+  let results = (outs AnyType:$lhs, AnyType:$rhs);
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::TwoResultsOp declarations
+// CHECK-NOT: SmallVectorImpl<::mlir::OpFoldResult>
+// CHECK: ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
+
+def Results_VariadicResultOp : Results_Op<"variadic_result"> {
+  let results = (outs Variadic<AnyType>:$res);
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::VariadicResultOp declarations
+// CHECK-NOT: SmallVectorImpl<::mlir::OpFoldResult>
+// CHECK: ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
+
+def Results_ZeroResultsOp : Results_Op<"zero_results"> {
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::ZeroResultsOp declarations
+// CHECK-NOT: SmallVectorImpl<::mlir::OpFoldResult>
+// CHECK: ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
diff --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index 48158b6b148aa..0f7665df45ecf 100644
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
@@ -3430,6 +3430,8 @@ void OpEmitter::genFolderDecls() {
       op.getNumResults() == 1 && op.getNumVariableLengthResults() == 0;
   if (hasSingleResult) {
     retType = "::mlir::OpFoldResult";
+  } else if (op.getDialect().useOpFoldResults()) {
+    retType = "::mlir::OpFoldResults";
   } else {
     paramList.emplace_back("::llvm::SmallVectorImpl<::mlir::OpFoldResult> &",
                            "results");

``````````

</details>


https://github.com/llvm/llvm-project/pull/228770