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,92 @@ 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 mlirAffineMapGet(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 mlirAffineMapDimAndSymbolGet(MlirContext ctx, unsigned dimCount, + unsigned symbolCount); + +/** Creates a single constant result affine map in the context. The affine map + * is owned by the context. */ +MlirAffineMap mlirAffineMapConstantMapGet(MlirContext ctx, int64_t val); + +/** Creates an affine map with 'numDims' identity in the context. The affine map + * is owned by the context. */ +MlirAffineMap mlirAffineMapMultiDimIdentityMapGet(MlirContext ctx, + unsigned 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 mlirAffineMapMinorIdentityMapGet(MlirContext ctx, unsigned dims, + unsigned results); + +/** Creates an affine map representing a permutation in the context. The affine + * map is owned by the context. */ +MlirAffineMap mlirAffineMapPermutationMapGet(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, + unsigned 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, + unsigned 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/include/mlir/CAPI/IR.h b/mlir/include/mlir/CAPI/IR.h --- a/mlir/include/mlir/CAPI/IR.h +++ b/mlir/include/mlir/CAPI/IR.h @@ -15,6 +15,7 @@ #ifndef MLIR_INCLUDE_MLIR_CAPI_IR_H #define MLIR_INCLUDE_MLIR_CAPI_IR_H +#include "mlir/CAPI/AffineMap.h" #include "mlir/CAPI/Wrap.h" #include "mlir/IR/MLIRContext.h" #include "mlir/IR/Module.h" 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 mlirAffineMapGet(MlirContext ctx) { + return wrap(AffineMap::get(unwrap(ctx))); +} + +MlirAffineMap mlirAffineMapDimAndSymbolGet(MlirContext ctx, unsigned dimCount, + unsigned symbolCount) { + return wrap(AffineMap::get(dimCount, symbolCount, unwrap(ctx))); +} + +MlirAffineMap mlirAffineMapConstantMapGet(MlirContext ctx, int64_t val) { + return wrap(AffineMap::getConstantMap(val, unwrap(ctx))); +} + +MlirAffineMap mlirAffineMapMultiDimIdentityMapGet(MlirContext ctx, + unsigned numDims) { + return wrap(AffineMap::getMultiDimIdentityMap(numDims, unwrap(ctx))); +} + +MlirAffineMap mlirAffineMapMinorIdentityMapGet(MlirContext ctx, unsigned dims, + unsigned results) { + return wrap(AffineMap::getMinorIdentityMap(dims, results, unwrap(ctx))); +} + +MlirAffineMap mlirAffineMapPermutationMapGet(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, + unsigned numResults) { + return wrap(unwrap(affineMap).getMajorSubMap(numResults)); +} + +MlirAffineMap mlirAffineMapGetMinorSubMap(MlirAffineMap affineMap, + unsigned 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 @@ -392,3 +392,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,123 @@ return 0; } +int printAffineMap(MlirContext ctx) { + MlirAffineMap affineMap = mlirAffineMapGet(ctx); + MlirAffineMap dimAndSymbolAffineMap = mlirAffineMapDimAndSymbolGet(ctx, 3, 2); + MlirAffineMap constAffineMap = mlirAffineMapConstantMapGet(ctx, 2); + MlirAffineMap multiDimIdentityAffineMap = + mlirAffineMapMultiDimIdentityMapGet(ctx, 3); + MlirAffineMap minorIdentityAffineMap = + mlirAffineMapMinorIdentityMapGet(ctx, 3, 2); + unsigned permutation[] = {1, 2, 0}; + MlirAffineMap permutationAffineMap = + mlirAffineMapPermutationMapGet( + ctx, sizeof(permutation) / sizeof(unsigned), permutation); + + mlirAffineMapDump(affineMap); + mlirAffineMapDump(dimAndSymbolAffineMap); + mlirAffineMapDump(constAffineMap); + mlirAffineMapDump(multiDimIdentityAffineMap); + mlirAffineMapDump(minorIdentityAffineMap); + mlirAffineMapDump(permutationAffineMap); + + if (!mlirAffineMapIsIdentity(affineMap) || + mlirAffineMapIsIdentity(dimAndSymbolAffineMap) || + mlirAffineMapIsIdentity(constAffineMap) || + !mlirAffineMapIsIdentity(multiDimIdentityAffineMap) || + mlirAffineMapIsIdentity(minorIdentityAffineMap) || + mlirAffineMapIsIdentity(permutationAffineMap)) + return 1; + + if (!mlirAffineMapIsMinorIdentity(affineMap) || + mlirAffineMapIsMinorIdentity(dimAndSymbolAffineMap) || + !mlirAffineMapIsMinorIdentity(multiDimIdentityAffineMap) || + !mlirAffineMapIsMinorIdentity(minorIdentityAffineMap) || + mlirAffineMapIsMinorIdentity(permutationAffineMap)) + return 2; + + if (!mlirAffineMapIsEmpty(affineMap) || + mlirAffineMapIsEmpty(dimAndSymbolAffineMap) || + mlirAffineMapIsEmpty(constAffineMap) || + mlirAffineMapIsEmpty(multiDimIdentityAffineMap) || + mlirAffineMapIsEmpty(minorIdentityAffineMap) || + mlirAffineMapIsEmpty(permutationAffineMap)) + return 3; + + if (mlirAffineMapIsSingleConstant(affineMap) || + mlirAffineMapIsSingleConstant(dimAndSymbolAffineMap) || + !mlirAffineMapIsSingleConstant(constAffineMap) || + mlirAffineMapIsSingleConstant(multiDimIdentityAffineMap) || + mlirAffineMapIsSingleConstant(minorIdentityAffineMap) || + mlirAffineMapIsSingleConstant(permutationAffineMap)) + return 4; + + if (mlirAffineMapGetSingleConstantResult(constAffineMap) != 2 ) + return 5; + + if (mlirAffineMapGetNumDims(affineMap) != 0 || + mlirAffineMapGetNumDims(dimAndSymbolAffineMap) != 3 || + mlirAffineMapGetNumDims(constAffineMap) != 0 || + mlirAffineMapGetNumDims(multiDimIdentityAffineMap) != 3 || + mlirAffineMapGetNumDims(minorIdentityAffineMap) != 3 || + mlirAffineMapGetNumDims(permutationAffineMap) != 3) + return 6; + + if (mlirAffineMapGetNumSymbols(affineMap) != 0 || + mlirAffineMapGetNumSymbols(dimAndSymbolAffineMap) != 2 || + mlirAffineMapGetNumSymbols(constAffineMap) != 0 || + mlirAffineMapGetNumSymbols(multiDimIdentityAffineMap) != 0 || + mlirAffineMapGetNumSymbols(minorIdentityAffineMap) != 0 || + mlirAffineMapGetNumSymbols(permutationAffineMap) != 0) + return 7; + + if (mlirAffineMapGetNumResults(affineMap) != 0 || + mlirAffineMapGetNumResults(dimAndSymbolAffineMap) != 0 || + mlirAffineMapGetNumResults(constAffineMap) != 1 || + mlirAffineMapGetNumResults(multiDimIdentityAffineMap) != 3 || + mlirAffineMapGetNumResults(minorIdentityAffineMap) != 2 || + mlirAffineMapGetNumResults(permutationAffineMap) != 3) + return 8; + + if (mlirAffineMapGetNumInputs(affineMap) != 0 || + mlirAffineMapGetNumInputs(dimAndSymbolAffineMap) != 5 || + mlirAffineMapGetNumInputs(constAffineMap) != 0 || + mlirAffineMapGetNumInputs(multiDimIdentityAffineMap) != 3 || + mlirAffineMapGetNumInputs(minorIdentityAffineMap) != 3 || + mlirAffineMapGetNumInputs(permutationAffineMap) != 3) + return 9; + + if (!mlirAffineMapIsProjectedPermutation(affineMap) || + !mlirAffineMapIsPermutation(affineMap) || + mlirAffineMapIsProjectedPermutation(dimAndSymbolAffineMap) || + mlirAffineMapIsPermutation(dimAndSymbolAffineMap) || + 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 +816,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;