[Mlir-commits] [mlir] [mlir-c] Add mlirValueReplaceUsesWithIf (PR #206544)

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/206544

>From 18161011c5fb8884ed05331b698d797ed6d88298 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Mon, 29 Jun 2026 10:31:27 -0700
Subject: [PATCH] [mlir-c] Add mlirValueReplaceUsesWithIf

---
 mlir/include/mlir-c/IR.h | 14 +++++++++++
 mlir/lib/CAPI/IR/IR.cpp  | 10 ++++++++
 mlir/test/CAPI/ir.c      | 52 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 311234e42df60..a81621ba4fef7 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -1123,6 +1123,20 @@ mlirValueReplaceAllUsesExcept(MlirValue of, MlirValue with,
                               intptr_t numExceptions,
                               MlirOperation *exceptions);
 
+/// Callback deciding whether a particular use should be replaced. It is passed
+/// the use as an MlirOpOperand (from which the owner operation, operand number
+/// and value can be queried) and the user-provided `userData`. Returns true to
+/// replace this use.
+typedef bool (*MlirOpOperandReplaceFilterCallback)(MlirOpOperand opOperand,
+                                                   void *userData);
+
+/// Replace uses of 'of' value with 'with' value, but only for the uses for
+/// which the `filter` callback returns true. `filter` must not be NULL.
+MLIR_CAPI_EXPORTED void
+mlirValueReplaceUsesWithIf(MlirValue of, MlirValue with,
+                           MlirOpOperandReplaceFilterCallback filter,
+                           void *userData);
+
 /// Gets the location of the value.
 MLIR_CAPI_EXPORTED MlirLocation mlirValueGetLocation(MlirValue v);
 
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index 94442e2be19a4..782e9d0372cb7 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -1218,6 +1218,16 @@ void mlirValueReplaceAllUsesExcept(MlirValue oldValue, MlirValue newValue,
   oldValueCpp.replaceAllUsesExcept(newValueCpp, exceptionSet);
 }
 
+void mlirValueReplaceUsesWithIf(MlirValue of, MlirValue with,
+                                MlirOpOperandReplaceFilterCallback filter,
+                                void *userData) {
+  assert(filter && "expected non-null filter callback");
+  unwrap(of).replaceUsesWithIf(unwrap(with),
+                               [filter, userData](OpOperand &operand) -> bool {
+                                 return filter(wrap(&operand), userData);
+                               });
+}
+
 MlirLocation mlirValueGetLocation(MlirValue v) {
   return wrap(unwrap(v).getLoc());
 }
diff --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c
index 57ae8b9a2819b..d77fbe4b63a85 100644
--- a/mlir/test/CAPI/ir.c
+++ b/mlir/test/CAPI/ir.c
@@ -2900,6 +2900,56 @@ int testDominanceInfo(MlirContext ctx) {
   return 0;
 }
 
+// Replace-uses filter that accepts only uses whose owner is an `arith.addi`.
+static bool useOwnerIsAddi(MlirOpOperand opOperand, void *userData) {
+  (void)userData;
+  MlirStringRef name =
+      mlirIdentifierStr(mlirOperationGetName(mlirOpOperandGetOwner(opOperand)));
+  return mlirStringRefEqual(name, mlirStringRefCreateFromCString("arith.addi"));
+}
+
+int testReplaceUsesWithIf(MlirContext ctx) {
+  fprintf(stderr, "@testReplaceUsesWithIf\n");
+  // CHECK-LABEL: @testReplaceUsesWithIf
+
+  mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("arith"));
+
+  const char *moduleStr = "func.func @f(%arg0: i32, %arg1: i32) -> i32 {\n"
+                          "  %0 = arith.addi %arg0, %arg0 : i32\n"
+                          "  %1 = arith.muli %arg0, %arg0 : i32\n"
+                          "  %2 = arith.addi %0, %1 : i32\n"
+                          "  return %2 : i32\n"
+                          "}\n";
+  MlirModule module =
+      mlirModuleCreateParse(ctx, mlirStringRefCreateFromCString(moduleStr));
+
+  MlirBlock moduleBody = mlirModuleGetBody(module);
+  MlirOperation funcOp = mlirBlockGetFirstOperation(moduleBody);
+  MlirRegion funcRegion = mlirOperationGetRegion(funcOp, 0);
+  MlirBlock funcBody = mlirRegionGetFirstBlock(funcRegion);
+  MlirValue arg0 = mlirBlockGetArgument(funcBody, 0);
+  MlirValue arg1 = mlirBlockGetArgument(funcBody, 1);
+  MlirOperation addOp = mlirBlockGetFirstOperation(funcBody);
+  MlirOperation mulOp = mlirOperationGetNextInBlock(addOp);
+
+  // Replace arg0 with arg1, but only in arith.addi ops.
+  mlirValueReplaceUsesWithIf(arg0, arg1, useOwnerIsAddi, NULL);
+
+  // The first addi now uses arg1 for both operands.
+  assert(mlirValueEqual(mlirOperationGetOperand(addOp, 0), arg1));
+  assert(mlirValueEqual(mlirOperationGetOperand(addOp, 1), arg1));
+
+  // The muli is not an addi, so its operands are untouched (still arg0).
+  assert(mlirValueEqual(mlirOperationGetOperand(mulOp, 0), arg0));
+  assert(mlirValueEqual(mlirOperationGetOperand(mulOp, 1), arg0));
+
+  mlirModuleDestroy(module);
+
+  // CHECK: testReplaceUsesWithIf: PASSED
+  fprintf(stderr, "testReplaceUsesWithIf: PASSED\n");
+  return 0;
+}
+
 int main(void) {
   MlirContext ctx = mlirContextCreate();
   registerAllUpstreamDialects(ctx);
@@ -2955,6 +3005,8 @@ int main(void) {
     return 19;
   if (testDominanceInfo(ctx))
     return 20;
+  if (testReplaceUsesWithIf(ctx))
+    return 21;
 
   // CHECK: DESTROY MAIN CONTEXT
   // CHECK: reportResourceDelete: resource_i64_blob