[Mlir-commits] [mlir] [mlir] Add option to cloning for different results (PR #184202)

[llvmlistbot at llvm.org]

https://github.com/tdegioanni-nvidia updated https://github.com/llvm/llvm-project/pull/184202

>From 6b9eba81423cac15991bd0317e869c4d258f47f2 Mon Sep 17 00:00:00 2001
From: Theo Degioanni <tdegioanni at nvidia.com>
Date: Mon, 2 Mar 2026 19:34:06 +0100
Subject: [PATCH 1/2] [mlir] Add option to cloning for different results

Co-authored-by: zero9178 <markus.boeck02 at gmail.com>
---
 mlir/include/mlir/IR/Operation.h           | 37 +++++++++++++++++-----
 mlir/lib/IR/Operation.cpp                  | 31 ++++++++++++------
 mlir/unittests/IR/OperationSupportTest.cpp | 16 ++++++++++
 3 files changed, 67 insertions(+), 17 deletions(-)

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index b2019574a820d..5e40b7b865fe1 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -140,18 +140,20 @@ class alignas(8) Operation final
   /// * Whether cloning should recursively traverse into the regions of the
   ///   operation or not.
   /// * Whether cloning should also clone the operands of the operation.
+  /// * Whether to use different result types or clone them.
   class CloneOptions {
   public:
     /// Default constructs an option with all flags set to false. That means all
     /// parts of an operation that may optionally not be cloned, are not cloned.
     CloneOptions();
 
-    /// Constructs an instance with the clone regions and clone operands flags
-    /// set accordingly.
-    CloneOptions(bool cloneRegions, bool cloneOperands);
+    /// Constructs an instance with the options set accordingly.
+    CloneOptions(bool cloneRegions, bool cloneOperands,
+                 std::optional<SmallVector<Type>> resultTypes);
 
-    /// Returns an instance with all flags set to true. This is the default
-    /// when using the clone method and clones all parts of the operation.
+    /// Returns an instance such that all elements of the operation are cloned.
+    /// This is the default when using the clone method and clones all parts of
+    /// the operation.
     static CloneOptions all();
 
     /// Configures whether cloning should traverse into any of the regions of
@@ -172,11 +174,29 @@ class alignas(8) Operation final
     /// Returns whether operands should be cloned as well.
     bool shouldCloneOperands() const { return cloneOperandsFlag; }
 
+    /// Configures different result types to use for the cloned operation.
+    /// If an empty optional, the result types are cloned from the original
+    /// operation.
+    CloneOptions &withResultTypes(std::optional<SmallVector<Type>> resultTypes);
+
+    /// Returns true if the results are cloned from the operation.
+    bool shouldCloneResults() const { return !resultTypes.has_value(); }
+
+    /// Returns the result types that should be used for the created operation
+    /// or `defaultResultTypes` if none were set.
+    TypeRange resultTypesOr(TypeRange defaultResultTypes) const {
+      if (resultTypes)
+        return *resultTypes;
+      return defaultResultTypes;
+    }
+
   private:
     /// Whether regions should be cloned.
     bool cloneRegionsFlag : 1;
     /// Whether operands should be cloned.
     bool cloneOperandsFlag : 1;
+    /// New result types to use in the cloned operation.
+    std::optional<SmallVector<Type>> resultTypes;
   };
 
   /// Create a deep copy of this operation, remapping any operands that use
@@ -185,7 +205,8 @@ class alignas(8) Operation final
   /// sub-operations to the corresponding operation that is copied, and adds
   /// those mappings to the map.
   /// Optionally, one may configure what parts of the operation to clone using
-  /// the options parameter.
+  /// the options parameter. If parts of the operation (e.g. results or regions)
+  /// are not cloned, they will not appear in the mapper.
   ///
   /// Calling this method from multiple threads is generally safe if through the
   /// process of cloning no new uses of 'Value's from outside the operation are
@@ -194,8 +215,8 @@ class alignas(8) Operation final
   /// mapper, it is possible to avoid adding uses to outside operands by
   /// remapping them to 'Value's owned by the caller thread.
   Operation *clone(IRMapping &mapper,
-                   CloneOptions options = CloneOptions::all());
-  Operation *clone(CloneOptions options = CloneOptions::all());
+                   const CloneOptions &options = CloneOptions::all());
+  Operation *clone(const CloneOptions &options = CloneOptions::all());
 
   /// Create a partial copy of this operation without traversing into attached
   /// regions. The new operation will have the same number of regions as the
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index bf8a918641dfb..46030936a3910 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -674,13 +674,18 @@ InFlightDiagnostic Operation::emitOpError(const Twine &message) {
 //===----------------------------------------------------------------------===//
 
 Operation::CloneOptions::CloneOptions()
-    : cloneRegionsFlag(false), cloneOperandsFlag(false) {}
+    : cloneRegionsFlag(false), cloneOperandsFlag(false),
+      resultTypes(std::nullopt) {}
 
-Operation::CloneOptions::CloneOptions(bool cloneRegions, bool cloneOperands)
-    : cloneRegionsFlag(cloneRegions), cloneOperandsFlag(cloneOperands) {}
+Operation::CloneOptions::CloneOptions(
+    bool cloneRegions, bool cloneOperands,
+    std::optional<SmallVector<Type>> resultTypes)
+    : cloneRegionsFlag(cloneRegions), cloneOperandsFlag(cloneOperands),
+      resultTypes(resultTypes) {}
 
 Operation::CloneOptions Operation::CloneOptions::all() {
-  return CloneOptions().cloneRegions().cloneOperands();
+  return CloneOptions().cloneRegions().cloneOperands().withResultTypes(
+      std::nullopt);
 }
 
 Operation::CloneOptions &Operation::CloneOptions::cloneRegions(bool enable) {
@@ -693,6 +698,12 @@ Operation::CloneOptions &Operation::CloneOptions::cloneOperands(bool enable) {
   return *this;
 }
 
+Operation::CloneOptions &Operation::CloneOptions::withResultTypes(
+    std::optional<SmallVector<Type>> resultTypes) {
+  this->resultTypes = std::move(resultTypes);
+  return *this;
+}
+
 /// Create a deep copy of this operation but keep the operation regions empty.
 /// Operands are remapped using `mapper` (if present), and `mapper` is updated
 /// to contain the results. The `mapResults` flag specifies whether the results
@@ -711,7 +722,7 @@ Operation *Operation::cloneWithoutRegions() {
 /// them alone if no entry is present).  Replaces references to cloned
 /// sub-operations to the corresponding operation that is copied, and adds
 /// those mappings to the map.
-Operation *Operation::clone(IRMapping &mapper, CloneOptions options) {
+Operation *Operation::clone(IRMapping &mapper, const CloneOptions &options) {
   SmallVector<Value, 8> operands;
   SmallVector<Block *, 2> successors;
 
@@ -728,7 +739,8 @@ Operation *Operation::clone(IRMapping &mapper, CloneOptions options) {
     successors.push_back(mapper.lookupOrDefault(successor));
 
   // Create the new operation.
-  auto *newOp = create(getLoc(), getName(), getResultTypes(), operands, attrs,
+  auto *newOp = create(getLoc(), getName(),
+                       options.resultTypesOr(getResultTypes()), operands, attrs,
                        getPropertiesStorage(), successors, getNumRegions());
   mapper.map(this, newOp);
 
@@ -739,13 +751,14 @@ Operation *Operation::clone(IRMapping &mapper, CloneOptions options) {
   }
 
   // Remember the mapping of any results.
-  for (unsigned i = 0, e = getNumResults(); i != e; ++i)
-    mapper.map(getResult(i), newOp->getResult(i));
+  if (options.shouldCloneResults())
+    for (unsigned i = 0, e = getNumResults(); i != e; ++i)
+      mapper.map(getResult(i), newOp->getResult(i));
 
   return newOp;
 }
 
-Operation *Operation::clone(CloneOptions options) {
+Operation *Operation::clone(const CloneOptions &options) {
   IRMapping mapper;
   return clone(mapper, options);
 }
diff --git a/mlir/unittests/IR/OperationSupportTest.cpp b/mlir/unittests/IR/OperationSupportTest.cpp
index 9f3e7ed34a27d..73babc4c4fbb7 100644
--- a/mlir/unittests/IR/OperationSupportTest.cpp
+++ b/mlir/unittests/IR/OperationSupportTest.cpp
@@ -356,4 +356,20 @@ TEST(OperationEquivalenceTest, HashWorksWithFlags) {
   opWithProperty2->destroy();
 }
 
+TEST(OperationCloneTest, CloneWithDifferentResults) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  Operation *useOp = createOp(&context, std::nullopt, builder.getI32Type());
+  IRMapping map;
+  Operation *cloneOp = useOp->clone(
+      map, Operation::CloneOptions::all().withResultTypes(
+               SmallVector<Type>{builder.getI32Type(), builder.getI16Type()}));
+
+  ASSERT_EQ(cloneOp->getNumResults(), 2u);
+  EXPECT_EQ(cloneOp->getResult(0).getType(), builder.getI32Type());
+  EXPECT_EQ(cloneOp->getResult(1).getType(), builder.getI16Type());
+  EXPECT_FALSE(map.contains(useOp->getResult(0)));
+}
+
 } // namespace

>From 84164758fb50cf000f668a83a99eda1528acbf44 Mon Sep 17 00:00:00 2001
From: Theo Degioanni <tdegioanni at nvidia.com>
Date: Mon, 2 Mar 2026 19:43:46 +0100
Subject: [PATCH 2/2] update test

---
 mlir/unittests/IR/OperationSupportTest.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/unittests/IR/OperationSupportTest.cpp b/mlir/unittests/IR/OperationSupportTest.cpp
index 73babc4c4fbb7..fc567bc4a8048 100644
--- a/mlir/unittests/IR/OperationSupportTest.cpp
+++ b/mlir/unittests/IR/OperationSupportTest.cpp
@@ -360,7 +360,7 @@ TEST(OperationCloneTest, CloneWithDifferentResults) {
   MLIRContext context;
   Builder builder(&context);
 
-  Operation *useOp = createOp(&context, std::nullopt, builder.getI32Type());
+  Operation *useOp = createOp(&context, {}, builder.getI32Type());
   IRMapping map;
   Operation *cloneOp = useOp->clone(
       map, Operation::CloneOptions::all().withResultTypes(