maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][sparse] remove singleton dimension level type (for now) 

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
This commit is contained in:
Aart Bik 2022-08-02 10:36:05 -07:00
parent f712775daf
commit 9921ef73c8
11 changed files with 10 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mlir/include/mlir-c/Dialect/SparseTensor.h
View File

@ -20,12 +20,10 @@ extern "C" {
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
/// - 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.
/// only nonzeros are stored (no duplicates).
///
/// These correspond to SparseTensorEncodingAttr::DimLevelType in the C++ API.
/// If updating, keep them in sync and update the static_assert in the impl
@ -33,7 +31,6 @@ MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor);
enum MlirSparseTensorDimLevelType {
MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE,
MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED,
MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON,
};
//===----------------------------------------------------------------------===//

7
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
View File

@ -84,11 +84,10 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
let extraClassDeclaration = [{
// Dimension level types that define sparse tensors:
// Dense - dimension is dense, every entry is stored
// Compressed - dimension is sparse, only nonzeros are stored
// Singleton - dimension contains single coordinate, no siblings
// Dense - dimension is dense, every entry is stored
// Compressed (unique) - dimension is sparse, only nonzeros are stored (no duplicates)
enum class DimLevelType {
Dense, Compressed, Singleton
Dense, Compressed
};
}];
}

2
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
View File

@ -21,7 +21,7 @@ namespace mlir {
namespace sparse_tensor {
/// Dimension level type for a tensor (undef means index does not appear).
enum Dim { kSparse, kDense, kSingle, kUndef };
enum Dim { kSparse, kDense, kUndef };
/// Tensor expression kind.
enum Kind {

3
mlir/lib/Bindings/Python/DialectSparseTensor.cpp
View File

@ -18,8 +18,7 @@ using namespace mlir::python::adaptors;
static void populateDialectSparseTensorSubmodule(const py::module &m) {
py::enum_<MlirSparseTensorDimLevelType>(m, "DimLevelType", py::module_local())
.value("dense", MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE)
.value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED)
.value("singleton", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON);
.value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED);
mlir_attribute_subclass(m, "EncodingAttr",
mlirAttributeIsASparseTensorEncodingAttr)

4
mlir/lib/CAPI/Dialect/SparseTensor.cpp
View File

@ -25,9 +25,7 @@ static_assert(
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),
SparseTensorEncodingAttr::DimLevelType::Compressed),
"MlirSparseTensorDimLevelType (C-API) and DimLevelType (C++) mismatch");
bool mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr) {

5
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
View File

@ -71,8 +71,6 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Dense);
} else if (strVal == "compressed") {
dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Compressed);
} else if (strVal == "singleton") {
dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Singleton);
} else {
parser.emitError(parser.getNameLoc(),
"unexpected dimension level type: ")
@ -126,9 +124,6 @@ void SparseTensorEncodingAttr::print(AsmPrinter &printer) const {
case DimLevelType::Compressed:
printer << "\"compressed\"";
break;
case DimLevelType::Singleton:
printer << "\"singleton\"";
break;
}
if (i != e - 1)
printer << ", ";

2
mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
View File

@ -148,8 +148,6 @@ DimLevelType mlir::sparse_tensor::dimLevelTypeEncoding(
return DimLevelType::kDense;
case SparseTensorEncodingAttr::DimLevelType::Compressed:
return DimLevelType::kCompressed;
case SparseTensorEncodingAttr::DimLevelType::Singleton:
return DimLevelType::kSingleton;
}
llvm_unreachable("Unknown SparseTensorEncodingAttr::DimLevelType");
}

5
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
View File

@ -378,11 +378,6 @@ static bool canUseDirectConversion(
if (alreadyCompressed)
return false; // Dense after Compressed not yet supported.
break;
case SparseTensorEncodingAttr::DimLevelType::Singleton:
// Although Singleton isn't generally supported yet, the direct
// conversion method doesn't have any particular problems with
// singleton after compressed.
break;
}
}
return true;

2
mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
View File

@ -147,8 +147,6 @@ static Dim toDim(const SparseTensorEncodingAttr &enc, unsigned d) {
SparseTensorEncodingAttr::DimLevelType tp = enc.getDimLevelType()[d];
if (tp == SparseTensorEncodingAttr::DimLevelType::Compressed)
return Dim::kSparse;
if (tp == SparseTensorEncodingAttr::DimLevelType::Singleton)
return Dim::kSingle;
}
return Dim::kDense;
}

3
mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
View File

@ -591,9 +591,6 @@ void Merger::dumpBits(const BitVector &bits) const {
case kDense:
llvm::dbgs() << "D";
break;
case kSingle:
llvm::dbgs() << "T";
break;
case kUndef:
llvm::dbgs() << "U";
break;

4
mlir/test/CAPI/sparse_tensor.c
View File

@ -25,7 +25,7 @@ static int testRoundtripEncoding(MlirContext ctx) {
// clang-format off
const char *originalAsm =
"#sparse_tensor.encoding<{ "
"dimLevelType = [ \"dense\", \"compressed\", \"singleton\"], "
"dimLevelType = [ \"dense\", \"compressed\", \"compressed\"], "
"dimOrdering = affine_map<(d0, d1, d2) -> (d0, d1, d2)>, "
"pointerBitWidth = 32, indexBitWidth = 64 }>";
// clang-format on
@ -40,7 +40,7 @@ static int testRoundtripEncoding(MlirContext ctx) {
mlirAffineMapDump(dimOrdering);
// CHECK: level_type: 0
// CHECK: level_type: 1
// CHECK: level_type: 2
// CHECK: level_type: 1
int numLevelTypes = mlirSparseTensorEncodingGetNumDimLevelTypes(originalAttr);
enum MlirSparseTensorDimLevelType *levelTypes =
malloc(sizeof(enum MlirSparseTensorDimLevelType) * numLevelTypes);