Add test for fix to tablegen for custom folders for ops that return a single

variadic result.

Add missing test for single line fix to `void OpEmitter::genFolderDecls()`
entitled "Fold away reduction over 0 dimensions."

PiperOrigin-RevId: 273880337

diff --git a/mlir/test/Transforms/test-canonicalize.mlir b/mlir/test/Transforms/test-canonicalize.mlir
index 7bae125..fdbfd59 100644 (file)
--- a/mlir/test/Transforms/test-canonicalize.mlir
+++ b/mlir/test/Transforms/test-canonicalize.mlir
@@ -27,3 +27,11 @@ func @remove_op_with_inner_ops_fold(%arg0 : i32) -> (i32) {
   }) : (i32) -> (i32)
   return %0 : i32
 }
+
+// CHECK-LABEL: func @remove_op_with_variadic_results_and_folder
+// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32, %[[ARG_1:[a-z0-9]*]]: i32)
+func @remove_op_with_variadic_results_and_folder(%arg0 : i32, %arg1 : i32) -> (i32, i32) {
+  // CHECK-NEXT: return %[[ARG_0]], %[[ARG_1]]
+  %0, %1 = "test.op_with_variadic_results_and_folder"(%arg0, %arg1) : (i32, i32) -> (i32, i32)
+  return %0, %1 : i32, i32
+}

diff --git a/mlir/test/lib/TestDialect/TestDialect.cpp b/mlir/test/lib/TestDialect/TestDialect.cpp
index 78a75f0..c8db796 100644 (file)
--- a/mlir/test/lib/TestDialect/TestDialect.cpp
+++ b/mlir/test/lib/TestDialect/TestDialect.cpp
@@ -235,6 +235,14 @@ OpFoldResult TestOpWithRegionFold::fold(ArrayRef<Attribute> operands) {
   return operand();
 }
 
+LogicalResult TestOpWithVariadicResultsAndFolder::fold(
+    ArrayRef<Attribute> operands, SmallVectorImpl<OpFoldResult> &results) {
+  for (Value *input : this->operands()) {
+    results.push_back(input);
+  }
+  return success();
+}
+
 SmallVector<Type, 2> mlir::OpWithInferTypeInterfaceOp::inferReturnTypes(
     llvm::Optional<Location> location, ArrayRef<Value *> operands,
     ArrayRef<NamedAttribute> attributes, ArrayRef<Region> regions) {

diff --git a/mlir/test/lib/TestDialect/TestOps.td b/mlir/test/lib/TestDialect/TestOps.td
index 68bae70..1122bea 100644 (file)
--- a/mlir/test/lib/TestDialect/TestOps.td
+++ b/mlir/test/lib/TestDialect/TestOps.td
@@ -485,6 +485,12 @@ def TestOpWithRegionFold : TEST_Op<"op_with_region_fold"> {
   let hasFolder = 1;
 }
 
+def TestOpWithVariadicResultsAndFolder: TEST_Op<"op_with_variadic_results_and_folder"> {
+  let arguments = (ins Variadic<I32>:$operands);
+  let results = (outs Variadic<I32>);
+  let hasFolder = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Test Patterns (Symbol Binding)