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

[makslevental]

Summary

This reverts commit 98fbd9d3f940928e68566c2d3794d89e361fbf88 because it breaks sanitizers. Will reland after testing sanitizers locally.

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(),
-                                   mlirValueGetType(self.get()));
-                             })
-      .def(
-          "replace_all_uses_with",
-          [](PyValue &self, PyValue &with) {
-            mlirValueReplaceAllUsesOfWith(self.get(), with.get());
-          },
-          kValueReplaceAllUsesWithDocstring);
+      .def_property_readonly("type", [](PyValue &self) {
+        return PyType(self.getParentOperation()->getContext(),
+      });
   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);
 
-  MlirValue constZeroValue = makeConstantLiteral(ctx, "0", "index");
-  MlirOperation constZero = mlirOpResultGetOwner(constZeroValue);
+  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);
   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);
-  MlirValue constOneValue = makeConstantLiteral(ctx, "1", "index");
-  MlirOperation constOne = mlirOpResultGetOwner(constOneValue);
+  MlirOperation constOne = mlirOperationCreate(&constOneState);
   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.
-  MlirValue constZeroValue = makeConstantLiteral(ctx, "0", "index");
-  MlirValue constOneValue = makeConstantLiteral(ctx, "1", "index");
+  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);
 
   // 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);
-  MlirOperation op2 = 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(op2);
-  mlirOperationDestroy(mlirOpResultGetOwner(constZeroValue));
-  mlirOperationDestroy(mlirOpResultGetOwner(constOneValue));
-  mlirOperationDestroy(mlirOpResultGetOwner(constTwoValue));
+  mlirOperationDestroy(constZero);
+  mlirOperationDestroy(constOne);
   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");
 
-  MlirValue constZeroValue = makeConstantLiteral(ctx, "0", "index");
-  MlirOperation constZero = mlirOpResultGetOwner(constZeroValue);
+  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);
 
   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.
-  MlirValue constZeroValue = makeConstantLiteral(ctx, "0", "index");
-  MlirOperation constZero = mlirOpResultGetOwner(constZeroValue);
+  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);
 
   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}")