[Mlir-commits] [mlir] 269c82d - [mlir][sparse] introduce new 2:4 block sparsity level type.
Peiming Liu
llvmlistbot at llvm.org
Wed Jul 12 16:33:58 PDT 2023
Author: Peiming Liu
Date: 2023-07-12T23:33:53Z
New Revision: 269c82d389f1f174b46002ec641dc358d73bc88c
URL: https://github.com/llvm/llvm-project/commit/269c82d389f1f174b46002ec641dc358d73bc88c
DIFF: https://github.com/llvm/llvm-project/commit/269c82d389f1f174b46002ec641dc358d73bc88c.diff
LOG: [mlir][sparse] introduce new 2:4 block sparsity level type.
Reviewed By: aartbik
Differential Revision: https://reviews.llvm.org/D155128
Added:
Modified:
mlir/include/mlir-c/Dialect/SparseTensor.h
mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
mlir/lib/Bindings/Python/DialectSparseTensor.cpp
mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h
index 0ad1a315e4c149..b2e4b96c65019c 100644
--- a/mlir/include/mlir-c/Dialect/SparseTensor.h
+++ b/mlir/include/mlir-c/Dialect/SparseTensor.h
@@ -26,19 +26,20 @@ MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor);
/// If updating, keep them in sync and update the static_assert in the impl
/// file.
enum MlirSparseTensorDimLevelType {
- MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE = 4, // 0b0001_00
- MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED = 8, // 0b0010_00
- MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU = 9, // 0b0010_01
- MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO = 10, // 0b0010_10
- MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO = 11, // 0b0010_11
- MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON = 16, // 0b0100_00
- MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU = 17, // 0b0100_01
- MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO = 18, // 0b0100_10
- MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO = 19, // 0b0100_11
- MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI = 32, // 0b1000_00
- MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NU = 33, // 0b1000_01
- MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NO = 34, // 0b1000_10
- MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NU_NO = 35, // 0b1000_11
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE = 4, // 0b00001_00
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED = 8, // 0b00010_00
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU = 9, // 0b00010_01
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO = 10, // 0b00010_10
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU_NO = 11, // 0b00010_11
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON = 16, // 0b00100_00
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU = 17, // 0b00100_01
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NO = 18, // 0b00100_10
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON_NU_NO = 19, // 0b00100_11
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI = 32, // 0b01000_00
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NU = 33, // 0b01000_01
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NO = 34, // 0b01000_10
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_WITH_HI_NU_NO = 35, // 0b01000_11
+ MLIR_SPARSE_TENSOR_DIM_LEVEL_TWO_OUT_OF_FOUR = 64, // 0b10000_00
};
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
index 522f2cd99b8657..b0348ed8eedcf1 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
@@ -171,30 +171,35 @@ enum class Action : uint32_t {
/// where we need to store an undefined or indeterminate `DimLevelType`.
/// It should not be used externally, since it does not indicate an
/// actual/representable format.
+///
+// TODO: We should generalize TwoOutOfFour to N out of M and use property to
+// encode the value of N and M.
enum class DimLevelType : uint8_t {
- Undef = 0, // 0b0000_00
- Dense = 4, // 0b0001_00
- Compressed = 8, // 0b0010_00
- CompressedNu = 9, // 0b0010_01
- CompressedNo = 10, // 0b0010_10
- CompressedNuNo = 11, // 0b0010_11
- Singleton = 16, // 0b0100_00
- SingletonNu = 17, // 0b0100_01
- SingletonNo = 18, // 0b0100_10
- SingletonNuNo = 19, // 0b0100_11
- CompressedWithHi = 32, // 0b1000_00
- CompressedWithHiNu = 33, // 0b1000_01
- CompressedWithHiNo = 34, // 0b1000_10
- CompressedWithHiNuNo = 35, // 0b1000_11
+ Undef = 0, // 0b00000_00
+ Dense = 4, // 0b00001_00
+ Compressed = 8, // 0b00010_00
+ CompressedNu = 9, // 0b00010_01
+ CompressedNo = 10, // 0b00010_10
+ CompressedNuNo = 11, // 0b00010_11
+ Singleton = 16, // 0b00100_00
+ SingletonNu = 17, // 0b00100_01
+ SingletonNo = 18, // 0b00100_10
+ SingletonNuNo = 19, // 0b00100_11
+ CompressedWithHi = 32, // 0b01000_00
+ CompressedWithHiNu = 33, // 0b01000_01
+ CompressedWithHiNo = 34, // 0b01000_10
+ CompressedWithHiNuNo = 35, // 0b01000_11
+ TwoOutOfFour = 64, // 0b10000_00
};
/// This enum defines all the storage formats supported by the sparse compiler,
/// without the level properties.
enum class LevelFormat : uint8_t {
- Dense = 4, // 0b0001_00
- Compressed = 8, // 0b0010_00
- Singleton = 16, // 0b0100_00
- CompressedWithHi = 32, // 0b1000_00
+ Dense = 4, // 0b00001_00
+ Compressed = 8, // 0b00010_00
+ Singleton = 16, // 0b00100_00
+ CompressedWithHi = 32, // 0b01000_00
+ TwoOutOfFour = 64, // 0b10000_00
};
/// Returns string representation of the given dimension level type.
@@ -229,6 +234,8 @@ constexpr const char *toMLIRString(DimLevelType dlt) {
return "compressed-hi-no";
case DimLevelType::CompressedWithHiNuNo:
return "compressed-hi-nu-no";
+ case DimLevelType::TwoOutOfFour:
+ return "compressed24";
}
return "";
}
@@ -239,7 +246,7 @@ constexpr bool isValidDLT(DimLevelType dlt) {
const uint8_t propertyBits = static_cast<uint8_t>(dlt) & 3;
// If undefined or dense, then must be unique and ordered.
// Otherwise, the format must be one of the known ones.
- return (formatBits <= 1)
+ return (formatBits <= 1 || formatBits == 16)
? (propertyBits == 0)
: (formatBits == 2 || formatBits == 4 || formatBits == 8);
}
@@ -254,6 +261,11 @@ constexpr bool isDenseDLT(DimLevelType dlt) {
return dlt == DimLevelType::Dense;
}
+/// Check if the `DimLevelType` is 2:4
+constexpr bool isTwoOutOfFourDLT(DimLevelType dlt) {
+ return dlt == DimLevelType::TwoOutOfFour;
+}
+
// We use the idiom `(dlt & ~3) == format` in order to only return true
// for valid DLTs. Whereas the `dlt & format` idiom is a bit faster but
// can return false-positives on invalid DLTs.
@@ -325,6 +337,11 @@ static_assert(
buildLevelType(LevelFormat::Dense, true, false) == std::nullopt &&
buildLevelType(LevelFormat::Dense, false, false) == std::nullopt &&
*buildLevelType(LevelFormat::Dense, true, true) == DimLevelType::Dense &&
+ buildLevelType(LevelFormat::TwoOutOfFour, false, true) == std::nullopt &&
+ buildLevelType(LevelFormat::TwoOutOfFour, true, false) == std::nullopt &&
+ buildLevelType(LevelFormat::TwoOutOfFour, false, false) == std::nullopt &&
+ *buildLevelType(LevelFormat::TwoOutOfFour, true, true) ==
+ DimLevelType::TwoOutOfFour &&
*buildLevelType(LevelFormat::Compressed, true, true) ==
DimLevelType::Compressed &&
*buildLevelType(LevelFormat::Compressed, true, false) ==
@@ -357,7 +374,8 @@ static_assert((isValidDLT(DimLevelType::Undef) &&
isValidDLT(DimLevelType::CompressedWithHi) &&
isValidDLT(DimLevelType::CompressedWithHiNu) &&
isValidDLT(DimLevelType::CompressedWithHiNo) &&
- isValidDLT(DimLevelType::CompressedWithHiNuNo)),
+ isValidDLT(DimLevelType::CompressedWithHiNuNo) &&
+ isValidDLT(DimLevelType::TwoOutOfFour)),
"isValidDLT definition is broken");
static_assert((!isCompressedDLT(DimLevelType::Dense) &&
@@ -394,6 +412,7 @@ static_assert((!isSingletonDLT(DimLevelType::Dense) &&
"isSingletonDLT definition is broken");
static_assert((isOrderedDLT(DimLevelType::Dense) &&
+ isOrderedDLT(DimLevelType::TwoOutOfFour) &&
isOrderedDLT(DimLevelType::Compressed) &&
isOrderedDLT(DimLevelType::CompressedNu) &&
!isOrderedDLT(DimLevelType::CompressedNo) &&
@@ -409,6 +428,7 @@ static_assert((isOrderedDLT(DimLevelType::Dense) &&
"isOrderedDLT definition is broken");
static_assert((isUniqueDLT(DimLevelType::Dense) &&
+ isUniqueDLT(DimLevelType::TwoOutOfFour) &&
isUniqueDLT(DimLevelType::Compressed) &&
!isUniqueDLT(DimLevelType::CompressedNu) &&
isUniqueDLT(DimLevelType::CompressedNo) &&
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
index df37ceee1cb85d..ad42d6467fa120 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
@@ -319,6 +319,7 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
::mlir::sparse_tensor::DimLevelType getLvlType(::mlir::sparse_tensor::Level l) const;
bool isDenseLvl(::mlir::sparse_tensor::Level l) const { return isDenseDLT(getLvlType(l)); }
+ bool isTwoOutOfFourLvl(::mlir::sparse_tensor::Level l) const { return isTwoOutOfFourDLT(getLvlType(l)); }
bool isCompressedLvl(::mlir::sparse_tensor::Level l) const { return isCompressedDLT(getLvlType(l)); }
bool isCompressedWithHiLvl(::mlir::sparse_tensor::Level l) const { return isCompressedWithHiDLT(getLvlType(l)); }
bool isSingletonLvl(::mlir::sparse_tensor::Level l) const { return isSingletonDLT(getLvlType(l)); }
diff --git a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp
index 2e8d535455a34f..d030883412640e 100644
--- a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp
+++ b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp
@@ -19,6 +19,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("compressed24", MLIR_SPARSE_TENSOR_DIM_LEVEL_TWO_OUT_OF_FOUR)
.value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED)
.value("compressed-nu", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NU)
.value("compressed-no", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED_NO)
diff --git a/mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h b/mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h
index 2d6a45a4409945..b1fb3a42e41fe2 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h
+++ b/mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h
@@ -34,7 +34,8 @@ namespace ir_detail {
DO(DimLevelType::CompressedWithHi) \
DO(DimLevelType::CompressedWithHiNu) \
DO(DimLevelType::CompressedWithHiNo) \
- DO(DimLevelType::CompressedWithHiNuNo)
+ DO(DimLevelType::CompressedWithHiNuNo) \
+ DO(DimLevelType::TwoOutOfFour)
#define LEVELTYPE_INITLIST_ELEMENT(lvlType) \
std::make_pair(StringRef(toMLIRString(lvlType)), lvlType),
#define LEVELTYPE_INITLIST \
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 3037429847dbf2..7dd15b96b30656 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -409,6 +409,7 @@ SparseTensorEncodingAttr::getStaticLvlSliceStride(Level lvl) const {
}
const static DimLevelType validDLTs[] = {DimLevelType::Dense,
+ DimLevelType::TwoOutOfFour,
DimLevelType::Compressed,
DimLevelType::CompressedNu,
DimLevelType::CompressedNo,
diff --git a/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir b/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
index 5a7c5a641a653c..99b1ea1afd2d61 100644
--- a/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
+++ b/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
@@ -130,6 +130,19 @@ func.func private @sparse_slice(tensor<?x?xf64, #CSR_SLICE>)
func.func private @sparse_slice(tensor<?x?xf64, #CSR_SLICE>)
+// -----
+
+// TODO: It is probably better to use [dense, dense, 2:4] (see NV_24 defined using new syntax
+// below) to encode a 2D matrix, but it would require dim2lvl mapping which is not ready yet.
+// So we take the simple path for now.
+#NV_24= #sparse_tensor.encoding<{
+ lvlTypes = [ "dense", "compressed24" ],
+}>
+
+// CHECK-LABEL: func private @sparse_2_out_of_4(
+// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed24" ] }>>
+func.func private @sparse_2_out_of_4(tensor<?x?xf64, #NV_24>)
+
///////////////////////////////////////////////////////////////////////////////
// Migration plan for new STEA surface syntax,
// use the NEW_SYNTAX on selected examples
@@ -204,3 +217,21 @@ func.func private @foo(%arg0: tensor<?x?x?x?xf64, #BCSR_implicit>) {
func.func private @foo(%arg0: tensor<?x?x?x?xf64, #BCSR_explicit>) {
return
}
+
+// -----
+
+#NV_24 = #sparse_tensor.encoding<{
+ NEW_SYNTAX =
+ ( i, j ) ->
+ ( i : dense,
+ j floordiv 4 : dense,
+ j mod 4 : compressed24
+ )
+}>
+
+//
+// CHECK-LABEL: func private @foo_2_out_of_4(
+// CHECK-SAME: tensor<?x?x?xf64, #sparse_tensor.encoding<{ lvlTypes = [ "dense", "dense", "compressed24" ] }>>
+func.func private @foo_2_out_of_4(%arg0: tensor<?x?x?xf64, #NV_24>) {
+ return
+}
More information about the Mlir-commits
mailing list