diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h --- a/mlir/include/mlir-c/Dialect/SparseTensor.h +++ b/mlir/include/mlir-c/Dialect/SparseTensor.h @@ -26,19 +26,20 @@ /// 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_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 --- a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h +++ b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h @@ -171,30 +171,35 @@ /// 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 @@ return "compressed-hi-no"; case DimLevelType::CompressedWithHiNuNo: return "compressed-hi-nu-no"; + case DimLevelType::TwoOutOfFour: + return "compressed24"; } return ""; } @@ -239,7 +246,7 @@ const uint8_t propertyBits = static_cast(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 @@ 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 @@ 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 @@ 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 @@ "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 @@ "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 --- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td +++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td @@ -319,6 +319,7 @@ ::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 --- a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp +++ b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp @@ -19,6 +19,7 @@ static void populateDialectSparseTensorSubmodule(const py::module &m) { py::enum_(m, "DimLevelType", py::module_local()) .value("dense", MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE) + .value("compressed24", MLIR_SPARSE_TENSOR_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 --- a/mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h +++ b/mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h @@ -34,7 +34,8 @@ 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 --- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp +++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp @@ -409,6 +409,7 @@ } 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 --- 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) +// ----- + +// 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> +func.func private @sparse_2_out_of_4(tensor) + /////////////////////////////////////////////////////////////////////////////// // Migration plan for new STEA surface syntax, // use the NEW_SYNTAX on selected examples @@ -204,3 +217,21 @@ func.func private @foo(%arg0: tensor) { 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> +func.func private @foo_2_out_of_4(%arg0: tensor) { + return +}