[mlir] Add FuncOp::eraseResults

[silvas]

Summary

I just found I needed this in an upcoming patch, and it seems generally
useful to have.

[silvas]

fix whitespace

[mehdi_amini]

mlir/include/mlir/IR/Function.h

71

C

mlir/include/mlir/IR/Function.h

public:
   LogicalResult verify();
 
   /// Erase a single argument at `argIndex`.
   void eraseArgument(unsigned argIndex) { eraseArguments({argIndex}); }
   /// Erases the arguments listed in `argIndices`.
   /// `argIndices` is allowed to have duplicates and can be in any order.
   void eraseArguments(ArrayRef<unsigned> argIndices);
 
+  /// Erase a single result at `resultIndex`.
+  void eraseResult(unsigned resultIndex) { eraseResults({resultIndex}); }
+  /// Erases the results listed in `resultIndices`.
+  /// `resultIndices` is allowed to have duplicates and can be in any order.
+  void eraseResults(ArrayRef<unsigned> resultIndices);
+
   /// Create a deep copy of this function and all of its blocks, remapping
   /// any operands that use values outside of the function using the map that is
   /// provided (leaving them alone if no entry is present). If the mapper
   /// contains entries for function arguments, these arguments are not included
   /// in the new function. Replaces references to cloned sub-values with the
   /// corresponding value that is copied, and adds those mappings to the mapper.
   FuncOp clone(BlockAndValueMapping &mapper);
   FuncOp clone();

mlir/lib/IR/Function.cpp

void FuncOp::eraseArguments(ArrayRef<unsigned> argIndices) {
   // We do this in reverse so that we erase later indices before earlier
   // indices, to avoid shifting the later indices.
   Block &entry = front();
   for (int i = 0; i < originalNumArgs; i++)
     if (shouldEraseArg(originalNumArgs - i - 1))
       entry.eraseArgument(originalNumArgs - i - 1);
 }
 
+void FuncOp::eraseResults(ArrayRef<unsigned> resultIndices) {
+  auto oldType = getType();
+  int originalNumResults = oldType.getNumResults();
+  llvm::BitVector eraseIndices(originalNumResults);
+  for (auto index : resultIndices)
+    eraseIndices.set(index);
+  auto shouldEraseResult = [&](int i) { return eraseIndices.test(i); };
+
+  // There are 2 things that need to be updated:
+  // - Function type.
+  // - Result attrs.
+
+  // Update the function type and result attrs.
+  SmallVector<Type, 4> newResultTypes;
+  SmallVector<MutableDictionaryAttr, 4> newResultAttrs;
+  for (int i = 0; i < originalNumResults; i++) {
+    if (shouldEraseResult(i))
+      continue;
+    newResultTypes.emplace_back(oldType.getResult(i));
+    newResultAttrs.emplace_back(getResultAttrDict(i));
+  }
+  setType(FunctionType::get(oldType.getInputs(), newResultTypes, getContext()));
+  setAllResultAttrs(newResultAttrs);
+}
+
 /// Clone the internal blocks from this function into dest and all attributes
 /// from this function to dest.
 void FuncOp::cloneInto(FuncOp dest, BlockAndValueMapping &mapper) {
   // Add the attributes of this function to dest.
   llvm::MapVector<Identifier, Attribute> newAttrs;
   for (auto &attr : dest.getAttrs())
     newAttrs.insert(attr);
   for (auto &attr : getAttrs())

mlir/test/IR/test-func-erase-result.mlir

+// RUN: mlir-opt %s -test-func-erase-result -split-input-file | FileCheck %s
+
+// CHECK: func @f(){{$}}
+// CHECK-NOT: attributes{{.*}}result
+func @f() -> (f32 {test.erase_this_result})
+
+// -----
+
+// CHECK: func @f() -> (f32 {test.A})
+// CHECK-NOT: attributes{{.*}}result
+func @f() -> (
+  f32 {test.erase_this_result},
+  f32 {test.A}
+)
+
+// -----
+
+// CHECK: func @f() -> (f32 {test.A})
+// CHECK-NOT: attributes{{.*}}result
+func @f() -> (
+  f32 {test.A},
+  f32 {test.erase_this_result}
+)
+
+// -----
+
+// CHECK: func @f() -> (f32 {test.A}, f32 {test.B})
+// CHECK-NOT: attributes{{.*}}result
+func @f() -> (
+  f32 {test.A},
+  f32 {test.erase_this_result},
+  f32 {test.B}
+)
+// -----
+
+// CHECK: func @f() -> (f32 {test.A}, f32 {test.B})
+// CHECK-NOT: attributes{{.*}}result
+func @f() -> (
+  f32 {test.A},
+  f32 {test.erase_this_result},
+  f32 {test.erase_this_result},
+  f32 {test.B}
+)
+
+// -----
+
+// CHECK: func @f() -> (f32 {test.A}, f32 {test.B}, f32 {test.C})
+// CHECK-NOT: attributes{{.*}}result
+func @f() -> (
+  f32 {test.A},
+  f32 {test.erase_this_result},
+  f32 {test.B},
+  f32 {test.erase_this_result},
+  f32 {test.C}
+)
+
+// -----
+
+// CHECK: func @f() -> (tensor<1xf32>, tensor<2xf32>, tensor<3xf32>)
+// CHECK-NOT: attributes{{.*}}result
+func @f() -> (
+  tensor<1xf32>,
+  f32 {test.erase_this_result},
+  tensor<2xf32>,
+  f32 {test.erase_this_result},
+  tensor<3xf32>
+)

mlir/test/lib/IR/TestFunc.cpp

for (FuncOp func : module.getOps<FuncOp>()) {
       // Reverse the order to test that unsorted index lists are handled
       // correctly.
       std::reverse(indicesToErase.begin(), indicesToErase.end());
       func.eraseArguments(indicesToErase);
     }
   }
 };
 
+/// This is a test pass for verifying FuncOp's eraseResult method.
+struct TestFuncEraseResult
+    : public PassWrapper<TestFuncEraseResult, OperationPass<ModuleOp>> {
+  void runOnOperation() override {
+    auto module = getOperation();
+
+    for (FuncOp func : module.getOps<FuncOp>()) {
+      SmallVector<unsigned, 4> indicesToErase;
+      for (auto resultIndex : llvm::seq<int>(0, func.getNumResults())) {
+        if (func.getResultAttr(resultIndex, "test.erase_this_result")) {
+          // Push back twice to test that duplicate indices are handled
+          // correctly.
+          indicesToErase.push_back(resultIndex);
+          indicesToErase.push_back(resultIndex);
+        }
+      }
+      // Reverse the order to test that unsorted index lists are handled
+      // correctly.
+      std::reverse(indicesToErase.begin(), indicesToErase.end());
+      func.eraseResults(indicesToErase);
+    }
+  }
+};
+
 /// This is a test pass for verifying FuncOp's setType method.
 struct TestFuncSetType
     : public PassWrapper<TestFuncSetType, OperationPass<ModuleOp>> {
   void runOnOperation() override {
     auto module = getOperation();
     SymbolTable symbolTable(module);
 
     for (FuncOp func : module.getOps<FuncOp>()) {
       auto sym = func.getAttrOfType<FlatSymbolRefAttr>("test.set_type_from");
       if (!sym)
         continue;
       func.setType(symbolTable.lookup<FuncOp>(sym.getValue()).getType());
     }
   }
 };
 } // end anonymous namespace
 
 namespace mlir {
 void registerTestFunc() {
-  PassRegistration<TestFuncEraseArg> pass("test-func-erase-arg",
-                                          "Test erasing func args.");
+  PassRegistration<TestFuncEraseArg>("test-func-erase-arg",
 
-  PassRegistration<TestFuncSetType> pass2("test-func-set-type",
-                                          "Test FuncOp::setType.");
+  PassRegistration<TestFuncEraseResult>("test-func-erase-result",
+                                        "Test erasing func results.");
+
+  PassRegistration<TestFuncSetType>("test-func-set-type",
 }
 } // namespace mlir