diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir-c/Dialect/SparseTensor.h @@ -0,0 +1,77 @@ +//===-- mlir-c/Dialect/SparseTensor.h - C API for SparseTensor ----*- 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_DIALECT_SPARSE_TENSOR_H +#define MLIR_C_DIALECT_SPARSE_TENSOR_H + +#include "mlir-c/AffineMap.h" +#include "mlir-c/Registration.h" + +#ifdef __cplusplus +extern "C" { +#endif + +MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor); + +/// Dimension level types that define sparse tensors: +/// - MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE - dimension is dense, every +/// entry is stored +/// - MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED - dimension is sparse, +/// only nonzeros are stored. +/// - MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON - dimension contains single +/// coordinate, no siblings. +/// +/// These correspond to SparseTensorEncodingAttr::DimLevelType in the C++ API. +/// If updating, keep them in sync and update the static_assert in the impl +/// file. +enum MlirSparseTensorDimLevelType { + MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE, + MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED, + MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON, +}; + +//===----------------------------------------------------------------------===// +// SparseTensorEncodingAttr +//===----------------------------------------------------------------------===// + +/// Checks whether the given attribute is a sparse_tensor.encoding attribute. +MLIR_CAPI_EXPORTED bool +mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr); + +/// Creates a sparse_tensor.encoding attribute with the given parameters. +MLIR_CAPI_EXPORTED MlirAttribute mlirSparseTensorEncodingAttrGet( + MlirContext ctx, intptr_t numDimLevelTypes, + enum MlirSparseTensorDimLevelType const *dimLevelTypes, + MlirAffineMap dimOrdering, int pointerBitWidth, int indexBitWidth); + +/// Returns the number of dim level types in a sparse_tensor.encoding attribute. +MLIR_CAPI_EXPORTED intptr_t +mlirSparseTensorEncodingGetNumDimLevelTypes(MlirAttribute attr); + +/// Returns a specified dim level type in a sparse_tensor.encoding attribute. +MLIR_CAPI_EXPORTED enum MlirSparseTensorDimLevelType +mlirSparseTensorEncodingAttrGetDimLevelType(MlirAttribute attr, intptr_t pos); + +/// Returns the dimension ordering in a sparse_tensor.encoding attribute. +MLIR_CAPI_EXPORTED MlirAffineMap +mlirSparseTensorEncodingAttrGetDimOrdering(MlirAttribute attr); + +/// Returns the pointer bit width in a sparse_tensor.encoding attribute. +MLIR_CAPI_EXPORTED int +mlirSparseTensorEncodingAttrGetPointerBitWidth(MlirAttribute attr); + +/// Returns the index bit width in a sparse_tensor.encoding attribute. +MLIR_CAPI_EXPORTED int +mlirSparseTensorEncodingAttrGetIndexBitWidth(MlirAttribute attr); + +#ifdef __cplusplus +} +#endif + +#endif // MLIR_C_DIALECT_SPARSE_TENSOR_H diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt --- a/mlir/lib/CAPI/Dialect/CMakeLists.txt +++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt @@ -1,21 +1,8 @@ -# TODO: Make the check source feature optional as an argument on *_add_library. -set(LLVM_OPTIONAL_SOURCES - Async.cpp - AsyncPasses.cpp - GPU.cpp - GPUPasses.cpp - Linalg.cpp - LinalgPasses.cpp - SCF.cpp - Shape.cpp - Standard.cpp - Tensor.cpp -) - add_mlir_public_c_api_library(MLIRCAPIAsync Async.cpp AsyncPasses.cpp + PARTIAL_SOURCES_INTENDED DEPENDS MLIRAsyncPassIncGen @@ -30,6 +17,7 @@ GPU.cpp GPUPasses.cpp + PARTIAL_SOURCES_INTENDED DEPENDS MLIRGPUPassIncGen @@ -43,6 +31,7 @@ Linalg.cpp LinalgPasses.cpp + PARTIAL_SOURCES_INTENDED DEPENDS MLIRLinalgPassIncGen @@ -56,6 +45,7 @@ add_mlir_public_c_api_library(MLIRCAPISCF SCF.cpp + PARTIAL_SOURCES_INTENDED LINK_LIBS PUBLIC MLIRCAPIIR MLIRSCF @@ -64,14 +54,25 @@ add_mlir_public_c_api_library(MLIRCAPIShape Shape.cpp + PARTIAL_SOURCES_INTENDED LINK_LIBS PUBLIC MLIRCAPIIR MLIRShape ) +add_mlir_public_c_api_library(MLIRCAPISparseTensor + SparseTensor.cpp + + PARTIAL_SOURCES_INTENDED + LINK_LIBS PUBLIC + MLIRCAPIIR + MLIRSparseTensor +) + add_mlir_public_c_api_library(MLIRCAPIStandard Standard.cpp + PARTIAL_SOURCES_INTENDED LINK_LIBS PUBLIC MLIRCAPIIR MLIRStandard @@ -80,6 +81,7 @@ add_mlir_public_c_api_library(MLIRCAPITensor Tensor.cpp + PARTIAL_SOURCES_INTENDED LINK_LIBS PUBLIC MLIRCAPIIR MLIRTensor diff --git a/mlir/lib/CAPI/Dialect/SparseTensor.cpp b/mlir/lib/CAPI/Dialect/SparseTensor.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/CAPI/Dialect/SparseTensor.cpp @@ -0,0 +1,71 @@ +//===- Tensor.cpp - C API for SparseTensor dialect ------------------------===// +// +// 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/Dialect/SparseTensor.h" +#include "mlir-c/IR.h" +#include "mlir/CAPI/AffineMap.h" +#include "mlir/CAPI/Registration.h" +#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" +#include "mlir/Support/LLVM.h" + +using namespace llvm; +using namespace mlir::sparse_tensor; + +MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor, + mlir::sparse_tensor::SparseTensorDialect) + +// Ensure the C-API enums are int-castable to C++ equivalents. +static_assert( + static_cast(MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE) == + static_cast(SparseTensorEncodingAttr::DimLevelType::Dense) && + static_cast(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED) == + static_cast( + SparseTensorEncodingAttr::DimLevelType::Compressed) && + static_cast(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON) == + static_cast(SparseTensorEncodingAttr::DimLevelType::Singleton), + "MlirSparseTensorDimLevelType (C-API) and DimLevelType (C++) mismatch"); + +bool mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr) { + return unwrap(attr).isa(); +} + +MlirAttribute mlirSparseTensorEncodingAttrGet( + MlirContext ctx, intptr_t numDimLevelTypes, + MlirSparseTensorDimLevelType const *dimLevelTypes, + MlirAffineMap dimOrdering, int pointerBitWidth, int indexBitWidth) { + SmallVector cppDimLevelTypes; + cppDimLevelTypes.resize(numDimLevelTypes); + for (intptr_t i = 0; i < numDimLevelTypes; ++i) + cppDimLevelTypes[i] = + static_cast(dimLevelTypes[i]); + return wrap(SparseTensorEncodingAttr::get(unwrap(ctx), cppDimLevelTypes, + unwrap(dimOrdering), + pointerBitWidth, indexBitWidth)); +} + +MlirAffineMap mlirSparseTensorEncodingAttrGetDimOrdering(MlirAttribute attr) { + return wrap(unwrap(attr).cast().getDimOrdering()); +} + +intptr_t mlirSparseTensorEncodingGetNumDimLevelTypes(MlirAttribute attr) { + return unwrap(attr).cast().getDimLevelType().size(); +} + +MlirSparseTensorDimLevelType +mlirSparseTensorEncodingAttrGetDimLevelType(MlirAttribute attr, intptr_t pos) { + return static_cast( + unwrap(attr).cast().getDimLevelType()[pos]); +} + +int mlirSparseTensorEncodingAttrGetPointerBitWidth(MlirAttribute attr) { + return unwrap(attr).cast().getPointerBitWidth(); +} + +int mlirSparseTensorEncodingAttrGetIndexBitWidth(MlirAttribute attr) { + return unwrap(attr).cast().getIndexBitWidth(); +} diff --git a/mlir/test/CAPI/CMakeLists.txt b/mlir/test/CAPI/CMakeLists.txt --- a/mlir/test/CAPI/CMakeLists.txt +++ b/mlir/test/CAPI/CMakeLists.txt @@ -1,44 +1,34 @@ -set(LLVM_OPTIONAL_SOURCES +function(_add_capi_test_executable name) + cmake_parse_arguments(ARG + "" + "" + "" + ${ARGN}) + set(LLVM_LINK_COMPONENTS + ) + add_llvm_executable(${name} + PARTIAL_SOURCES_INTENDED + ${ARG_UNPARSED_ARGUMENTS}) + llvm_update_compile_flags(${name}) + target_link_libraries(${name} + PRIVATE + MLIRPublicAPI) +endfunction(_add_capi_test_executable) + +_add_capi_test_executable(mlir-capi-execution-engine-test execution_engine.c - ir.c - pass.c +DEPENDS + MLIRConversionPassIncGen ) -set(LLVM_LINK_COMPONENTS - ) - -add_llvm_executable(mlir-capi-ir-test +_add_capi_test_executable(mlir-capi-ir-test ir.c - ) -llvm_update_compile_flags(mlir-capi-ir-test) - -get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS) -target_link_libraries(mlir-capi-ir-test - PRIVATE - MLIRPublicAPI - ) - +) -add_llvm_executable(mlir-capi-pass-test +_add_capi_test_executable(mlir-capi-pass-test pass.c - ) -llvm_update_compile_flags(mlir-capi-pass-test) - -get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS) -target_link_libraries(mlir-capi-pass-test - PRIVATE - MLIRPublicAPI - ) - -add_llvm_executable(mlir-capi-execution-engine-test - execution_engine.c -DEPENDS - MLIRConversionPassIncGen - ) -llvm_update_compile_flags(mlir-capi-execution-engine-test) +) -get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS) -target_link_libraries(mlir-capi-execution-engine-test - PRIVATE - MLIRPublicAPI - ) +_add_capi_test_executable(mlir-capi-sparse-tensor-test + sparse_tensor.c +) diff --git a/mlir/test/CAPI/sparse_tensor.c b/mlir/test/CAPI/sparse_tensor.c new file mode 100644 --- /dev/null +++ b/mlir/test/CAPI/sparse_tensor.c @@ -0,0 +1,79 @@ +//===- sparse_tensor.c - Test of sparse_tensor 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-sparse-tensor-test 2>&1 | FileCheck %s + +#include "mlir-c/Dialect/SparseTensor.h" +#include "mlir-c/IR.h" +#include "mlir-c/Registration.h" + +#include +#include +#include +#include +#include + +// CHECK-LABEL: testRoundtripEncoding() +static int testRoundtripEncoding(MlirContext ctx) { + fprintf(stderr, "testRoundtripEncoding()\n"); + // clang-format off + const char *originalAsm = + "#sparse_tensor.encoding<{ " + "dimLevelType = [ \"dense\", \"compressed\", \"singleton\"], " + "dimOrdering = affine_map<(d0, d1, d2) -> (d0, d1, d2)>, " + "pointerBitWidth = 32, indexBitWidth = 64 }>"; + // clang-format on + MlirAttribute originalAttr = + mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString(originalAsm)); + // CHECK: isa: 1 + fprintf(stderr, "isa: %d\n", + mlirAttributeIsASparseTensorEncodingAttr(originalAttr)); + MlirAffineMap dimOrdering = + mlirSparseTensorEncodingAttrGetDimOrdering(originalAttr); + // CHECK: (d0, d1, d2) -> (d0, d1, d2) + mlirAffineMapDump(dimOrdering); + // CHECK: level_type: 0 + // CHECK: level_type: 1 + // CHECK: level_type: 2 + int numLevelTypes = mlirSparseTensorEncodingGetNumDimLevelTypes(originalAttr); + enum MlirSparseTensorDimLevelType *levelTypes = + alloca(sizeof(enum MlirSparseTensorDimLevelType) * numLevelTypes); + for (int i = 0; i < numLevelTypes; ++i) { + levelTypes[i] = + mlirSparseTensorEncodingAttrGetDimLevelType(originalAttr, i); + fprintf(stderr, "level_type: %d\n", levelTypes[i]); + } + // CHECK: pointer: 32 + int pointerBitWidth = + mlirSparseTensorEncodingAttrGetPointerBitWidth(originalAttr); + fprintf(stderr, "pointer: %d\n", pointerBitWidth); + // CHECK: index: 64 + int indexBitWidth = + mlirSparseTensorEncodingAttrGetIndexBitWidth(originalAttr); + fprintf(stderr, "index: %d\n", indexBitWidth); + + MlirAttribute newAttr = mlirSparseTensorEncodingAttrGet( + ctx, numLevelTypes, levelTypes, dimOrdering, pointerBitWidth, + indexBitWidth); + mlirAttributeDump(newAttr); // For debugging filecheck output. + // CHECK: equal: 1 + fprintf(stderr, "equal: %d\n", mlirAttributeEqual(originalAttr, newAttr)); + return 0; +} + +int main() { + MlirContext ctx = mlirContextCreate(); + mlirDialectHandleRegisterDialect(mlirGetDialectHandle__sparse_tensor__(), + ctx); + if (testRoundtripEncoding(ctx)) + return 1; + + mlirContextDestroy(ctx); + return 0; +} diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt --- a/mlir/test/CMakeLists.txt +++ b/mlir/test/CMakeLists.txt @@ -61,6 +61,7 @@ mlir-capi-execution-engine-test mlir-capi-ir-test mlir-capi-pass-test + mlir-capi-sparse-tensor-test mlir-cpu-runner mlir-edsc-builder-api-test mlir-linalg-ods-gen