9921ef73c8
Although we have plans to support this, and many other, dimension level type(s), currently the tag is not supported. It will be easy to add this back once support is added. NOTE: based on discussion in https://discourse.llvm.org/t/overcoming-sparsification-limitation-on-level-types/62585 https://github.com/llvm/llvm-project/issues/51658 Reviewed By: Peiming Differential Revision: https://reviews.llvm.org/D131002
70 lines
2.9 KiB
C++
70 lines
2.9 KiB
C++
//===- 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),
|
|
"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();
|
|
}
|