[mlir][CAPI] Attribute set/remove on operations.

[stellaraccident]

Summary

• New functions: mlirOperationSetAttributeByName, mlirOperationRemoveAttributeByName
• Also adds some *IsNull checks and standardizes the rest to use "static inline" form, which makes them all non-opaque and not part of the ABI (which is desirable).
• Changes needed to resolve TODOs in npcomp PyTorch capture.

[stellaraccident]

Rebase

mlir/include/mlir-c/IR.h

 
 /** Creates an MLIR context and transfers its ownership to the caller. */
 MlirContext mlirContextCreate();
 
 /** Checks if two contexts are equal. */
 int mlirContextEqual(MlirContext ctx1, MlirContext ctx2);
 
 /** Checks whether a context is null. */
-inline int mlirContextIsNull(MlirContext context) { return !context.ptr; }
+static inline int mlirContextIsNull(MlirContext context) {
+  return !context.ptr;
+}
 
 /** Takes an MLIR context owned by the caller and destroys it. */
 void mlirContextDestroy(MlirContext context);
 
 /** Sets whether unregistered dialects are allowed in this context. */
 void mlirContextSetAllowUnregisteredDialects(MlirContext context, int allow);
 
 /** Returns whether the context allows unregistered dialects. */

 /*============================================================================*/
 /* Dialect API.                                                               */
 /*============================================================================*/
 
 /** Returns the context that owns the dialect. */
 MlirContext mlirDialectGetContext(MlirDialect dialect);
 
 /** Checks if the dialect is null. */
-int mlirDialectIsNull(MlirDialect dialect);
+static inline int mlirDialectIsNull(MlirDialect dialect) {
+  return !dialect.ptr;
+}
 
 /** Checks if two dialects that belong to the same context are equal. Dialects
  * from different contexts will not compare equal. */
 int mlirDialectEqual(MlirDialect dialect1, MlirDialect dialect2);
 
 /** Returns the namespace of the given dialect. */
 MlirStringRef mlirDialectGetNamespace(MlirDialect dialect);
 

 
 /** Parses a module from the string and transfers ownership to the caller. */
 MlirModule mlirModuleCreateParse(MlirContext context, const char *module);
 
 /** Gets the context that a module was created with. */
 MlirContext mlirModuleGetContext(MlirModule module);
 
 /** Checks whether a module is null. */
-inline int mlirModuleIsNull(MlirModule module) { return !module.ptr; }
+static inline int mlirModuleIsNull(MlirModule module) { return !module.ptr; }
 
 /** Takes a module owned by the caller and deletes it. */
 void mlirModuleDestroy(MlirModule module);
 
 /** Views the module as a generic operation. */
 MlirOperation mlirModuleGetOperation(MlirModule module);
 
 /*============================================================================*/

 
 /** Creates an operation and transfers ownership to the caller. */
 MlirOperation mlirOperationCreate(const MlirOperationState *state);
 
 /** Takes an operation owned by the caller and destroys it. */
 void mlirOperationDestroy(MlirOperation op);
 
 /** Checks whether the underlying operation is null. */
-int mlirOperationIsNull(MlirOperation op);
+static inline int mlirOperationIsNull(MlirOperation op) { return !op.ptr; }
 
 /** Returns the number of regions attached to the given operation. */
 intptr_t mlirOperationGetNumRegions(MlirOperation op);
 
 /** Returns `pos`-th region attached to the operation. */
 MlirRegion mlirOperationGetRegion(MlirOperation op, intptr_t pos);
 
 /** Returns an operation immediately following the given operation it its

 
 /** Return `pos`-th attribute of the operation. */
 MlirNamedAttribute mlirOperationGetAttribute(MlirOperation op, intptr_t pos);
 
 /** Returns an attribute attached to the operation given its name. */
 MlirAttribute mlirOperationGetAttributeByName(MlirOperation op,
                                               const char *name);
 
+/** Sets an attribute by name, replacing the existing if it exists or
+ * adding a new one otherwise. */
+void mlirOperationSetAttributeByName(MlirOperation op, const char *name,
+                                     MlirAttribute attr);
+
+/** Removes an attribute by name. Returns 0 if the attribute was not found
+ * and !0 if removed. */
+int mlirOperationRemoveAttributeByName(MlirOperation op, const char *name);
+
 /** Prints an operation by sending chunks of the string representation and
  * forwarding `userData to `callback`. Note that the callback may be called
  * several times with consecutive chunks of the string. */
 void mlirOperationPrint(MlirOperation op, MlirStringCallback callback,
                         void *userData);
 
 /** Prints an operation to stderr. */
 void mlirOperationDump(MlirOperation op);
 
 /*============================================================================*/
 /* Region API.                                                                */
 /*============================================================================*/
 
 /** Creates a new empty region and transfers ownership to the caller. */
 MlirRegion mlirRegionCreate();
 
 /** Takes a region owned by the caller and destroys it. */
 void mlirRegionDestroy(MlirRegion region);
 
 /** Checks whether a region is null. */
-int mlirRegionIsNull(MlirRegion region);
+static inline int mlirRegionIsNull(MlirRegion region) { return !region.ptr; }
 
 /** Gets the first block in the region. */
 MlirBlock mlirRegionGetFirstBlock(MlirRegion region);
 
 /** Takes a block owned by the caller and appends it to the given region. */
 void mlirRegionAppendOwnedBlock(MlirRegion region, MlirBlock block);
 
 /** Takes a block owned by the caller and inserts it at `pos` to the given

 /** Creates a new empty block with the given argument types and transfers
  * ownership to the caller. */
 MlirBlock mlirBlockCreate(intptr_t nArgs, MlirType *args);
 
 /** Takes a block owned by the caller and destroys it. */
 void mlirBlockDestroy(MlirBlock block);
 
 /** Checks whether a block is null. */
-int mlirBlockIsNull(MlirBlock block);
+static inline int mlirBlockIsNull(MlirBlock block) { return !block.ptr; }
 
 /** Returns the block immediately following the given block in its parent
  * region. */
 MlirBlock mlirBlockGetNextInRegion(MlirBlock block);
 
 /** Returns the first operation in the block. */
 MlirOperation mlirBlockGetFirstOperation(MlirBlock block);
 

  * several times with consecutive chunks of the string. */
 void mlirBlockPrint(MlirBlock block, MlirStringCallback callback,
                     void *userData);
 
 /*============================================================================*/
 /* Value API.                                                                 */
 /*============================================================================*/
 
+/** Returns whether the value is null. */
+static inline int mlirValueIsNull(MlirValue value) { return !value.ptr; }
+
 /** Returns the type of the value. */
 MlirType mlirValueGetType(MlirValue value);
 
 /** Prints a value by sending chunks of the string representation and
  * forwarding `userData to `callback`. Note that the callback may be called
  * several times with consecutive chunks of the string. */
 void mlirValuePrint(MlirValue value, MlirStringCallback callback,
                     void *userData);
 
 /*============================================================================*/
 /* Type API.                                                                  */
 /*============================================================================*/
 
 /** Parses a type. The type is owned by the context. */
 MlirType mlirTypeParseGet(MlirContext context, const char *type);
 
 /** Gets the context that a type was created with. */
 MlirContext mlirTypeGetContext(MlirType type);
 
 /** Checks whether a type is null. */
-inline int mlirTypeIsNull(MlirType type) { return !type.ptr; }
+static inline int mlirTypeIsNull(MlirType type) { return !type.ptr; }
 
 /** Checks if two types are equal. */
 int mlirTypeEqual(MlirType t1, MlirType t2);
 
 /** Prints a location by sending chunks of the string representation and
  * forwarding `userData to `callback`. Note that the callback may be called
  * several times with consecutive chunks of the string. */
 void mlirTypePrint(MlirType type, MlirStringCallback callback, void *userData);
 
 /** Prints the type to the standard error stream. */
 void mlirTypeDump(MlirType type);
 
 /*============================================================================*/
 /* Attribute API.                                                             */
 /*============================================================================*/
 
 /** Parses an attribute. The attribute is owned by the context. */
 MlirAttribute mlirAttributeParseGet(MlirContext context, const char *attr);
 
 /** Gets the context that an attribute was created with. */
 MlirContext mlirAttributeGetContext(MlirAttribute attribute);
 
 /** Checks whether an attribute is null. */
-inline int mlirAttributeIsNull(MlirAttribute attr) { return !attr.ptr; }
+static inline int mlirAttributeIsNull(MlirAttribute attr) { return !attr.ptr; }
 
 /** Checks if two attributes are equal. */
 int mlirAttributeEqual(MlirAttribute a1, MlirAttribute a2);
 
 /** Prints an attribute by sending chunks of the string representation and
  * forwarding `userData to `callback`. Note that the callback may be called
  * several times with consecutive chunks of the string. */
 void mlirAttributePrint(MlirAttribute attr, MlirStringCallback callback,

mlir/lib/CAPI/IR/IR.cpp

 /* ========================================================================== */
 /* Dialect API.                                                               */
 /* ========================================================================== */
 
 MlirContext mlirDialectGetContext(MlirDialect dialect) {
   return wrap(unwrap(dialect)->getContext());
 }
 
-int mlirDialectIsNull(MlirDialect dialect) {
-  return unwrap(dialect) == nullptr;
-}
-
 int mlirDialectEqual(MlirDialect dialect1, MlirDialect dialect2) {
   return unwrap(dialect1) == unwrap(dialect2);
 }
 
 MlirStringRef mlirDialectGetNamespace(MlirDialect dialect) {
   return wrap(unwrap(dialect)->getNamespace());
 }
 
MlirOperation mlirOperationCreate(const MlirOperationState *state) {
   free(state->successors);
   free(state->regions);
   free(state->attributes);
   return result;
 }
 
 void mlirOperationDestroy(MlirOperation op) { unwrap(op)->erase(); }
 
-int mlirOperationIsNull(MlirOperation op) { return unwrap(op) == nullptr; }
-
 intptr_t mlirOperationGetNumRegions(MlirOperation op) {
   return static_cast<intptr_t>(unwrap(op)->getNumRegions());
 }
 
 MlirRegion mlirOperationGetRegion(MlirOperation op, intptr_t pos) {
   return wrap(&unwrap(op)->getRegion(static_cast<unsigned>(pos)));
 }
 
MlirNamedAttribute mlirOperationGetAttribute(MlirOperation op, intptr_t pos) {
   return MlirNamedAttribute{attr.first.c_str(), wrap(attr.second)};
 }
 
 MlirAttribute mlirOperationGetAttributeByName(MlirOperation op,
                                               const char *name) {
   return wrap(unwrap(op)->getAttr(name));
 }
 
+void mlirOperationSetAttributeByName(MlirOperation op, const char *name,
+                                     MlirAttribute attr) {
+  unwrap(op)->setAttr(name, unwrap(attr));
+}
+
+int mlirOperationRemoveAttributeByName(MlirOperation op, const char *name) {
+  auto removeResult = unwrap(op)->removeAttr(name);
+  return removeResult == MutableDictionaryAttr::RemoveResult::Removed;
+}
+
 void mlirOperationPrint(MlirOperation op, MlirStringCallback callback,
                         void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(op)->print(stream);
   stream.flush();
 }
 
 void mlirOperationDump(MlirOperation op) { return unwrap(op)->dump(); }
void mlirRegionInsertOwnedBlockBefore(MlirRegion region, MlirBlock reference,
   unwrap(region)->getBlocks().insert(Region::iterator(unwrap(reference)),
                                      unwrap(block));
 }
 
 void mlirRegionDestroy(MlirRegion region) {
   delete static_cast<Region *>(region.ptr);
 }
 
-int mlirRegionIsNull(MlirRegion region) { return unwrap(region) == nullptr; }
-
 /* ========================================================================== */
 /* Block API.                                                                 */
 /* ========================================================================== */
 
 MlirBlock mlirBlockCreate(intptr_t nArgs, MlirType *args) {
   Block *b = new Block;
   for (intptr_t i = 0; i < nArgs; ++i)
     b->addArgument(unwrap(args[i]));
void mlirBlockInsertOwnedOperationBefore(MlirBlock block,
   assert(unwrap(reference)->getBlock() == unwrap(block) &&
          "expected reference operation to belong to the block");
   unwrap(block)->getOperations().insert(Block::iterator(unwrap(reference)),
                                         unwrap(operation));
 }
 
 void mlirBlockDestroy(MlirBlock block) { delete unwrap(block); }
 
-int mlirBlockIsNull(MlirBlock block) { return unwrap(block) == nullptr; }
-
 intptr_t mlirBlockGetNumArguments(MlirBlock block) {
   return static_cast<intptr_t>(unwrap(block)->getNumArguments());
 }
 
 MlirValue mlirBlockGetArgument(MlirBlock block, intptr_t pos) {
   return wrap(unwrap(block)->getArgument(static_cast<unsigned>(pos)));
 }
 

mlir/test/CAPI/ir.c

void collectStats(MlirOperation operation) {
   fprintf(stderr, "Number of values: %u\n", stats.numValues);
 }
 
 static void printToStderr(const char *str, intptr_t len, void *userData) {
   (void)userData;
   fwrite(str, 1, len, stderr);
 }
 
-static void printFirstOfEach(MlirOperation operation) {
+static void printFirstOfEach(MlirContext ctx, MlirOperation operation) {
   // Assuming we are given a module, go to the first operation of the first
   // function.
   MlirRegion region = mlirOperationGetRegion(operation, 0);
   MlirBlock block = mlirRegionGetFirstBlock(region);
   operation = mlirBlockGetFirstOperation(block);
   region = mlirOperationGetRegion(operation, 0);
   block = mlirRegionGetFirstBlock(region);
   operation = mlirBlockGetFirstOperation(block);
 
   // In the module we created, the first operation of the first function is an
-  // "std.dim", which has an attribute an a single result that we can use to
+  // "std.dim", which has an attribute and a single result that we can use to
   // test the printing mechanism.
   mlirBlockPrint(block, printToStderr, NULL);
   fprintf(stderr, "\n");
+  fprintf(stderr, "First operation: ");
   mlirOperationPrint(operation, printToStderr, NULL);
   fprintf(stderr, "\n");
 
-  MlirNamedAttribute namedAttr = mlirOperationGetAttribute(operation, 0);
-  mlirAttributePrint(namedAttr.attribute, printToStderr, NULL);
+  // Get the attribute by index.
+  MlirNamedAttribute namedAttr0 = mlirOperationGetAttribute(operation, 0);
+  fprintf(stderr, "Get attr 0: ");
+  mlirAttributePrint(namedAttr0.attribute, printToStderr, NULL);
   fprintf(stderr, "\n");
 
+  // Now re-get the attribute by name.
+  MlirAttribute attr0ByName =
+      mlirOperationGetAttributeByName(operation, namedAttr0.name);
+  fprintf(stderr, "Get attr 0 by name: ");
+  mlirAttributePrint(attr0ByName, printToStderr, NULL);
+  fprintf(stderr, "\n");
+
+  // Get a non-existing attribute and assert that it is null (sanity).
+  fprintf(stderr, "does_not_exist is null: %d\n",
+          mlirAttributeIsNull(
+              mlirOperationGetAttributeByName(operation, "does_not_exist")));
+
+  // Get result 0 and its type.
   MlirValue value = mlirOperationGetResult(operation, 0);
+  fprintf(stderr, "Result 0: ");
   mlirValuePrint(value, printToStderr, NULL);
   fprintf(stderr, "\n");
+  fprintf(stderr, "Value is null: %d\n", mlirValueIsNull(value));
 
   MlirType type = mlirValueGetType(value);
+  fprintf(stderr, "Result 0 type: ");
   mlirTypePrint(type, printToStderr, NULL);
   fprintf(stderr, "\n");
+
+  // Set a custom attribute.
+  mlirOperationSetAttributeByName(operation, "custom_attr",
+                                  mlirBoolAttrGet(ctx, 1));
+  fprintf(stderr, "Op with set attr: ");
+  mlirOperationPrint(operation, printToStderr, NULL);
+  fprintf(stderr, "\n");
+
+  // Remove the attribute.
+  fprintf(stderr, "Remove attr: %d\n",
+          mlirOperationRemoveAttributeByName(operation, "custom_attr"));
+  fprintf(stderr, "Remove attr again: %d\n",
+          mlirOperationRemoveAttributeByName(operation, "custom_attr"));
+  fprintf(stderr, "Removed attr is null: %d\n",
+          mlirAttributeIsNull(
+              mlirOperationGetAttributeByName(operation, "custom_attr")));
 }
 
 /// Creates an operation with a region containing multiple blocks with
 /// operations and dumps it. The blocks and operations are inserted using
 /// block/operation-relative API and their final order is checked.
 static void buildWithInsertionsAndPrint(MlirContext ctx) {
   MlirLocation loc = mlirLocationUnknownGet(ctx);
 
int main() {
   // clang-format off
   // CHECK: Number of operations: 13
   // CHECK: Number of attributes: 4
   // CHECK: Number of blocks: 3
   // CHECK: Number of regions: 3
   // CHECK: Number of values: 9
   // clang-format on
 
-  printFirstOfEach(module);
+  printFirstOfEach(ctx, module);
   // clang-format off
   // CHECK:   %[[C0:.*]] = constant 0 : index
   // CHECK:   %[[DIM:.*]] = dim %{{.*}}, %[[C0]] : memref<?xf32>
   // CHECK:   %[[C1:.*]] = constant 1 : index
   // CHECK:   scf.for %[[I:.*]] = %[[C0]] to %[[DIM]] step %[[C1]] {
   // CHECK:     %[[LHS:.*]] = load %{{.*}}[%[[I]]] : memref<?xf32>
   // CHECK:     %[[RHS:.*]] = load %{{.*}}[%[[I]]] : memref<?xf32>
   // CHECK:     %[[SUM:.*]] = addf %[[LHS]], %[[RHS]] : f32
   // CHECK:     store %[[SUM]], %{{.*}}[%[[I]]] : memref<?xf32>
   // CHECK:   }
   // CHECK: return
-  // CHECK: constant 0 : index
-  // CHECK: 0 : index
-  // CHECK: constant 0 : index
-  // CHECK: index
+  // CHECK: First operation: {{.*}} = constant 0 : index
+  // CHECK: Get attr 0: 0 : index
+  // CHECK: Get attr 0 by name: 0 : index
+  // CHECK: does_not_exist is null: 1
+  // CHECK: Result 0: {{.*}} = constant 0 : index
+  // CHECK: Value is null: 0
+  // CHECK: Result 0 type: index
+  // CHECK: Op with set attr: {{.*}} {custom_attr = true}
+  // CHECK: Remove attr: 1
+  // CHECK: Remove attr again: 0
+  // CHECK: Removed attr is null: 1
   // clang-format on
 
   mlirModuleDestroy(moduleOp);
 
   buildWithInsertionsAndPrint(ctx);
   // clang-format off
   // CHECK-LABEL:  "insertion.order.test"
   // CHECK:      ^{{.*}}(%{{.*}}: i1