[mlir] Support setting operand values in C and Python APIs.

[mikeurbach]

Summary

This adds mlirOperationSetOperand to the IR C API, similar to the
function to get an operand.

In the Python API, this adds operands[index] = value syntax, similar
to the syntax to get an operand with operands[index].

[mikeurbach]

This is something @jdd and I will use in circt, and seemed generally useful upstream in mlir. However, we weren't sure about the addition of mlirOperationSetOperand to the C API. Are there deeper considerations preventing this from being added? Or is it just something that was demand-driven and has yet to be demanded? I've been testing this patch locally with some small programs for about a week and haven't found any issues with setting operands from Python via the proposed C API.

[mehdi_amini]

mlir/lib/CAPI/IR/IR.cpp

357

Can you add coverage to this function through the C unitest?

[mikeurbach]

Add C unittest for operands API

[mikeurbach]

Add cleanup code to new C unittest.

mlir/include/mlir-c/IR.h

 
 /// Returns the number of operands of the operation.
 MLIR_CAPI_EXPORTED intptr_t mlirOperationGetNumOperands(MlirOperation op);
 
 /// Returns `pos`-th operand of the operation.
 MLIR_CAPI_EXPORTED MlirValue mlirOperationGetOperand(MlirOperation op,
                                                      intptr_t pos);
 
+/// Sets the `pos`-th operand of the operation.
+MLIR_CAPI_EXPORTED void mlirOperationSetOperand(MlirOperation op, intptr_t pos,
+                                                MlirValue newValue);
+
 /// Returns the number of results of the operation.
 MLIR_CAPI_EXPORTED intptr_t mlirOperationGetNumResults(MlirOperation op);
 
 /// Returns `pos`-th result of the operation.
 MLIR_CAPI_EXPORTED MlirValue mlirOperationGetResult(MlirOperation op,
                                                     intptr_t pos);
 
 /// Returns the number of successor blocks of the operation.

mlir/lib/Bindings/Python/IRCore.cpp

public:
   PyValue getElement(intptr_t pos) {
     return PyValue(operation, mlirOperationGetOperand(operation->get(), pos));
   }
 
   PyOpOperandList slice(intptr_t startIndex, intptr_t length, intptr_t step) {
     return PyOpOperandList(operation, startIndex, length, step);
   }
 
+  void dunderSetItem(intptr_t index, PyValue value) {
+    index = wrapIndex(index);
+    mlirOperationSetOperand(operation->get(), index, value.get());
+  }
+
+  static void bindDerived(ClassTy &c) {
+    c.def("__setitem__", &PyOpOperandList::dunderSetItem);
+  }
+
 private:
   PyOperationRef operation;
 };
 
 /// A list of operation results. Internally, these are stored as consecutive
 /// elements, random access is cheap. The result list is associated with the
 /// operation whose results these are, and extends the lifetime of this
 /// operation.

mlir/lib/Bindings/Python/PybindUtils.h

 /// A derived class may additionally define:
 ///   - a `static void bindDerived(ClassTy &)` method to bind additional methods
 ///     the python class.
 template <typename Derived, typename ElementTy>
 class Sliceable {
 protected:
   using ClassTy = pybind11::class_<Derived>;
 
+  intptr_t wrapIndex(intptr_t index) {
+    if (index < 0)
+      index = length + index;
+    if (index < 0 || index >= length) {
+      throw python::SetPyError(PyExc_IndexError,
+                               "attempt to access out of bounds");
+    }
+    return index;
+  }
+
 public:
   explicit Sliceable(intptr_t startIndex, intptr_t length, intptr_t step)
       : startIndex(startIndex), length(length), step(step) {
     assert(length >= 0 && "expected non-negative slice length");
   }
 
   /// Returns the length of the slice.
   intptr_t dunderLen() const { return length; }
 
   /// Returns the element at the given slice index. Supports negative indices
   /// by taking elements in inverse order. Throws if the index is out of bounds.
   ElementTy dunderGetItem(intptr_t index) {
     // Negative indices mean we count from the end.
-    if (index < 0)
-      index = length + index;
-    if (index < 0 || index >= length) {
-      throw python::SetPyError(PyExc_IndexError,
-                               "attempt to access out of bounds");
-    }
+    index = wrapIndex(index);
 
     // Compute the linear index given the current slice properties.
     int linearIndex = index * step + startIndex;
     assert(linearIndex >= 0 &&
            linearIndex < static_cast<Derived *>(this)->getNumElements() &&
            "linear index out of bounds, the slice is ill-formed");
     return static_cast<Derived *>(this)->getElement(linearIndex);
   }

mlir/lib/CAPI/IR/IR.cpp

 intptr_t mlirOperationGetNumOperands(MlirOperation op) {
   return static_cast<intptr_t>(unwrap(op)->getNumOperands());
 }
 
 MlirValue mlirOperationGetOperand(MlirOperation op, intptr_t pos) {
   return wrap(unwrap(op)->getOperand(static_cast<unsigned>(pos)));
 }
 
+void mlirOperationSetOperand(MlirOperation op, intptr_t pos,
+                             MlirValue newValue) {
+  unwrap(op)->setOperand(static_cast<unsigned>(pos), unwrap(newValue));
+}
+
 intptr_t mlirOperationGetNumResults(MlirOperation op) {
   return static_cast<intptr_t>(unwrap(op)->getNumResults());
 }
 
 MlirValue mlirOperationGetResult(MlirOperation op, intptr_t pos) {
   return wrap(unwrap(op)->getResult(static_cast<unsigned>(pos)));
 }
 

mlir/test/Bindings/Python/ir_operation.py

with Context() as ctx:
     fourth = consumer.operands[::2][1::2]
     for operand in fourth:
       print(operand)
 
 
 run(testOperationOperandsSlice)
 
 
+# CHECK-LABEL: TEST: testOperationOperandsSet
+def testOperationOperandsSet():
+  with Context() as ctx, Location.unknown(ctx):
+    ctx.allow_unregistered_dialects = True
+    module = Module.parse(r"""
+      func @f1() {
+        %0 = "test.producer0"() : () -> i64
+        %1 = "test.producer1"() : () -> i64
+        %2 = "test.producer2"() : () -> i64
+        "test.consumer"(%0) : (i64) -> ()
+        return
+      }""")
+    func = module.body.operations[0]
+    entry_block = func.regions[0].blocks[0]
+    producer1 = entry_block.operations[1]
+    producer2 = entry_block.operations[2]
+    consumer = entry_block.operations[3]
+    assert len(consumer.operands) == 1
+    type = consumer.operands[0].type
+
+    # CHECK: test.producer1
+    consumer.operands[0] = producer1.result
+    print(consumer.operands[0])
+
+    # CHECK: test.producer2
+    consumer.operands[-1] = producer2.result
+    print(consumer.operands[0])
+
+
+run(testOperationOperandsSet)
+
+
 # CHECK-LABEL: TEST: testDetachedOperation
 def testDetachedOperation():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     i32 = IntegerType.get_signed(32)
     op1 = Operation.create(
         "custom.op1", results=[i32, i32], regions=1, attributes={

mlir/test/CAPI/ir.c

static int testBackreferences() {
 
   mlirOperationDestroy(op);
   mlirContextDestroy(ctx);
 
   // CHECK-LABEL: @test_backreferences
   return 0;
 }
 
+/// Tests operand APIs.
+int testOperands() {
+  fprintf(stderr, "@testOperands\n");
+  // CHECK-LABEL: @testOperands
+
+  MlirContext ctx = mlirContextCreate();
+  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("std.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("std.constant"), loc);
+  mlirOperationStateAddResults(&constOneState, 1, &indexType);
+  mlirOperationStateAddAttributes(&constOneState, 1, &indexOneValueAttr);
+  MlirOperation constOne = mlirOperationCreate(&constOneState);
+  MlirValue constOneValue = mlirOperationGetResult(constOne, 0);
+
+  // Create the operation under test.
+  MlirOperationState opState =
+      mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op"), loc);
+  MlirValue initialOperands[] = {constZeroValue};
+  mlirOperationStateAddOperands(&opState, 1, initialOperands);
+  MlirOperation op = mlirOperationCreate(&opState);
+
+  // Test operand APIs.
+  intptr_t numOperands = mlirOperationGetNumOperands(op);
+  fprintf(stderr, "Num Operands: %ld\n", numOperands);
+  // CHECK: Num Operands: 1
+
+  MlirValue opOperand = mlirOperationGetOperand(op, 0);
+  fprintf(stderr, "Original operand: ");
+  mlirValuePrint(opOperand, printToStderr, NULL);
+  // CHECK: Original operand: {{.+}} {value = 0 : index}
+
+  mlirOperationSetOperand(op, 0, constOneValue);
+  opOperand = mlirOperationGetOperand(op, 0);
+  fprintf(stderr, "Updated operand: ");
+  mlirValuePrint(opOperand, printToStderr, NULL);
+  // CHECK: Updated operand: {{.+}} {value = 1 : index}
+
+  mlirOperationDestroy(op);
+  mlirOperationDestroy(constZero);
+  mlirOperationDestroy(constOne);
+  mlirContextDestroy(ctx);
+
+  return 0;
+}
+
 // Wraps a diagnostic into additional text we can match against.
 MlirLogicalResult errorHandler(MlirDiagnostic diagnostic, void *userData) {
   fprintf(stderr, "processing diagnostic (userData: %ld) <<\n", (long)userData);
   mlirDiagnosticPrint(diagnostic, printToStderr, NULL);
   fprintf(stderr, "\n");
   MlirLocation loc = mlirDiagnosticGetLocation(diagnostic);
   mlirLocationPrint(loc, printToStderr, NULL);
   assert(mlirDiagnosticGetNumNotes(diagnostic) == 0);
int main() {
   if (affineMapFromExprs(ctx))
     return 7;
   if (printIntegerSet(ctx))
     return 8;
   if (registerOnlyStd())
     return 9;
   if (testBackreferences())
     return 10;
+  if (testOperands())
+    return 11;
 
   mlirContextDestroy(ctx);
 
   testDiagnostics();
   return 0;
 }