[Mlir-commits] [mlir] bcfa7ba - [mlir][CAPI] Add CAPI bindings for the sparse_tensor dialect.
Stella Laurenzo
llvmlistbot at llvm.org
Mon May 10 09:55:19 PDT 2021
Author: Stella Laurenzo
Date: 2021-05-10T16:54:56Z
New Revision: bcfa7baec8bbf45b98bcde60305efa23df7399e6
URL: https://github.com/llvm/llvm-project/commit/bcfa7baec8bbf45b98bcde60305efa23df7399e6
DIFF: https://github.com/llvm/llvm-project/commit/bcfa7baec8bbf45b98bcde60305efa23df7399e6.diff
LOG: [mlir][CAPI] Add CAPI bindings for the sparse_tensor dialect.
* Adds dialect registration, hand coded 'encoding' attribute and test.
* An MLIR CAPI tablegen backend for attributes does not exist, and this is a relatively complicated case. I opted to hand code it in a canonical way for now, which will provide a reasonable blueprint for building out the tablegen version in the future.
* Also added a (local) CMake function for declaring new CAPI tests, since it was getting repetitive/buggy.
Differential Revision: https://reviews.llvm.org/D102141
Added:
mlir/include/mlir-c/Dialect/SparseTensor.h
mlir/lib/CAPI/Dialect/SparseTensor.cpp
mlir/test/CAPI/sparse_tensor.c
Modified:
mlir/lib/CAPI/Dialect/CMakeLists.txt
mlir/test/CAPI/CMakeLists.txt
mlir/test/CMakeLists.txt
Removed:
################################################################################
diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h
new file mode 100644
index 0000000000000..2615a1655c15b
--- /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
index dd9bd6f67881d..69371fd5718f9 100644
--- 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 @@ add_mlir_public_c_api_library(MLIRCAPIGPU
GPU.cpp
GPUPasses.cpp
+ PARTIAL_SOURCES_INTENDED
DEPENDS
MLIRGPUPassIncGen
@@ -43,6 +31,7 @@ add_mlir_public_c_api_library(MLIRCAPILinalg
Linalg.cpp
LinalgPasses.cpp
+ PARTIAL_SOURCES_INTENDED
DEPENDS
MLIRLinalgPassIncGen
@@ -56,6 +45,7 @@ add_mlir_public_c_api_library(MLIRCAPILinalg
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(MLIRCAPISCF
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(MLIRCAPIStandard
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
index 0000000000000..f35c14af209b7
--- /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<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE) ==
+ static_cast<int>(SparseTensorEncodingAttr::DimLevelType::Dense) &&
+ static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED) ==
+ static_cast<int>(
+ SparseTensorEncodingAttr::DimLevelType::Compressed) &&
+ static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON) ==
+ static_cast<int>(SparseTensorEncodingAttr::DimLevelType::Singleton),
+ "MlirSparseTensorDimLevelType (C-API) and DimLevelType (C++) mismatch");
+
+bool mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr) {
+ return unwrap(attr).isa<SparseTensorEncodingAttr>();
+}
+
+MlirAttribute mlirSparseTensorEncodingAttrGet(
+ MlirContext ctx, intptr_t numDimLevelTypes,
+ MlirSparseTensorDimLevelType const *dimLevelTypes,
+ MlirAffineMap dimOrdering, int pointerBitWidth, int indexBitWidth) {
+ SmallVector<SparseTensorEncodingAttr::DimLevelType> cppDimLevelTypes;
+ cppDimLevelTypes.resize(numDimLevelTypes);
+ for (intptr_t i = 0; i < numDimLevelTypes; ++i)
+ cppDimLevelTypes[i] =
+ static_cast<SparseTensorEncodingAttr::DimLevelType>(dimLevelTypes[i]);
+ return wrap(SparseTensorEncodingAttr::get(unwrap(ctx), cppDimLevelTypes,
+ unwrap(dimOrdering),
+ pointerBitWidth, indexBitWidth));
+}
+
+MlirAffineMap mlirSparseTensorEncodingAttrGetDimOrdering(MlirAttribute attr) {
+ return wrap(unwrap(attr).cast<SparseTensorEncodingAttr>().getDimOrdering());
+}
+
+intptr_t mlirSparseTensorEncodingGetNumDimLevelTypes(MlirAttribute attr) {
+ return unwrap(attr).cast<SparseTensorEncodingAttr>().getDimLevelType().size();
+}
+
+MlirSparseTensorDimLevelType
+mlirSparseTensorEncodingAttrGetDimLevelType(MlirAttribute attr, intptr_t pos) {
+ return static_cast<MlirSparseTensorDimLevelType>(
+ unwrap(attr).cast<SparseTensorEncodingAttr>().getDimLevelType()[pos]);
+}
+
+int mlirSparseTensorEncodingAttrGetPointerBitWidth(MlirAttribute attr) {
+ return unwrap(attr).cast<SparseTensorEncodingAttr>().getPointerBitWidth();
+}
+
+int mlirSparseTensorEncodingAttrGetIndexBitWidth(MlirAttribute attr) {
+ return unwrap(attr).cast<SparseTensorEncodingAttr>().getIndexBitWidth();
+}
diff --git a/mlir/test/CAPI/CMakeLists.txt b/mlir/test/CAPI/CMakeLists.txt
index d01cb0220b2c4..6e377e5b987c5 100644
--- 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
index 0000000000000..549622e275bee
--- /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 <assert.h>
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+// 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
index bf05922976ae0..f8e349b5b606d 100644
--- a/mlir/test/CMakeLists.txt
+++ b/mlir/test/CMakeLists.txt
@@ -61,6 +61,7 @@ set(MLIR_TEST_DEPENDS
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
More information about the Mlir-commits
mailing list