[mlir] Expose affine expression to C API

[zhanghb97]

Summary

This patch provides C API for MLIR affine expression.

• Implement C API for methods of AffineExpr class.
• Implement C API for methods of derived classes (AffineBinaryOpExpr, AffineDimExpr, AffineSymbolExpr, and AffineConstantExpr).
• Expose IsA and Get methods for each affine binary operation expression.

[ftynse]

mlir/include/mlir-c/AffineExpr.h

22

I am strongly interested in changing this and not exposing the enum here, even though the C++ API does. To me, this feels like an early design iteration that was never migrated to the typeid/isa scheme. Could we remove this, the GetKind method and instead provide: mlirAffineExprIsADim, mlirAffineExprIsAFloorDiv for checking (including the top-level mlirAffineExprIsABinaryExpr that matches all binary expressions) and mlirAffineFoorDivExprGet etc. that match the already existing mlirAffineDimExprGet?

WDYT? @stellaraccident?

mlir/lib/CAPI/IR/AffineExpr.cpp

122

I'd suggest adding a helper function that takes MlirAffineExprKind and returns an mlir::AffineExprKind. Then this function reduces to return wrap(getAffineBinaryOpExpr(convertKind(kind), unwrap(lhs), unwrap(rhs));

[zhanghb97]

mlir/include/mlir-c/AffineExpr.h

22

I noticed that Diagnostics C API exposes the enum MlirDiagnosticSeverity so I used the enum here because these two are similar. Now I think the isa scheme is more consistent with other APIs. Looking forward to Stella's opinion.
I also want to confirm that if we use the scheme, we should

• remove the enum
• change the getKind to IsA for each kind
• change mlirAffineBinaryOpExprGet to mlirAffine***ExprGet for each binary operation.

Is that correct?

[ftynse]

mlir/include/mlir-c/AffineExpr.h

22

Yes, I guessed that you followed the diagnostics. The reason why I don't like it for the AffineExpr hierarchy is that it does have subclasses, but only for some expressions (e.g., AffineDimExpr is a proper subclass, but AffineAddExpr is not). And worse, the kind enum serves both for subclassing purposes and for further subdivision within a class. This creates all sorts of problems about what should be handled when, leading to non-exhaustive switch conditions like you probably noticed given the number of llvm_unreachable you had to include.

And yes, your understanding of my proposal is correct.

[zhanghb97]

Rebase.

Use the typeid/isa scheme, expose IsA and Get methods for each affine binary operation expression.

mlir/include/mlir-c/AffineExpr.h

+/*===-- mlir-c/AffineExpr.h - C API for MLIR Affine Expressions ---*- C -*-===*\
+|*                                                                            *|
+|* 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                    *|
+|*                                                                            *|
+\*===----------------------------------------------------------------------===*/
+
+#ifndef MLIR_C_AFFINEEXPR_H
+#define MLIR_C_AFFINEEXPR_H
+
+#include "mlir-c/AffineMap.h"
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+DEFINE_C_API_STRUCT(MlirAffineExpr, const void);
+
+/** Kind of a affine expression. */
+enum MlirAffineExprKind {
+  MlirAffineExprAdd,
+  MlirAffineExprMul,
+  MlirAffineExprMod,
+  MlirAffineExprFloorDiv,
+  MlirAffineExprCeilDiv,
+  MlirAffineExprConstant,
+  MlirAffineExprDimId,
+  MlirAffineExprSymbolId,
+};
+typedef enum MlirAffineExprKind MlirAffineExprKind;
+
+/** Gets the context that owns the affine expression. */
+MlirContext mlirAffineExprGetContext(MlirAffineExpr affineExpr);
+
+/** Gets the kind of the given affine expression. */
+MlirAffineExprKind mlirAffineExprGetKind(MlirAffineExpr affineExpr);
+
+/** Prints an affine expression 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 mlirAffineExprPrint(MlirAffineExpr affineExpr, MlirStringCallback callback,
+                         void *userData);
+
+/** Prints the affine expression to the standard error stream. */
+void mlirAffineExprDump(MlirAffineExpr affineExpr);
+
+/** Checks whether the given affine expression is made out of only symbols and
+ * constants. */
+int mlirAffineExprIsSymbolicOrConstant(MlirAffineExpr affineExpr);
+
+/** Checks whether the given affine expression is a pure affine expression, i.e.
+ * mul, floordiv, ceildic, and mod is only allowed w.r.t constants. */
+int mlirAffineExprIsPureAffine(MlirAffineExpr affineExpr);
+
+/** Returns the greatest known integral divisor of this affine expression. The
+ * result is always positive. */
+int64_t mlirAffineExprGetLargestKnownDivisor(MlirAffineExpr affineExpr);
+
+/** Checks whether the given affine expression is a multiple of 'factor'. */
+int mlirAffineExprIsMultipleOf(MlirAffineExpr affineExpr, int64_t factor);
+
+/** Checks whether the given affine expression involves AffineDimExpr
+ * 'position'. */
+int mlirAffineExprIsFunctionOfDim(MlirAffineExpr affineExpr, intptr_t position);
+
+/*============================================================================*/
+/* Affine Dimension Expression.                                               */
+/*============================================================================*/
+
+/** Creates an affine dimension expression with 'position' in the context. */
+MlirAffineExpr mlirAffineDimExprGet(MlirContext ctx, intptr_t position);
+
+/** Returns the position of the given affine dimension expression. */
+intptr_t mlirAffineDimExprGetPosition(MlirAffineExpr affineExpr);
+
+/*============================================================================*/
+/* Affine Symbol Expression.                                                  */
+/*============================================================================*/
+
+/** Creates an affine symbol expression with 'position' in the context. */
+MlirAffineExpr mlirAffineSymbolExprGet(MlirContext ctx, intptr_t position);
+
+/** Returns the position of the given affine symbol expression. */
+intptr_t mlirAffineSymbolExprGetPosition(MlirAffineExpr affineExpr);
+
+/*============================================================================*/
+/* Affine Constant Expression.                                                */
+/*============================================================================*/
+
+/** Creates an affine constant expression with 'constant' in the context. */
+MlirAffineExpr mlirAffineConstantExprGet(MlirContext ctx, int64_t constant);
+
+/** Returns the value of the given affine constant expression. */
+int64_t mlirAffineConstantExprGetValue(MlirAffineExpr affineExpr);
+
+/*============================================================================*/
+/* Affine Binary Operation Expression.                                        */
+/*============================================================================*/
+
+/** Creates a specific 'kind' affine binary operation expression with its 'lhs'
+ * and 'rhs'. */
+MlirAffineExpr mlirAffineBinaryOpExprGet(MlirAffineExprKind kind,
+                                         MlirAffineExpr lhs,
+                                         MlirAffineExpr rhs);
+
+/** Returns the left hand side affine expression of the given affine binary
+ * operation expression. */
+MlirAffineExpr mlirAffineBinaryOpExprGetLHS(MlirAffineExpr affineExpr);
+
+/** Returns the right hand side affine expression of the given affine binary
+ * operation expression. */
+MlirAffineExpr mlirAffineBinaryOpExprGetRHS(MlirAffineExpr affineExpr);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_AFFINEEXPR_H

mlir/include/mlir/CAPI/AffineExpr.h

+//===- AffineExpr.h - C API Utils for Affine Expressions --------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains declarations of implementation details of the C API for
+// MLIR Affine Expression. This file should not be included from C++ code other
+// than C API implementation nor from C code.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CAPI_AFFINEEXPR_H
+#define MLIR_CAPI_AFFINEEXPR_H
+
+#include "mlir-c/AffineExpr.h"
+#include "mlir/CAPI/Wrap.h"
+#include "mlir/IR/AffineExpr.h"
+
+DEFINE_C_API_METHODS(MlirAffineExpr, mlir::AffineExpr)
+
+#endif // MLIR_CAPI_AFFINEEXPR_H

mlir/include/mlir/IR/AffineExpr.h

public:
   /// Example:
   ///   expr: `d0 + d2`
   ///   map:  `(d0, d1, d2)[s0, s1] -> (d0 + s1, d1 + s0, d0 + d1 + d2)`
   ///   returned expr: `d0 * 2 + d1 + d2 + s1`
   AffineExpr compose(AffineMap map) const;
 
   friend ::llvm::hash_code hash_value(AffineExpr arg);
 
+  /// Methods supporting C API.
+  const void *getAsOpaquePointer() const {
+    return static_cast<const void *>(expr);
+  }
+  static AffineExpr getFromOpaquePointer(const void *pointer) {
+    return AffineExpr(
+        reinterpret_cast<ImplType *>(const_cast<void *>(pointer)));
+  }
+
 protected:
   ImplType *expr;
 };
 
 /// Affine binary operation expression. An affine binary operation could be an
 /// add, mul, floordiv, ceildiv, or a modulo operation. (Subtraction is
 /// represented through a multiply by -1 and add.) These expressions are always
 /// constructed in a simplified form. For eg., the LHS and RHS operands can't

mlir/lib/CAPI/IR/AffineExpr.cpp

+//===- AffineExpr.cpp - C API for MLIR Affine Expressions -----------------===//
+//
+// 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/AffineExpr.h"
+#include "mlir-c/AffineMap.h"
+#include "mlir-c/IR.h"
+#include "mlir/CAPI/AffineExpr.h"
+#include "mlir/CAPI/AffineMap.h"
+#include "mlir/CAPI/IR.h"
+#include "mlir/CAPI/Utils.h"
+#include "mlir/IR/AffineExpr.h"
+
+using namespace mlir;
+
+MlirContext mlirAffineExprGetContext(MlirAffineExpr affineExpr) {
+  return wrap(unwrap(affineExpr).getContext());
+}
+
+MlirAffineExprKind mlirAffineExprGetKind(MlirAffineExpr affineExpr) {
+  switch (unwrap(affineExpr).getKind()) {
+  case mlir::AffineExprKind::Add:
+    return MlirAffineExprAdd;
+  case mlir::AffineExprKind::Mul:
+    return MlirAffineExprMul;
+  case mlir::AffineExprKind::Mod:
+    return MlirAffineExprMod;
+  case mlir::AffineExprKind::FloorDiv:
+    return MlirAffineExprFloorDiv;
+  case mlir::AffineExprKind::CeilDiv:
+    return MlirAffineExprCeilDiv;
+  case mlir::AffineExprKind::Constant:
+    return MlirAffineExprConstant;
+  case mlir::AffineExprKind::DimId:
+    return MlirAffineExprDimId;
+  case mlir::AffineExprKind::SymbolId:
+    return MlirAffineExprSymbolId;
+  }
+  llvm_unreachable("unhandled affine expression kind");
+}
+
+void mlirAffineExprPrint(MlirAffineExpr affineExpr, MlirStringCallback callback,
+                         void *userData) {
+  mlir::detail::CallbackOstream stream(callback, userData);
+  unwrap(affineExpr).print(stream);
+  stream.flush();
+}
+
+void mlirAffineExprDump(MlirAffineExpr affineExpr) {
+  unwrap(affineExpr).dump();
+}
+
+int mlirAffineExprIsSymbolicOrConstant(MlirAffineExpr affineExpr) {
+  return unwrap(affineExpr).isSymbolicOrConstant();
+}
+
+int mlirAffineExprIsPureAffine(MlirAffineExpr affineExpr) {
+  return unwrap(affineExpr).isPureAffine();
+}
+
+int64_t mlirAffineExprGetLargestKnownDivisor(MlirAffineExpr affineExpr) {
+  return unwrap(affineExpr).getLargestKnownDivisor();
+}
+
+int mlirAffineExprIsMultipleOf(MlirAffineExpr affineExpr, int64_t factor) {
+  return unwrap(affineExpr).isMultipleOf(factor);
+}
+
+int mlirAffineExprIsFunctionOfDim(MlirAffineExpr affineExpr,
+                                  intptr_t position) {
+  return unwrap(affineExpr).isFunctionOfDim(position);
+}
+
+/*============================================================================*/
+/* Affine Dimension Expression.                                               */
+/*============================================================================*/
+
+MlirAffineExpr mlirAffineDimExprGet(MlirContext ctx, intptr_t position) {
+  return wrap(getAffineDimExpr(position, unwrap(ctx)));
+}
+
+intptr_t mlirAffineDimExprGetPosition(MlirAffineExpr affineExpr) {
+  return unwrap(affineExpr).cast<AffineDimExpr>().getPosition();
+}
+
+/*============================================================================*/
+/* Affine Symbol Expression.                                                  */
+/*============================================================================*/
+
+MlirAffineExpr mlirAffineSymbolExprGet(MlirContext ctx, intptr_t position) {
+  return wrap(getAffineSymbolExpr(position, unwrap(ctx)));
+}
+
+intptr_t mlirAffineSymbolExprGetPosition(MlirAffineExpr affineExpr) {
+  return unwrap(affineExpr).cast<AffineSymbolExpr>().getPosition();
+}
+
+/*============================================================================*/
+/* Affine Constant Expression.                                                */
+/*============================================================================*/
+
+MlirAffineExpr mlirAffineConstantExprGet(MlirContext ctx, int64_t constant) {
+  return wrap(getAffineConstantExpr(constant, unwrap(ctx)));
+}
+
+int64_t mlirAffineConstantExprGetValue(MlirAffineExpr affineExpr) {
+  return unwrap(affineExpr).cast<AffineConstantExpr>().getValue();
+}
+
+/*============================================================================*/
+/* Affine Binary Operation Expression.                                        */
+/*============================================================================*/
+
+MlirAffineExpr mlirAffineBinaryOpExprGet(MlirAffineExprKind kind,
+                                         MlirAffineExpr lhs,
+                                         MlirAffineExpr rhs) {
+  switch (kind) {
+  case MlirAffineExprAdd:
+    return wrap(getAffineBinaryOpExpr(mlir::AffineExprKind::Add, unwrap(lhs),
+                                      unwrap(rhs)));
+  case MlirAffineExprMul:
+    return wrap(getAffineBinaryOpExpr(mlir::AffineExprKind::Mul, unwrap(lhs),
+                                      unwrap(rhs)));
+  case MlirAffineExprMod:
+    return wrap(getAffineBinaryOpExpr(mlir::AffineExprKind::Mod, unwrap(lhs),
+                                      unwrap(rhs)));
+  case MlirAffineExprFloorDiv:
+    return wrap(getAffineBinaryOpExpr(mlir::AffineExprKind::FloorDiv,
+                                      unwrap(lhs), unwrap(rhs)));
+  case MlirAffineExprCeilDiv:
+    return wrap(getAffineBinaryOpExpr(mlir::AffineExprKind::CeilDiv,
+                                      unwrap(lhs), unwrap(rhs)));
+  case MlirAffineExprConstant:
+    return wrap(getAffineBinaryOpExpr(mlir::AffineExprKind::Constant,
+                                      unwrap(lhs), unwrap(rhs)));
+  case MlirAffineExprDimId:
+    return wrap(getAffineBinaryOpExpr(mlir::AffineExprKind::DimId, unwrap(lhs),
+                                      unwrap(rhs)));
+  case MlirAffineExprSymbolId:
+    return wrap(getAffineBinaryOpExpr(mlir::AffineExprKind::SymbolId,
+                                      unwrap(lhs), unwrap(rhs)));
+  }
+  llvm_unreachable("unhandled affine expression kind");
+}
+
+MlirAffineExpr mlirAffineBinaryOpExprGetLHS(MlirAffineExpr affineExpr) {
+  return wrap(unwrap(affineExpr).cast<AffineBinaryOpExpr>().getLHS());
+}
+
+MlirAffineExpr mlirAffineBinaryOpExprGetRHS(MlirAffineExpr affineExpr) {
+  return wrap(unwrap(affineExpr).cast<AffineBinaryOpExpr>().getRHS());
+}

mlir/lib/CAPI/IR/CMakeLists.txt

 # Main API.
 add_mlir_library(MLIRCAPIIR
+  AffineExpr.cpp
   AffineMap.cpp
   Diagnostics.cpp
   IR.cpp
   StandardAttributes.cpp
   StandardTypes.cpp
   Support.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir-c
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRParser
   MLIRSupport
   )

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/AffineExpr.h"
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Registration.h"
 #include "mlir-c/StandardAttributes.h"
 #include "mlir-c/StandardDialect.h"
 #include "mlir-c/StandardTypes.h"
 
int printAffineMap(MlirContext ctx) {
 
   mlirAffineMapDump(subMap);
   mlirAffineMapDump(majorSubMap);
   mlirAffineMapDump(minorSubMap);
 
   return 0;
 }
 
+int printAffineExpr(MlirContext ctx) {
+  MlirAffineExpr affineDimExpr = mlirAffineDimExprGet(ctx, 5);
+  MlirAffineExpr affineSymbolExpr = mlirAffineSymbolExprGet(ctx, 5);
+  MlirAffineExpr affineConstantExpr = mlirAffineConstantExprGet(ctx, 5);
+  MlirAffineExpr affineBinaryOpExpr = mlirAffineBinaryOpExprGet(
+      MlirAffineExprAdd, affineDimExpr, affineSymbolExpr);
+
+  // Tests mlirAffineExprDump.
+  mlirAffineExprDump(affineDimExpr);
+  mlirAffineExprDump(affineSymbolExpr);
+  mlirAffineExprDump(affineConstantExpr);
+  mlirAffineExprDump(affineBinaryOpExpr);
+
+  // Tests methods of affine binary operation expression.
+  mlirAffineExprDump(mlirAffineBinaryOpExprGetLHS(affineBinaryOpExpr));
+  mlirAffineExprDump(mlirAffineBinaryOpExprGetRHS(affineBinaryOpExpr));
+
+  // Tests methods of affine dimension expression.
+  if (mlirAffineDimExprGetPosition(affineDimExpr) != 5)
+    return 1;
+
+  // Tests methods of affine symbol expression.
+  if (mlirAffineSymbolExprGetPosition(affineSymbolExpr) != 5)
+    return 2;
+
+  // Tests methods of affine constant expression.
+  if (mlirAffineConstantExprGetValue(affineConstantExpr) != 5)
+    return 3;
+
+  // Tests methods of affine expression.
+  if (mlirAffineExprGetKind(affineDimExpr) != MlirAffineExprDimId ||
+      mlirAffineExprGetKind(affineSymbolExpr) != MlirAffineExprSymbolId ||
+      mlirAffineExprGetKind(affineConstantExpr) != MlirAffineExprConstant ||
+      mlirAffineExprGetKind(affineBinaryOpExpr) != MlirAffineExprAdd)
+    return 4;
+
+  if (mlirAffineExprIsSymbolicOrConstant(affineDimExpr) ||
+      !mlirAffineExprIsSymbolicOrConstant(affineSymbolExpr) ||
+      !mlirAffineExprIsSymbolicOrConstant(affineConstantExpr) ||
+      mlirAffineExprIsSymbolicOrConstant(affineBinaryOpExpr))
+    return 5;
+
+  if (!mlirAffineExprIsPureAffine(affineDimExpr) ||
+      !mlirAffineExprIsPureAffine(affineSymbolExpr) ||
+      !mlirAffineExprIsPureAffine(affineConstantExpr) ||
+      !mlirAffineExprIsPureAffine(affineBinaryOpExpr))
+    return 6;
+
+  if (mlirAffineExprGetLargestKnownDivisor(affineDimExpr) != 1 ||
+      mlirAffineExprGetLargestKnownDivisor(affineSymbolExpr) != 1 ||
+      mlirAffineExprGetLargestKnownDivisor(affineConstantExpr) != 5 ||
+      mlirAffineExprGetLargestKnownDivisor(affineBinaryOpExpr) != 1)
+    return 7;
+
+  if (!mlirAffineExprIsMultipleOf(affineDimExpr, 1) ||
+      !mlirAffineExprIsMultipleOf(affineSymbolExpr, 1) ||
+      !mlirAffineExprIsMultipleOf(affineConstantExpr, 5) ||
+      !mlirAffineExprIsMultipleOf(affineBinaryOpExpr, 1))
+    return 8;
+
+  if (!mlirAffineExprIsFunctionOfDim(affineDimExpr, 5) ||
+      mlirAffineExprIsFunctionOfDim(affineSymbolExpr, 5) ||
+      mlirAffineExprIsFunctionOfDim(affineConstantExpr, 5) ||
+      !mlirAffineExprIsFunctionOfDim(affineBinaryOpExpr, 5))
+    return 9;
+
+  return 0;
+}
+
 int registerOnlyStd() {
   MlirContext ctx = mlirContextCreate();
   // The built-in dialect is always loaded.
   if (mlirContextGetNumLoadedDialects(ctx) != 1)
     return 1;
 
   MlirDialect std =
       mlirContextGetOrLoadDialect(ctx, mlirStandardDialectGetNamespace());
int main() {
   // CHECK: (d0, d1, d2) -> (d0)
   // CHECK: (d0, d1, d2) -> (d2)
   // CHECK: 0
   // clang-format on
   fprintf(stderr, "@affineMap\n");
   errcode = printAffineMap(ctx);
   fprintf(stderr, "%d\n", errcode);
 
+  // clang-format off
+  // CHECK: d5
+  // CHECK: s5
+  // CHECK: 5
+  // CHECK: d5 + s5
+  // CHECK: d5
+  // CHECK: s5
+  // CHECK: 0
+  // clang-format on
+  fprintf(stderr, "@affineExpr\n");
+  errcode = printAffineExpr(ctx);
+  fprintf(stderr, "%d\n", errcode);
+
   fprintf(stderr, "@registration\n");
   errcode = registerOnlyStd();
   fprintf(stderr, "%d\n", errcode);
   // clang-format off
   // CHECK-LABEL: @registration
   // CHECK: 0
   // clang-format on