diff --git a/mlir/include/mlir-c/AffineMap.h b/mlir/include/mlir-c/AffineMap.h --- a/mlir/include/mlir-c/AffineMap.h +++ b/mlir/include/mlir-c/AffineMap.h @@ -16,7 +16,95 @@ extern "C" { #endif -DEFINE_C_API_STRUCT(MlirAffineMap, const void); +/** Creates a zero result affine map with no dimensions or symbols in the + * context. The affine map is owned by the context. */ +MlirAffineMap mlirAffineMapEmptyGet(MlirContext ctx); + +/** Creates a zero result affine map of the given dimensions and symbols in the + * context. The affine map is owned by the context. */ +MlirAffineMap mlirAffineMapGet(MlirContext ctx, intptr_t dimCount, + intptr_t symbolCount); + +/** Creates a single constant result affine map in the context. The affine map + * is owned by the context. */ +MlirAffineMap mlirAffineMapConstantGet(MlirContext ctx, int64_t val); + +/** Creates an affine map with 'numDims' identity in the context. The affine map + * is owned by the context. */ +MlirAffineMap mlirAffineMapMultiDimIdentityGet(MlirContext ctx, + intptr_t numDims); + +/** Creates an identity affine map on the most minor dimensions in the context. + * The affine map is owned by the context. The function asserts that the number + * of dimensions is greater or equal to the number of results. */ +MlirAffineMap mlirAffineMapMinorIdentityGet(MlirContext ctx, intptr_t dims, + intptr_t results); + +/** Creates an affine map with a permutation expression and its size in the + * context. The permutation expression is a non-empty vector of integers. + * The elements of the permutation vector must be continuous from 0 and cannot + * be repeated (i.e. `[1,2,0]` is a valid permutation. `[2,0]` or `[1,1,2]` is + * an invalid invalid permutation.) The affine map is owned by the context. */ +MlirAffineMap mlirAffineMapPermutationGet(MlirContext ctx, intptr_t size, + unsigned *permutation); + +/** Checks whether the given affine map is an identity affine map. The function + * asserts that the number of dimensions is greater or equal to the number of + * results. */ +int mlirAffineMapIsIdentity(MlirAffineMap affineMap); + +/** Checks whether the given affine map is a minor identity affine map. */ +int mlirAffineMapIsMinorIdentity(MlirAffineMap affineMap); + +/** Checks whether the given affine map is an empty affine map. */ +int mlirAffineMapIsEmpty(MlirAffineMap affineMap); + +/** Checks whether the given affine map is a single result constant affine + * map. */ +int mlirAffineMapIsSingleConstant(MlirAffineMap affineMap); + +/** Returns the constant result of the given affine map. The function asserts + * that the map has a single constant result. */ +int64_t mlirAffineMapGetSingleConstantResult(MlirAffineMap affineMap); + +/** Returns the number of dimensions of the given affine map. */ +unsigned mlirAffineMapGetNumDims(MlirAffineMap affineMap); + +/** Returns the number of symbols of the given affine map. */ +unsigned mlirAffineMapGetNumSymbols(MlirAffineMap affineMap); + +/** Returns the number of results of the given affine map. */ +unsigned mlirAffineMapGetNumResults(MlirAffineMap affineMap); + +/** Returns the number of inputs (dimensions + symbols) of the given affine + * map. */ +unsigned mlirAffineMapGetNumInputs(MlirAffineMap affineMap); + +/** Checks whether the given affine map represents a subset of a symbol-less + * permutation map. */ +int mlirAffineMapIsProjectedPermutation(MlirAffineMap affineMap); + +/** Checks whether the given affine map represents a symbol-less permutation + * map. */ +int mlirAffineMapIsPermutation(MlirAffineMap affineMap); + +/** Returns the affine map consisting of the `resultPos` subset. */ +MlirAffineMap mlirAffineMapGetSubMap(MlirAffineMap affineMap, intptr_t size, + unsigned *resultPos); + +/** Returns the affine map consisting of the most major `numResults` results. + * Returns the null AffineMap if the `numResults` is equal to zero. + * Returns the `affineMap` if `numResults` is greater or equals to number of + * results of the given affine map. */ +MlirAffineMap mlirAffineMapGetMajorSubMap(MlirAffineMap affineMap, + intptr_t numResults); + +/** Returns the affine map consisting of the most minor `numResults` results. + * Returns the null AffineMap if the `numResults` is equal to zero. + * Returns the `affineMap` if `numResults` is greater or equals to number of + * results of the given affine map. */ +MlirAffineMap mlirAffineMapGetMinorSubMap(MlirAffineMap affineMap, + intptr_t numResults); #ifdef __cplusplus } diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h --- a/mlir/include/mlir-c/IR.h +++ b/mlir/include/mlir-c/IR.h @@ -55,6 +55,7 @@ DEFINE_C_API_STRUCT(MlirType, const void); DEFINE_C_API_STRUCT(MlirLocation, const void); DEFINE_C_API_STRUCT(MlirModule, const void); +DEFINE_C_API_STRUCT(MlirAffineMap, const void); /** Named MLIR attribute. * @@ -360,6 +361,30 @@ /** Associates an attribute with the name. Takes ownership of neither. */ MlirNamedAttribute mlirNamedAttributeGet(const char *name, MlirAttribute attr); +/*============================================================================*/ +/* Affine Map API. */ +/*============================================================================*/ + +/** Gets the context that the given affine map was created with*/ +MlirContext mlirAffineMapGetContext(MlirAffineMap affineMap); + +/** Checks whether an affine map is null. */ +inline int mlirAffineMapIsNull(MlirAffineMap affineMap) { + return !affineMap.ptr; +} + +/** Checks if two affine maps are equal. */ +int mlirAffineMapEqual(MlirAffineMap a1, MlirAffineMap a2); + +/** Prints an affine map 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 mlirAffineMapPrint(MlirAffineMap affineMap, MlirStringCallback callback, + void *userData); + +/** Prints the affine map to the standard error stream. */ +void mlirAffineMapDump(MlirAffineMap affineMap); + #ifdef __cplusplus } #endif diff --git a/mlir/lib/CAPI/IR/AffineMap.cpp b/mlir/lib/CAPI/IR/AffineMap.cpp --- a/mlir/lib/CAPI/IR/AffineMap.cpp +++ b/mlir/lib/CAPI/IR/AffineMap.cpp @@ -9,7 +9,98 @@ #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" -// This is a placeholder for affine map bindings. The file is here to serve as a -// compilation unit that includes the headers. +// TODO: expose the C API related to `AffineExpr` and mutable affine map. + +using namespace mlir; + +MlirAffineMap mlirAffineMapEmptyGet(MlirContext ctx) { + return wrap(AffineMap::get(unwrap(ctx))); +} + +MlirAffineMap mlirAffineMapGet(MlirContext ctx, intptr_t dimCount, + intptr_t symbolCount) { + return wrap(AffineMap::get(dimCount, symbolCount, unwrap(ctx))); +} + +MlirAffineMap mlirAffineMapConstantGet(MlirContext ctx, int64_t val) { + return wrap(AffineMap::getConstantMap(val, unwrap(ctx))); +} + +MlirAffineMap mlirAffineMapMultiDimIdentityGet(MlirContext ctx, + intptr_t numDims) { + return wrap(AffineMap::getMultiDimIdentityMap(numDims, unwrap(ctx))); +} + +MlirAffineMap mlirAffineMapMinorIdentityGet(MlirContext ctx, intptr_t dims, + intptr_t results) { + return wrap(AffineMap::getMinorIdentityMap(dims, results, unwrap(ctx))); +} + +MlirAffineMap mlirAffineMapPermutationGet(MlirContext ctx, intptr_t size, + unsigned *permutation) { + return wrap(AffineMap::getPermutationMap( + llvm::makeArrayRef(permutation, static_cast(size)), unwrap(ctx))); +} + +int mlirAffineMapIsIdentity(MlirAffineMap affineMap) { + return unwrap(affineMap).isIdentity(); +} + +int mlirAffineMapIsMinorIdentity(MlirAffineMap affineMap) { + return unwrap(affineMap).isMinorIdentity(); +} + +int mlirAffineMapIsEmpty(MlirAffineMap affineMap) { + return unwrap(affineMap).isEmpty(); +} + +int mlirAffineMapIsSingleConstant(MlirAffineMap affineMap) { + return unwrap(affineMap).isSingleConstant(); +} + +int64_t mlirAffineMapGetSingleConstantResult(MlirAffineMap affineMap) { + return unwrap(affineMap).getSingleConstantResult(); +} + +unsigned mlirAffineMapGetNumDims(MlirAffineMap affineMap) { + return unwrap(affineMap).getNumDims(); +} + +unsigned mlirAffineMapGetNumSymbols(MlirAffineMap affineMap) { + return unwrap(affineMap).getNumSymbols(); +} + +unsigned mlirAffineMapGetNumResults(MlirAffineMap affineMap) { + return unwrap(affineMap).getNumResults(); +} + +unsigned mlirAffineMapGetNumInputs(MlirAffineMap affineMap) { + return unwrap(affineMap).getNumInputs(); +} + +int mlirAffineMapIsProjectedPermutation(MlirAffineMap affineMap) { + return unwrap(affineMap).isProjectedPermutation(); +} + +int mlirAffineMapIsPermutation(MlirAffineMap affineMap) { + return unwrap(affineMap).isPermutation(); +} + +MlirAffineMap mlirAffineMapGetSubMap(MlirAffineMap affineMap, intptr_t size, + unsigned *resultPos) { + return wrap(unwrap(affineMap).getSubMap( + llvm::makeArrayRef(resultPos, static_cast(size)))); +} + +MlirAffineMap mlirAffineMapGetMajorSubMap(MlirAffineMap affineMap, + intptr_t numResults) { + return wrap(unwrap(affineMap).getMajorSubMap(numResults)); +} + +MlirAffineMap mlirAffineMapGetMinorSubMap(MlirAffineMap affineMap, + intptr_t numResults) { + return wrap(unwrap(affineMap).getMinorSubMap(numResults)); +} diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp --- a/mlir/lib/CAPI/IR/IR.cpp +++ b/mlir/lib/CAPI/IR/IR.cpp @@ -8,6 +8,7 @@ #include "mlir-c/IR.h" +#include "mlir/CAPI/AffineMap.h" #include "mlir/CAPI/IR.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Dialect.h" @@ -392,3 +393,24 @@ MlirNamedAttribute mlirNamedAttributeGet(const char *name, MlirAttribute attr) { return MlirNamedAttribute{name, attr}; } + +/*============================================================================*/ +/* Affine Map API. */ +/*============================================================================*/ + +MlirContext mlirAffineMapGetContext(MlirAffineMap affineMap) { + return wrap(unwrap(affineMap).getContext()); +} + +int mlirAffineMapEqual(MlirAffineMap a1, MlirAffineMap a2) { + return unwrap(a1) == unwrap(a2); +} + +void mlirAffineMapPrint(MlirAffineMap affineMap, MlirStringCallback callback, + void *userData) { + CallbackOstream stream(callback, userData); + unwrap(affineMap).print(stream); + stream.flush(); +} + +void mlirAffineMapDump(MlirAffineMap affineMap) { unwrap(affineMap).dump(); } diff --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c --- a/mlir/test/CAPI/ir.c +++ b/mlir/test/CAPI/ir.c @@ -14,6 +14,7 @@ #include "mlir-c/Registration.h" #include "mlir-c/StandardAttributes.h" #include "mlir-c/StandardTypes.h" +#include "mlir-c/AffineMap.h" #include #include @@ -587,6 +588,121 @@ return 0; } +int printAffineMap(MlirContext ctx) { + MlirAffineMap emptyAffineMap = mlirAffineMapEmptyGet(ctx); + MlirAffineMap affineMap = mlirAffineMapGet(ctx, 3, 2); + MlirAffineMap constAffineMap = mlirAffineMapConstantGet(ctx, 2); + MlirAffineMap multiDimIdentityAffineMap = + mlirAffineMapMultiDimIdentityGet(ctx, 3); + MlirAffineMap minorIdentityAffineMap = + mlirAffineMapMinorIdentityGet(ctx, 3, 2); + unsigned permutation[] = {1, 2, 0}; + MlirAffineMap permutationAffineMap = mlirAffineMapPermutationGet( + ctx, sizeof(permutation) / sizeof(unsigned), permutation); + + mlirAffineMapDump(emptyAffineMap); + mlirAffineMapDump(affineMap); + mlirAffineMapDump(constAffineMap); + mlirAffineMapDump(multiDimIdentityAffineMap); + mlirAffineMapDump(minorIdentityAffineMap); + mlirAffineMapDump(permutationAffineMap); + + if (!mlirAffineMapIsIdentity(emptyAffineMap) || + mlirAffineMapIsIdentity(affineMap) || + mlirAffineMapIsIdentity(constAffineMap) || + !mlirAffineMapIsIdentity(multiDimIdentityAffineMap) || + mlirAffineMapIsIdentity(minorIdentityAffineMap) || + mlirAffineMapIsIdentity(permutationAffineMap)) + return 1; + + if (!mlirAffineMapIsMinorIdentity(emptyAffineMap) || + mlirAffineMapIsMinorIdentity(affineMap) || + !mlirAffineMapIsMinorIdentity(multiDimIdentityAffineMap) || + !mlirAffineMapIsMinorIdentity(minorIdentityAffineMap) || + mlirAffineMapIsMinorIdentity(permutationAffineMap)) + return 2; + + if (!mlirAffineMapIsEmpty(emptyAffineMap) || + mlirAffineMapIsEmpty(affineMap) || + mlirAffineMapIsEmpty(constAffineMap) || + mlirAffineMapIsEmpty(multiDimIdentityAffineMap) || + mlirAffineMapIsEmpty(minorIdentityAffineMap) || + mlirAffineMapIsEmpty(permutationAffineMap)) + return 3; + + if (mlirAffineMapIsSingleConstant(emptyAffineMap) || + mlirAffineMapIsSingleConstant(affineMap) || + !mlirAffineMapIsSingleConstant(constAffineMap) || + mlirAffineMapIsSingleConstant(multiDimIdentityAffineMap) || + mlirAffineMapIsSingleConstant(minorIdentityAffineMap) || + mlirAffineMapIsSingleConstant(permutationAffineMap)) + return 4; + + if (mlirAffineMapGetSingleConstantResult(constAffineMap) != 2) + return 5; + + if (mlirAffineMapGetNumDims(emptyAffineMap) != 0 || + mlirAffineMapGetNumDims(affineMap) != 3 || + mlirAffineMapGetNumDims(constAffineMap) != 0 || + mlirAffineMapGetNumDims(multiDimIdentityAffineMap) != 3 || + mlirAffineMapGetNumDims(minorIdentityAffineMap) != 3 || + mlirAffineMapGetNumDims(permutationAffineMap) != 3) + return 6; + + if (mlirAffineMapGetNumSymbols(emptyAffineMap) != 0 || + mlirAffineMapGetNumSymbols(affineMap) != 2 || + mlirAffineMapGetNumSymbols(constAffineMap) != 0 || + mlirAffineMapGetNumSymbols(multiDimIdentityAffineMap) != 0 || + mlirAffineMapGetNumSymbols(minorIdentityAffineMap) != 0 || + mlirAffineMapGetNumSymbols(permutationAffineMap) != 0) + return 7; + + if (mlirAffineMapGetNumResults(emptyAffineMap) != 0 || + mlirAffineMapGetNumResults(affineMap) != 0 || + mlirAffineMapGetNumResults(constAffineMap) != 1 || + mlirAffineMapGetNumResults(multiDimIdentityAffineMap) != 3 || + mlirAffineMapGetNumResults(minorIdentityAffineMap) != 2 || + mlirAffineMapGetNumResults(permutationAffineMap) != 3) + return 8; + + if (mlirAffineMapGetNumInputs(emptyAffineMap) != 0 || + mlirAffineMapGetNumInputs(affineMap) != 5 || + mlirAffineMapGetNumInputs(constAffineMap) != 0 || + mlirAffineMapGetNumInputs(multiDimIdentityAffineMap) != 3 || + mlirAffineMapGetNumInputs(minorIdentityAffineMap) != 3 || + mlirAffineMapGetNumInputs(permutationAffineMap) != 3) + return 9; + + if (!mlirAffineMapIsProjectedPermutation(emptyAffineMap) || + !mlirAffineMapIsPermutation(emptyAffineMap) || + mlirAffineMapIsProjectedPermutation(affineMap) || + mlirAffineMapIsPermutation(affineMap) || + mlirAffineMapIsProjectedPermutation(constAffineMap) || + mlirAffineMapIsPermutation(constAffineMap) || + !mlirAffineMapIsProjectedPermutation(multiDimIdentityAffineMap) || + !mlirAffineMapIsPermutation(multiDimIdentityAffineMap) || + !mlirAffineMapIsProjectedPermutation(minorIdentityAffineMap) || + mlirAffineMapIsPermutation(minorIdentityAffineMap) || + !mlirAffineMapIsProjectedPermutation(permutationAffineMap) || + !mlirAffineMapIsPermutation(permutationAffineMap)) + return 10; + + unsigned sub[] = {1}; + + MlirAffineMap subMap = mlirAffineMapGetSubMap( + multiDimIdentityAffineMap, sizeof(sub) / sizeof(unsigned), sub); + MlirAffineMap majorSubMap = mlirAffineMapGetMajorSubMap( + multiDimIdentityAffineMap, 1); + MlirAffineMap minorSubMap = mlirAffineMapGetMinorSubMap( + multiDimIdentityAffineMap, 1); + + mlirAffineMapDump(subMap); + mlirAffineMapDump(majorSubMap); + mlirAffineMapDump(minorSubMap); + + return 0; +} + int main() { MlirContext ctx = mlirContextCreate(); mlirRegisterAllDialects(ctx); @@ -698,6 +814,22 @@ errcode = printStandardAttributes(ctx); fprintf(stderr, "%d\n", errcode); + // clang-format off + // CHECK-LABEL: @affineMap + // CHECK: () -> () + // CHECK: (d0, d1, d2)[s0, s1] -> () + // CHECK: () -> (2) + // CHECK: (d0, d1, d2) -> (d0, d1, d2) + // CHECK: (d0, d1, d2) -> (d1, d2) + // CHECK: (d0, d1, d2) -> (d1, d2, d0) + // CHECK: (d0, d1, d2) -> (d1) + // CHECK: (d0, d1, d2) -> (d0) + // CHECK: (d0, d1, d2) -> (d2) + // CHECK: 0 + fprintf(stderr, "@affineMap\n"); + errcode = printAffineMap(ctx); + fprintf(stderr, "%d\n", errcode); + mlirContextDestroy(ctx); return 0;