[mlir][CAPI] Add result type inference to the CAPI.

[stellaraccident]

Summary

• Adds a flag to MlirOperationState to enable result type inference using the InferTypeOpInterface.
• I chose this level of implementation for a couple of reasons: a) In the creation flow is naturally where generated and custom builder code will be invoking such a thing b) it is a bit more efficient to share the data structure and unpacking vs having a standalone entry-point c) we can always decide to expose more of these interfaces with first-class APIs, but that doesn't preclude that we will always want to use this one in this way (and less API surface area for common things is better for API stability and evolution).
• I struggled to find an appropriate way to test it since we don't link the test dialect into anything CAPI accessible at present. I opted instead for one of the simplest ops I found in a regular dialect which implements the interface.
• This does not do any trait-based type selection. That will be left to generated tablegen wrappers.

[mehdi_amini]

mlir/lib/CAPI/IR/IR.cpp

260

(spurious braces)

[stellaraccident]

Comments

mlir/include/mlir-c/IR.h

struct MlirOperationState {
   intptr_t nOperands;
   MlirValue *operands;
   intptr_t nRegions;
   MlirRegion *regions;
   intptr_t nSuccessors;
   MlirBlock *successors;
   intptr_t nAttributes;
   MlirNamedAttribute *attributes;
+  bool enableResultTypeInference;
 };
 typedef struct MlirOperationState MlirOperationState;
 
 /// Constructs an operation state from a name and a location.
 MLIR_CAPI_EXPORTED MlirOperationState mlirOperationStateGet(MlirStringRef name,
                                                             MlirLocation loc);
 
 /// Adds a list of components to the operation state.
 MLIR_CAPI_EXPORTED void mlirOperationStateAddResults(MlirOperationState *state,
                                                      intptr_t n,
                                                      MlirType const *results);
 MLIR_CAPI_EXPORTED void
 mlirOperationStateAddOperands(MlirOperationState *state, intptr_t n,
                               MlirValue const *operands);
 MLIR_CAPI_EXPORTED void
 mlirOperationStateAddOwnedRegions(MlirOperationState *state, intptr_t n,
                                   MlirRegion const *regions);
 MLIR_CAPI_EXPORTED void
 mlirOperationStateAddSuccessors(MlirOperationState *state, intptr_t n,
                                 MlirBlock const *successors);
 MLIR_CAPI_EXPORTED void
 mlirOperationStateAddAttributes(MlirOperationState *state, intptr_t n,
                                 MlirNamedAttribute const *attributes);
 
+/// Enables result type inference for the operation under construction. If
+/// enabled, then the caller must not have called
+/// mlirOperationStateAddResults(). Note that if enabled, the
+/// mlirOperationCreate() call is failable: it will return a null operation
+/// on inference failure and will emit diagnostics.
+MLIR_CAPI_EXPORTED void
+mlirOperationStateEnableResultTypeInference(MlirOperationState *state);
+
 //===----------------------------------------------------------------------===//
 // Op Printing flags API.
 // While many of these are simple settings that could be represented in a
 // struct, they are wrapped in a heap allocated object and accessed via
 // functions to maximize the possibility of compatibility over time.
 //===----------------------------------------------------------------------===//
 
 /** Creates new printing flags with defaults, intended for customization.

 MLIR_CAPI_EXPORTED void
 mlirOpPrintingFlagsUseLocalScope(MlirOpPrintingFlags flags);
 
 //===----------------------------------------------------------------------===//
 // Operation API.
 //===----------------------------------------------------------------------===//
 
 /// Creates an operation and transfers ownership to the caller.
-MLIR_CAPI_EXPORTED MlirOperation
-mlirOperationCreate(const MlirOperationState *state);
+/// Note that caller owned child objects are transferred in this call and must
+/// not be further used. Particularly, this applies to any regions added to
+/// the state (the implementation may invalidate any such pointers).
+///
+/// This call can fail under the following conditions, in which case, it will
+/// return a null operation and emit diagnostics:
+///   - Result type inference is enabled and cannot be performed.
+MLIR_CAPI_EXPORTED MlirOperation mlirOperationCreate(MlirOperationState *state);
 
 /// Takes an operation owned by the caller and destroys it.
 MLIR_CAPI_EXPORTED void mlirOperationDestroy(MlirOperation op);
 
 /// Checks whether the underlying operation is null.
 static inline bool mlirOperationIsNull(MlirOperation op) { return !op.ptr; }
 
 /** Checks whether two operation handles point to the same operation. This does

mlir/lib/CAPI/IR/IR.cpp

 #include "mlir/CAPI/Support.h"
 #include "mlir/CAPI/Utils.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Verifier.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Parser.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Context API.
 //===----------------------------------------------------------------------===//
 
MlirOperationState mlirOperationStateGet(MlirStringRef name, MlirLocation loc) {
   state.nOperands = 0;
   state.operands = nullptr;
   state.nRegions = 0;
   state.regions = nullptr;
   state.nSuccessors = 0;
   state.successors = nullptr;
   state.nAttributes = 0;
   state.attributes = nullptr;
+  state.enableResultTypeInference = false;
   return state;
 }
 
 #define APPEND_ELEMS(type, sizeName, elemName)                                 \
   state->elemName =                                                            \
       (type *)realloc(state->elemName, (state->sizeName + n) * sizeof(type));  \
   memcpy(state->elemName + state->sizeName, elemName, n * sizeof(type));       \
   state->sizeName += n;
void mlirOperationStateAddSuccessors(MlirOperationState *state, intptr_t n,
                                      MlirBlock const *successors) {
   APPEND_ELEMS(MlirBlock, nSuccessors, successors);
 }
 void mlirOperationStateAddAttributes(MlirOperationState *state, intptr_t n,
                                      MlirNamedAttribute const *attributes) {
   APPEND_ELEMS(MlirNamedAttribute, nAttributes, attributes);
 }
 
+void mlirOperationStateEnableResultTypeInference(MlirOperationState *state) {
+  state->enableResultTypeInference = true;
+}
+
 //===----------------------------------------------------------------------===//
 // Operation API.
 //===----------------------------------------------------------------------===//
 
-MlirOperation mlirOperationCreate(const MlirOperationState *state) {
+static LogicalResult inferOperationTypes(OperationState &state) {
+  MLIRContext *context = state.getContext();
+  const AbstractOperation *abstractOp =
+      AbstractOperation::lookup(state.name.getStringRef(), context);
+  if (!abstractOp) {
+    emitError(state.location)
+        << "type inference was requested for the operation " << state.name
+        << ", but the operation was not registered. Ensure that the dialect "
+           "containing the operation is linked into MLIR and registered with "
+           "the context";
+    return failure();
+  }
+
+  // Fallback to inference via an op interface.
+  auto *inferInterface = abstractOp->getInterface<InferTypeOpInterface>();
+  if (!inferInterface) {
+    emitError(state.location)
+        << "type inference was requested for the operation " << state.name
+        << ", but the operation does not support type inference. Result "
+           "types must be specified explicitly.";
+    return failure();
+  }
+
+  if (succeeded(inferInterface->inferReturnTypes(
+          context, state.location, state.operands,
+          state.attributes.getDictionary(context), state.regions, state.types)))
+    return success();
+
+  // Diagnostic emitted by interface.
+  return failure();
+}
+
+MlirOperation mlirOperationCreate(MlirOperationState *state) {
   assert(state);
   OperationState cppState(unwrap(state->location), unwrap(state->name));
   SmallVector<Type, 4> resultStorage;
   SmallVector<Value, 8> operandStorage;
   SmallVector<Block *, 2> successorStorage;
   cppState.addTypes(unwrapList(state->nResults, state->results, resultStorage));
   cppState.addOperands(
       unwrapList(state->nOperands, state->operands, operandStorage));
   cppState.addSuccessors(
       unwrapList(state->nSuccessors, state->successors, successorStorage));
 
   cppState.attributes.reserve(state->nAttributes);
   for (intptr_t i = 0; i < state->nAttributes; ++i)
     cppState.addAttribute(unwrap(state->attributes[i].name),
                           unwrap(state->attributes[i].attribute));
 
   for (intptr_t i = 0; i < state->nRegions; ++i)
     cppState.addRegion(std::unique_ptr<Region>(unwrap(state->regions[i])));
 
-  MlirOperation result = wrap(Operation::create(cppState));
   free(state->results);
   free(state->operands);
   free(state->successors);
   free(state->regions);
   free(state->attributes);
+
+  // Infer result types.
+  if (state->enableResultTypeInference) {
+    assert(cppState.types.empty() &&
+           "result type inference enabled and result types provided");
+    if (failed(inferOperationTypes(cppState)))
+      return {nullptr};
+  }
+
+  MlirOperation result = wrap(Operation::create(cppState));
   return result;
 }
 
 void mlirOperationDestroy(MlirOperation op) { unwrap(op)->erase(); }
 
 bool mlirOperationEqual(MlirOperation op, MlirOperation other) {
   return unwrap(op) == unwrap(other);
 }

mlir/test/CAPI/ir.c

static void buildWithInsertionsAndPrint(MlirContext ctx) {
   // CHECK:        "dummy.op5"
   // CHECK:      ^{{.*}}(%{{.*}}: i3
   // CHECK:        "dummy.op6"
   // CHECK:      ^{{.*}}(%{{.*}}: i4
   // CHECK:        "dummy.op7"
   // clang-format on
 }
 
+/// Creates operations with type inference and tests various failure modes.
+static int createOperationWithTypeInference(MlirContext ctx) {
+  MlirLocation loc = mlirLocationUnknownGet(ctx);
+  MlirAttribute iAttr = mlirIntegerAttrGet(mlirIntegerTypeGet(ctx, 32), 4);
+
+  // The shape.const_size op implements result type inference and is only used
+  // for that reason.
+  MlirOperationState state = mlirOperationStateGet(
+      mlirStringRefCreateFromCString("shape.const_size"), loc);
+  MlirNamedAttribute valueAttr = mlirNamedAttributeGet(
+      mlirIdentifierGet(ctx, mlirStringRefCreateFromCString("value")), iAttr);
+  mlirOperationStateAddAttributes(&state, 1, &valueAttr);
+  mlirOperationStateEnableResultTypeInference(&state);
+
+  // Expect result type inference to succeed.
+  MlirOperation op = mlirOperationCreate(&state);
+  if (mlirOperationIsNull(op)) {
+    fprintf(stderr, "ERROR: Result type inference unexpectedly failed");
+    return 1;
+  }
+
+  // CHECK: RESULT_TYPE_INFERENCE: !shape.size
+  fprintf(stderr, "RESULT_TYPE_INFERENCE: ");
+  mlirTypeDump(mlirValueGetType(mlirOperationGetResult(op, 0)));
+  fprintf(stderr, "\n");
+  mlirOperationDestroy(op);
+  return 0;
+}
+
 /// Dumps instances of all builtin types to check that C API works correctly.
 /// Additionally, performs simple identity checks that a builtin type
 /// constructed with C API can be inspected and has the expected type. The
 /// latter achieves full coverage of C API for builtin types. Returns 0 on
 /// success and a non-zero error code on failure.
 static int printBuiltinTypes(MlirContext ctx) {
   // Integer types.
   MlirType i32 = mlirIntegerTypeGet(ctx, 32);
int printBuiltinAttributes(MlirContext ctx) {
   uint32_t *uint32RawData =
       (uint32_t *)mlirDenseElementsAttrGetRawData(uint32Elements);
   int32_t *int32RawData =
       (int32_t *)mlirDenseElementsAttrGetRawData(int32Elements);
   uint64_t *uint64RawData =
       (uint64_t *)mlirDenseElementsAttrGetRawData(uint64Elements);
   int64_t *int64RawData =
       (int64_t *)mlirDenseElementsAttrGetRawData(int64Elements);
-  float *floatRawData =
-      (float *)mlirDenseElementsAttrGetRawData(floatElements);
+  float *floatRawData = (float *)mlirDenseElementsAttrGetRawData(floatElements);
   double *doubleRawData =
       (double *)mlirDenseElementsAttrGetRawData(doubleElements);
   if (uint32RawData[0] != 0u || uint32RawData[1] != 1u ||
-      int32RawData[0] != 0 || int32RawData[1] != 1 ||
-      uint64RawData[0] != 0u || uint64RawData[1] != 1u ||
-      int64RawData[0] != 0 || int64RawData[1] != 1 ||
+      int32RawData[0] != 0 || int32RawData[1] != 1 || uint64RawData[0] != 0u ||
+      uint64RawData[1] != 1u || int64RawData[0] != 0 || int64RawData[1] != 1 ||
       floatRawData[0] != 0.0f || floatRawData[1] != 1.0f ||
       doubleRawData[0] != 0.0 || doubleRawData[1] != 1.0)
     return 18;
 
   mlirAttributeDump(splatBool);
   mlirAttributeDump(splatUInt32);
   mlirAttributeDump(splatInt32);
   mlirAttributeDump(splatUInt64);

 }
 
 int main() {
   MlirContext ctx = mlirContextCreate();
   mlirRegisterAllDialects(ctx);
   if (constructAndTraverseIr(ctx))
     return 1;
   buildWithInsertionsAndPrint(ctx);
+  if (createOperationWithTypeInference(ctx))
+    return 2;
 
   if (printBuiltinTypes(ctx))
-    return 2;
-  if (printBuiltinAttributes(ctx))
     return 3;
-  if (printAffineMap(ctx))
+  if (printBuiltinAttributes(ctx))
     return 4;
-  if (printAffineExpr(ctx))
+  if (printAffineMap(ctx))
     return 5;
-  if (affineMapFromExprs(ctx))
+  if (printAffineExpr(ctx))
     return 6;
-  if (registerOnlyStd())
+  if (affineMapFromExprs(ctx))
     return 7;
+  if (registerOnlyStd())
+    return 8;
 
   mlirContextDestroy(ctx);
 
   testDiagnostics();
   return 0;
 }