forked from OSchip/llvm-project
[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
This commit is contained in:
Stella Laurenzo
parent
22f834210a
commit
bcfa7baec8
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
|
|
@ -1,44 +1,34 @@
|
|||
set(LLVM_OPTIONAL_SOURCES
|
||||
execution_engine.c
|
||||
ir.c
|
||||
pass.c
|
||||
)
|
||||
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)
|
||||
|
||||
set(LLVM_LINK_COMPONENTS
|
||||
)
|
||||
|
||||
add_llvm_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
|
||||
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
|
||||
_add_capi_test_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-ir-test
|
||||
ir.c
|
||||
)
|
||||
|
||||
_add_capi_test_executable(mlir-capi-pass-test
|
||||
pass.c
|
||||
)
|
||||
|
||||
_add_capi_test_executable(mlir-capi-sparse-tensor-test
|
||||
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;
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue