Diff 458518

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

Show First 20 Lines • Show All 618 Lines • ▼ Show 20 Lines	def SparseTensor_YieldOp : SparseTensor_Op<"yield", [NoSideEffect, Terminator]>,
}];		}];

let assemblyFormat = [{		let assemblyFormat = [{
$result attr-dict `:` type($result)		$result attr-dict `:` type($result)
}];		}];
let hasVerifier = 1;		let hasVerifier = 1;
}		}

//===----------------------------------------------------------------------===//
// Sparse Tensor Storage Operation. These operations are used internally by
// sparse tensor codegen to progressively lower sparse tensors.
//===----------------------------------------------------------------------===//

def SparseTensor_StorageOp : SparseTensor_Op<"storage", []>,
Arguments<(ins Variadic<AnyType>:$inputs)>,
Results<(outs AnyTuple:$result)> {
let summary = "Pack a list of value into one sparse tensor storage value";
let description = [{
Pack a list of value into one sparse tensor storage value (represented as
a tuple) at the given index.

The result tuple elements' type should match the corresponding type in the
input array.

Example:

```mlir
%0 = sparse_tensor.storage(%1, %2): memref<?xf64>, memref<?xf64>
to tuple<memref<?xf64>, memref<?xf64>>
```
}];

let assemblyFormat = " attr-dict `(` $inputs `)``:` type($inputs) `to` type($result)";
let hasVerifier = 1;
}

def SparseTensor_StorageGetOp : SparseTensor_Op<"storage_get", []>,
Arguments<(ins AnyTuple:$storage,
IndexAttr:$idx)>,
Results<(outs AnyType:$result)> {
let summary = "Get the data stored in the sparse tensor storage at the given index";
let description = [{
Get the data stored in the sparse tensor storage (represented as a tuple)
at the given index.

The result type should match the corresponding element type in the tuple.

Example:

```mlir
%0 = sparse_tensor.storage_get %arg0[0] : tuple<memref<?xf64>, memref<?xf64>, f64> to memref<?xf64>
```
}];

let assemblyFormat = " $storage attr-dict `[`$idx`]` `:` type($storage) `to` type($result)";
let hasVerifier = 1;
}

def SparseTensor_StorageSetOp : SparseTensor_Op<"storage_set", []>,
Arguments<(ins AnyTuple:$storage,
AnyType:$value,
IndexAttr:$idx)>,
Results<(outs AnyTuple:$result)> {
let summary = "Set the data stored in the sparse tensor storage at given index";
let description = [{
Set the data stored in the sparse tensor storage (represented as a tuple)
at the given index. Return a new SSA value with the corresponding element
updated (others remain unchanged).

The result type should match the original tuple type with only the updated
element type changed accordingly.

Example:

```mlir
%0 = sparse_tensor.storage_set %arg0, %arg1 at 0 : tuple<memref<?xf64>, memref<?xf64>, f64>, memref<?xf64> to tuple<memref<?xf64>, memref<?xf64>, f64>
```
}];

let assemblyFormat = " $storage attr-dict `[`$idx`]``,` $value `:` type($storage) `,` type($value) `to` type($result)";
let hasVerifier = 1;
}

#endif // SPARSETENSOR_OPS		#endif // SPARSETENSOR_OPS

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

	Show First 20 Lines • Show All 150 Lines • ▼ Show 20 Lines

	/// Sets up sparse tensor conversion rules.			/// Sets up sparse tensor conversion rules.
	void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,			void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
	RewritePatternSet &patterns);			RewritePatternSet &patterns);

	std::unique_ptr<Pass> createSparseTensorCodegenPass();			std::unique_ptr<Pass> createSparseTensorCodegenPass();

	//===----------------------------------------------------------------------===//			//===----------------------------------------------------------------------===//
	// The SparseTensorStorageExpansion pass.
	//===----------------------------------------------------------------------===//

	/// Sparse tensor storage type converter from compound to expanded form.
	class SparseTensorStorageTupleExpander : public TypeConverter {
	public:
	SparseTensorStorageTupleExpander();
	};

	/// Sets up sparse tensor storage expansion rules.
	void populateSparseTensorStorageExpansionPatterns(TypeConverter &typeConverter,
	RewritePatternSet &patterns);

	std::unique_ptr<Pass> createSparseTensorStorageExpansionPass();

	//===----------------------------------------------------------------------===//
	// Other rewriting rules and passes.			// Other rewriting rules and passes.
	//===----------------------------------------------------------------------===//			//===----------------------------------------------------------------------===//

	void populateSparseTensorRewriting(RewritePatternSet &patterns);			void populateSparseTensorRewriting(RewritePatternSet &patterns);

	std::unique_ptr<Pass> createDenseBufferizationPass(			std::unique_ptr<Pass> createDenseBufferizationPass(
	const bufferization::OneShotBufferizationOptions &options);			const bufferization::OneShotBufferizationOptions &options);

	Show All 11 Lines

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

Show First 20 Lines • Show All 169 Lines • ▼ Show 20 Lines	let dependentDialects = [
"arith::ArithmeticDialect",		"arith::ArithmeticDialect",
"bufferization::BufferizationDialect",		"bufferization::BufferizationDialect",
"memref::MemRefDialect",		"memref::MemRefDialect",
"scf::SCFDialect",		"scf::SCFDialect",
"sparse_tensor::SparseTensorDialect",		"sparse_tensor::SparseTensorDialect",
];		];
}		}

def SparseTensorStorageExpansion : Pass<"sparse-tensor-storage-expansion", "ModuleOp"> {
let summary = "Expand compounded sparse tensor storage into individual SSA values";
let description = [{
A pass that expands sparse tensor storage (aggregated by tuple) into
individual SSA values. It also lowers sparse tensor storage operations,
e.g., sparse_tensor.storage_get and sparse_tensor.storage_set.

Example of the conversion:

```mlir
Before:
func.func @sparse_storage_set(%arg0: tuple<memref<?xf64>,
memref<?xf64>,
f64>)
-> tuple<memref<?xf64>,
memref<?xf64>,
f64> {
return %arg0 : tuple<memref<?xf64>, memref<?xf64>, f64>
}
After:
func.func @sparse_storage_set(%arg0: memref<?xf64>,
%arg1: memref<?xf64>,
%arg2: f64)
-> (memref<?xf64>, memref<?xf64>, f64) {
return %arg0, %arg1, %arg2 : memref<?xf64>, memref<?xf64>, f64
}
```
}];
let constructor = "mlir::createSparseTensorStorageExpansionPass()";
let dependentDialects = [
"sparse_tensor::SparseTensorDialect",
];
}


#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES		#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES

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

Show First 20 Lines • Show All 477 Lines • ▼ Show 20 Lines	if (isa<BinaryOp>(parentOp) \|\| isa<UnaryOp>(parentOp) \|\|
isa<ReduceOp>(parentOp))		isa<ReduceOp>(parentOp))
return success();		return success();

return emitOpError(		return emitOpError(
"expected parent op to be sparse_tensor unary, binary, or reduce");		"expected parent op to be sparse_tensor unary, binary, or reduce");
}		}

//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//
// Sparse Tensor Storage Operation.
//===----------------------------------------------------------------------===//

LogicalResult StorageOp::verify() {
auto retTypes = getResult().getType().getTypes();
if (retTypes.size() != getInputs().size())
return emitError("The number of inputs is inconsistent with output tuple");

for (auto pair : llvm::zip(getInputs(), retTypes)) {
auto input = std::get<0>(pair);
auto retTy = std::get<1>(pair);

if (input.getType() != retTy)
return emitError(llvm::formatv("Type mismatch between input (type={0}) "
"and output tuple element (type={1})",
input.getType(), retTy));
}
return success();
}

LogicalResult StorageGetOp::verify() {
uint64_t extractIdx = getIdx().getZExtValue();
auto innerTypeArray = getStorage().getType().getTypes();
if (extractIdx >= innerTypeArray.size())
return emitError(llvm::formatv(
"Out-of-bound access with index={0} on tuple with length={1}",
extractIdx, innerTypeArray.size()));

auto expectedTy = getStorage().getType().getType(extractIdx);
auto returnTy = getResult().getType();
if (expectedTy != returnTy)
return emitError(llvm::formatv(
"Type mismatch between the returning type (type={0}) and the "
"corresponding element type at index {1} (type={2})",
expectedTy, extractIdx, returnTy));
return success();
}

LogicalResult StorageSetOp::verify() {
uint64_t setIdx = getIdx().getZExtValue();
SmallVector<Type, 8> expectedElemTy(getStorage().getType().getTypes());
if (setIdx >= expectedElemTy.size())
return emitError(llvm::formatv(
"Out-of-bound access with index = {0} on tuple with length={1}", setIdx,
expectedElemTy.size()));

// Updates the element type after storage_set.
expectedElemTy[setIdx] = getValue().getType();
auto expectedTy = TupleType::get(getContext(), expectedElemTy);
auto returnTy = getResult().getType();
if (expectedTy != returnTy)
return emitError(
llvm::formatv("Type mismatch between the returning type "
"(type={0}) and the expected type (type={1})",
returnTy, expectedTy));
return success();
}

//===----------------------------------------------------------------------===//
// TensorDialect Methods.		// TensorDialect Methods.
//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//

void SparseTensorDialect::initialize() {		void SparseTensorDialect::initialize() {
addAttributes<		addAttributes<
#define GET_ATTRDEF_LIST		#define GET_ATTRDEF_LIST
#include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc"		#include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc"
>();		>();
Show All 10 Lines

mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt

	add_mlir_dialect_library(MLIRSparseTensorTransforms			add_mlir_dialect_library(MLIRSparseTensorTransforms
	BufferizableOpInterfaceImpl.cpp			BufferizableOpInterfaceImpl.cpp
	CodegenUtils.cpp			CodegenUtils.cpp
	DenseBufferizationPass.cpp			DenseBufferizationPass.cpp
	Sparsification.cpp			Sparsification.cpp
	SparseTensorCodegen.cpp			SparseTensorCodegen.cpp
	SparseTensorConversion.cpp			SparseTensorConversion.cpp
	SparseTensorPasses.cpp			SparseTensorPasses.cpp
	SparseTensorRewriting.cpp			SparseTensorRewriting.cpp
	SparseTensorStorageExpansion.cpp

	ADDITIONAL_HEADER_DIRS			ADDITIONAL_HEADER_DIRS
	${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor			${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor

	DEPENDS			DEPENDS
	MLIRSparseTensorPassIncGen			MLIRSparseTensorPassIncGen

	LINK_LIBS PUBLIC			LINK_LIBS PUBLIC
	Show All 20 Lines

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

Show First 20 Lines • Show All 48 Lines • ▼ Show 20 Lines	static unsigned toStored(const SparseTensorEncodingAttr &enc, unsigned i) {
auto order = enc.getDimOrdering();		auto order = enc.getDimOrdering();
if (order) {		if (order) {
assert(order.isPermutation());		assert(order.isPermutation());
return order.getPermutedPosition(i);		return order.getPermutedPosition(i);
}		}
return i;		return i;
}		}

		/// Flatten a list of operands that may contain sparse tensors.
		aartbikUnsubmitted Done Reply Inline Actions tenors -> tensors aartbik: tenors -> tensors
		PeimingAuthorUnsubmitted Done Reply Inline Actions Nice catch! Peiming: Nice catch!
		static void flattenOperands(ValueRange operands,
		SmallVectorImpl<Value> &flattened) {
		// In case of
		// sparse_tensor, c, sparse_tensor
		// ==>
		// memref ..., c, memref ...
		for (auto operand : operands) {
		if (auto cast =
		dyn_cast<UnrealizedConversionCastOp>(operand.getDefiningOp());
		cast && getSparseTensorEncoding(cast->getResultTypes()[0]))
		// An unrealized_conversion_cast will be inserted by type converter to
		// inter-mix the gap between 1:N conversion between sparse tensors and
		// fields. In this case, take the operands in the cast and replace the
		// sparse tensor output with the flattened type array.
		flattened.append(cast.getOperands().begin(), cast.getOperands().end());
		else
		flattened.push_back(operand);
		}
		}

/// Maps a sparse tensor type to the appropriate compounded buffers.		/// Maps a sparse tensor type to the appropriate compounded buffers.
static Optional<Type> convertSparseTensorType(Type type) {		static Optional<LogicalResult>
		convertSparseTensorType(Type type, SmallVectorImpl<Type> &fields) {
		aartbikUnsubmitted Done Reply Inline Actions this is no longer Optional, but just LogicalResult, since you have success() at L139 aartbik: this is no longer Optional, but just LogicalResult, since you have success() at L139
		PeimingAuthorUnsubmitted Done Reply Inline Actions This is a callback to the type converter. Return llvm::None tells the converter to try next converting rule. Peiming: This is a callback to the type converter. Return llvm::None tells the converter to try next…
		aartbikUnsubmitted Done Reply Inline Actions you are right! aartbik: you are right!
auto enc = getSparseTensorEncoding(type);		auto enc = getSparseTensorEncoding(type);
if (!enc)		if (!enc)
return llvm::None;		return llvm::None;
		aartbikUnsubmitted Done Reply Inline Actions return failure() aartbik: return failure()
// Construct the basic types.		// Construct the basic types.
auto context = type.getContext();		auto context = type.getContext();
unsigned idxWidth = enc.getIndexBitWidth();		unsigned idxWidth = enc.getIndexBitWidth();
unsigned ptrWidth = enc.getPointerBitWidth();		unsigned ptrWidth = enc.getPointerBitWidth();
RankedTensorType rType = type.cast<RankedTensorType>();		RankedTensorType rType = type.cast<RankedTensorType>();
Type indexType = IndexType::get(context);		Type indexType = IndexType::get(context);
Type idxType = idxWidth ? IntegerType::get(context, idxWidth) : indexType;		Type idxType = idxWidth ? IntegerType::get(context, idxWidth) : indexType;
Type ptrType = ptrWidth ? IntegerType::get(context, ptrWidth) : indexType;		Type ptrType = ptrWidth ? IntegerType::get(context, ptrWidth) : indexType;
Show All 11 Lines	convertSparseTensorType(Type type, SmallVectorImpl<Type> &fields) {
// memref<? x ptr> pointers-d ; pointers for sparse dim d		// memref<? x ptr> pointers-d ; pointers for sparse dim d
// memref<? x idx> indices-d ; indices for sparse dim d		// memref<? x idx> indices-d ; indices for sparse dim d
// ; if singleton:		// ; if singleton:
// memref<? x idx> indices-d ; indices for singleton dim d		// memref<? x idx> indices-d ; indices for singleton dim d
// memref<? x eltType> values ; values		// memref<? x eltType> values ; values
// };		// };
//		//
unsigned rank = rType.getShape().size();		unsigned rank = rType.getShape().size();
SmallVector<Type, 8> fields;
// The dimSizes array.		// The dimSizes array.
fields.push_back(MemRefType::get({rank}, indexType));		fields.push_back(MemRefType::get({rank}, indexType));
// Per-dimension storage.		// Per-dimension storage.
for (unsigned r = 0; r < rank; r++) {		for (unsigned r = 0; r < rank; r++) {
// Dimension level types apply in order to the reordered dimension.		// Dimension level types apply in order to the reordered dimension.
// As a result, the compound type can be constructed directly in the given		// As a result, the compound type can be constructed directly in the given
// order. Clients of this type know what field is what from the sparse		// order. Clients of this type know what field is what from the sparse
// tensor type.		// tensor type.
Show All 12 Lines	for (unsigned r = 0; r < rank; r++) {
case SparseTensorEncodingAttr::DimLevelType::SingletonNo:		case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:		case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, idxType));		fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, idxType));
break;		break;
}		}
}		}
// The values array.		// The values array.
fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, eltType));		fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, eltType));
// Sparse tensor storage (temporarily) lives in a tuple. This allows a		return success();
// simple 1:1 type conversion during codegen. A subsequent pass uses
// a 1:N type conversion to expand the tuple into its fields.
return TupleType::get(context, fields);
}		}

// Returns field index of sparse tensor type for pointers/indices, when set.		// Returns field index of sparse tensor type for pointers/indices, when set.
static unsigned getFieldIndex(Type type, unsigned ptrDim, unsigned idxDim) {		static unsigned getFieldIndex(Type type, unsigned ptrDim, unsigned idxDim) {
auto enc = getSparseTensorEncoding(type);		auto enc = getSparseTensorEncoding(type);
assert(enc);		assert(enc);
RankedTensorType rType = type.cast<RankedTensorType>();		RankedTensorType rType = type.cast<RankedTensorType>();
unsigned field = 1; // start at DimSizes;		unsigned field = 1; // start at DimSizes;
Show All 23 Lines	case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
field++;		field++;
break;		break;
}		}
}		}
llvm_unreachable("failed to find ptr/idx field index");		llvm_unreachable("failed to find ptr/idx field index");
return -1;		return -1;
}		}

/// Returns field type in tuple at given index.
static Type getFieldType(Value tuple, unsigned field) {
return tuple.getType().cast<TupleType>().getType(field);
}

/// Creates tuple get operation at given index.
static Value createTupleGet(OpBuilder &builder, Location loc, Value tuple,
unsigned field) {
Type indexType = builder.getIndexType();
return builder.create<StorageGetOp>(loc, getFieldType(tuple, field), tuple,
builder.getIntegerAttr(indexType, field));
}

/// Creates tuple.
static Value createTupleMake(OpBuilder &builder, Location loc, Type type,
ValueRange values) {
return builder.create<StorageOp>(loc, type, values);
}

/// Create allocation operation.		/// Create allocation operation.
static Value createAllocation(OpBuilder &builder, Location loc, Type type,		static Value createAllocation(OpBuilder &builder, Location loc, Type type,
Value sz) {		Value sz) {
auto memType = MemRefType::get({ShapedType::kDynamicSize}, type);		auto memType = MemRefType::get({ShapedType::kDynamicSize}, type);
return builder.create<memref::AllocOp>(loc, memType, sz);		return builder.create<memref::AllocOp>(loc, memType, sz);
}		}

/// Creates allocation tuple for sparse tensor type.		/// Creates allocation for each field in sparse tensor type.
///		///
/// TODO: for efficiency, we will need heuristis to make educated guesses		/// TODO: for efficiency, we will need heuristis to make educated guesses
/// on the required final sizes; also, we will need an improved		/// on the required final sizes; also, we will need an improved
/// memory allocation scheme with capacity and reallocation		/// memory allocation scheme with capacity and reallocation
///		///
static Value createAllocTuple(OpBuilder &builder, Location loc, Type type,		static void createAllocFields(OpBuilder &builder, Location loc, Type type,
ValueRange dynSizes) {		ValueRange dynSizes,
		SmallVectorImpl<Value> &fields) {
auto enc = getSparseTensorEncoding(type);		auto enc = getSparseTensorEncoding(type);
assert(enc);		assert(enc);
// Construct the basic types.		// Construct the basic types.
unsigned idxWidth = enc.getIndexBitWidth();		unsigned idxWidth = enc.getIndexBitWidth();
unsigned ptrWidth = enc.getPointerBitWidth();		unsigned ptrWidth = enc.getPointerBitWidth();
RankedTensorType rType = type.cast<RankedTensorType>();		RankedTensorType rType = type.cast<RankedTensorType>();
Type indexType = builder.getIndexType();		Type indexType = builder.getIndexType();
Type idxType = idxWidth ? builder.getIntegerType(idxWidth) : indexType;		Type idxType = idxWidth ? builder.getIntegerType(idxWidth) : indexType;
Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;		Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;
Type eltType = rType.getElementType();		Type eltType = rType.getElementType();
// Build the allocation tuple, using heuristics for pre-allocation.
auto shape = rType.getShape();		auto shape = rType.getShape();
unsigned rank = shape.size();		unsigned rank = shape.size();
SmallVector<Value, 8> fields;
bool allDense = true;		bool allDense = true;
Value one = constantIndex(builder, loc, 1);		Value one = constantIndex(builder, loc, 1);
Value linear = one;		Value linear = one;
Value heuristic = one; // FIX, see TODO above		Value heuristic = one; // FIX, see TODO above
// Build original sizes.		// Build original sizes.
SmallVector<Value, 8> sizes;		SmallVector<Value, 8> sizes;
for (unsigned r = 0, o = 0; r < rank; r++) {		for (unsigned r = 0, o = 0; r < rank; r++) {
if (ShapedType::isDynamic(shape[r]))		if (ShapedType::isDynamic(shape[r]))
Show All 32 Lines	case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
allDense = false;		allDense = false;
break;		break;
}		}
}		}
// The values array. For all-dense, the full length is required.		// The values array. For all-dense, the full length is required.
// In all other case, we resort to the heuristical initial value.		// In all other case, we resort to the heuristical initial value.
Value valuesSz = allDense ? linear : heuristic;		Value valuesSz = allDense ? linear : heuristic;
fields.push_back(createAllocation(builder, loc, eltType, valuesSz));		fields.push_back(createAllocation(builder, loc, eltType, valuesSz));
// Construct tuple allocation.
Type tupleType = *convertSparseTensorType(type);
return createTupleMake(builder, loc, tupleType, fields);
}		}

/// Returns integral constant, if defined.		/// Returns integral constant, if defined.
static Optional<int64_t> getConstantInt(Value val) {		static Optional<int64_t> getConstantInt(Value val) {
if (auto constantOp = val.getDefiningOp<arith::ConstantOp>())		if (auto constantOp = val.getDefiningOp<arith::ConstantOp>())
return constantOp.getValue().cast<IntegerAttr>().getInt();		return constantOp.getValue().cast<IntegerAttr>().getInt();
return {};		return {};
}		}

//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//
// Codegen rules.		// Codegen rules.
//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//

/// Sparse codegen rule for returns.		/// Sparse tensor storage conversion rule for returns.
class SparseReturnConverter : public OpConversionPattern<func::ReturnOp> {		class SparseReturnConverter : public OpConversionPattern<func::ReturnOp> {
public:		public:
using OpConversionPattern::OpConversionPattern;		using OpConversionPattern::OpConversionPattern;
LogicalResult		LogicalResult
matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor,		matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {		ConversionPatternRewriter &rewriter) const override {
rewriter.replaceOpWithNewOp<func::ReturnOp>(op, adaptor.getOperands());		SmallVector<Value, 8> flattened;
		flattenOperands(adaptor.getOperands(), flattened);
		// Create a return with the flattened value extracted from sparse tensors.
		rewriter.replaceOpWithNewOp<func::ReturnOp>(op, flattened);
		return success();
		}
		};

		/// Sparse tensor storage conversion rule for calls.
		class SparseCallConverter : public OpConversionPattern<func::CallOp> {
		public:
		// The default CallOp converter can not handle 1:N type conversion.
		using OpConversionPattern::OpConversionPattern;
		LogicalResult
		matchAndRewrite(func::CallOp op, OpAdaptor adaptor,
		ConversionPatternRewriter &rewriter) const override {
		Location loc = op.getLoc();
		// In case of:
		// sparse_tensor, f, sparse_tensor = call @foo(...)
		// ==>
		// memref..., f, memref = call @foo(...) replace with
		// cast(memref...)->sparse_tensor, f, cast(memref...)->sparse_tensor
		SmallVector<Type, 8> finalRetTy;
		if (failed(typeConverter->convertTypes(op.getResultTypes(), finalRetTy)))
		return failure();

		// (1) Genereates new call with flattened return value.
		SmallVector<Value, 8> flattened;
		flattenOperands(adaptor.getOperands(), flattened);
		auto newCall = rewriter.create<func::CallOp>(loc, op.getCallee(),
		finalRetTy, flattened);
		// (2) Create cast operation for sparse tensor returns.
		SmallVector<Value, 4> castedRet;
		// Tracks the offset of current return value (of the orignal call)
		// relative to the new call (after sparse tensor flattening);
		unsigned retOffset = 0;
		// Temporal buffer to hold the flattened list of type for
		// a sparse tensor.
		SmallVector<Type, 8> sparseFlat;
		for (auto ret : op.getResults()) {
		assert(retOffset < newCall.getNumResults());
		auto retType = ret.getType();
		if (failed(typeConverter->convertType(retType, sparseFlat)))
		// This should never happen.
		llvm_unreachable("Failed to convert type in sparse tensor codegen");

		// Converted types can not be empty when the type conversion succeed.
		assert(!sparseFlat.empty());
		if (sparseFlat.size() > 1) {
		auto flatSize = sparseFlat.size();
		ValueRange sparseElem(iterator_range<ResultRange::iterator>(
		newCall.result_begin() + retOffset,
		newCall.result_begin() + retOffset + flatSize));
		auto castOp = rewriter.create<UnrealizedConversionCastOp>(
		loc, TypeRange({retType}), sparseElem);
		castedRet.push_back(castOp.getResult(0));
		retOffset += flatSize;
		} else {
		// If this is an 1:1 conversion, no need for casting.
		castedRet.push_back(newCall.getResult(retOffset));
		retOffset++;
		}
		sparseFlat.clear();
		}

		assert(castedRet.size() == op.getNumResults());
		rewriter.replaceOp(op, castedRet);
return success();		return success();
}		}
};		};

/// Sparse codegen rule for dimension accesses.		/// Sparse codegen rule for dimension accesses.
class SparseDimOpConverter : public OpConversionPattern<tensor::DimOp> {		class SparseDimOpConverter : public OpConversionPattern<tensor::DimOp> {
public:		public:
using OpConversionPattern::OpConversionPattern;		using OpConversionPattern::OpConversionPattern;
Show All 12 Lines	matchAndRewrite(tensor::DimOp op, OpAdaptor adaptor,
Location loc = op->getLoc();		Location loc = op->getLoc();
auto shape = op.getSource().getType().cast<RankedTensorType>().getShape();		auto shape = op.getSource().getType().cast<RankedTensorType>().getShape();
if (!ShapedType::isDynamic(shape[*index])) {		if (!ShapedType::isDynamic(shape[*index])) {
rewriter.replaceOp(op, constantIndex(rewriter, loc, shape[*index]));		rewriter.replaceOp(op, constantIndex(rewriter, loc, shape[*index]));
return success();		return success();
}		}
// Any other query can consult the dimSizes array at field 0 using,		// Any other query can consult the dimSizes array at field 0 using,
// accounting for the reordering applied to the sparse storage.		// accounting for the reordering applied to the sparse storage.
Value tuple = adaptor.getSource();		auto tuple = llvm::cast<UnrealizedConversionCastOp>(
Value dimSizes = createTupleGet(rewriter, loc, tuple, 0);		adaptor.getSource().getDefiningOp());
		aartbikUnsubmitted Done Reply Inline Actions can we give this a better name then cast? I think the original tuple would still make a lot of sense here (as in a,b,c tuple, not tuple(a,b,c)). aartbik: can we give this a better name then cast? I think the original tuple would still make a lot of…
		PeimingAuthorUnsubmitted Done Reply Inline Actions I also abstracted the to_value/point/indices to share a SparseGetterOpConverter as the base class to reduce code repetition here. Peiming: I also abstracted the to_value/point/indices to share a SparseGetterOpConverter as the base…
rewriter.replaceOpWithNewOp<memref::LoadOp>(		rewriter.replaceOpWithNewOp<memref::LoadOp>(
op, dimSizes, constantIndex(rewriter, loc, toStored(enc, *index)));		op, tuple.getInputs().front(),
		constantIndex(rewriter, loc, toStored(enc, *index)));
return success();		return success();
}		}
};		};

/// Sparse codegen rule for trivial tensor casts.		/// Sparse codegen rule for trivial tensor casts.
class SparseCastConverter : public OpConversionPattern<tensor::CastOp> {		class SparseCastConverter : public OpConversionPattern<tensor::CastOp> {
public:		public:
using OpConversionPattern::OpConversionPattern;		using OpConversionPattern::OpConversionPattern;
Show All 19 Lines	public:
matchAndRewrite(bufferization::AllocTensorOp op, OpAdaptor adaptor,		matchAndRewrite(bufferization::AllocTensorOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {		ConversionPatternRewriter &rewriter) const override {
RankedTensorType resType = op.getType();		RankedTensorType resType = op.getType();
auto enc = getSparseTensorEncoding(resType);		auto enc = getSparseTensorEncoding(resType);
if (!enc)		if (!enc)
return failure();		return failure();
if (op.getCopy())		if (op.getCopy())
return rewriter.notifyMatchFailure(op, "tensor copy not implemented");		return rewriter.notifyMatchFailure(op, "tensor copy not implemented");
// Construct allocation tuple.
Value tuple = createAllocTuple(rewriter, op->getLoc(), resType,		// Construct allocation for each field.
		aartbikUnsubmitted Done Reply Inline Actions move below // Construct aartbik: move below // Construct
adaptor.getOperands());		Location loc = op.getLoc();
rewriter.replaceOp(op, tuple);		SmallVector<Value, 8> fields;
		createAllocFields(rewriter, loc, resType, adaptor.getOperands(), fields);
		rewriter.replaceOpWithNewOp<UnrealizedConversionCastOp>(
		op, TypeRange{resType}, fields);
return success();		return success();
}		}
};		};

/// Sparse codegen rule for the dealloc operator.		/// Sparse codegen rule for the dealloc operator.
class SparseTensorDeallocConverter		class SparseTensorDeallocConverter
: public OpConversionPattern<bufferization::DeallocTensorOp> {		: public OpConversionPattern<bufferization::DeallocTensorOp> {
public:		public:
using OpConversionPattern::OpConversionPattern;		using OpConversionPattern::OpConversionPattern;
LogicalResult		LogicalResult
matchAndRewrite(bufferization::DeallocTensorOp op, OpAdaptor adaptor,		matchAndRewrite(bufferization::DeallocTensorOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {		ConversionPatternRewriter &rewriter) const override {
auto enc = getSparseTensorEncoding(op.getTensor().getType());		auto enc = getSparseTensorEncoding(op.getTensor().getType());
if (!enc)		if (!enc)
return failure();		return failure();
// Replace the tuple deallocation with field deallocations.
Location loc = op->getLoc();		// Replace the sparse tensor deallocation with field deallocations.
aartbikUnsubmitted Done Reply Inline Actions please keep Location and also the place of this decl (after the comment on // Replace) aartbik: please keep Location and also the place of this decl (after the comment on // Replace)
		aartbikUnsubmitted Done Reply Inline Actions below // Replace aartbik: below // Replace
Value tuple = adaptor.getTensor();		Location loc = op.getLoc();
for (unsigned i = 0, sz = tuple.getType().cast<TupleType>().size(); i < sz;		auto tuple = llvm::cast<UnrealizedConversionCastOp>(
i++) {		adaptor.getTensor().getDefiningOp());
Value mem = createTupleGet(rewriter, loc, tuple, i);		for (auto input : tuple.getInputs())
rewriter.create<memref::DeallocOp>(loc, mem);		// Deallocate every buffer used to store the sparse tensor handler.
}		rewriter.create<memref::DeallocOp>(loc, input);

rewriter.eraseOp(op);		rewriter.eraseOp(op);
return success();		return success();
}		}
};		};

/// Sparse codegen rule for pointer accesses.		/// Sparse codegen rule for tensor rematerialization.
class SparseToPointersConverter : public OpConversionPattern<ToPointersOp> {		class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> {
public:		public:
using OpConversionPattern::OpConversionPattern;		using OpConversionPattern::OpConversionPattern;
LogicalResult		LogicalResult
matchAndRewrite(ToPointersOp op, OpAdaptor adaptor,		matchAndRewrite(LoadOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {		ConversionPatternRewriter &rewriter) const override {
Optional<int64_t> index = getConstantInt(adaptor.getOperands()[1]);		if (op.getHasInserts()) {
		aartbikUnsubmitted Done Reply Inline Actions tuple? aartbik: tuple?
if (!index)		// Finalize any pending insertions.
return failure();		// TODO: implement
// Replace the requested pointer access with corresponding field.		}
Location loc = op->getLoc();		rewriter.replaceOp(op, adaptor.getOperands());
Value tuple = adaptor.getTensor();
unsigned i = getFieldIndex(op.getTensor().getType(), /ptrDim=/*index, -1);
rewriter.replaceOp(op, createTupleGet(rewriter, loc, tuple, i));
return success();		return success();
}		}
};		};

/// Sparse codegen rule for index accesses.		/// Base class for getter-like operations, e.g., to_indices, to_pointers.
class SparseToIndicesConverter : public OpConversionPattern<ToIndicesOp> {		template <typename SourceOp, typename Base>
		class SparseGetterOpConverter : public OpConversionPattern<SourceOp> {
public:		public:
using OpConversionPattern::OpConversionPattern;		using OpAdaptor = typename SourceOp::Adaptor;
		using OpConversionPattern<SourceOp>::OpConversionPattern;
LogicalResult		LogicalResult
matchAndRewrite(ToIndicesOp op, OpAdaptor adaptor,		matchAndRewrite(SourceOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {		ConversionPatternRewriter &rewriter) const override {
Optional<int64_t> index = getConstantInt(adaptor.getOperands()[1]);		// Replace the requested pointer access with corresponding field.
if (!index)		// The cast_op is inserted by type converter to intermix 1:N type
		// conversion.
		auto tuple = llvm::cast<UnrealizedConversionCastOp>(
		adaptor.getTensor().getDefiningOp());
		auto idx = Base::getIndexForOp(tuple, op);
		if (!idx)
		// Failed to get the index.
return failure();		return failure();
// Replace the requested indices access with corresponding field.		auto fields = tuple.getInputs();
		aartbikUnsubmitted Done Reply Inline Actions tuple? aartbik: tuple?
Location loc = op->getLoc();		assert(*idx < fields.size());
Value tuple = adaptor.getTensor();		rewriter.replaceOp(op, fields[*idx]);
unsigned i = getFieldIndex(op.getTensor().getType(), -1, /idxDim=/*index);
rewriter.replaceOp(op, createTupleGet(rewriter, loc, tuple, i));
return success();		return success();
}		}
};		};

/// Sparse codegen rule for value accesses.		/// Sparse codegen rule for pointer accesses.
class SparseToValuesConverter : public OpConversionPattern<ToValuesOp> {		class SparseToPointersConverter
		: public SparseGetterOpConverter<ToPointersOp, SparseToPointersConverter> {
public:		public:
using OpConversionPattern::OpConversionPattern;		using SparseGetterOpConverter::SparseGetterOpConverter;
		aartbikUnsubmitted Done Reply Inline Actions tuple? aartbik: tuple?
LogicalResult		// Callback for SparseGetterOpConverter.
matchAndRewrite(ToValuesOp op, OpAdaptor adaptor,		static Optional<unsigned> getIndexForOp(UnrealizedConversionCastOp /tuple/,
ConversionPatternRewriter &rewriter) const override {		ToPointersOp op) {
// Replace the requested values access with corresponding field.		Optional<int64_t> dim = getConstantInt(op.getDim());
		aartbikUnsubmitted Done Reply Inline Actions Note that in the original rule, I returned failure(), i.e. the rewriting does not occur (which is the idiom used elsehwere too for stuff the sparse compiler cannot handle). We typically generate these ops ourselves, so they are always constant index. But just in case they are not, it seems safer to not touch them. An assert only happens in debug mode, in prod, we will crash on L491 While writing these rules, I was thinking that we should really put in the restriction in the op itself. The variable case can be handled, but requires setting up some indexing array to get the field at runtime. Does not seem worth it. aartbik: Note that in the original rule, I returned failure(), i.e. the rewriting does not occur (which…
Location loc = op->getLoc();		if (!dim)
Value tuple = adaptor.getTensor();		return llvm::None; // variable dim
unsigned i = tuple.getType().cast<TupleType>().size() - 1; // last		return getFieldIndex(op.getTensor().getType(), /ptrDim=/*dim, -1);
rewriter.replaceOp(op, createTupleGet(rewriter, loc, tuple, i));
return success();
}		}
};		};

/// Sparse codegen rule for tensor rematerialization.		/// Sparse codegen rule for index accesses.
class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> {		class SparseToIndicesConverter
		: public SparseGetterOpConverter<ToIndicesOp, SparseToIndicesConverter> {
public:		public:
using OpConversionPattern::OpConversionPattern;		using SparseGetterOpConverter::SparseGetterOpConverter;
LogicalResult		// Callback for SparseGetterOpConverter.
matchAndRewrite(LoadOp op, OpAdaptor adaptor,		static Optional<unsigned> getIndexForOp(UnrealizedConversionCastOp /tuple/,
ConversionPatternRewriter &rewriter) const override {		ToIndicesOp op) {
if (op.getHasInserts()) {		Optional<int64_t> dim = getConstantInt(op.getDim());
// Finalize any pending insertions.		if (!dim)
		aartbikUnsubmitted Done Reply Inline Actions Is it possible to let Base::getFieldForOp() return unsigned field. If so, you can push assert(i < fields.size()); return fields[i]; into the base class, and make all derived classes even shorter? aartbik: Is it possible to let Base::getFieldForOp() return unsigned field. If so, you can push…
		PeimingAuthorUnsubmitted Done Reply Inline Actions Good suggestion! Peiming: Good suggestion!
// TODO: implement		return llvm::None; // variable dim
		return getFieldIndex(op.getTensor().getType(), -1, /idxDim=/*dim);
}		}
rewriter.replaceOp(op, adaptor.getOperands());		};
return success();
		/// Sparse codegen rule for value accesses.
		class SparseToValuesConverter
		: public SparseGetterOpConverter<ToValuesOp, SparseToValuesConverter> {
		public:
		using SparseGetterOpConverter::SparseGetterOpConverter;
		// Callback for SparseGetterOpConverter.
		static Optional<unsigned> getIndexForOp(UnrealizedConversionCastOp tuple,
		ToValuesOp /op/) {
		// The last field holds the value buffer.
		return tuple.getInputs().size() - 1;
}		}
};		};

} // namespace		} // namespace

//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//
// Sparse tensor type conversion into an actual buffer.		// Sparse tensor type conversion into an actual buffer.
//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//

mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {		mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
addConversion([](Type type) { return type; });		addConversion([](Type type) { return type; });
addConversion(convertSparseTensorType);		addConversion(convertSparseTensorType);
}		}

//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//
// Public method for populating conversion rules.		// Public method for populating conversion rules.
//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//

/// Populates the given patterns list with conversion rules required for		/// Populates the given patterns list with conversion rules required for
/// the sparsification of linear algebra operations.		/// the sparsification of linear algebra operations.
void mlir::populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,		void mlir::populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
RewritePatternSet &patterns) {		RewritePatternSet &patterns) {
patterns.add<SparseReturnConverter, SparseDimOpConverter, SparseCastConverter,		patterns.add<SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
SparseTensorAllocConverter, SparseTensorDeallocConverter,		SparseCastConverter, SparseTensorAllocConverter,
SparseToPointersConverter, SparseToIndicesConverter,		SparseTensorDeallocConverter, SparseToPointersConverter,
SparseToValuesConverter, SparseTensorLoadConverter>(		SparseToIndicesConverter, SparseToValuesConverter,
typeConverter, patterns.getContext());		SparseTensorLoadConverter>(typeConverter, patterns.getContext());
}		}

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

Show All 18 Lines
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"		#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"		#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"		#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

namespace mlir {		namespace mlir {
#define GEN_PASS_DEF_SPARSIFICATIONPASS		#define GEN_PASS_DEF_SPARSIFICATIONPASS
#define GEN_PASS_DEF_SPARSETENSORCONVERSIONPASS		#define GEN_PASS_DEF_SPARSETENSORCONVERSIONPASS
#define GEN_PASS_DEF_SPARSETENSORCODEGEN		#define GEN_PASS_DEF_SPARSETENSORCODEGEN
#define GEN_PASS_DEF_SPARSETENSORSTORAGEEXPANSION
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"		#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
} // namespace mlir		} // namespace mlir

using namespace mlir;		using namespace mlir;
using namespace mlir::sparse_tensor;		using namespace mlir::sparse_tensor;

namespace {		namespace {

▲ Show 20 Lines • Show All 113 Lines • ▼ Show 20 Lines	struct SparseTensorCodegenPass
SparseTensorCodegenPass() = default;		SparseTensorCodegenPass() = default;
SparseTensorCodegenPass(const SparseTensorCodegenPass &pass) = default;		SparseTensorCodegenPass(const SparseTensorCodegenPass &pass) = default;

void runOnOperation() override {		void runOnOperation() override {
auto *ctx = &getContext();		auto *ctx = &getContext();
RewritePatternSet patterns(ctx);		RewritePatternSet patterns(ctx);
SparseTensorTypeToBufferConverter converter;		SparseTensorTypeToBufferConverter converter;
ConversionTarget target(*ctx);		ConversionTarget target(*ctx);
// Almost everything in the sparse dialect must go!		// Everything in the sparse dialect must go!
		aartbikUnsubmitted Done Reply Inline Actions This is again, Everything! aartbik: This is again, Everything!
target.addIllegalDialect<SparseTensorDialect>();		target.addIllegalDialect<SparseTensorDialect>();
target.addLegalOp<StorageGetOp, StorageSetOp, StorageOp>();
// All dynamic rules below accept new function, call, return, and various		// All dynamic rules below accept new function, call, return, and various
// tensor and bufferization operations as legal output of the rewriting		// tensor and bufferization operations as legal output of the rewriting
// provided that all sparse tensor types have been fully rewritten.		// provided that all sparse tensor types have been fully rewritten.
target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {		target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
return converter.isSignatureLegal(op.getFunctionType());		return converter.isSignatureLegal(op.getFunctionType());
});		});
target.addDynamicallyLegalOp<func::CallOp>([&](func::CallOp op) {		target.addDynamicallyLegalOp<func::CallOp>([&](func::CallOp op) {
return converter.isSignatureLegal(op.getCalleeType());		return converter.isSignatureLegal(op.getCalleeType());
});		});
target.addDynamicallyLegalOp<func::ReturnOp>([&](func::ReturnOp op) {		target.addDynamicallyLegalOp<func::ReturnOp>([&](func::ReturnOp op) {
return converter.isLegal(op.getOperandTypes());		return converter.isLegal(op.getOperandTypes());
});		});
target.addDynamicallyLegalOp<bufferization::AllocTensorOp>(		target.addDynamicallyLegalOp<bufferization::AllocTensorOp>(
[&](bufferization::AllocTensorOp op) {		[&](bufferization::AllocTensorOp op) {
return converter.isLegal(op.getType());		return converter.isLegal(op.getType());
});		});
target.addDynamicallyLegalOp<bufferization::DeallocTensorOp>(		target.addDynamicallyLegalOp<bufferization::DeallocTensorOp>(
[&](bufferization::DeallocTensorOp op) {		[&](bufferization::DeallocTensorOp op) {
return converter.isLegal(op.getTensor().getType());		return converter.isLegal(op.getTensor().getType());
});		});
// Legal dialects may occur in generated code.		// Legal dialects may occur in generated code.
target.addLegalDialect<arith::ArithmeticDialect,		target.addLegalDialect<arith::ArithmeticDialect,
bufferization::BufferizationDialect,		bufferization::BufferizationDialect,
memref::MemRefDialect, scf::SCFDialect>();		memref::MemRefDialect, scf::SCFDialect>();
// Populate with rules and apply rewriting rules.
populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
converter);
populateCallOpTypeConversionPattern(patterns, converter);
scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
target);
populateSparseTensorCodegenPatterns(converter, patterns);
if (failed(applyPartialConversion(getOperation(), target,
std::move(patterns))))
signalPassFailure();
}
};

struct SparseTensorStorageExpansionPass
: public impl::SparseTensorStorageExpansionBase<
SparseTensorStorageExpansionPass> {

SparseTensorStorageExpansionPass() = default;
SparseTensorStorageExpansionPass(
const SparseTensorStorageExpansionPass &pass) = default;

void runOnOperation() override {
auto *ctx = &getContext();
RewritePatternSet patterns(ctx);
SparseTensorStorageTupleExpander converter;
ConversionTarget target(*ctx);
// Now, everything in the sparse dialect must go!
target.addIllegalDialect<SparseTensorDialect>();
// All dynamic rules below accept new function, call, return.
target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
return converter.isSignatureLegal(op.getFunctionType());
});
target.addDynamicallyLegalOp<func::CallOp>([&](func::CallOp op) {
return converter.isSignatureLegal(op.getCalleeType());
});
target.addDynamicallyLegalOp<func::ReturnOp>([&](func::ReturnOp op) {
return converter.isLegal(op.getOperandTypes());
});
// We generate UnrealizedConversionCastOp to intermix tuples and a
// list of types.
target.addLegalOp<UnrealizedConversionCastOp>();		target.addLegalOp<UnrealizedConversionCastOp>();
// Populate with rules and apply rewriting rules.		// Populate with rules and apply rewriting rules.
populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,		populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
converter);		converter);
scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,		scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
		aartbikUnsubmitted Done Reply Inline Actions commented out? aartbik: commented out?
target);		target);
populateSparseTensorStorageExpansionPatterns(converter, patterns);		populateSparseTensorCodegenPatterns(converter, patterns);
if (failed(applyPartialConversion(getOperation(), target,		if (failed(applyPartialConversion(getOperation(), target,
std::move(patterns))))		std::move(patterns))))
signalPassFailure();		signalPassFailure();
}		}
};		};

} // namespace		} // namespace

Show All 33 Lines
std::unique_ptr<Pass> mlir::createSparseTensorConversionPass(		std::unique_ptr<Pass> mlir::createSparseTensorConversionPass(
const SparseTensorConversionOptions &options) {		const SparseTensorConversionOptions &options) {
return std::make_unique<SparseTensorConversionPass>(options);		return std::make_unique<SparseTensorConversionPass>(options);
}		}

std::unique_ptr<Pass> mlir::createSparseTensorCodegenPass() {		std::unique_ptr<Pass> mlir::createSparseTensorCodegenPass() {
return std::make_unique<SparseTensorCodegenPass>();		return std::make_unique<SparseTensorCodegenPass>();
}		}

std::unique_ptr<Pass> mlir::createSparseTensorStorageExpansionPass() {
return std::make_unique<SparseTensorStorageExpansionPass>();
}

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageExpansion.cpp

This file was deleted.

	//===- SparseTensorStorageExpansion.cpp - Sparse tensor storage expansion ===//
	//
	// 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
	//
	//===----------------------------------------------------------------------===//
	//
	// The sparse tensor storage expansion pass expands the compound storage for
	// sparse tensors (using tuple) to flattened SSA values.
	//
	//===----------------------------------------------------------------------===//

	#include "CodegenUtils.h"

	#include "mlir/Dialect/Func/IR/FuncOps.h"
	#include "mlir/Dialect/MemRef/IR/MemRef.h"
	#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
	#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
	#include "mlir/Dialect/Tensor/IR/Tensor.h"
	#include "mlir/Transforms/DialectConversion.h"

	using namespace mlir;
	using namespace mlir::sparse_tensor;

	namespace {

	//===----------------------------------------------------------------------===//
	// Helper methods.
	//===----------------------------------------------------------------------===//

	/// Expands sparse tensor storage tuple.
	static Optional<LogicalResult>
	convertSparseTensorStorageTuple(Type t, SmallVectorImpl<Type> &result) {
	if (auto tuple = t.dyn_cast<TupleType>()) {
	// Note that it does not handle nest tuples, but it is fine
	// for sparse compiler as they will not be generated.
	result.append(tuple.getTypes().begin(), tuple.getTypes().end());
	return success();
	}
	return llvm::None;
	}

	/// Flatten a list of operands that may contain tuples.
	static void flattenOperands(ValueRange operands,
	SmallVectorImpl<Value> &flattened) {
	// In case of
	// tuple<a, b>, c, tuple<d, e>
	// ==>
	// a, b, c, d, e
	for (auto operand : operands) {
	if (auto cast =
	dyn_cast<UnrealizedConversionCastOp>(operand.getDefiningOp());
	cast && cast->getResultTypes()[0].isa<TupleType>())
	// An unrealized_conversion_cast will be inserted by type converter to
	// inter-mix the gap between 1:N conversion between tuple and types.
	// In this case, take the operands in the cast and replace the tuple
	// output with the flattened type array.
	flattened.append(cast.getOperands().begin(), cast.getOperands().end());
	else
	flattened.push_back(operand);
	}
	}
	//===----------------------------------------------------------------------===//
	// Conversion rules.
	//===----------------------------------------------------------------------===//

	/// Sparse tensor storage conversion rule for sparse_tensor::storage.
	class SparseStorageConversion : public OpConversionPattern<StorageOp> {
	public:
	using OpConversionPattern::OpConversionPattern;
	LogicalResult
	matchAndRewrite(StorageOp op, OpAdaptor adaptor,
	ConversionPatternRewriter &rewriter) const override {
	// Simply convert it to a unrealize_conversion_cast.
	// We should guarantee that all uses of sparse_tensor.storage op will
	// be eventually eliminated by accessing the flattened SSA values directly.
	rewriter.replaceOpWithNewOp<UnrealizedConversionCastOp>(
	op, TypeRange{op.getType()}, adaptor.getInputs());
	return success();
	}
	};

	/// Sparse tensor storage conversion rule for sparse_tensor::storage_get.
	class SparseStorageGetConverter : public OpConversionPattern<StorageGetOp> {
	public:
	using OpConversionPattern::OpConversionPattern;
	LogicalResult
	matchAndRewrite(StorageGetOp op, OpAdaptor adaptor,
	ConversionPatternRewriter &rewriter) const override {
	auto castOp =
	cast<UnrealizedConversionCastOp>(adaptor.getStorage().getDefiningOp());
	uint64_t idx = op.getIdx().getZExtValue();
	assert(idx < castOp.getOperands().size());

	rewriter.replaceOp(op, castOp.getOperand(idx));
	return success();
	}
	};

	/// Sparse tensor storage conversion rule for sparse_tensor::storage_set.
	class SparseStorageSetConverter : public OpConversionPattern<StorageSetOp> {
	public:
	using OpConversionPattern::OpConversionPattern;
	LogicalResult
	matchAndRewrite(StorageSetOp op, OpAdaptor adaptor,
	ConversionPatternRewriter &rewriter) const override {
	auto castOp =
	cast<UnrealizedConversionCastOp>(adaptor.getStorage().getDefiningOp());
	uint64_t idx = op.getIdx().getZExtValue();

	SmallVector<Value, 8> values(castOp.getOperands());
	assert(idx < values.size());

	// Updates the corresponding element.
	values[idx] = adaptor.getValue();
	rewriter.replaceOpWithNewOp<UnrealizedConversionCastOp>(
	op, TypeRange{op.getType()}, values);
	return success();
	}
	};

	/// Sparse tensor storage conversion rule for returns.
	class SparseStorageReturnConverter
	: public OpConversionPattern<func::ReturnOp> {
	public:
	using OpConversionPattern::OpConversionPattern;
	LogicalResult
	matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor,
	ConversionPatternRewriter &rewriter) const override {
	SmallVector<Value, 8> flattened;
	flattenOperands(adaptor.getOperands(), flattened);
	// Create a return with the flattened value extracted from tuple.
	rewriter.replaceOpWithNewOp<func::ReturnOp>(op, flattened);
	return success();
	}
	};

	/// Sparse tensor storage conversion rule for calls.
	class SparseStorageCallConverter : public OpConversionPattern<func::CallOp> {
	public:
	// The default CallOp converter can not handle 1:N type conversion properly
	using OpConversionPattern::OpConversionPattern;
	LogicalResult
	matchAndRewrite(func::CallOp op, OpAdaptor adaptor,
	ConversionPatternRewriter &rewriter) const override {
	Location loc = op.getLoc();
	// In case of:
	// tuple(a, b), f, tuple(c, d) = call @foo(...)
	// ==>
	// a, b, f, c, d = call @foo(...)
	// cast(a, b)->tuple, f, cast(c,d)->tuple
	SmallVector<Type, 8> finalRetTy;
	if (failed(typeConverter->convertTypes(op.getResultTypes(), finalRetTy)))
	return failure();

	// (1) Genereates new call with flattened return value.
	SmallVector<Value, 8> flattened;
	flattenOperands(adaptor.getOperands(), flattened);
	auto newCall = rewriter.create<func::CallOp>(loc, op.getCallee(),
	finalRetTy, flattened);

	// (2) Create cast operation for tuple returns.
	SmallVector<Value, 4> castedRet;
	// Tracks the offset of current return value (of the orignal call)
	// relative to the new call (after tuple flattening);
	unsigned retOffset = 0;
	for (auto ret : op.getResults()) {
	assert(retOffset < newCall.getNumResults());
	auto tupleRet = ret.getType().dyn_cast<TupleType>();
	if (tupleRet) {
	auto tupleSize = tupleRet.size();
	// NOTE: The range is computed under the assumption of non-recursive
	// tuple type.
	ValueRange tupleElem(iterator_range<ResultRange::iterator>(
	newCall.result_begin() + retOffset,
	newCall.result_begin() + retOffset + tupleSize));
	auto castOp = rewriter.create<UnrealizedConversionCastOp>(
	loc, TypeRange({tupleRet}), tupleElem);
	castedRet.push_back(castOp.getResult(0));
	retOffset += tupleSize;
	} else {
	// If this not a tuple, simply add it into returned values.
	castedRet.push_back(ret);
	retOffset++;
	}
	}

	assert(castedRet.size() == op.getNumResults());
	rewriter.replaceOp(op, castedRet);
	return success();
	}
	};

	} // namespace

	//===----------------------------------------------------------------------===//
	// Sparse tensor storage expansion
	//===----------------------------------------------------------------------===//

	mlir::SparseTensorStorageTupleExpander::SparseTensorStorageTupleExpander() {
	addConversion([](Type type) { return type; });
	addConversion(convertSparseTensorStorageTuple);
	}

	//===----------------------------------------------------------------------===//
	// Public method for populating conversion rules.
	//===----------------------------------------------------------------------===//

	/// Populates the given patterns list with conversion rules required
	/// to expand compounded sparse tensor tuples.
	void mlir::populateSparseTensorStorageExpansionPatterns(
	TypeConverter &typeConverter, RewritePatternSet &patterns) {
	patterns.add<SparseStorageConversion, SparseStorageGetConverter,
	SparseStorageSetConverter, SparseStorageReturnConverter,
	SparseStorageCallConverter>(typeConverter,
	patterns.getContext());
	}

mlir/test/Dialect/SparseTensor/codegen.mlir

// RUN: mlir-opt %s --sparse-tensor-codegen --canonicalize --cse \| FileCheck %s --check-prefixes=CHECK,CHECK-CODEGEN		// RUN: mlir-opt %s --sparse-tensor-codegen --canonicalize --cse \| FileCheck %s
// RUN: mlir-opt %s --sparse-tensor-codegen --sparse-tensor-storage-expansion --canonicalize --cse \| FileCheck %s --check-prefixes=CHECK,CHECK-STORAGE

#SparseVector = #sparse_tensor.encoding<{		#SparseVector = #sparse_tensor.encoding<{
dimLevelType = [ "compressed" ],		dimLevelType = [ "compressed" ],
indexBitWidth = 64,		indexBitWidth = 64,
pointerBitWidth = 32		pointerBitWidth = 32
}>		}>

#Dense2D = #sparse_tensor.encoding<{		#Dense2D = #sparse_tensor.encoding<{
Show All 25 Lines	#DCSR = #sparse_tensor.encoding<{
pointerBitWidth = 32		pointerBitWidth = 32
}>		}>

#Dense3D = #sparse_tensor.encoding<{		#Dense3D = #sparse_tensor.encoding<{
dimLevelType = [ "dense", "dense", "dense" ],		dimLevelType = [ "dense", "dense", "dense" ],
dimOrdering = affine_map<(i, j, k) -> (k, i, j)>		dimOrdering = affine_map<(i, j, k) -> (k, i, j)>
}>		}>

// CHECK-CODEGEN-LABEL: func @sparse_nop(		// CHECK-LABEL: func @sparse_nop(
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<1xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>>)		// CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
// CHECK-CODEGEN: return %[[A]] : tuple<memref<1xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>>		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
//		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-STORAGE-LABEL: func @sparse_nop(		// CHECK-SAME: %[[A3:.*3]]: memref<?xf64>)
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<1xindex>,		// CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] : memref<1xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref<?xf64>)
// CHECK-STORAGE: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] : memref<1xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
func.func @sparse_nop(%arg0: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {		func.func @sparse_nop(%arg0: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
return %arg0 : tensor<?xf64, #SparseVector>		return %arg0 : tensor<?xf64, #SparseVector>
}		}

// CHECK-CODEGEN-LABEL: func @sparse_nop_cast(		// CHECK-LABEL: func @sparse_nop_multi_ret(
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<1xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>>)		// CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
// CHECK-CODEGEN: return %[[A]] : tuple<memref<1xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>>		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
//		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-STORAGE-LABEL: func @sparse_nop_cast(		// CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<1xindex>,		// CHECK-SAME: %[[A4:.*4]]: memref<1xindex>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xi32>,		// CHECK-SAME: %[[A5:.*5]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref<?xi64>,		// CHECK-SAME: %[[A6:.*6]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref<?xf32>)		// CHECK-SAME: %[[A7:.*7]]: memref<?xf64>) ->
// CHECK-STORAGE: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] : memref<1xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>		// CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A7]]
		func.func @sparse_nop_multi_ret(%arg0: tensor<?xf64, #SparseVector>,
		%arg1: tensor<?xf64, #SparseVector>) ->
		(tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) {
		return %arg0, %arg1 : tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>
		}

		// CHECK-LABEL: func @sparse_nop_call(
		// CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
		// CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
		// CHECK-SAME: %[[A4:.*4]]: memref<1xindex>,
		// CHECK-SAME: %[[A5:.*5]]: memref<?xi32>,
		// CHECK-SAME: %[[A6:.*6]]: memref<?xi64>,
		// CHECK-SAME: %[[A7:.*7]]: memref<?xf64>)
		// CHECK: %[[T0:.*]]:8 = call @sparse_nop_multi_ret(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A7]])
		// CHECK: return %[[T0]]#0, %[[T0]]#1, %[[T0]]#2, %[[T0]]#3, %[[T0]]#4, %[[T0]]#5, %[[T0]]#6, %[[T0]]#7
		func.func @sparse_nop_call(%arg0: tensor<?xf64, #SparseVector>,
		%arg1: tensor<?xf64, #SparseVector>) ->
		(tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) {
		%1, %2 = call @sparse_nop_multi_ret(%arg0, %arg1) :
		(tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) ->
		(tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>)
		return %1, %2: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>
		}

		//
		// CHECK-LABEL: func @sparse_nop_cast(
		// CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
		// CHECK-SAME: %[[A3:.*3]]: memref<?xf32>)
		// CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] : memref<1xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>
func.func @sparse_nop_cast(%arg0: tensor<64xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {		func.func @sparse_nop_cast(%arg0: tensor<64xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {
%0 = tensor.cast %arg0 : tensor<64xf32, #SparseVector> to tensor<?xf32, #SparseVector>		%0 = tensor.cast %arg0 : tensor<64xf32, #SparseVector> to tensor<?xf32, #SparseVector>
return %0 : tensor<?xf32, #SparseVector>		return %0 : tensor<?xf32, #SparseVector>
}		}

// CHECK-CODEGEN-LABEL: func @sparse_nop_cast_3d(		//
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<3xindex>, memref<?xf32>>)		// CHECK-LABEL: func @sparse_nop_cast_3d(
// CHECK-CODEGEN: return %[[A]] : tuple<memref<3xindex>, memref<?xf32>>		// CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
//		// CHECK-SAME: %[[A1:.*1]]: memref<?xf32>)
// CHECK-STORAGE-LABEL: func @sparse_nop_cast_3d(		// CHECK: return %[[A0]], %[[A1]] : memref<3xindex>, memref<?xf32>
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<3xindex>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xf32>)
// CHECK-STORAGE: return %[[A0]], %[[A1]] : memref<3xindex>, memref<?xf32>
func.func @sparse_nop_cast_3d(%arg0: tensor<10x20x30xf32, #Dense3D>) -> tensor<?x?x?xf32, #Dense3D> {		func.func @sparse_nop_cast_3d(%arg0: tensor<10x20x30xf32, #Dense3D>) -> tensor<?x?x?xf32, #Dense3D> {
%0 = tensor.cast %arg0 : tensor<10x20x30xf32, #Dense3D> to tensor<?x?x?xf32, #Dense3D>		%0 = tensor.cast %arg0 : tensor<10x20x30xf32, #Dense3D> to tensor<?x?x?xf32, #Dense3D>
return %0 : tensor<?x?x?xf32, #Dense3D>		return %0 : tensor<?x?x?xf32, #Dense3D>
}		}

// CHECK-CODEGEN-LABEL: func @sparse_dense_2d(
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xf64>>)
//		//
// CHECK-STORAGE-LABEL: func @sparse_dense_2d(		// CHECK-LABEL: func @sparse_dense_2d(
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>,		// CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xf64>) {		// CHECK-SAME: %[[A1:.*1]]: memref<?xf64>) {
// CHECK-STORAGE: return		// CHECK: return
func.func @sparse_dense_2d(%arg0: tensor<?x?xf64, #Dense2D>) {		func.func @sparse_dense_2d(%arg0: tensor<?x?xf64, #Dense2D>) {
return		return
}		}

// CHECK-CODEGEN-LABEL: func @sparse_row(
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>>)
//		//
// CHECK-STORAGE-LABEL: func @sparse_row(		// CHECK-LABEL: func @sparse_row(
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>,		// CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xi32>,		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref<?xi64>,		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref<?xf64>) {		// CHECK-SAME: %[[A3:.*3]]: memref<?xf64>) {
// CHECK-STORAGE: return		// CHECK: return
func.func @sparse_row(%arg0: tensor<?x?xf64, #Row>) {		func.func @sparse_row(%arg0: tensor<?x?xf64, #Row>) {
return		return
}		}

// CHECK-CODEGEN-LABEL: func @sparse_csr(
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>>)
//		//
// CHECK-STORAGE-LABEL: func @sparse_csr(		// CHECK-LABEL: func @sparse_csr(
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>,		// CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xi32>,		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref<?xi64>,		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref<?xf64>) {		// CHECK-SAME: %[[A3:.*3]]: memref<?xf64>) {
// CHECK-STORAGE: return		// CHECK: return
func.func @sparse_csr(%arg0: tensor<?x?xf64, #CSR>) {		func.func @sparse_csr(%arg0: tensor<?x?xf64, #CSR>) {
return		return
}		}

// CHECK-CODEGEN-LABEL: func @sparse_dcsr(
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xf64>>)
//		//
// CHECK-STORAGE-LABEL: func @sparse_dcsr(		// CHECK-LABEL: func @sparse_dcsr(
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>,		// CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xi32>,		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref<?xi64>,		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref<?xi32>,		// CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A4:.*4]]: memref<?xi64>,		// CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A5:.*5]]: memref<?xf64>) {		// CHECK-SAME: %[[A5:.*5]]: memref<?xf64>) {
// CHECK-STORAGE: return		// CHECK: return
func.func @sparse_dcsr(%arg0: tensor<?x?xf64, #DCSR>) {		func.func @sparse_dcsr(%arg0: tensor<?x?xf64, #DCSR>) {
return		return
}		}

//		//
// Querying for dimension 1 in the tensor type can immediately		// Querying for dimension 1 in the tensor type can immediately
// fold using the original static dimension sizes.		// fold using the original static dimension sizes.
//		//
// CHECK-CODEGEN-LABEL: func @sparse_dense_3d(		//
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<3xindex>, memref<?xf64>>)		// CHECK-LABEL: func @sparse_dense_3d(
// CHECK-CODEGEN: %[[C:.*]] = arith.constant 20 : index		// CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
// CHECK-CODEGEN: return %[[C]] : index		// CHECK-SAME: %[[A1:.*1]]: memref<?xf64>)
//		// CHECK: %[[C:.*]] = arith.constant 20 : index
// CHECK-STORAGE-LABEL: func @sparse_dense_3d(		// CHECK: return %[[C]] : index
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<3xindex>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xf64>)
// CHECK-STORAGE: %[[C:.*]] = arith.constant 20 : index
// CHECK-STORAGE: return %[[C]] : index
func.func @sparse_dense_3d(%arg0: tensor<10x20x30xf64, #Dense3D>) -> index {		func.func @sparse_dense_3d(%arg0: tensor<10x20x30xf64, #Dense3D>) -> index {
%c = arith.constant 1 : index		%c = arith.constant 1 : index
%0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #Dense3D>		%0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #Dense3D>
return %0 : index		return %0 : index
}		}

//		//
// Querying for dimension 1 in the tensor type needs to be permuted		// Querying for dimension 1 in the tensor type needs to be permuted
// into querying for dimension 2 in the stored sparse tensor scheme,		// into querying for dimension 2 in the stored sparse tensor scheme,
// since the latter honors the dimOrdering.		// since the latter honors the dimOrdering.
//		//
// CHECK-CODEGEN-LABEL: func @sparse_dense_3d_dyn(		//
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<3xindex>, memref<?xf64>>)		// CHECK-LABEL: func @sparse_dense_3d_dyn(
// CHECK-CODEGEN: %[[C:.*]] = arith.constant 2 : index		// CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
// CHECK-CODEGEN: %[[F:.*]] = sparse_tensor.storage_get %[[A]][0] : tuple<memref<3xindex>, memref<?xf64>> to memref<3xindex>		// CHECK-SAME: %[[A1:.*1]]: memref<?xf64>)
// CHECK-CODEGEN: %[[L:.*]] = memref.load %[[F]][%[[C]]] : memref<3xindex>		// CHECK: %[[C:.*]] = arith.constant 2 : index
// CHECK-CODEGEN: return %[[L]] : index		// CHECK: %[[L:.*]] = memref.load %[[A0]][%[[C]]] : memref<3xindex>
//		// CHECK: return %[[L]] : index
// CHECK-STORAGE-LABEL: func @sparse_dense_3d_dyn(
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<3xindex>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xf64>)
// CHECK-STORAGE: %[[C:.*]] = arith.constant 2 : index
// CHECK-STORAGE: %[[L:.*]] = memref.load %[[A0]][%[[C]]] : memref<3xindex>
// CHECK-STORAGE: return %[[L]] : index
func.func @sparse_dense_3d_dyn(%arg0: tensor<?x?x?xf64, #Dense3D>) -> index {		func.func @sparse_dense_3d_dyn(%arg0: tensor<?x?x?xf64, #Dense3D>) -> index {
%c = arith.constant 1 : index		%c = arith.constant 1 : index
%0 = tensor.dim %arg0, %c : tensor<?x?x?xf64, #Dense3D>		%0 = tensor.dim %arg0, %c : tensor<?x?x?xf64, #Dense3D>
return %0 : index		return %0 : index
}		}

// CHECK-CODEGEN-LABEL: func @sparse_pointers_dcsr(		//
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xf64>>)		// CHECK-LABEL: func @sparse_pointers_dcsr(
// CHECK-CODEGEN: %[[F:.*]] = sparse_tensor.storage_get %[[A]][3] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<?xi32>		// CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
// CHECK-CODEGEN: return %[[F]] : memref<?xi32>		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
//		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-STORAGE-LABEL: func @sparse_pointers_dcsr(		// CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>,		// CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xi32>,		// CHECK-SAME: %[[A5:.*5]]: memref<?xf64>)
// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref<?xi64>,		// CHECK: return %[[A3]] : memref<?xi32>
// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A4:.*4]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A5:.*5]]: memref<?xf64>)
// CHECK-STORAGE: return %[[A3]] : memref<?xi32>
func.func @sparse_pointers_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi32> {		func.func @sparse_pointers_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi32> {
%c = arith.constant 1 : index		%c = arith.constant 1 : index
%0 = sparse_tensor.pointers %arg0, %c : tensor<?x?xf64, #DCSR> to memref<?xi32>		%0 = sparse_tensor.pointers %arg0, %c : tensor<?x?xf64, #DCSR> to memref<?xi32>
return %0 : memref<?xi32>		return %0 : memref<?xi32>
}		}

// CHECK-CODEGEN-LABEL: func @sparse_indices_dcsr(		//
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xf64>>)		// CHECK-LABEL: func @sparse_indices_dcsr(
// CHECK-CODEGEN: %[[F:.*]] = sparse_tensor.storage_get %[[A]][4] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<?xi64>		// CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
// CHECK-CODEGEN: return %[[F]] : memref<?xi64>		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
//		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-STORAGE-LABEL: func @sparse_indices_dcsr(		// CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>,		// CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xi32>,		// CHECK-SAME: %[[A5:.*5]]: memref<?xf64>)
// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref<?xi64>,		// CHECK: return %[[A4]] : memref<?xi64>
// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A4:.*4]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A5:.*5]]: memref<?xf64>)
// CHECK-STORAGE: return %[[A4]] : memref<?xi64>
func.func @sparse_indices_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi64> {		func.func @sparse_indices_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi64> {
%c = arith.constant 1 : index		%c = arith.constant 1 : index
%0 = sparse_tensor.indices %arg0, %c : tensor<?x?xf64, #DCSR> to memref<?xi64>		%0 = sparse_tensor.indices %arg0, %c : tensor<?x?xf64, #DCSR> to memref<?xi64>
return %0 : memref<?xi64>		return %0 : memref<?xi64>
}		}

// CHECK-CODEGEN-LABEL: func @sparse_values_dcsr(		//
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xf64>>)		// CHECK-LABEL: func @sparse_values_dcsr(
// CHECK-CODEGEN: %[[F:.*]] = sparse_tensor.storage_get %[[A]][5] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<?xf64>		// CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
// CHECK-CODEGEN: return %[[F]] : memref<?xf64>		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
//		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-STORAGE-LABEL: func @sparse_values_dcsr(		// CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>,		// CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xi32>,		// CHECK-SAME: %[[A5:.*5]]: memref<?xf64>)
// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref<?xi64>,		// CHECK: return %[[A5]] : memref<?xf64>
// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A4:.*4]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A5:.*5]]: memref<?xf64>)
// CHECK-STORAGE: return %[[A5]] : memref<?xf64>
func.func @sparse_values_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xf64> {		func.func @sparse_values_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xf64> {
%0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64>		%0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64>
return %0 : memref<?xf64>		return %0 : memref<?xf64>
}		}

// CHECK-CODEGEN-LABEL: func @sparse_dealloc_csr(		//
// CHECK-CODEGEN-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>>)		// CHECK-LABEL: func @sparse_dealloc_csr(
// CHECK-CODEGEN: %[[F0:.*]] = sparse_tensor.storage_get %[[A]][0] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<2xindex>		// CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
// CHECK-CODEGEN: memref.dealloc %[[F0]] : memref<2xindex>		// CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
// CHECK-CODEGEN: %[[F1:.*]] = sparse_tensor.storage_get %[[A]][1] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<?xi32>		// CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-CODEGEN: memref.dealloc %[[F1]] : memref<?xi32>		// CHECK-SAME: %[[A3:.*3]]: memref<?xf64>) {
// CHECK-CODEGEN: %[[F2:.*]] = sparse_tensor.storage_get %[[A]][2] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<?xi64>		// CHECK: memref.dealloc %[[A0]] : memref<2xindex>
// CHECK-CODEGEN: memref.dealloc %[[F2]] : memref<?xi64>		// CHECK: memref.dealloc %[[A1]] : memref<?xi32>
// CHECK-CODEGEN: %[[F3:.*]] = sparse_tensor.storage_get %[[A]][3] : tuple<memref<2xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>> to memref<?xf64>		// CHECK: memref.dealloc %[[A2]] : memref<?xi64>
// CHECK-CODEGEN: memref.dealloc %[[F3]] : memref<?xf64>		// CHECK: memref.dealloc %[[A3]] : memref<?xf64>
// CHECK-CODEGEN: return		// CHECK: return
//
// CHECK-STORAGE-LABEL: func @sparse_dealloc_csr(
// CHECK-STORAGE-SAME: %[[A0:.*0]]: memref<2xindex>,
// CHECK-STORAGE-SAME: %[[A1:.*1]]: memref<?xi32>,
// CHECK-STORAGE-SAME: %[[A2:.*2]]: memref<?xi64>,
// CHECK-STORAGE-SAME: %[[A3:.*3]]: memref<?xf64>) {
// CHECK-STORAGE: memref.dealloc %[[A0]] : memref<2xindex>
// CHECK-STORAGE: memref.dealloc %[[A1]] : memref<?xi32>
// CHECK-STORAGE: memref.dealloc %[[A2]] : memref<?xi64>
// CHECK-STORAGE: memref.dealloc %[[A3]] : memref<?xf64>
// CHECK-STORAGE: return
func.func @sparse_dealloc_csr(%arg0: tensor<?x?xf64, #CSR>) {		func.func @sparse_dealloc_csr(%arg0: tensor<?x?xf64, #CSR>) {
bufferization.dealloc_tensor %arg0 : tensor<?x?xf64, #CSR>		bufferization.dealloc_tensor %arg0 : tensor<?x?xf64, #CSR>
return		return
}		}

// CHECK-LABEL: func @sparse_alloc_csc(		// CHECK-LABEL: func @sparse_alloc_csc(
// CHECK-SAME: %[[A:.*]]: index) ->		// CHECK-SAME: %[[A:.*]]: index) ->
// CHECK-CODEGEN-SAME: tuple<memref<2xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>>		// CHECK-SAME: memref<2xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
		aartbikUnsubmitted Done Reply Inline Actions make sure to indent the same, probably by moving all others to the left aartbik: make sure to indent the same, probably by moving all others to the left
// CHECK-STORAGE-SAME: memref<2xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index		// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index		// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index		// CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
// CHECK: %[[T0:.*]] = memref.alloc() : memref<2xindex>		// CHECK: %[[T0:.*]] = memref.alloc() : memref<2xindex>
// CHECK: memref.store %[[A]], %[[T0]][%[[C0]]] : memref<2xindex>		// CHECK: memref.store %[[A]], %[[T0]][%[[C0]]] : memref<2xindex>
// CHECK: memref.store %[[C10]], %[[T0]][%[[C1]]] : memref<2xindex>		// CHECK: memref.store %[[C10]], %[[T0]][%[[C1]]] : memref<2xindex>
// CHECK: %[[T1:.*]] = memref.alloc() : memref<1xindex>		// CHECK: %[[T1:.*]] = memref.alloc() : memref<1xindex>
// CHECK: %[[T2:.*]] = memref.cast %[[T1]] : memref<1xindex> to memref<?xindex>		// CHECK: %[[T2:.*]] = memref.cast %[[T1]] : memref<1xindex> to memref<?xindex>
// CHECK: %[[T3:.*]] = memref.alloc() : memref<1xindex>		// CHECK: %[[T3:.*]] = memref.alloc() : memref<1xindex>
// CHECK: %[[T4:.*]] = memref.cast %[[T3]] : memref<1xindex> to memref<?xindex>		// CHECK: %[[T4:.*]] = memref.cast %[[T3]] : memref<1xindex> to memref<?xindex>
// CHECK: %[[T5:.*]] = memref.alloc() : memref<1xf64>		// CHECK: %[[T5:.*]] = memref.alloc() : memref<1xf64>
// CHECK: %[[T6:.*]] = memref.cast %[[T5]] : memref<1xf64> to memref<?xf64>		// CHECK: %[[T6:.*]] = memref.cast %[[T5]] : memref<1xf64> to memref<?xf64>
// CHECK-CODEGEN: %[[T:.*]] = sparse_tensor.storage(%[[T0]], %[[T2]], %[[T4]], %[[T6]])		// CHECK: return %[[T0]], %[[T2]], %[[T4]], %[[T6]]
// CHECK-CODEGEN: return %[[T]]
// CHECK-STORAGE: return %[[T0]], %[[T2]], %[[T4]], %[[T6]]
func.func @sparse_alloc_csc(%arg0: index) -> tensor<10x?xf64, #CSC> {		func.func @sparse_alloc_csc(%arg0: index) -> tensor<10x?xf64, #CSC> {
%0 = bufferization.alloc_tensor(%arg0) : tensor<10x?xf64, #CSC>		%0 = bufferization.alloc_tensor(%arg0) : tensor<10x?xf64, #CSC>
%1 = sparse_tensor.load %0 : tensor<10x?xf64, #CSC>		%1 = sparse_tensor.load %0 : tensor<10x?xf64, #CSC>
return %1 : tensor<10x?xf64, #CSC>		return %1 : tensor<10x?xf64, #CSC>
}		}

// CHECK-LABEL: func @sparse_alloc_3d() ->		// CHECK-LABEL: func @sparse_alloc_3d() ->
// CHECK-CODEGEN-SAME: tuple<memref<3xindex>, memref<?xf64>>		// CHECK-SAME: memref<3xindex>, memref<?xf64>
		aartbikUnsubmitted Done Reply Inline Actions same, indent by moving others left aartbik: same, indent by moving others left
// CHECK-STORAGE-SAME: memref<3xindex>, memref<?xf64>
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index		// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index		// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index		// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index		// CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
// CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index		// CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index
// CHECK-DAG: %[[C30:.*]] = arith.constant 30 : index		// CHECK-DAG: %[[C30:.*]] = arith.constant 30 : index
// CHECK: %[[A0:.*]] = memref.alloc() : memref<3xindex>		// CHECK: %[[A0:.*]] = memref.alloc() : memref<3xindex>
// CHECK: memref.store %[[C30]], %[[A0]][%[[C0]]] : memref<3xindex>		// CHECK: memref.store %[[C30]], %[[A0]][%[[C0]]] : memref<3xindex>
// CHECK: memref.store %[[C10]], %[[A0]][%[[C1]]] : memref<3xindex>		// CHECK: memref.store %[[C10]], %[[A0]][%[[C1]]] : memref<3xindex>
// CHECK: memref.store %[[C20]], %[[A0]][%[[C2]]] : memref<3xindex>		// CHECK: memref.store %[[C20]], %[[A0]][%[[C2]]] : memref<3xindex>
// CHECK: %[[A:.*]] = memref.alloc() : memref<6000xf64>		// CHECK: %[[A:.*]] = memref.alloc() : memref<6000xf64>
// CHECK: %[[A1:.*]] = memref.cast %[[A]] : memref<6000xf64> to memref<?xf64>		// CHECK: %[[A1:.*]] = memref.cast %[[A]] : memref<6000xf64> to memref<?xf64>
// CHECK-CODEGEN: %[[T:.*]] = sparse_tensor.storage(%[[A0]], %[[A1]])		// CHECK: return %[[A0]], %[[A1]] : memref<3xindex>, memref<?xf64>
// CHECK-CODEGEN: return %[[T]] : tuple<memref<3xindex>, memref<?xf64>>
// CHECK-STORAGE: return %[[A0]], %[[A1]] : memref<3xindex>, memref<?xf64>
func.func @sparse_alloc_3d() -> tensor<10x20x30xf64, #Dense3D> {		func.func @sparse_alloc_3d() -> tensor<10x20x30xf64, #Dense3D> {
%0 = bufferization.alloc_tensor() : tensor<10x20x30xf64, #Dense3D>		%0 = bufferization.alloc_tensor() : tensor<10x20x30xf64, #Dense3D>
%1 = sparse_tensor.load %0 : tensor<10x20x30xf64, #Dense3D>		%1 = sparse_tensor.load %0 : tensor<10x20x30xf64, #Dense3D>
return %1 : tensor<10x20x30xf64, #Dense3D>		return %1 : tensor<10x20x30xf64, #Dense3D>
}		}

mlir/test/Dialect/SparseTensor/invalid.mlir

Show First 20 Lines • Show All 436 Lines • ▼ Show 20 Lines	func.func @invalid_concat_size_mismatch(%arg0: tensor<2x4xf64, #DC>,
%arg2: tensor<4x4xf64, #DC>) -> tensor<9x4xf64, #DC> {		%arg2: tensor<4x4xf64, #DC>) -> tensor<9x4xf64, #DC> {
// expected-error@+1 {{All dimensions (expect for the concatenating one) should be equal}}		// expected-error@+1 {{All dimensions (expect for the concatenating one) should be equal}}
%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}		%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
: tensor<2x4xf64, #DC>,		: tensor<2x4xf64, #DC>,
tensor<3x3xf64, #DC>,		tensor<3x3xf64, #DC>,
tensor<4x4xf64, #DC> to tensor<9x4xf64, #DC>		tensor<4x4xf64, #DC> to tensor<9x4xf64, #DC>
return %0 : tensor<9x4xf64, #DC>		return %0 : tensor<9x4xf64, #DC>
}		}

// -----

func.func @sparse_storage_new(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: f64) ->
tuple<memref<?xf64>, memref<?xf64>> {
// expected-error@+1{{The number of inputs is inconsistent with output}}
%0 = sparse_tensor.storage(%arg0, %arg1, %arg2)
: memref<?xf64>, memref<?xf64>, f64 to tuple<memref<?xf64>, memref<?xf64>>
return %0 : tuple<memref<?xf64>, memref<?xf64>>
}

// -----

func.func @sparse_storage_new(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: f64) ->
tuple<memref<?xi64>, memref<?xf64>, f64> {
// expected-error@+1{{Type mismatch between}}
%0 = sparse_tensor.storage(%arg0, %arg1, %arg2)
: memref<?xf64>, memref<?xf64>, f64 to tuple<memref<?xi64>, memref<?xf64>, f64>
return %0 : tuple<memref<?xi64>, memref<?xf64>, f64>
}

// -----

func.func @sparse_storage_get(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>) -> memref<?xf64> {
// expected-error@+1{{Out-of-bound access}}
%0 = sparse_tensor.storage_get %arg0[3]
: tuple<memref<?xf64>, memref<?xf64>, f64> to
memref<?xf64>
return %0 : memref<?xf64>
}

// -----

func.func @sparse_storage_get(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>) -> memref<?xf64> {
// expected-error@+1{{Type mismatch}}
%0 = sparse_tensor.storage_get %arg0[2]
: tuple<memref<?xf64>, memref<?xf64>, f64> to
memref<?xf64>
return %0 : memref<?xf64>
}

// -----

func.func @sparse_storage_set(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>, %arg1: memref<?xf64>) -> tuple<memref<?xf64>, memref<?xf64>, f64> {
// expected-error@+1{{Out-of-bound access}}
%0 = sparse_tensor.storage_set %arg0[3], %arg1
: tuple<memref<?xf64>, memref<?xf64>, f64>, memref<?xf64> to
tuple<memref<?xf64>, memref<?xf64>, f64>
return %0 : tuple<memref<?xf64>, memref<?xf64>, f64>
}

// -----

func.func @sparse_storage_set(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>, %arg1: memref<?xf64>) -> tuple<memref<?xf64>, memref<?xf64>, f64> {
// expected-error@+1{{Type mismatch}}
%0 = sparse_tensor.storage_set %arg0[2], %arg1
: tuple<memref<?xf64>, memref<?xf64>, f64>, memref<?xf64> to
tuple<memref<?xf64>, memref<?xf64>, f64>
return %0 : tuple<memref<?xf64>, memref<?xf64>, f64>
}

mlir/test/Dialect/SparseTensor/roundtrip.mlir

Show First 20 Lines • Show All 308 Lines • ▼ Show 20 Lines	func.func @concat_sparse_sparse(%arg0: tensor<2x4xf64, #SparseMatrix>,
%arg1: tensor<3x4xf64, #SparseMatrix>,		%arg1: tensor<3x4xf64, #SparseMatrix>,
%arg2: tensor<4x4xf64, #SparseMatrix>) -> tensor<9x4xf64, #SparseMatrix> {		%arg2: tensor<4x4xf64, #SparseMatrix>) -> tensor<9x4xf64, #SparseMatrix> {
%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}		%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
: tensor<2x4xf64, #SparseMatrix>,		: tensor<2x4xf64, #SparseMatrix>,
tensor<3x4xf64, #SparseMatrix>,		tensor<3x4xf64, #SparseMatrix>,
tensor<4x4xf64, #SparseMatrix> to tensor<9x4xf64, #SparseMatrix>		tensor<4x4xf64, #SparseMatrix> to tensor<9x4xf64, #SparseMatrix>
return %0 : tensor<9x4xf64, #SparseMatrix>		return %0 : tensor<9x4xf64, #SparseMatrix>
}		}

// -----


// CHECK: func @sparse_storage_new(
// CHECK-SAME: %[[A0:.*0]]: memref<?xf64>,
// CHECK-SAME: %[[A1:.*1]]: memref<?xf64>,
// CHECK-SAME: %[[A2:.*]]: f64
// CHECK: %[[TMP_0:.*]] = sparse_tensor.storage(%[[A0]], %[[A1]], %[[A2]])
// CHECK: return %[[TMP_0]] : tuple<memref<?xf64>, memref<?xf64>, f64>
func.func @sparse_storage_new(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: f64) ->
tuple<memref<?xf64>, memref<?xf64>, f64> {
%0 = sparse_tensor.storage(%arg0, %arg1, %arg2)
: memref<?xf64>, memref<?xf64>, f64 to tuple<memref<?xf64>, memref<?xf64>, f64>
return %0 : tuple<memref<?xf64>, memref<?xf64>, f64>
}

// -----

// CHECK-LABEL: func @sparse_storage_get(
// CHECK-SAME: %[[A0:.*]]: tuple<memref<?xf64>, memref<?xf64>, f64>
// CHECK: %[[TMP0:.*]] = sparse_tensor.storage_get %[[A0]][0] :
// CHECK-SAME: tuple<memref<?xf64>, memref<?xf64>, f64>
// CHECK-SAME: to memref<?xf64>
// CHECK: return %[[TMP0]] : memref<?xf64>
func.func @sparse_storage_get(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>) -> memref<?xf64> {
%0 = sparse_tensor.storage_get %arg0[0]
: tuple<memref<?xf64>, memref<?xf64>, f64> to memref<?xf64>
return %0 : memref<?xf64>
}

// -----

// CHECK-LABEL: func @sparse_storage_set(
// CHECK-SAME: %[[A0:.*]]: tuple<memref<?xf64>, memref<?xf64>, f64>,
// CHECK-SAME: %[[A1:.*]]: memref<?xf64>
// CHECK: %[[TMP0:.*]] = sparse_tensor.storage_set %[[A0]][0], %[[A1]] :
// CHECK-SAME: tuple<memref<?xf64>, memref<?xf64>, f64>,
// CHECK-SAME: memref<?xf64>
// CHECK-SAME: to tuple<memref<?xf64>, memref<?xf64>, f64>
// CHECK: return %0 : tuple<memref<?xf64>, memref<?xf64>, f64>
func.func @sparse_storage_set(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>, %arg1: memref<?xf64>) -> tuple<memref<?xf64>, memref<?xf64>, f64> {
%0 = sparse_tensor.storage_set %arg0[0], %arg1
: tuple<memref<?xf64>, memref<?xf64>, f64>, memref<?xf64> to
tuple<memref<?xf64>, memref<?xf64>, f64>
return %0 : tuple<memref<?xf64>, memref<?xf64>, f64>
}

mlir/test/Dialect/SparseTensor/sparse_tensor_storage.mlir

This file was deleted.

	// RUN: mlir-opt %s -sparse-tensor-storage-expansion -cse \| FileCheck %s

	// CHECK-LABEL: func @sparse_storage_expand(
	// CHECK-SAME: %[[TMP_arg0:.*0]]: memref<?xf64>,
	// CHECK-SAME: %[[TMP_arg1:.*1]]: memref<?xf64>,
	// CHECK-SAME: %[[TMP_arg2:.*]]: f64
	// CHECK return %[[TMP_arg0]], %[[TMP_arg1]], %[[TMP_arg2]]
	func.func @sparse_storage_expand(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>)
	-> tuple<memref<?xf64>, memref<?xf64>, f64> {
	return %arg0 : tuple<memref<?xf64>, memref<?xf64>, f64>
	}

	// CHECK-LABEL: func @call_sparse_storage_expand(
	// CHECK-SAME: %[[TMP_arg0:.*0]]: memref<?xf64>,
	// CHECK-SAME: %[[TMP_arg1:.*1]]: memref<?xf64>,
	// CHECK-SAME: %[[TMP_arg2:.*]]: f64)
	// CHECK: %[[TMP_0:.*]]:3 = call @sparse_storage_expand(%[[TMP_arg0]], %[[TMP_arg1]], %[[TMP_arg2]])
	// CHECK: return %[[TMP_0]]#0, %[[TMP_0]]#1, %[[TMP_0]]#2 : memref<?xf64>, memref<?xf64>, f64
	func.func @call_sparse_storage_expand(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>)
	-> tuple<memref<?xf64>, memref<?xf64>, f64> {
	%1 = call @sparse_storage_expand(%arg0) : (tuple<memref<?xf64>, memref<?xf64>, f64>) ->
	tuple<memref<?xf64>, memref<?xf64>, f64>
	return %1 : tuple<memref<?xf64>, memref<?xf64>, f64>
	}

	// CHECK-LABEL: func @sparse_storage(
	// CHECK-SAME: %[[TMP_arg0:.*0]]: memref<?xf64>,
	// CHECK-SAME: %[[TMP_arg1:.*1]]: memref<?xf64>,
	// CHECK-SAME: %[[TMP_arg2:.*2]]: memref<?xf64>)
	// CHECK: return %[[TMP_arg0]], %[[TMP_arg1]], %[[TMP_arg2]]
	func.func @sparse_storage(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: memref<?xf64>)
	-> tuple<memref<?xf64>, memref<?xf64>, memref<?xf64>> {
	%1 = sparse_tensor.storage(%arg0, %arg1, %arg2) : memref<?xf64>, memref<?xf64>, memref<?xf64> to tuple<memref<?xf64>, memref<?xf64>, memref<?xf64>>
	return %1 : tuple<memref<?xf64>, memref<?xf64>, memref<?xf64>>
	}

	// CHECK-LABEL: func @sparse_storage_get(
	// CHECK-SAME: %[[TMP_arg0:.*0]]: memref<?xf64>,
	// CHECK-SAME: %[[TMP_arg1:.*1]]: memref<?xf64>,
	// CHECK-SAME: %[[TMP_arg2:.*]]: f64)
	// CHECK: return %[[TMP_arg0]] : memref<?xf64>
	func.func @sparse_storage_get(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>) -> memref<?xf64> {
	%0 = sparse_tensor.storage_get %arg0[0]
	: tuple<memref<?xf64>, memref<?xf64>, f64> to memref<?xf64>
	return %0 : memref<?xf64>
	}

	// CHECK-LABEL: func @sparse_storage_set(
	// CHECK-SAME: %[[TMP_arg0:.*0]]: memref<?xf64>,
	// CHECK-SAME: %[[TMP_arg1:.*1]]: memref<?xf64>,
	// CHECK-SAME: %[[TMP_arg2:.*]]: f64,
	// CHECK-SAME: %[[TMP_arg3:.*]]: memref<?xf64>)
	// CHECK: return %[[TMP_arg3]], %[[TMP_arg1]], %[[TMP_arg2]] : memref<?xf64>, memref<?xf64>, f64
	func.func @sparse_storage_set(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>,
	%arg1: memref<?xf64>) -> tuple<memref<?xf64>, memref<?xf64>, f64> {
	%0 = sparse_tensor.storage_set %arg0[0], %arg1
	: tuple<memref<?xf64>, memref<?xf64>, f64>, memref<?xf64> to
	tuple<memref<?xf64>, memref<?xf64>, f64>
	return %0 : tuple<memref<?xf64>, memref<?xf64>, f64>
	}

This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] Refactoring: remove dependence on tuple type when lowering sparse tensors.
ClosedPublic

Details

Diff Detail

Event Timeline

Revision Contents

Diff 458518

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

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

mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageExpansion.cpp

mlir/test/Dialect/SparseTensor/codegen.mlir

mlir/test/Dialect/SparseTensor/invalid.mlir

mlir/test/Dialect/SparseTensor/roundtrip.mlir

mlir/test/Dialect/SparseTensor/sparse_tensor_storage.mlir

This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] Refactoring: remove dependence on tuple type when lowering sparse tensors.ClosedPublic

Details

Diff Detail

Event Timeline

Revision Contents

Diff 458518

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

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

mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageExpansion.cpp

mlir/test/Dialect/SparseTensor/codegen.mlir

mlir/test/Dialect/SparseTensor/invalid.mlir

mlir/test/Dialect/SparseTensor/roundtrip.mlir

mlir/test/Dialect/SparseTensor/sparse_tensor_storage.mlir

[mlir][sparse] Refactoring: remove dependence on tuple type when lowering sparse tensors.
ClosedPublic