f708a549b8
This change adds a new `SparseTensorType` class for making the "dim" vs "lvl" distinction more overt, and for abstracting over the differences between sparse-tensors and dense-tensors. In addition, this change also adds new type aliases `Dimension`, `Level`, and `FieldIndex` to make code more self-documenting. Although the diff is very large, the majority of the changes are mechanical in nature (e.g., changing types to use the new aliases, updating variable names to match, etc). Along the way I also made many variables `const` when they could be; the majority of which required only adding the keyword. A few places had conditional definitions of these variables, requiring actual code changes; however, that was only done when the overall change was extremely local and easy to extract. All these changes are included in the current patch only because it would be too onerous to split them off into a separate patch. Reviewed By: aartbik Differential Revision: https://reviews.llvm.org/D143800
89 lines
3.8 KiB
C++
89 lines
3.8 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>(DimLevelType::Dense) &&
|
|
static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED) ==
|
|
static_cast<int>(DimLevelType::Compressed) &&
|
|
static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU) ==
|
|
static_cast<int>(DimLevelType::CompressedNu) &&
|
|
static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO) ==
|
|
static_cast<int>(DimLevelType::CompressedNo) &&
|
|
static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO) ==
|
|
static_cast<int>(DimLevelType::CompressedNuNo) &&
|
|
static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON) ==
|
|
static_cast<int>(DimLevelType::Singleton) &&
|
|
static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU) ==
|
|
static_cast<int>(DimLevelType::SingletonNu) &&
|
|
static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO) ==
|
|
static_cast<int>(DimLevelType::SingletonNo) &&
|
|
static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO) ==
|
|
static_cast<int>(DimLevelType::SingletonNuNo),
|
|
"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, MlirAffineMap higherOrdering,
|
|
int pointerBitWidth, int indexBitWidth) {
|
|
SmallVector<DimLevelType> cppDimLevelTypes;
|
|
cppDimLevelTypes.resize(numDimLevelTypes);
|
|
for (intptr_t i = 0; i < numDimLevelTypes; ++i)
|
|
cppDimLevelTypes[i] = static_cast<DimLevelType>(dimLevelTypes[i]);
|
|
return wrap(SparseTensorEncodingAttr::get(
|
|
unwrap(ctx), cppDimLevelTypes, unwrap(dimOrdering),
|
|
unwrap(higherOrdering), pointerBitWidth, indexBitWidth));
|
|
}
|
|
|
|
MlirAffineMap mlirSparseTensorEncodingAttrGetDimOrdering(MlirAttribute attr) {
|
|
return wrap(unwrap(attr).cast<SparseTensorEncodingAttr>().getDimOrdering());
|
|
}
|
|
|
|
MlirAffineMap
|
|
mlirSparseTensorEncodingAttrGetHigherOrdering(MlirAttribute attr) {
|
|
return wrap(
|
|
unwrap(attr).cast<SparseTensorEncodingAttr>().getHigherOrdering());
|
|
}
|
|
|
|
intptr_t mlirSparseTensorEncodingGetNumDimLevelTypes(MlirAttribute attr) {
|
|
return unwrap(attr).cast<SparseTensorEncodingAttr>().getLvlRank();
|
|
}
|
|
|
|
MlirSparseTensorDimLevelType
|
|
mlirSparseTensorEncodingAttrGetDimLevelType(MlirAttribute attr, intptr_t lvl) {
|
|
return static_cast<MlirSparseTensorDimLevelType>(
|
|
unwrap(attr).cast<SparseTensorEncodingAttr>().getLvlType(lvl));
|
|
}
|
|
|
|
int mlirSparseTensorEncodingAttrGetPointerBitWidth(MlirAttribute attr) {
|
|
return unwrap(attr).cast<SparseTensorEncodingAttr>().getPointerBitWidth();
|
|
}
|
|
|
|
int mlirSparseTensorEncodingAttrGetIndexBitWidth(MlirAttribute attr) {
|
|
return unwrap(attr).cast<SparseTensorEncodingAttr>().getIndexBitWidth();
|
|
}
|