[MLIR][python bindings] implement `replace_all_uses_with` on `PyValue`

[makslevental]

Summary

This diff adds replace_all_uses_with (replaceAllUsesWith) to the python api for PyValue, through the CAPI (mlirValueReplaceAllUsesOfWith). Plus house-keeping (dead code). Originally there was some discussion about whether the LiveOperationMap needs to be updated but I think we can skirt issue that by just "replacing" and then leaving elimination for a final dce pass (via the PassManager).

[makslevental]

lint

[makslevental]

lint

[makslevental]

add newline

[makslevental]

docs plus more dce

[makslevental]

mlir/test/CAPI/ir.c

31–49

small factorization

1867–1875

I don't know why the order here is swapped - constOneValue is used in the previous block of tests (where dummy.op is created) so its first use should be dummy.op. Either the use_iterator is "confused" (not likely) or I am (likely).

[makslevental]

Ping

[mehdi_amini]

mlir/test/CAPI/ir.c

1867–1875

The use-list ordering is supposed to be an opaque implementation detail that can't be relied on.
In general they are added push_front to the use-list I think.

[makslevental]

mlir/test/CAPI/ir.c

1867–1875

So then mlirValueGetFirstUse should be something like mlirValueGetLatestUse?

[kstoimenov]

Looks like this broke our ASAN bot: https://lab.llvm.org/buildbot/#/builders/168/builds/13141. Could you please revert or fix forward?

Thanks!

[makslevental]

In D148816#4296272, @kstoimenov wrote:

Looks like this broke our ASAN bot: https://lab.llvm.org/buildbot/#/builders/168/builds/13141. Could you please revert or fix forward?

Thanks!

Sorry just seeing this now - will revert.

mlir/include/mlir-c/IR.h

 /// several times with consecutive chunks of the string.
 MLIR_CAPI_EXPORTED void
 mlirValuePrint(MlirValue value, MlirStringCallback callback, void *userData);
 
 /// Returns an op operand representing the first use of the value, or a null op
 /// operand if there are no uses.
 MLIR_CAPI_EXPORTED MlirOpOperand mlirValueGetFirstUse(MlirValue value);
 
+/// Replace all uses of 'of' value with the 'with' value, updating anything in
+/// the IR that uses 'of' to use the other value instead.  When this returns
+/// there are zero uses of 'of'.
+MLIR_CAPI_EXPORTED void mlirValueReplaceAllUsesOfWith(MlirValue of,
+                                                      MlirValue with);
+
 //===----------------------------------------------------------------------===//
 // OpOperand API.
 //===----------------------------------------------------------------------===//
 
 /// Returns whether the op operand is null.
 MLIR_CAPI_EXPORTED bool mlirOpOperandIsNull(MlirOpOperand opOperand);
 
 /// Returns the owner operation of an op operand.

mlir/lib/Bindings/Python/IRCore.cpp

 //===- IRModules.cpp - IR Submodules of pybind module ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "IRModule.h"
 
 #include "Globals.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
-#include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/Debug.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
-//#include "mlir-c/Registration.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <optional>
 #include <utility>
 
 namespace py = pybind11;
 using namespace mlir;

 static const char kValueDunderStrDocstring[] =
     R"(Returns the string form of the value.
 
 If the value is a block argument, this is the assembly form of its type and the
 position in the argument list. If the value is an operation result, this is
 equivalent to printing the operation that produced it.
 )";
 
+static const char kValueReplaceAllUsesWithDocstring[] =
+    R"(Replace all uses of value with the new value, updating anything in
+the IR that uses 'self' to use the other value instead.
+)";
+
 //------------------------------------------------------------------------------
 // Utilities.
 //------------------------------------------------------------------------------
 
 /// Helper for creating an @classmethod.
 template <class Func, typename... Args>
 py::object classmethod(Func f, Args... args) {
   py::object cf = py::cpp_function(f, args...);
py::class_<PyValue>(m, "Value", py::module_local())
             PyPrintAccumulator printAccum;
             printAccum.parts.append("Value(");
             mlirValuePrint(self.get(), printAccum.getCallback(),
                            printAccum.getUserData());
             printAccum.parts.append(")");
             return printAccum.join();
           },
           kValueDunderStrDocstring)
-      .def_property_readonly("type", [](PyValue &self) {
-        return PyType(self.getParentOperation()->getContext(),
+      .def_property_readonly("type",
+                             [](PyValue &self) {
+                               return PyType(
+                                   self.getParentOperation()->getContext(),
                                    mlirValueGetType(self.get()));
-      });
+                             })
+      .def(
+          "replace_all_uses_with",
+          [](PyValue &self, PyValue &with) {
+            mlirValueReplaceAllUsesOfWith(self.get(), with.get());
+          },
+          kValueReplaceAllUsesWithDocstring);
   PyBlockArgument::bind(m);
   PyOpResult::bind(m);
   PyOpOperand::bind(m);
 
   //----------------------------------------------------------------------------
   // Mapping of SymbolTable.
   //----------------------------------------------------------------------------
   py::class_<PySymbolTable>(m, "SymbolTable", py::module_local())

mlir/lib/CAPI/IR/IR.cpp

MlirOpOperand mlirValueGetFirstUse(MlirValue value) {
   if (cppValue.use_empty())
     return {};
 
   OpOperand *opOperand = cppValue.use_begin().getOperand();
 
   return wrap(opOperand);
 }
 
+void mlirValueReplaceAllUsesOfWith(MlirValue oldValue, MlirValue newValue) {
+  unwrap(oldValue).replaceAllUsesWith(unwrap(newValue));
+}
+
 //===----------------------------------------------------------------------===//
 // OpOperand API.
 //===----------------------------------------------------------------------===//
 
 bool mlirOpOperandIsNull(MlirOpOperand opOperand) { return !opOperand.ptr; }
 
 MlirOperation mlirOpOperandGetOwner(MlirOpOperand opOperand) {
   return wrap(unwrap(opOperand)->getOwner());

mlir/test/CAPI/ir.c

 
 #include <assert.h>
 #include <inttypes.h>
 #include <math.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 
+MlirValue makeConstantLiteral(MlirContext ctx, const char *literalStr,
+                              const char *typeStr) {
+  MlirLocation loc = mlirLocationUnknownGet(ctx);
+  char attrStr[50];
+  sprintf(attrStr, "%s : %s", literalStr, typeStr);
+  MlirAttribute literal =
+      mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString(attrStr));
+  MlirNamedAttribute valueAttr = mlirNamedAttributeGet(
+      mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("value")), literal);
+  MlirOperationState constState = mlirOperationStateGet(
+      mlirStringRefCreateFromCString("arith.constant"), loc);
+  MlirType type =
+      mlirTypeParseGet(ctx, mlirStringRefCreateFromCString(typeStr));
+  mlirOperationStateAddResults(&constState, 1, &type);
+  mlirOperationStateAddAttributes(&constState, 1, &valueAttr);
+  MlirOperation constOp = mlirOperationCreate(&constState);
+  return mlirOperationGetResult(constOp, 0);
+}
+
 static void registerAllUpstreamDialects(MlirContext ctx) {
   MlirDialectRegistry registry = mlirDialectRegistryCreate();
   mlirRegisterAllDialects(registry);
   mlirContextAppendDialectRegistry(ctx, registry);
   mlirDialectRegistryDestroy(registry);
 }
 
 void populateLoopBody(MlirContext ctx, MlirBlock loopBody,
MlirNamedAttribute funcAttrs[] = {
           funcNameAttr)};
   MlirOperationState funcState = mlirOperationStateGet(
       mlirStringRefCreateFromCString("func.func"), location);
   mlirOperationStateAddAttributes(&funcState, 2, funcAttrs);
   mlirOperationStateAddOwnedRegions(&funcState, 1, &funcBodyRegion);
   MlirOperation func = mlirOperationCreate(&funcState);
   mlirBlockInsertOwnedOperation(moduleBody, 0, func);
 
-  MlirType indexType =
-      mlirTypeParseGet(ctx, mlirStringRefCreateFromCString("index"));
-  MlirAttribute indexZeroLiteral =
-      mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("0 : index"));
-  MlirNamedAttribute indexZeroValueAttr = mlirNamedAttributeGet(
-      mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("value")),
-      indexZeroLiteral);
-  MlirOperationState constZeroState = mlirOperationStateGet(
-      mlirStringRefCreateFromCString("arith.constant"), location);
-  mlirOperationStateAddResults(&constZeroState, 1, &indexType);
-  mlirOperationStateAddAttributes(&constZeroState, 1, &indexZeroValueAttr);
-  MlirOperation constZero = mlirOperationCreate(&constZeroState);
+  MlirValue constZeroValue = makeConstantLiteral(ctx, "0", "index");
+  MlirOperation constZero = mlirOpResultGetOwner(constZeroValue);
   mlirBlockAppendOwnedOperation(funcBody, constZero);
 
   MlirValue funcArg0 = mlirBlockGetArgument(funcBody, 0);
-  MlirValue constZeroValue = mlirOperationGetResult(constZero, 0);
   MlirValue dimOperands[] = {funcArg0, constZeroValue};
   MlirOperationState dimState = mlirOperationStateGet(
       mlirStringRefCreateFromCString("memref.dim"), location);
   mlirOperationStateAddOperands(&dimState, 2, dimOperands);
+  MlirType indexType =
+      mlirTypeParseGet(ctx, mlirStringRefCreateFromCString("index"));
   mlirOperationStateAddResults(&dimState, 1, &indexType);
   MlirOperation dim = mlirOperationCreate(&dimState);
   mlirBlockAppendOwnedOperation(funcBody, dim);
 
   MlirRegion loopBodyRegion = mlirRegionCreate();
   MlirBlock loopBody = mlirBlockCreate(0, NULL, NULL);
   mlirBlockAddArgument(loopBody, indexType, location);
   mlirRegionAppendOwnedBlock(loopBodyRegion, loopBody);
 
   MlirAttribute indexOneLiteral =
       mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("1 : index"));
   MlirNamedAttribute indexOneValueAttr = mlirNamedAttributeGet(
       mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("value")),
       indexOneLiteral);
   MlirOperationState constOneState = mlirOperationStateGet(
       mlirStringRefCreateFromCString("arith.constant"), location);
   mlirOperationStateAddResults(&constOneState, 1, &indexType);
   mlirOperationStateAddAttributes(&constOneState, 1, &indexOneValueAttr);
-  MlirOperation constOne = mlirOperationCreate(&constOneState);
+  MlirValue constOneValue = makeConstantLiteral(ctx, "1", "index");
+  MlirOperation constOne = mlirOpResultGetOwner(constOneValue);
   mlirBlockAppendOwnedOperation(funcBody, constOne);
 
   MlirValue dimValue = mlirOperationGetResult(dim, 0);
-  MlirValue constOneValue = mlirOperationGetResult(constOne, 0);
   MlirValue loopOperands[] = {constZeroValue, dimValue, constOneValue};
   MlirOperationState loopState = mlirOperationStateGet(
       mlirStringRefCreateFromCString("scf.for"), location);
   mlirOperationStateAddOperands(&loopState, 3, loopOperands);
   mlirOperationStateAddOwnedRegions(&loopState, 1, &loopBodyRegion);
   MlirOperation loop = mlirOperationCreate(&loopState);
   mlirBlockAppendOwnedOperation(funcBody, loop);
 
if (!mlirTypeIsAOpaque(opaque) ||
     return 25;
   mlirTypeDump(opaque);
   fprintf(stderr, "\n");
   // CHECK: !dialect.type
 
   return 0;
 }
 
-void callbackSetFixedLengthString(const char *data, intptr_t len,
-                                  void *userData) {
-  strncpy(userData, data, len);
-}
-
 bool stringIsEqual(const char *lhs, MlirStringRef rhs) {
   if (strlen(lhs) != rhs.length) {
     return false;
   }
   return !strncmp(lhs, rhs.data, rhs.length);
 }
 
 int printBuiltinAttributes(MlirContext ctx) {
int testOperands(void) {
   // CHECK-LABEL: @testOperands
 
   MlirContext ctx = mlirContextCreate();
   registerAllUpstreamDialects(ctx);
 
   mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("arith"));
   mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("test"));
   MlirLocation loc = mlirLocationUnknownGet(ctx);
-  MlirType indexType = mlirIndexTypeGet(ctx);
 
   // Create some constants to use as operands.
-  MlirAttribute indexZeroLiteral =
-      mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("0 : index"));
-  MlirNamedAttribute indexZeroValueAttr = mlirNamedAttributeGet(
-      mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("value")),
-      indexZeroLiteral);
-  MlirOperationState constZeroState = mlirOperationStateGet(
-      mlirStringRefCreateFromCString("arith.constant"), loc);
-  mlirOperationStateAddResults(&constZeroState, 1, &indexType);
-  mlirOperationStateAddAttributes(&constZeroState, 1, &indexZeroValueAttr);
-  MlirOperation constZero = mlirOperationCreate(&constZeroState);
-  MlirValue constZeroValue = mlirOperationGetResult(constZero, 0);
-
-  MlirAttribute indexOneLiteral =
-      mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("1 : index"));
-  MlirNamedAttribute indexOneValueAttr = mlirNamedAttributeGet(
-      mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("value")),
-      indexOneLiteral);
-  MlirOperationState constOneState = mlirOperationStateGet(
-      mlirStringRefCreateFromCString("arith.constant"), loc);
-  mlirOperationStateAddResults(&constOneState, 1, &indexType);
-  mlirOperationStateAddAttributes(&constOneState, 1, &indexOneValueAttr);
-  MlirOperation constOne = mlirOperationCreate(&constOneState);
-  MlirValue constOneValue = mlirOperationGetResult(constOne, 0);
+  MlirValue constZeroValue = makeConstantLiteral(ctx, "0", "index");
+  MlirValue constOneValue = makeConstantLiteral(ctx, "1", "index");
 
   // Create the operation under test.
   mlirContextSetAllowUnregisteredDialects(ctx, true);
   MlirOperationState opState =
       mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op"), loc);
   MlirValue initialOperands[] = {constZeroValue};
   mlirOperationStateAddOperands(&opState, 1, initialOperands);
   MlirOperation op = mlirOperationCreate(&opState);
int testOperands(void) {
   // CHECK: Use operandNumber: 0
 
   use2 = mlirOpOperandGetNextUse(use2);
   if (!mlirOpOperandIsNull(use2)) {
     fprintf(stderr, "ERROR: Next use should be null\n");
     return 3;
   }
 
+  MlirOperationState op2State =
+      mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op2"), loc);
+  MlirValue initialOperands2[] = {constOneValue};
+  mlirOperationStateAddOperands(&op2State, 1, initialOperands2);
+  (void)mlirOperationCreate(&op2State);
+
+  MlirOpOperand use3 = mlirValueGetFirstUse(constOneValue);
+  fprintf(stderr, "First use owner: ");
+  mlirOperationPrint(mlirOpOperandGetOwner(use3), printToStderr, NULL);
+  fprintf(stderr, "\n");
+  // CHECK: First use owner: "dummy.op2"
+
+  use3 = mlirOpOperandGetNextUse(mlirValueGetFirstUse(constOneValue));
+  fprintf(stderr, "Second use owner: ");
+  mlirOperationPrint(mlirOpOperandGetOwner(use3), printToStderr, NULL);
+  fprintf(stderr, "\n");
+  // CHECK: Second use owner: "dummy.op"
+
+  MlirValue constTwoValue = makeConstantLiteral(ctx, "2", "index");
+  mlirValueReplaceAllUsesOfWith(constOneValue, constTwoValue);
+
+  use3 = mlirValueGetFirstUse(constOneValue);
+  if (!mlirOpOperandIsNull(use3)) {
+    fprintf(stderr, "ERROR: Use should be null\n");
+    return 4;
+  }
+
+  MlirOpOperand use4 = mlirValueGetFirstUse(constTwoValue);
+  fprintf(stderr, "First replacement use owner: ");
+  mlirOperationPrint(mlirOpOperandGetOwner(use4), printToStderr, NULL);
+  fprintf(stderr, "\n");
+  // CHECK: First replacement use owner: "dummy.op"
+
+  use4 = mlirOpOperandGetNextUse(mlirValueGetFirstUse(constTwoValue));
+  fprintf(stderr, "Second replacement use owner: ");
+  mlirOperationPrint(mlirOpOperandGetOwner(use4), printToStderr, NULL);
+  fprintf(stderr, "\n");
+  // CHECK: Second replacement use owner: "dummy.op2"
+
   mlirOperationDestroy(op);
-  mlirOperationDestroy(constZero);
-  mlirOperationDestroy(constOne);
+  mlirOperationDestroy(mlirOpResultGetOwner(constZeroValue));
+  mlirOperationDestroy(mlirOpResultGetOwner(constOneValue));
+  mlirOperationDestroy(mlirOpResultGetOwner(constTwoValue));
   mlirContextDestroy(ctx);
 
   return 0;
 }
 
 /// Tests clone APIs.
 int testClone(void) {
   fprintf(stderr, "@testClone\n");
   // CHECK-LABEL: @testClone
 
   MlirContext ctx = mlirContextCreate();
   registerAllUpstreamDialects(ctx);
 
   mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("func"));
   MlirLocation loc = mlirLocationUnknownGet(ctx);
-  MlirType indexType = mlirIndexTypeGet(ctx);
   MlirStringRef valueStringRef = mlirStringRefCreateFromCString("value");
 
-  MlirAttribute indexZeroLiteral =
-      mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("0 : index"));
-  MlirNamedAttribute indexZeroValueAttr = mlirNamedAttributeGet(
-      mlirIdentifierGet(ctx, valueStringRef), indexZeroLiteral);
-  MlirOperationState constZeroState = mlirOperationStateGet(
-      mlirStringRefCreateFromCString("arith.constant"), loc);
-  mlirOperationStateAddResults(&constZeroState, 1, &indexType);
-  mlirOperationStateAddAttributes(&constZeroState, 1, &indexZeroValueAttr);
-  MlirOperation constZero = mlirOperationCreate(&constZeroState);
+  MlirValue constZeroValue = makeConstantLiteral(ctx, "0", "index");
+  MlirOperation constZero = mlirOpResultGetOwner(constZeroValue);
 
   MlirAttribute indexOneLiteral =
       mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("1 : index"));
   MlirOperation constOne = mlirOperationClone(constZero);
   mlirOperationSetAttributeByName(constOne, valueStringRef, indexOneLiteral);
 
   mlirOperationPrint(constZero, printToStderr, NULL);
   mlirOperationPrint(constOne, printToStderr, NULL);
int testTypeID(MlirContext ctx) {
 
   if (mlirTypeIDEqual(i32ID, i32AttrID)) {
     fprintf(stderr, "ERROR: Expected integer type id to not equal integer "
                     "attribute type id\n");
     return 4;
   }
 
   MlirLocation loc = mlirLocationUnknownGet(ctx);
-  MlirType indexType = mlirIndexTypeGet(ctx);
-  MlirStringRef valueStringRef = mlirStringRefCreateFromCString("value");
 
   // Create a registered operation, which should have a type id.
-  MlirAttribute indexZeroLiteral =
-      mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString("0 : index"));
-  MlirNamedAttribute indexZeroValueAttr = mlirNamedAttributeGet(
-      mlirIdentifierGet(ctx, valueStringRef), indexZeroLiteral);
-  MlirOperationState constZeroState = mlirOperationStateGet(
-      mlirStringRefCreateFromCString("arith.constant"), loc);
-  mlirOperationStateAddResults(&constZeroState, 1, &indexType);
-  mlirOperationStateAddAttributes(&constZeroState, 1, &indexZeroValueAttr);
-  MlirOperation constZero = mlirOperationCreate(&constZeroState);
+  MlirValue constZeroValue = makeConstantLiteral(ctx, "0", "index");
+  MlirOperation constZero = mlirOpResultGetOwner(constZeroValue);
 
   if (!mlirOperationVerify(constZero)) {
     fprintf(stderr, "ERROR: Expected operation to verify correctly\n");
     return 5;
   }
 
   if (mlirOperationIsNull(constZero)) {
     fprintf(stderr, "ERROR: Expected registered operation to be present\n");

mlir/test/python/ir/value.py

def testValueUses():
   # CHECK: Use owner: "custom.op2"
   # CHECK: Use operand_number: 0
   # CHECK: Use owner: "custom.op2"
   # CHECK: Use operand_number: 0
   for use in value.uses:
     assert use.owner in [op1, op2]
     print(f"Use owner: {use.owner}")
     print(f"Use operand_number: {use.operand_number}")
+
+# CHECK-LABEL: TEST: testValueReplaceAllUsesWith
+@run
+def testValueReplaceAllUsesWith():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with Location.unknown(ctx):
+    i32 = IntegerType.get_signless(32)
+    module = Module.create()
+    with InsertionPoint(module.body):
+      value = Operation.create("custom.op1", results=[i32]).results[0]
+      op1 = Operation.create("custom.op2", operands=[value])
+      op2 = Operation.create("custom.op2", operands=[value])
+      value2 = Operation.create("custom.op3", results=[i32]).results[0]
+      value.replace_all_uses_with(value2)
+
+  assert len(list(value.uses)) == 0
+
+  # CHECK: Use owner: "custom.op2"
+  # CHECK: Use operand_number: 0
+  # CHECK: Use owner: "custom.op2"
+  # CHECK: Use operand_number: 0
+  for use in value2.uses:
+    assert use.owner in [op1, op2]
+    print(f"Use owner: {use.owner}")
+    print(f"Use operand_number: {use.operand_number}")