Add FunctionType to MLIR C and Python bindings.

[ftynse]

mlir/lib/CAPI/IR/StandardTypes.cpp

321–335

Nit: consider adding explicit casts between signed (intptr_t) and unsigned integers to avoid warnings.

[stellaraccident]

Address comments.

[zhanghb97]

mlir/lib/Bindings/Python/IRModules.cpp

1300

I notice that some def/def_static take py::arg, while others don't take. When do we need to use py::arg? What is the purpose of using py::arg?

[stellaraccident]

Rebase.

[stellaraccident]

mlir/lib/Bindings/Python/IRModules.cpp

1300

It is generally a good idea to include py::arg wherever there is an API where people may want to invoke things with keywords or have other hints to that effect (i.e. most of them). Any that are missing should not be considered intentional unless if it is a purely positional/internal/dunder API.

mlir/include/mlir-c/StandardTypes.h

MlirType mlirTupleTypeGet(MlirContext ctx, intptr_t numElements,
                           MlirType *elements);
 
 /** Returns the number of types contained in a tuple. */
 intptr_t mlirTupleTypeGetNumTypes(MlirType type);
 
 /** Returns the pos-th type in the tuple type. */
 MlirType mlirTupleTypeGetType(MlirType type, intptr_t pos);
 
+/*============================================================================*/
+/* Function type.                                                             */
+/*============================================================================*/
+
+/** Checks whether the given type is a function type. */
+int mlirTypeIsAFunction(MlirType type);
+
+/** Creates a function type, mapping a list of input types to result types. */
+MlirType mlirFunctionTypeGet(MlirContext ctx, intptr_t numInputs,
+                             MlirType *inputs, intptr_t numResults,
+                             MlirType *results);
+
+/** Returns the number of input types. */
+intptr_t mlirFunctionTypeGetNumInputs(MlirType type);
+
+/** Returns the number of result types. */
+intptr_t mlirFunctionTypeGetNumResults(MlirType type);
+
+/** Returns the pos-th input type. */
+MlirType mlirFunctionTypeGetInput(MlirType type, intptr_t pos);
+
+/** Returns the pos-th result type. */
+MlirType mlirFunctionTypeGetResult(MlirType type, intptr_t pos);
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_STANDARDTYPES_H

mlir/lib/Bindings/Python/IRModules.cpp

c.def_property_readonly(
         "num_types",
         [](PyTupleType &self) -> intptr_t {
           return mlirTupleTypeGetNumTypes(self.type);
         },
         "Returns the number of types contained in a tuple.");
   }
 };
 
+/// Function type.
+class PyFunctionType : public PyConcreteType<PyFunctionType> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFunction;
+  static constexpr const char *pyClassName = "FunctionType";
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](PyMlirContext &context, std::vector<PyType> inputs,
+           std::vector<PyType> results) {
+          SmallVector<MlirType, 4> inputsRaw(inputs.begin(), inputs.end());
+          SmallVector<MlirType, 4> resultsRaw(results.begin(), results.end());
+          MlirType t = mlirFunctionTypeGet(context.get(), inputsRaw.size(),
+                                           inputsRaw.data(), resultsRaw.size(),
+                                           resultsRaw.data());
+          return PyFunctionType(context.getRef(), t);
+        },
+        py::arg("context"), py::arg("inputs"), py::arg("results"),
+        "Gets a FunctionType from a list of input and result types");
+    c.def_property_readonly(
+        "inputs",
+        [](PyFunctionType &self) {
+          MlirType t = self.type;
+          auto contextRef = self.getContext();
+          py::list types;
+          for (intptr_t i = 0, e = mlirFunctionTypeGetNumInputs(self.type);
+               i < e; ++i) {
+            types.append(PyType(contextRef, mlirFunctionTypeGetInput(t, i)));
+          }
+          return types;
+        },
+        "Returns the list of input types in the FunctionType.");
+    c.def_property_readonly(
+        "results",
+        [](PyFunctionType &self) {
+          MlirType t = self.type;
+          auto contextRef = self.getContext();
+          py::list types;
+          for (intptr_t i = 0, e = mlirFunctionTypeGetNumResults(self.type);
+               i < e; ++i) {
+            types.append(PyType(contextRef, mlirFunctionTypeGetResult(t, i)));
+          }
+          return types;
+        },
+        "Returns the list of result types in the FunctionType.");
+  }
+};
+
 } // namespace
 
 //------------------------------------------------------------------------------
 // Populates the pybind11 IR submodule.
 //------------------------------------------------------------------------------
 
 void mlir::python::populateIRSubmodule(py::module &m) {
   // Mapping of MlirContext
void mlir::python::populateIRSubmodule(py::module &m) {
   PyComplexType::bind(m);
   PyShapedType::bind(m);
   PyVectorType::bind(m);
   PyRankedTensorType::bind(m);
   PyUnrankedTensorType::bind(m);
   PyMemRefType::bind(m);
   PyUnrankedMemRefType::bind(m);
   PyTupleType::bind(m);
+  PyFunctionType::bind(m);
 
   // Container bindings.
   PyBlockIterator::bind(m);
   PyBlockList::bind(m);
   PyOperationIterator::bind(m);
   PyOperationList::bind(m);
   PyRegionIterator::bind(m);
   PyRegionList::bind(m);
 }

mlir/lib/CAPI/IR/StandardTypes.cpp

 //===- StandardTypes.cpp - C Interface to MLIR Standard Types -------------===//
 //
 // 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 "mlir-c/StandardTypes.h"
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/IR.h"
 #include "mlir/CAPI/AffineMap.h"
 #include "mlir/CAPI/IR.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/StandardTypes.h"
+#include "mlir/IR/Types.h"
 
 using namespace mlir;
 
 /* ========================================================================== */
 /* Integer types.                                                             */
 /* ========================================================================== */
 
 int mlirTypeIsAInteger(MlirType type) {

 
 intptr_t mlirTupleTypeGetNumTypes(MlirType type) {
   return unwrap(type).cast<TupleType>().size();
 }
 
 MlirType mlirTupleTypeGetType(MlirType type, intptr_t pos) {
   return wrap(unwrap(type).cast<TupleType>().getType(static_cast<size_t>(pos)));
 }
+
+/*============================================================================*/
+/* Function type.                                                             */
+/*============================================================================*/
+
+int mlirTypeIsAFunction(MlirType type) {
+  return unwrap(type).isa<FunctionType>();
+}
+
+MlirType mlirFunctionTypeGet(MlirContext ctx, intptr_t numInputs,
+                             MlirType *inputs, intptr_t numResults,
+                             MlirType *results) {
+  SmallVector<Type, 4> inputsList;
+  SmallVector<Type, 4> resultsList;
+  (void)unwrapList(numInputs, inputs, inputsList);
+  (void)unwrapList(numResults, results, resultsList);
+  return wrap(FunctionType::get(inputsList, resultsList, unwrap(ctx)));
+}
+
+intptr_t mlirFunctionTypeGetNumInputs(MlirType type) {
+  return unwrap(type).cast<FunctionType>().getNumInputs();
+}
+
+intptr_t mlirFunctionTypeGetNumResults(MlirType type) {
+  return unwrap(type).cast<FunctionType>().getNumResults();
+}
+
+MlirType mlirFunctionTypeGetInput(MlirType type, intptr_t pos) {
+  assert(pos >= 0 && "pos in array must be positive");
+  return wrap(
+      unwrap(type).cast<FunctionType>().getInput(static_cast<unsigned>(pos)));
+}
+
+MlirType mlirFunctionTypeGetResult(MlirType type, intptr_t pos) {
+  assert(pos >= 0 && "pos in array must be positive");
+  return wrap(
+      unwrap(type).cast<FunctionType>().getResult(static_cast<unsigned>(pos)));
+}

mlir/test/Bindings/Python/ir_types.py

def testTupleType():
   # CHECK: tuple type: tuple<i32, f32, vector<2x3xf32>>
   print("tuple type:", tuple_type)
   # CHECK: number of types: 3
   print("number of types:", tuple_type.num_types)
   # CHECK: pos-th type in the tuple type: f32
   print("pos-th type in the tuple type:", tuple_type.get_type(1))
 
 run(testTupleType)
+
+
+# CHECK-LABEL: TEST: testFunctionType
+def testFunctionType():
+  ctx = mlir.ir.Context()
+  input_types = [mlir.ir.IntegerType.get_signless(ctx, 32),
+                 mlir.ir.IntegerType.get_signless(ctx, 16)]
+  result_types = [mlir.ir.IndexType(ctx)]
+  func = mlir.ir.FunctionType.get(ctx, input_types, result_types)
+  # CHECK: INPUTS: [Type(i32), Type(i16)]
+  print("INPUTS:", func.inputs)
+  # CHECK: RESULTS: [Type(index)]
+  print("RESULTS:", func.results)
+
+
+run(testFunctionType)

mlir/test/CAPI/ir.c

 /*===- ir.c - Simple test of C APIs ---------------------------------------===*\
 |*                                                                            *|
 |* 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                    *|
 |*                                                                            *|
 \*===----------------------------------------------------------------------===*/
 
 /* RUN: mlir-capi-ir-test 2>&1 | FileCheck %s
  */
 
-#include "mlir-c/AffineMap.h"
 #include "mlir-c/IR.h"
+#include "mlir-c/AffineMap.h"
 #include "mlir-c/Registration.h"
 #include "mlir-c/StandardAttributes.h"
 #include "mlir-c/StandardTypes.h"
 
 #include <assert.h>
 #include <math.h>
 #include <stdio.h>
 #include <stdlib.h>
static int printStandardTypes(MlirContext ctx) {
   MlirType tuple = mlirTupleTypeGet(ctx, 2, types);
   if (!mlirTypeIsATuple(tuple) || mlirTupleTypeGetNumTypes(tuple) != 2 ||
       !mlirTypeEqual(mlirTupleTypeGetType(tuple, 0), unrankedMemRef) ||
       !mlirTypeEqual(mlirTupleTypeGetType(tuple, 1), f32))
     return 20;
   mlirTypeDump(tuple);
   fprintf(stderr, "\n");
 
+  // Function type.
+  MlirType funcInputs[2] = {mlirIndexTypeGet(ctx), mlirIntegerTypeGet(ctx, 1)};
+  MlirType funcResults[3] = {mlirIntegerTypeGet(ctx, 16),
+                             mlirIntegerTypeGet(ctx, 32),
+                             mlirIntegerTypeGet(ctx, 64)};
+  MlirType funcType = mlirFunctionTypeGet(ctx, 2, funcInputs, 3, funcResults);
+  if (mlirFunctionTypeGetNumInputs(funcType) != 2)
+    return 21;
+  if (mlirFunctionTypeGetNumResults(funcType) != 3)
+    return 22;
+  if (!mlirTypeEqual(funcInputs[0], mlirFunctionTypeGetInput(funcType, 0)) ||
+      !mlirTypeEqual(funcInputs[1], mlirFunctionTypeGetInput(funcType, 1)))
+    return 23;
+  if (!mlirTypeEqual(funcResults[0], mlirFunctionTypeGetResult(funcType, 0)) ||
+      !mlirTypeEqual(funcResults[1], mlirFunctionTypeGetResult(funcType, 1)) ||
+      !mlirTypeEqual(funcResults[2], mlirFunctionTypeGetResult(funcType, 2)))
+    return 24;
+  mlirTypeDump(funcType);
+  fprintf(stderr, "\n");
+
   return 0;
 }
 
 void callbackSetFixedLengthString(const char *data, intptr_t len,
                                   void *userData) {
   strncpy(userData, data, len);
 }
 
int printAffineMap(MlirContext ctx) {
   if (!mlirAffineMapIsMinorIdentity(emptyAffineMap) ||
       mlirAffineMapIsMinorIdentity(affineMap) ||
       !mlirAffineMapIsMinorIdentity(multiDimIdentityAffineMap) ||
       !mlirAffineMapIsMinorIdentity(minorIdentityAffineMap) ||
       mlirAffineMapIsMinorIdentity(permutationAffineMap))
     return 2;
 
   if (!mlirAffineMapIsEmpty(emptyAffineMap) ||
-      mlirAffineMapIsEmpty(affineMap) ||
-      mlirAffineMapIsEmpty(constAffineMap) ||
+      mlirAffineMapIsEmpty(affineMap) || mlirAffineMapIsEmpty(constAffineMap) ||
       mlirAffineMapIsEmpty(multiDimIdentityAffineMap) ||
       mlirAffineMapIsEmpty(minorIdentityAffineMap) ||
       mlirAffineMapIsEmpty(permutationAffineMap))
     return 3;
 
   if (mlirAffineMapIsSingleConstant(emptyAffineMap) ||
       mlirAffineMapIsSingleConstant(affineMap) ||
       !mlirAffineMapIsSingleConstant(constAffineMap) ||
int main() {
   // CHECK: none
   // CHECK: complex<f32>
   // CHECK: vector<2x3xf32>
   // CHECK: tensor<2x3xf32>
   // CHECK: tensor<*xf32>
   // CHECK: memref<2x3xf32, 2>
   // CHECK: memref<*xf32, 4>
   // CHECK: tuple<memref<*xf32, 4>, f32>
+  // CHECK: (index, i1) -> (i16, i32, i64)
   // CHECK: 0
   // clang-format on
   fprintf(stderr, "@types\n");
   int errcode = printStandardTypes(ctx);
   fprintf(stderr, "%d\n", errcode);
 
   // clang-format off
   // CHECK-LABEL: @attrs