[mlir][sparse] Update verifier for block sparsity and singleton (#69389) (7b9fb1c2) · Commits · llvm-doe / llvm-project

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

+15 −0

Original line number	Diff line number	Diff line
		@@ -173,6 +173,21 @@ AffineMap inferLvlToDim(AffineMap dimToLvl, MLIRContext *context);
		/// Asserts on failure (so only use when known to succeed).
		AffineMap inverseBlockSparsity(AffineMap dimToLvl, MLIRContext *context);

		/// Given the dimToLvl map, returns the block sizes in a vector.
		/// For instance, a 2x3 block will return [2, 3]. Unblocked dimension i
		/// will return 0, and i floordiv 1, i mod 1 will return 1. Therefore,
		/// the example below will return [0, 1].
		/// map = ( i, j ) ->
		/// ( i : dense,
		/// j floordiv 1 : compressed,
		/// j mod 1 : dense
		/// )
		/// Only valid block sparsity will be accepted.
		SmallVector<unsigned> getBlockSize(AffineMap dimToLvl);

		/// Given the dimToLvl map, returns if it's block sparsity.
		bool isBlockSparsity(AffineMap dimToLvl);

		//
		// Reordering.
		//

mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp

+8 −3

Original line number	Diff line number	Diff line
		@@ -356,10 +356,15 @@ AffineMap DimLvlMap::getDimToLvlMap(MLIRContext *context) const {
		AffineMap DimLvlMap::getLvlToDimMap(MLIRContext *context) const {
		SmallVector<AffineExpr> dimAffines;
		dimAffines.reserve(getDimRank());
		for (const auto &dimSpec : dimSpecs)
		dimAffines.push_back(dimSpec.getExpr().getAffineExpr());
		for (const auto &dimSpec : dimSpecs) {
		auto expr = dimSpec.getExpr().getAffineExpr();
		if (expr) {
		dimAffines.push_back(expr);
		}
		}
		auto map = AffineMap::get(getLvlRank(), getSymRank(), dimAffines, context);
		if (map.isIdentity()) return AffineMap();
		if (dimAffines.empty() \|\| map.isIdentity())
		return AffineMap();
		return map;
		}

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

+70 −17

Original line number	Diff line number	Diff line
		@@ -455,6 +455,7 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
		SmallVector<DimLevelType> lvlTypes;
		SmallVector<SparseTensorDimSliceAttr> dimSlices;
		AffineMap dimToLvl = {};
		AffineMap lvlToDim = {};
		unsigned posWidth = 0;
		unsigned crdWidth = 0;
		StringRef attrName;
		@@ -568,6 +569,7 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {

		ERROR_IF(dimToLvl, "Cannot mix `dimToLvl` with `map`")
		dimToLvl = dlm.getDimToLvlMap(parser.getContext());
		lvlToDim = dlm.getLvlToDimMap(parser.getContext());
		break;
		}
		} // switch
		@@ -582,8 +584,9 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
		#undef RETURN_ON_FAIL

		// Construct struct-like storage for attribute.
		// TODO: Fetch lvlToDim if user provides one
		AffineMap lvlToDim = inferLvlToDim(dimToLvl, parser.getContext());
		if (!lvlToDim \|\| lvlToDim.isEmpty()) {
		lvlToDim = inferLvlToDim(dimToLvl, parser.getContext());
		}
		return parser.getChecked<SparseTensorEncodingAttr>(
		parser.getContext(), lvlTypes, dimToLvl, lvlToDim, posWidth, crdWidth,
		dimSlices);
		@@ -663,6 +666,17 @@ SparseTensorEncodingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
		return emitError() << "unexpected position bitwidth: " << posWidth;
		if (!acceptBitWidth(crdWidth))
		return emitError() << "unexpected coordinate bitwidth: " << crdWidth;
		if (auto it = std::find_if(lvlTypes.begin(), lvlTypes.end(), isSingletonDLT);
		it != std::end(lvlTypes)) {
		if (it == lvlTypes.begin() \|\|
		(!isCompressedDLT((it - 1)) && !isLooseCompressedDLT((it - 1))))
		return emitError() << "expected compressed or loose_compressed level "
		"before singleton level";
		if (!std::all_of(it, lvlTypes.end(),
		[](DimLevelType i) { return isSingletonDLT(i); }))
		return emitError() << "expected all singleton lvlTypes "
		"following a singleton level";
		}
		// Before we can check that the level-rank is consistent/coherent
		// across all fields, we need to define it. The source-of-truth for
		// the `getLvlRank` method is the length of the level-types array,
		@@ -678,19 +692,12 @@ SparseTensorEncodingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
		return emitError()
		<< "level-rank mismatch between dimToLvl and lvlTypes: "
		<< dimToLvl.getNumResults() << " != " << lvlRank;
		// TODO: The following is attempting to match the old error-conditions
		// from prior to merging dimOrdering and higherOrdering into dimToLvl.
		// That is, we currently require `dimToLvl` to be either a permutation
		// (as when higherOrdering is the identity) or expansive (as per the
		// constraints on higherOrdering). However, those constraints do
		// not match the intended semantics of `dimToLvl`. As we improve the
		// compiler to actually handle non-permutations, we need to update these
		// checks to match what is actually supported. In particular, this is
		// where we'll have to check that when `lvlToDim` is provided then it
		// is indeed an inverse of `dimToLvl`, and when it isn't provided then
		// it can be automatically inferred.
		if (dimRank == lvlRank && !dimToLvl.isPermutation())
		return emitError() << "expected a permutation affine map for dimToLvl";
		auto inferRes = inferLvlToDim(dimToLvl, dimToLvl.getContext());
		// Symbols can't be inferred but are acceptable.
		if (!inferRes && dimToLvl.getNumSymbols() == 0)
		return emitError() << "failed to infer lvlToDim from dimToLvl";
		if (lvlToDim && (inferRes != lvlToDim))
		return emitError() << "expected lvlToDim to be an inverse of dimToLvl";
		if (dimRank > lvlRank)
		return emitError() << "unexpected dimToLvl mapping from " << dimRank
		<< " to " << lvlRank;
		@@ -758,8 +765,7 @@ AffineMap mlir::sparse_tensor::inferLvlToDim(AffineMap dimToLvl,
		lvlToDim = AffineMap();
		} else if (map.isPermutation()) {
		lvlToDim = inversePermutation(map);
		} else {
		// TODO: check if it's block sparsity
		} else if (isBlockSparsity(map)) {
		lvlToDim = inverseBlockSparsity(map, context);
		}
		return lvlToDim;
		@@ -818,6 +824,53 @@ AffineMap mlir::sparse_tensor::inverseBlockSparsity(AffineMap dimToLvl,
		return dimToLvl.get(dimToLvl.getNumResults(), 0, lvlExprs, context);
		}

		SmallVector<unsigned> mlir::sparse_tensor::getBlockSize(AffineMap dimToLvl) {
		assert(isBlockSparsity(dimToLvl) &&
		"expected dimToLvl to be block sparsity for calling getBlockSize");
		SmallVector<unsigned> blockSize;
		for (auto result : dimToLvl.getResults()) {
		if (auto binOp = result.dyn_cast<AffineBinaryOpExpr>()) {
		if (result.getKind() == AffineExprKind::Mod) {
		blockSize.push_back(
		binOp.getRHS().dyn_cast<AffineConstantExpr>().getValue());
		}
		} else {
		blockSize.push_back(0);
		}
		}
		return blockSize;
		}

		bool mlir::sparse_tensor::isBlockSparsity(AffineMap dimToLvl) {
		if (!dimToLvl)
		return false;
		std::map<unsigned, int64_t> coeffientMap;
		for (auto result : dimToLvl.getResults()) {
		if (auto binOp = result.dyn_cast<AffineBinaryOpExpr>()) {
		auto pos = binOp.getLHS().dyn_cast<AffineDimExpr>().getPosition();
		if (result.getKind() == AffineExprKind::FloorDiv) {
		// Expect only one floordiv for each dimension.
		if (coeffientMap.find(pos) != coeffientMap.end())
		return false;
		coeffientMap[pos] =
		binOp.getRHS().dyn_cast<AffineConstantExpr>().getValue();
		} else if (result.getKind() == AffineExprKind::Mod) {
		// Expect floordiv before mod.
		if (coeffientMap.find(pos) == coeffientMap.end())
		return false;
		// Expect mod to have the same coefficient as floordiv.
		if (binOp.getRHS().dyn_cast<AffineConstantExpr>().getValue() !=
		coeffientMap[pos]) {
		return false;
		}
		} else {
		return false;
		}
		}
		}
		return !coeffientMap.empty();
		}

		bool mlir::sparse_tensor::isCOOType(SparseTensorEncodingAttr enc,
		Level startLvl, bool isUnique) {
		if (!enc \|\|

mlir/test/Dialect/SparseTensor/invalid_encoding.mlir

+21 −2

Original line number	Diff line number	Diff line
		@@ -60,7 +60,7 @@ func.func private @tensor_sizes_mismatch(%arg0: tensor<8xi32, #a>) -> ()

		// -----

		// expected-error@+1 {{unexpected dimToLvl mapping from 2 to 1}}
		// expected-error@+1 {{failed to infer lvlToDim from dimToLvl}}
		#a = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense)}>
		func.func private @tensor_sizes_mismatch(%arg0: tensor<8xi32, #a>) -> ()

		@@ -119,7 +119,7 @@ func.func private @tensor_dimtolvl_mismatch(%arg0: tensor<8xi32, #a>) -> ()

		// -----

		// expected-error@+1 {{expected a permutation affine map for dimToLvl}}
		// expected-error@+1 {{failed to infer lvlToDim from dimToLvl}}
		#a = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense, d0 : compressed)}>
		func.func private @tensor_no_permutation(%arg0: tensor<16x32xf32, #a>) -> ()

		@@ -251,3 +251,22 @@ func.func private @too_few_lvl_decl(%arg0: tensor<?x?xf64, #TooFewLvlDecl>) {
		func.func private @wrong_order_lvl_decl(%arg0: tensor<?x?xf64, #WrongOrderLvlDecl>) {
		return
		}

		// -----

		// expected-error@+1 {{expected lvlToDim to be an inverse of dimToLvl}}
		#BSR_explicit = #sparse_tensor.encoding<{
		map =
		{il, jl, ii, jj}
		( i = il * 3 + ii,
		j = jl * 2 + jj
		) ->
		( il = i floordiv 2 : dense,
		jl = j floordiv 3 : compressed,
		ii = i mod 2 : dense,
		jj = j mod 3 : dense
		)
		}>
		func.func private @BSR_explicit(%arg0: tensor<?x?xf64, #BSR_explicit>) {
		return
		}