This is an archive of the discontinued LLVM Phabricator instance.

Differential D138823

[mlir][sparse] Improve concatenate operator rewrite for annotated all dense dimensions results.
ClosedPublic

Authored by bixia on Nov 28 2022, 8:13 AM.

Details

Reviewers
aartbik
nicolasvasilache
Commits
rG101a0c84f7ca: [mlir][sparse] Improve concatenate operator rewrite for annotated all dense…

Diff Detail

Event Timeline

bixia created this revision.Nov 28 2022, 8:13 AM
Herald added a reviewer: aartbik. · View Herald TranscriptNov 28 2022, 8:13 AM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: jsetoain, Moerafaat, anlunx and 22 others. · View Herald Transcript
bixia requested review of this revision.Nov 28 2022, 8:13 AM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptNov 28 2022, 8:13 AM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
Harbormaster completed remote builds in B199769: Diff 478249.Nov 28 2022, 8:56 AM
aartbik accepted this revision.Nov 28 2022, 5:27 PM
aartbik added inline comments.
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
502

maybe make a note in the comment on the all dense case

mlir/test/Dialect/SparseTensor/sparse_concat_codegen.mlir
5

maybe in a next revision, can we also have an example that reorders (like CSC)?

This revision is now accepted and ready to land.Nov 28 2022, 5:27 PM
bixia updated this revision to Diff 478626.Nov 29 2022, 9:20 AM
bixia marked 2 inline comments as done.

Add a test with permutation. Add code comment.

bixia added inline comments.Nov 29 2022, 9:28 AM
mlir/test/Dialect/SparseTensor/sparse_concat_codegen.mlir
5

Add a test for reordering.

Harbormaster completed remote builds in B200054: Diff 478626.Nov 29 2022, 10:35 AM
Closed by commit rG101a0c84f7ca: [mlir][sparse] Improve concatenate operator rewrite for annotated all dense… (authored by bixia). · Explain WhyNov 29 2022, 12:26 PM
This revision was automatically updated to reflect the committed changes.
bixia added a commit: rG101a0c84f7ca: [mlir][sparse] Improve concatenate operator rewrite for annotated all dense….

Revision Contents

PathSize
mlir/
lib/
Dialect/
SparseTensor/
Transforms/
29 lines
test/
Dialect/
SparseTensor/
178 lines
Integration/
Dialect/
SparseTensor/
CPU/
33 lines

Diff 478678

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

Show First 20 Lines • Show All 484 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(ConcatenateOp op,
auto dstTp = op.getType().cast<RankedTensorType>(); auto dstTp = op.getType().cast<RankedTensorType>();
uint64_t conDim = op.getDimension().getZExtValue(); uint64_t conDim = op.getDimension().getZExtValue();
SmallVector<Value> sizes; SmallVector<Value> sizes;
concatSizesFromInputs(rewriter, sizes, loc, dstTp, op.getInputs(), conDim); concatSizesFromInputs(rewriter, sizes, loc, dstTp, op.getInputs(), conDim);
// %t = concatenate %s1, %s2, %s3 {dim = 1} // %t = concatenate %s1, %s2, %s3 {dim = 1}
// ==> // ==>
// if (isSparseDst) // if (isSparseDst)
// if (allDense)
// %tmp = bufferization.alloc_tensor dstTp
// else
// %tmp = bufferization.alloc_tensor : unordered COO // %tmp = bufferization.alloc_tensor : unordered COO
// else // else
// %tmp = memref.alloc : dense tensor // %tmp = memref.alloc : dense tensor
// foreach in %s1 : insert d0, d1, %tmp // foreach in %s1 : insert d0, d1, %tmp
// foreach in %s2 : insert d0, d1 + size(s1), %tmp // foreach in %s2 : insert d0, d1 + size(s1), %tmp
// foreach in %s3 : insert d0, d1 + size(s1) + size(s2), %tmp // foreach in %s3 : insert d0, d1 + size(s1) + size(s2), %tmp
// %t = convert_to_dest_tensor(%tmp) // %t = convert_to_dest_tensor(%tmp)
aartbikUnsubmitted
Done
ReplyInline Actions

maybe make a note in the comment on the all dense case

aartbik: maybe make a note in the comment on the all dense case
SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp); SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
Value dst; // Destination tensor for inserting source tensor values. Value dst; // Destination tensor for inserting source tensor values.
bool allDense = false;
if (encDst) { if (encDst) {
allDense = llvm::all_of(encDst.getDimLevelType(),
[](DimLevelType dlt) { return isDenseDLT(dlt); });
SmallVector<Value> dynSizes; SmallVector<Value> dynSizes;
getDynamicSizes(dstTp, sizes, dynSizes); getDynamicSizes(dstTp, sizes, dynSizes);
RankedTensorType cooTp = getUnorderedCOOFromType(dstTp); RankedTensorType tp = dstTp;
dst = rewriter.create<AllocTensorOp>(loc, cooTp, dynSizes).getResult(); if (!allDense) {
tp = getUnorderedCOOFromType(dstTp);
encDst = getSparseTensorEncoding(tp);
}
dst = rewriter.create<AllocTensorOp>(loc, tp, dynSizes).getResult();
} else { } else {
// TODO: Dense buffers should be allocated/deallocated via the callback // TODO: Dense buffers should be allocated/deallocated via the callback
// in BufferizationOptions. // in BufferizationOptions.
dst = allocDenseTensor(rewriter, loc, dstTp, sizes); dst = allocDenseTensor(rewriter, loc, dstTp, sizes);
} }
int64_t rank = dstTp.getRank(); int64_t rank = dstTp.getRank();
Value offset = constantIndex(rewriter, loc, 0); Value offset = constantIndex(rewriter, loc, 0);
SmallVector<Value> initArgs; SmallVector<Value> initArgs;
if (encDst) if (encDst)
initArgs.push_back(dst); initArgs.push_back(dst);
ForeachOp foreachOp; ForeachOp foreachOp;
for (Value input : op.getInputs()) { for (Value input : op.getInputs()) {
// Build a for op for each input tensor to append new values into the // Build a for op for each input tensor to append new values into the
// output tensor. // output tensor.
foreachOp = rewriter.create<ForeachOp>( foreachOp = rewriter.create<ForeachOp>(
loc, input, initArgs, loc, input, initArgs,
[&](OpBuilder &builder, Location loc, ValueRange args, Value v, [&](OpBuilder &builder, Location loc, ValueRange args, Value v,
ValueRange reduc) { ValueRange reduc) {
SmallVector<Value> indices; SmallVector<Value> indices(rank, Value());
for (int64_t i = 0; i < rank; i++) { for (int64_t i = 0; i < rank; i++) {
Value idx = args[i]; Value idx = args[i];
if (i == static_cast<int64_t>(conDim)) if (i == static_cast<int64_t>(conDim))
// Transform coordinates for the concatenating dim. // Transform coordinates for the concatenating dim.
idx = builder.create<arith::AddIOp>(loc, idx, offset); idx = builder.create<arith::AddIOp>(loc, idx, offset);
indices.push_back(idx); indices[toStoredDim(encDst, i)] = idx;
} }
if (encDst) { if (encDst) {
Value cond = genIsNonzero(rewriter, loc, v); Value cond = genIsNonzero(rewriter, loc, v);
scf::IfOp ifOp = builder.create<scf::IfOp>( scf::IfOp ifOp = builder.create<scf::IfOp>(
loc, TypeRange(reduc.front().getType()), cond, /*else*/ true); loc, TypeRange(reduc.front().getType()), cond, /*else*/ true);
builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
Value t = Value t =
builder.create<InsertOp>(loc, v, reduc.front(), indices); builder.create<InsertOp>(loc, v, reduc.front(), indices);
Show All 17 Lines for (Value input : op.getInputs()) {
if (encDst) { if (encDst) {
dst = foreachOp.getResult(0); dst = foreachOp.getResult(0);
initArgs[0] = dst; initArgs[0] = dst;
} }
} }
if (encDst) { if (encDst) {
dst = rewriter.create<LoadOp>(loc, dst, true); dst = rewriter.create<LoadOp>(loc, dst, true);
Value converted = rewriter.create<ConvertOp>(loc, dstTp, dst).getResult(); if (!allDense) {
rewriter.create<DeallocTensorOp>(loc, dst); Value tmpCoo = dst;
rewriter.replaceOp(op, converted); dst = rewriter.create<ConvertOp>(loc, dstTp, tmpCoo).getResult();
rewriter.create<DeallocTensorOp>(loc, tmpCoo);
}
rewriter.replaceOp(op, dst);
} else { } else {
rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstTp, dst); rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstTp, dst);
} }
return success(); return success();
} }
}; };
/// Sparse rewriting rule for the convert operator. /// Sparse rewriting rule for the convert operator.
▲ Show 20 Lines • Show All 532 Lines • Show Last 20 Lines

mlir/test/Dialect/SparseTensor/sparse_concat_codegen.mlir

// RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false enable-convert=false" \ // RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false enable-convert=false" \
// RUN: | FileCheck %s // RUN: | FileCheck %s
#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}> #DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
#DENSE = #sparse_tensor.encoding<{dimLevelType = ["dense", "dense"]}>
aartbikUnsubmitted
Done
ReplyInline Actions

maybe in a next revision, can we also have an example that reorders (like CSC)?

aartbik: maybe in a next revision, can we also have an example that reorders (like CSC)?
bixiaAuthorUnsubmitted
Done
ReplyInline Actions

Add a test for reordering.

bixia: Add a test for reordering.
#DENSE_P = #sparse_tensor.encoding<{
dimLevelType = ["dense", "dense"],
dimOrdering = affine_map<(i,j) -> (j,i)>
}>
// CHECK-LABEL: @concat_sparse_sparse( // CHECK-LABEL: @concat_sparse_sparse(
// CHECK-SAME: %[[TMP_arg0:.*]]: tensor<2x4xf64, #sparse_tensor // CHECK-SAME: %[[TMP_arg0:.*]]: tensor<2x4xf64, #sparse_tensor
// CHECK-SAME: %[[TMP_arg1:.*]]: tensor<3x4xf64, #sparse_tensor // CHECK-SAME: %[[TMP_arg1:.*]]: tensor<3x4xf64, #sparse_tensor
// CHECK-SAME: %[[TMP_arg2:.*]]: tensor<4x4xf64, #sparse_tensor // CHECK-SAME: %[[TMP_arg2:.*]]: tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_c0:.*]] = arith.constant 0 : index // CHECK: %[[TMP_c0:.*]] = arith.constant 0 : index
// CHECK: %[[TMP_c1:.*]] = arith.constant 1 : index // CHECK: %[[TMP_c1:.*]] = arith.constant 1 : index
// CHECK: %[[TMP_c5:.*]] = arith.constant 5 : index // CHECK: %[[TMP_c5:.*]] = arith.constant 5 : index
// CHECK: %[[TMP_c2:.*]] = arith.constant 2 : index // CHECK: %[[TMP_c2:.*]] = arith.constant 2 : index
▲ Show 20 Lines • Show All 237 Lines • ▼ Show 20 Linesfunc.func @concat_sparse_sparse_dense(%arg0: tensor<2x4xf64, #DCSR>,
%arg2: tensor<4x4xf64, #DCSR>) %arg2: tensor<4x4xf64, #DCSR>)
-> tensor<?x?xf64> { -> tensor<?x?xf64> {
%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index} %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
: tensor<2x4xf64, #DCSR>, : tensor<2x4xf64, #DCSR>,
tensor<3x4xf64, #DCSR>, tensor<3x4xf64, #DCSR>,
tensor<4x4xf64, #DCSR> to tensor<?x?xf64> tensor<4x4xf64, #DCSR> to tensor<?x?xf64>
return %0 : tensor<?x?xf64> return %0 : tensor<?x?xf64>
} }
// CHECK-LABEL: @concat_sparse_sparse_annotated_dense(
// CHECK-SAME: %[[TMP_arg0:.*]]: tensor<2x4xf64, #sparse_tensor
// CHECK-SAME: %[[TMP_arg1:.*]]: tensor<3x4xf64, #sparse_tensor
// CHECK-SAME: %[[TMP_arg2:.*]]: tensor<4x4xf64, #sparse_tensor
// CHECK-DAG: %[[TMP_c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[TMP_c1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[TMP_c5:.*]] = arith.constant 5 : index
// CHECK-DAG: %[[TMP_c2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[TMP_c9:.*]] = arith.constant 9 : index
// CHECK-DAG: %[[TMP_c4:.*]] = arith.constant 4 : index
// CHECK: %[[TMP_0:.*]] = bufferization.alloc_tensor(%[[TMP_c9]], %[[TMP_c4]]) : tensor<?x?xf64, #sparse_tensor
// CHECK: %[[TMP_1:.*]] = sparse_tensor.pointers %[[TMP_arg0]] {dimension = 0 : index} : tensor<2x4xf64, #sparse_tensor
// CHECK: %[[TMP_2:.*]] = sparse_tensor.indices %[[TMP_arg0]] {dimension = 0 : index} : tensor<2x4xf64, #sparse_tensor
// CHECK: %[[TMP_3:.*]] = sparse_tensor.pointers %[[TMP_arg0]] {dimension = 1 : index} : tensor<2x4xf64, #sparse_tensor
// CHECK: %[[TMP_4:.*]] = sparse_tensor.indices %[[TMP_arg0]] {dimension = 1 : index} : tensor<2x4xf64, #sparse_tensor
// CHECK: %[[TMP_5:.*]] = sparse_tensor.values %[[TMP_arg0]] : tensor<2x4xf64, #sparse_tensor
// CHECK: %[[TMP_6:.*]] = memref.load %[[TMP_1]][%[[TMP_c0]]] : memref<?xindex>
// CHECK: %[[TMP_7:.*]] = memref.load %[[TMP_1]][%[[TMP_c1]]] : memref<?xindex>
// CHECK: %[[RET_1:.*]] = scf.for %[[TMP_arg3:.*]] = %[[TMP_6]] to %[[TMP_7]] step %[[TMP_c1]] iter_args(%[[A0:.*]] = %[[TMP_0]])
// CHECK: %[[TMP_23:.*]] = memref.load %[[TMP_2]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK-DAG: %[[TMP_25:.*]] = memref.load %[[TMP_3]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK-DAG: %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
// CHECK: %[[TMP_26:.*]] = memref.load %[[TMP_3]][%[[TMP_24]]] : memref<?xindex>
// CHECK: %[[RET_4:.*]] = scf.for %[[TMP_arg4:.*]] = %[[TMP_25]] to %[[TMP_26]] step %[[TMP_c1]] iter_args(%[[A1:.*]] = %[[A0]])
// CHECK: %[[TMP_27:.*]] = memref.load %[[TMP_4]][%[[TMP_arg4]]] : memref<?xindex>
// CHECK: %[[TMP_28:.*]] = memref.load %[[TMP_5]][%[[TMP_arg4]]] : memref<?xf64>
// CHECK: %[[NEW_1:.*]] = sparse_tensor.insert %[[TMP_28]] into %[[A1]][%[[TMP_23]], %[[TMP_27]]] : tensor<?x?xf64, #sparse_tensor
// CHECK: scf.yield %[[NEW_1]]
// CHECK: }
// CHECK: scf.yield %[[RET_4]]
// CHECK: }
// CHECK: %[[TMP_8:.*]] = sparse_tensor.pointers %[[TMP_arg1]] {dimension = 0 : index} : tensor<3x4xf64, #sparse_tensor
// CHECK: %[[TMP_9:.*]] = sparse_tensor.indices %[[TMP_arg1]] {dimension = 0 : index} : tensor<3x4xf64, #sparse_tensor
// CHECK: %[[TMP_10:.*]] = sparse_tensor.pointers %[[TMP_arg1]] {dimension = 1 : index} : tensor<3x4xf64, #sparse_tensor
// CHECK: %[[TMP_11:.*]] = sparse_tensor.indices %[[TMP_arg1]] {dimension = 1 : index} : tensor<3x4xf64, #sparse_tensor
// CHECK: %[[TMP_12:.*]] = sparse_tensor.values %[[TMP_arg1]] : tensor<3x4xf64, #sparse_tensor
// CHECK: %[[TMP_13:.*]] = memref.load %[[TMP_8]][%[[TMP_c0]]] : memref<?xindex>
// CHECK: %[[TMP_14:.*]] = memref.load %[[TMP_8]][%[[TMP_c1]]] : memref<?xindex>
// CHECK: %[[RET_2:.*]] = scf.for %[[TMP_arg3:.*]] = %[[TMP_13]] to %[[TMP_14]] step %[[TMP_c1]] iter_args(%[[A2:.*]] = %[[RET_1]])
// CHECK: %[[TMP_23:.*]] = memref.load %[[TMP_9]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK-DAG: %[[TMP_25:.*]] = memref.load %[[TMP_10]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK-DAG: %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
// CHECK: %[[TMP_26:.*]] = memref.load %[[TMP_10]][%[[TMP_24]]] : memref<?xindex>
// CHECK: %[[RET_5:.*]] = scf.for %[[TMP_arg4:.*]] = %[[TMP_25]] to %[[TMP_26]] step %[[TMP_c1]] iter_args(%[[A3:.*]] = %[[A2]])
// CHECK: %[[TMP_27:.*]] = memref.load %[[TMP_11]][%[[TMP_arg4]]] : memref<?xindex>
// CHECK: %[[TMP_28:.*]] = memref.load %[[TMP_12]][%[[TMP_arg4]]] : memref<?xf64>
// CHECK: %[[TMP_29:.*]] = arith.addi %[[TMP_23]], %[[TMP_c2]] : index
// CHECK: %[[NEW_2:.*]] = sparse_tensor.insert %[[TMP_28]] into %[[A3]][%[[TMP_29]], %[[TMP_27]]] : tensor<?x?xf64, #sparse_tensor
// CHECK: scf.yield %[[NEW_2]]
// CHECK: }
// CHECK: scf.yield %[[RET_5]]
// CHECK: }
// CHECK: %[[TMP_15:.*]] = sparse_tensor.pointers %[[TMP_arg2]] {dimension = 0 : index} : tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_16:.*]] = sparse_tensor.indices %[[TMP_arg2]] {dimension = 0 : index} : tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_17:.*]] = sparse_tensor.pointers %[[TMP_arg2]] {dimension = 1 : index} : tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_18:.*]] = sparse_tensor.indices %[[TMP_arg2]] {dimension = 1 : index} : tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_19:.*]] = sparse_tensor.values %[[TMP_arg2]] : tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_20:.*]] = memref.load %[[TMP_15]][%[[TMP_c0]]] : memref<?xindex>
// CHECK: %[[TMP_21:.*]] = memref.load %[[TMP_15]][%[[TMP_c1]]] : memref<?xindex>
// CHECK: %[[RET_3:.*]] = scf.for %[[TMP_arg3:.*]] = %[[TMP_20]] to %[[TMP_21]] step %[[TMP_c1]] iter_args(%[[A4:.*]] = %[[RET_2]])
// CHECK: %[[TMP_23:.*]] = memref.load %[[TMP_16]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK: %[[TMP_25:.*]] = memref.load %[[TMP_17]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK: %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
// CHECK: %[[TMP_26:.*]] = memref.load %[[TMP_17]][%[[TMP_24]]] : memref<?xindex>
// CHECK: %[[RET_6:.*]] = scf.for %[[TMP_arg4:.*]] = %[[TMP_25]] to %[[TMP_26]] step %[[TMP_c1]] iter_args(%[[A5:.*]] = %[[A4]])
// CHECK: %[[TMP_27:.*]] = memref.load %[[TMP_18]][%[[TMP_arg4]]] : memref<?xindex>
// CHECK: %[[TMP_28:.*]] = memref.load %[[TMP_19]][%[[TMP_arg4]]] : memref<?xf64>
// CHECK: %[[TMP_29:.*]] = arith.addi %[[TMP_23]], %[[TMP_c5]] : index
// CHECK: %[[NEW_3:.*]] = sparse_tensor.insert %[[TMP_28]] into %[[A5]][%[[TMP_29]], %[[TMP_27]]] : tensor<?x?xf64, #sparse_tensor
// CHECK: scf.yield %[[NEW_3]]
// CHECK: }
// CHECK: scf.yield %[[RET_6]]
// CHECK: }
// CHECK: %[[R:.*]] = sparse_tensor.load %[[RET_3:.*]] hasInserts : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>>
// CHECK: return %[[R]] : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>>
func.func @concat_sparse_sparse_annotated_dense(%arg0: tensor<2x4xf64, #DCSR>,
%arg1: tensor<3x4xf64, #DCSR>,
%arg2: tensor<4x4xf64, #DCSR>)
-> tensor<?x?xf64, #DENSE> {
%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
: tensor<2x4xf64, #DCSR>,
tensor<3x4xf64, #DCSR>,
tensor<4x4xf64, #DCSR> to tensor<?x?xf64, #DENSE>
return %0 : tensor<?x?xf64, #DENSE>
}
// CHECK-LABEL: @concat_sparse_sparse_annotated_dense_permute(
// CHECK-SAME: %[[TMP_arg0:.*]]: tensor<2x4xf64, #sparse_tensor
// CHECK-SAME: %[[TMP_arg1:.*]]: tensor<3x4xf64, #sparse_tensor
// CHECK-SAME: %[[TMP_arg2:.*]]: tensor<4x4xf64, #sparse_tensor
// CHECK-DAG: %[[TMP_c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[TMP_c1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[TMP_c5:.*]] = arith.constant 5 : index
// CHECK-DAG: %[[TMP_c2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[TMP_c9:.*]] = arith.constant 9 : index
// CHECK-DAG: %[[TMP_c4:.*]] = arith.constant 4 : index
// CHECK: %[[TMP_0:.*]] = bufferization.alloc_tensor(%[[TMP_c9]], %[[TMP_c4]]) : tensor<?x?xf64, #sparse_tensor
// CHECK: %[[TMP_1:.*]] = sparse_tensor.pointers %[[TMP_arg0]] {dimension = 0 : index} : tensor<2x4xf64, #sparse_tensor
// CHECK: %[[TMP_2:.*]] = sparse_tensor.indices %[[TMP_arg0]] {dimension = 0 : index} : tensor<2x4xf64, #sparse_tensor
// CHECK: %[[TMP_3:.*]] = sparse_tensor.pointers %[[TMP_arg0]] {dimension = 1 : index} : tensor<2x4xf64, #sparse_tensor
// CHECK: %[[TMP_4:.*]] = sparse_tensor.indices %[[TMP_arg0]] {dimension = 1 : index} : tensor<2x4xf64, #sparse_tensor
// CHECK: %[[TMP_5:.*]] = sparse_tensor.values %[[TMP_arg0]] : tensor<2x4xf64, #sparse_tensor
// CHECK: %[[TMP_6:.*]] = memref.load %[[TMP_1]][%[[TMP_c0]]] : memref<?xindex>
// CHECK: %[[TMP_7:.*]] = memref.load %[[TMP_1]][%[[TMP_c1]]] : memref<?xindex>
// CHECK: %[[RET_1:.*]] = scf.for %[[TMP_arg3:.*]] = %[[TMP_6]] to %[[TMP_7]] step %[[TMP_c1]] iter_args(%[[A0:.*]] = %[[TMP_0]])
// CHECK: %[[TMP_23:.*]] = memref.load %[[TMP_2]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK-DAG: %[[TMP_25:.*]] = memref.load %[[TMP_3]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK-DAG: %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
// CHECK: %[[TMP_26:.*]] = memref.load %[[TMP_3]][%[[TMP_24]]] : memref<?xindex>
// CHECK: %[[RET_4:.*]] = scf.for %[[TMP_arg4:.*]] = %[[TMP_25]] to %[[TMP_26]] step %[[TMP_c1]] iter_args(%[[A1:.*]] = %[[A0]])
// CHECK: %[[TMP_27:.*]] = memref.load %[[TMP_4]][%[[TMP_arg4]]] : memref<?xindex>
// CHECK: %[[TMP_28:.*]] = memref.load %[[TMP_5]][%[[TMP_arg4]]] : memref<?xf64>
// CHECK: %[[NEW_1:.*]] = sparse_tensor.insert %[[TMP_28]] into %[[A1]][%[[TMP_27]], %[[TMP_23]]] : tensor<?x?xf64, #sparse_tensor
// CHECK: scf.yield %[[NEW_1]]
// CHECK: }
// CHECK: scf.yield %[[RET_4]]
// CHECK: }
// CHECK: %[[TMP_8:.*]] = sparse_tensor.pointers %[[TMP_arg1]] {dimension = 0 : index} : tensor<3x4xf64, #sparse_tensor
// CHECK: %[[TMP_9:.*]] = sparse_tensor.indices %[[TMP_arg1]] {dimension = 0 : index} : tensor<3x4xf64, #sparse_tensor
// CHECK: %[[TMP_10:.*]] = sparse_tensor.pointers %[[TMP_arg1]] {dimension = 1 : index} : tensor<3x4xf64, #sparse_tensor
// CHECK: %[[TMP_11:.*]] = sparse_tensor.indices %[[TMP_arg1]] {dimension = 1 : index} : tensor<3x4xf64, #sparse_tensor
// CHECK: %[[TMP_12:.*]] = sparse_tensor.values %[[TMP_arg1]] : tensor<3x4xf64, #sparse_tensor
// CHECK: %[[TMP_13:.*]] = memref.load %[[TMP_8]][%[[TMP_c0]]] : memref<?xindex>
// CHECK: %[[TMP_14:.*]] = memref.load %[[TMP_8]][%[[TMP_c1]]] : memref<?xindex>
// CHECK: %[[RET_2:.*]] = scf.for %[[TMP_arg3:.*]] = %[[TMP_13]] to %[[TMP_14]] step %[[TMP_c1]] iter_args(%[[A2:.*]] = %[[RET_1]])
// CHECK: %[[TMP_23:.*]] = memref.load %[[TMP_9]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK-DAG: %[[TMP_25:.*]] = memref.load %[[TMP_10]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK-DAG: %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
// CHECK: %[[TMP_26:.*]] = memref.load %[[TMP_10]][%[[TMP_24]]] : memref<?xindex>
// CHECK: %[[RET_5:.*]] = scf.for %[[TMP_arg4:.*]] = %[[TMP_25]] to %[[TMP_26]] step %[[TMP_c1]] iter_args(%[[A3:.*]] = %[[A2]])
// CHECK: %[[TMP_27:.*]] = memref.load %[[TMP_11]][%[[TMP_arg4]]] : memref<?xindex>
// CHECK: %[[TMP_28:.*]] = memref.load %[[TMP_12]][%[[TMP_arg4]]] : memref<?xf64>
// CHECK: %[[TMP_29:.*]] = arith.addi %[[TMP_23]], %[[TMP_c2]] : index
// CHECK: %[[NEW_2:.*]] = sparse_tensor.insert %[[TMP_28]] into %[[A3]][%[[TMP_27]], %[[TMP_29]]] : tensor<?x?xf64, #sparse_tensor
// CHECK: scf.yield %[[NEW_2]]
// CHECK: }
// CHECK: scf.yield %[[RET_5]]
// CHECK: }
// CHECK: %[[TMP_15:.*]] = sparse_tensor.pointers %[[TMP_arg2]] {dimension = 0 : index} : tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_16:.*]] = sparse_tensor.indices %[[TMP_arg2]] {dimension = 0 : index} : tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_17:.*]] = sparse_tensor.pointers %[[TMP_arg2]] {dimension = 1 : index} : tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_18:.*]] = sparse_tensor.indices %[[TMP_arg2]] {dimension = 1 : index} : tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_19:.*]] = sparse_tensor.values %[[TMP_arg2]] : tensor<4x4xf64, #sparse_tensor
// CHECK: %[[TMP_20:.*]] = memref.load %[[TMP_15]][%[[TMP_c0]]] : memref<?xindex>
// CHECK: %[[TMP_21:.*]] = memref.load %[[TMP_15]][%[[TMP_c1]]] : memref<?xindex>
// CHECK: %[[RET_3:.*]] = scf.for %[[TMP_arg3:.*]] = %[[TMP_20]] to %[[TMP_21]] step %[[TMP_c1]] iter_args(%[[A4:.*]] = %[[RET_2]])
// CHECK: %[[TMP_23:.*]] = memref.load %[[TMP_16]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK: %[[TMP_25:.*]] = memref.load %[[TMP_17]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK: %[[TMP_24:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
// CHECK: %[[TMP_26:.*]] = memref.load %[[TMP_17]][%[[TMP_24]]] : memref<?xindex>
// CHECK: %[[RET_6:.*]] = scf.for %[[TMP_arg4:.*]] = %[[TMP_25]] to %[[TMP_26]] step %[[TMP_c1]] iter_args(%[[A5:.*]] = %[[A4]])
// CHECK: %[[TMP_27:.*]] = memref.load %[[TMP_18]][%[[TMP_arg4]]] : memref<?xindex>
// CHECK: %[[TMP_28:.*]] = memref.load %[[TMP_19]][%[[TMP_arg4]]] : memref<?xf64>
// CHECK: %[[TMP_29:.*]] = arith.addi %[[TMP_23]], %[[TMP_c5]] : index
// CHECK: %[[NEW_3:.*]] = sparse_tensor.insert %[[TMP_28]] into %[[A5]][%[[TMP_27]], %[[TMP_29]]] : tensor<?x?xf64, #sparse_tensor
// CHECK: scf.yield %[[NEW_3]]
// CHECK: }
// CHECK: scf.yield %[[RET_6]]
// CHECK: }
// CHECK: %[[R:.*]] = sparse_tensor.load %[[RET_3:.*]] hasInserts : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>>
// CHECK: return %[[R]] : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>>
func.func @concat_sparse_sparse_annotated_dense_permute(%arg0: tensor<2x4xf64, #DCSR>,
%arg1: tensor<3x4xf64, #DCSR>,
%arg2: tensor<4x4xf64, #DCSR>)
-> tensor<?x?xf64, #DENSE_P> {
%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
: tensor<2x4xf64, #DCSR>,
tensor<3x4xf64, #DCSR>,
tensor<4x4xf64, #DCSR> to tensor<?x?xf64, #DENSE_P>
return %0 : tensor<?x?xf64, #DENSE_P>
}
No newline at end of file

mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate.mlir

// DEFINE: %{option} = enable-runtime-library=true // DEFINE: %{option} = enable-runtime-library=true
// DEFINE: %{command} = mlir-opt %s --sparse-compiler=%{option} | \ // DEFINE: %{command} = mlir-opt %s --sparse-compiler=%{option} | \
// DEFINE: mlir-cpu-runner \ // DEFINE: mlir-cpu-runner \
// DEFINE: -e entry -entry-point-result=void \ // DEFINE: -e entry -entry-point-result=void \
// DEFINE: -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \ // DEFINE: -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
// DEFINE: FileCheck %s // DEFINE: FileCheck %s
// //
// RUN: %{command} // RUN: %{command}
// //
// Do the same run, but now with direct IR generation. // Do the same run, but now with direct IR generation.
// REDEFINE: %{option} = enable-runtime-library=false // REDEFINE: %{option} = enable-runtime-library=false
// RUN: %{command} // RUN: %{command}
#MAT_C_C = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}> #MAT_C_C = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
#MAT_D_C = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}> #MAT_D_C = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
#MAT_C_D = #sparse_tensor.encoding<{dimLevelType = ["compressed", "dense"]}> #MAT_C_D = #sparse_tensor.encoding<{dimLevelType = ["compressed", "dense"]}>
#MAT_D_D = #sparse_tensor.encoding<{
dimLevelType = ["dense", "dense"],
dimOrdering = affine_map<(i,j) -> (j,i)>
}>
#MAT_C_C_P = #sparse_tensor.encoding<{ #MAT_C_C_P = #sparse_tensor.encoding<{
dimLevelType = [ "compressed", "compressed" ], dimLevelType = [ "compressed", "compressed" ],
dimOrdering = affine_map<(i,j) -> (j,i)> dimOrdering = affine_map<(i,j) -> (j,i)>
}> }>
#MAT_C_D_P = #sparse_tensor.encoding<{ #MAT_C_D_P = #sparse_tensor.encoding<{
dimLevelType = [ "compressed", "dense" ], dimLevelType = [ "compressed", "dense" ],
Show All 19 Linesmodule {
// Concats all sparse matrices (with different encodings) to a dense matrix. // Concats all sparse matrices (with different encodings) to a dense matrix.
func.func @concat_sparse_dense(%arg0: tensor<2x4xf64, #MAT_C_C>, %arg1: tensor<3x4xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64> { func.func @concat_sparse_dense(%arg0: tensor<2x4xf64, #MAT_C_C>, %arg1: tensor<3x4xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64> {
%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index} %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
: tensor<2x4xf64, #MAT_C_C>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<9x4xf64> : tensor<2x4xf64, #MAT_C_C>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<9x4xf64>
return %0 : tensor<9x4xf64> return %0 : tensor<9x4xf64>
} }
// Concats all sparse matrices (with different encodings) to a annotated all dense matrix.
func.func @concat_sparse_annotated_dense(%arg0: tensor<2x4xf64, #MAT_C_C>, %arg1: tensor<3x4xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64, #MAT_D_D> {
%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
: tensor<2x4xf64, #MAT_C_C>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<9x4xf64, #MAT_D_D>
return %0 : tensor<9x4xf64, #MAT_D_D>
}
// Concats mix sparse and dense matrices to a sparse matrix // Concats mix sparse and dense matrices to a sparse matrix
func.func @concat_mix_sparse(%arg0: tensor<2x4xf64>, %arg1: tensor<3x4xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64, #MAT_C_C> { func.func @concat_mix_sparse(%arg0: tensor<2x4xf64>, %arg1: tensor<3x4xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64, #MAT_C_C> {
%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index} %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
: tensor<2x4xf64>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<9x4xf64, #MAT_C_C> : tensor<2x4xf64>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C> to tensor<9x4xf64, #MAT_C_C>
return %0 : tensor<9x4xf64, #MAT_C_C> return %0 : tensor<9x4xf64, #MAT_C_C>
} }
// Concats mix sparse and dense matrices to a dense matrix // Concats mix sparse and dense matrices to a dense matrix
▲ Show 20 Lines • Show All 149 Lines • ▼ Show 20 Lines func.func @dump_mat_dense_9x4(%A: tensor<9x4xf64>) {
%du = arith.constant -1.0 : f64 %du = arith.constant -1.0 : f64
%v = vector.transfer_read %A[%c0, %c0], %du: tensor<9x4xf64>, vector<9x4xf64> %v = vector.transfer_read %A[%c0, %c0], %du: tensor<9x4xf64>, vector<9x4xf64>
vector.print %v : vector<9x4xf64> vector.print %v : vector<9x4xf64>
return return
} }
func.func @dump_mat_annotated_dense_9x4(%A: tensor<9x4xf64, #MAT_D_D>) {
%c0 = arith.constant 0 : index
%du = arith.constant -1.0 : f64
%n = sparse_tensor.number_of_entries %A : tensor<9x4xf64, #MAT_D_D>
vector.print %n : index
%1 = sparse_tensor.values %A : tensor<9x4xf64, #MAT_D_D> to memref<?xf64>
%2 = vector.transfer_read %1[%c0], %du: memref<?xf64>, vector<36xf64>
vector.print %2 : vector<36xf64>
return
}
func.func @dump_mat_4x9(%A: tensor<4x9xf64, #MAT_C_C>) { func.func @dump_mat_4x9(%A: tensor<4x9xf64, #MAT_C_C>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%du = arith.constant -1.0 : f64 %du = arith.constant -1.0 : f64
%c = sparse_tensor.convert %A : tensor<4x9xf64, #MAT_C_C> to tensor<4x9xf64> %c = sparse_tensor.convert %A : tensor<4x9xf64, #MAT_C_C> to tensor<4x9xf64>
%v = vector.transfer_read %c[%c0, %c0], %du: tensor<4x9xf64>, vector<4x9xf64> %v = vector.transfer_read %c[%c0, %c0], %du: tensor<4x9xf64>, vector<4x9xf64>
vector.print %v : vector<4x9xf64> vector.print %v : vector<4x9xf64>
▲ Show 20 Lines • Show All 191 Lines • ▼ Show 20 Lines func.func @entry() {
// CHECK-NEXT: ( ( 1, 0, 1, 0, 1, 0, 0, 1.5, 1 ), ( 3.1, 0, 1, 0, 0.5, 0, 3.5, 0, 0 ), ( 0, 2, 0, 0, 1, 1, 5, 2, 0 ), ( 0, 0, 5, 2, 0, 1, 0.5, 0, 0 ) ) // CHECK-NEXT: ( ( 1, 0, 1, 0, 1, 0, 0, 1.5, 1 ), ( 3.1, 0, 1, 0, 0.5, 0, 3.5, 0, 0 ), ( 0, 2, 0, 0, 1, 1, 5, 2, 0 ), ( 0, 0, 5, 2, 0, 1, 0.5, 0, 0 ) )
// CHECK-NEXT: 18 // CHECK-NEXT: 18
// CHECK-NEXT: ( 1, 1, 1, 1.5, 1, 3.1, 1, 0.5, 3.5, 2, 1, 1, 5, 2, 5, 2, 1, 0.5 ) // CHECK-NEXT: ( 1, 1, 1, 1.5, 1, 3.1, 1, 0.5, 3.5, 2, 1, 1, 5, 2, 5, 2, 1, 0.5 )
%16 = call @concat_mix_sparse_dyn(%m42, %sm43cd, %sm44dc) %16 = call @concat_mix_sparse_dyn(%m42, %sm43cd, %sm44dc)
: (tensor<4x2xf64>, tensor<4x3xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<?x?xf64, #MAT_C_C> : (tensor<4x2xf64>, tensor<4x3xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<?x?xf64, #MAT_C_C>
call @dump_mat_dyn(%16) : (tensor<?x?xf64, #MAT_C_C>) -> () call @dump_mat_dyn(%16) : (tensor<?x?xf64, #MAT_C_C>) -> ()
// CHECK-NEXT: 36
// CHECK-NEXT: ( 1, 0, 1, 0, 1, 0, 0, 1, 1, 0, 2, 0, 0.5, 5, 0, 3.5, 5, 0.5, 3, 0, 1, 0, 2, 1.5, 0, 2, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0 )
%17 = call @concat_sparse_annotated_dense(%sm24cc, %sm34cd, %sm44dc)
: (tensor<2x4xf64, #MAT_C_C>, tensor<3x4xf64, #MAT_C_D>, tensor<4x4xf64, #MAT_D_C>) -> tensor<9x4xf64, #MAT_D_D>
call @dump_mat_annotated_dense_9x4(%17) : (tensor<9x4xf64, #MAT_D_D>) -> ()
// Release resources. // Release resources.
bufferization.dealloc_tensor %sm24cc : tensor<2x4xf64, #MAT_C_C> bufferization.dealloc_tensor %sm24cc : tensor<2x4xf64, #MAT_C_C>
bufferization.dealloc_tensor %sm34cd : tensor<3x4xf64, #MAT_C_D> bufferization.dealloc_tensor %sm34cd : tensor<3x4xf64, #MAT_C_D>
bufferization.dealloc_tensor %sm42cc : tensor<4x2xf64, #MAT_C_C> bufferization.dealloc_tensor %sm42cc : tensor<4x2xf64, #MAT_C_C>
bufferization.dealloc_tensor %sm43cd : tensor<4x3xf64, #MAT_C_D> bufferization.dealloc_tensor %sm43cd : tensor<4x3xf64, #MAT_C_D>
bufferization.dealloc_tensor %sm44dc : tensor<4x4xf64, #MAT_D_C> bufferization.dealloc_tensor %sm44dc : tensor<4x4xf64, #MAT_D_C>
bufferization.dealloc_tensor %sm24ccp : tensor<2x4xf64, #MAT_C_C_P> bufferization.dealloc_tensor %sm24ccp : tensor<2x4xf64, #MAT_C_C_P>
bufferization.dealloc_tensor %sm34cdp : tensor<3x4xf64, #MAT_C_D_P> bufferization.dealloc_tensor %sm34cdp : tensor<3x4xf64, #MAT_C_D_P>
Show All 12 Lines func.func @entry() {
bufferization.dealloc_tensor %9 : tensor<4x9xf64> bufferization.dealloc_tensor %9 : tensor<4x9xf64>
bufferization.dealloc_tensor %10 : tensor<4x9xf64, #MAT_C_C> bufferization.dealloc_tensor %10 : tensor<4x9xf64, #MAT_C_C>
bufferization.dealloc_tensor %11 : tensor<4x9xf64> bufferization.dealloc_tensor %11 : tensor<4x9xf64>
bufferization.dealloc_tensor %12 : tensor<4x9xf64, #MAT_C_C_P> bufferization.dealloc_tensor %12 : tensor<4x9xf64, #MAT_C_C_P>
bufferization.dealloc_tensor %13 : tensor<4x9xf64> bufferization.dealloc_tensor %13 : tensor<4x9xf64>
bufferization.dealloc_tensor %14 : tensor<4x9xf64, #MAT_C_C> bufferization.dealloc_tensor %14 : tensor<4x9xf64, #MAT_C_C>
bufferization.dealloc_tensor %15 : tensor<4x9xf64> bufferization.dealloc_tensor %15 : tensor<4x9xf64>
bufferization.dealloc_tensor %16 : tensor<?x?xf64, #MAT_C_C> bufferization.dealloc_tensor %16 : tensor<?x?xf64, #MAT_C_C>
bufferization.dealloc_tensor %17 : tensor<9x4xf64, #MAT_D_D>
return return
} }
} }