This is an archive of the discontinued LLVM Phabricator instance.

Differential D138436

[mlir][sparse] Fix a bug in concatenate operator rewriting.
ClosedPublic

Authored by bixia on Nov 21 2022, 7:12 AM.

Details

Reviewers
aartbik
nicolasvasilache
wrengr
Peiming
Commits
rG6c01b5cdaddc: [mlir][sparse] Fix a bug in concatenate operator rewriting.
Summary

When calculating the dynamic dimensions for the concatenate result, we
shouldn't accumulate the sizes for the non-concatenating dimensions.

Diff Detail

Event Timeline

bixia created this revision.Nov 21 2022, 7:12 AM
Herald added a reviewer: aartbik. · View Herald TranscriptNov 21 2022, 7:12 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 21 2022, 7:12 AM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptNov 21 2022, 7:12 AM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
bixia added reviewers: wrengr, Peiming.Nov 21 2022, 7:13 AM
Peiming added inline comments.Nov 21 2022, 9:00 AM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
472–477

When was it deleted?

Peiming added inline comments.Nov 21 2022, 9:12 AM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
472–477

NVM, but are you sure there is no compiler warning?

e.g., conDim (size_t) == d.index() (int)?

bixia added inline comments.Nov 21 2022, 12:22 PM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
472–477

It was a bug when this code was added.

472–477

d.index() is size_t see here

Peiming accepted this revision.Nov 21 2022, 1:21 PM
Peiming added inline comments.
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
472–477

You are right!

This revision is now accepted and ready to land.Nov 21 2022, 1:21 PM
aartbik accepted this revision.Nov 21 2022, 1:58 PM
aartbik added inline comments.
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
472–477

Can you please add a comment to this test (so we will never delete it again)

wrengr added inline comments.Nov 21 2022, 2:02 PM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
472–477

If the issue is about signedness warnings, then I could move mlir/ExecutionEngine/SparseTensor/ArithmeticUtils.h into MLIRSparseTensorEnums so that safelyEQ can be used here

bixia updated this revision to Diff 477018.Nov 21 2022, 3:24 PM
bixia marked an inline comment as done.

Add a comment.

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
472–477

Thanks! But the two types are the same here.

Harbormaster completed remote builds in B198873: Diff 477018.Nov 22 2022, 12:33 AM
Closed by commit rG6c01b5cdaddc: [mlir][sparse] Fix a bug in concatenate operator rewriting. (authored by bixia). · Explain WhyNov 22 2022, 8:17 AM
This revision was automatically updated to reflect the committed changes.
bixia added a commit: rG6c01b5cdaddc: [mlir][sparse] Fix a bug in concatenate operator rewriting..

Revision Contents

PathSize
mlir/
lib/
Dialect/
SparseTensor/
Transforms/
9 lines
test/
Dialect/
SparseTensor/
89 lines

Diff 477207

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

Show First 20 Lines • Show All 463 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(ConcatenateOp op,
SmallVector<Value> dynSizes; SmallVector<Value> dynSizes;
if (!rtp.hasStaticShape()) { if (!rtp.hasStaticShape()) {
ArrayRef<int64_t> rShape = rtp.getShape(); ArrayRef<int64_t> rShape = rtp.getShape();
for (const auto &d : llvm::enumerate(rShape)) { for (const auto &d : llvm::enumerate(rShape)) {
if (d.value() == ShapedType::kDynamic) { if (d.value() == ShapedType::kDynamic) {
Value v = Value v =
createOrFoldDimOp(rewriter, loc, op.getOperand(0), d.index()); createOrFoldDimOp(rewriter, loc, op.getOperand(0), d.index());
rewriter.create<tensor::DimOp>(loc, op.getOperand(0), d.index()); rewriter.create<tensor::DimOp>(loc, op.getOperand(0), d.index());
if (conDim == d.index()) {
// Adding the size of the concatenating dimension.
for (const auto &opnd : op.getOperands().drop_front()) { for (const auto &opnd : op.getOperands().drop_front()) {
Value t = createOrFoldDimOp(rewriter, loc, opnd, d.index()); Value t = createOrFoldDimOp(rewriter, loc, opnd, d.index());
v = rewriter.create<arith::AddIOp>(loc, v, t); v = rewriter.create<arith::AddIOp>(loc, v, t);
} }
PeimingUnsubmitted
Not Done
ReplyInline Actions

When was it deleted?

Peiming: When was it deleted?
PeimingUnsubmitted
Not Done
ReplyInline Actions

NVM, but are you sure there is no compiler warning?

e.g., conDim (size_t) == d.index() (int)?

Peiming: NVM, but are you sure there is no compiler warning? e.g., `conDim (size_t) == d.index() (int)`?
bixiaAuthorUnsubmitted
Done
ReplyInline Actions

d.index() is size_t see here

bixia: d.index() is size_t see [here](https://github.com/llvm/llvm…
PeimingUnsubmitted
Not Done
ReplyInline Actions

You are right!

Peiming: You are right!
bixiaAuthorUnsubmitted
Done
ReplyInline Actions

It was a bug when this code was added.

bixia: It was a bug when this code was added.
aartbikUnsubmitted
Done
ReplyInline Actions

Can you please add a comment to this test (so we will never delete it again)

aartbik: Can you please add a comment to this test (so we will never delete it again)
wrengrUnsubmitted
Not Done
ReplyInline Actions

If the issue is about signedness warnings, then I could move mlir/ExecutionEngine/SparseTensor/ArithmeticUtils.h into MLIRSparseTensorEnums so that safelyEQ can be used here

wrengr: If the issue is about signedness warnings, then I could move…
bixiaAuthorUnsubmitted
Done
ReplyInline Actions

Thanks! But the two types are the same here.

bixia: Thanks! But the two types are the same here.
}
dynSizes.push_back(v); dynSizes.push_back(v);
} }
} }
} }
// %t = concatenate %s1, %s2, %s3 {dim = 1} // %t = concatenate %s1, %s2, %s3 {dim = 1}
// ==> // ==>
// %tmp = bufferization.alloc_tensor : unordered COO // %tmp = bufferization.alloc_tensor : unordered COO
▲ Show 20 Lines • Show All 588 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: --sparsification | FileCheck %s // RUN: | FileCheck %s
#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}> #DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
// 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
▲ Show 20 Lines • Show All 71 Lines • ▼ Show 20 Linesfunc.func @concat_sparse_sparse(%arg0: tensor<2x4xf64, #DCSR>,
%arg2: tensor<4x4xf64, #DCSR>) %arg2: tensor<4x4xf64, #DCSR>)
-> tensor<9x4xf64, #DCSR> { -> tensor<9x4xf64, #DCSR> {
%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<9x4xf64, #DCSR> tensor<4x4xf64, #DCSR> to tensor<9x4xf64, #DCSR>
return %0 : tensor<9x4xf64, #DCSR> return %0 : tensor<9x4xf64, #DCSR>
} }
// CHECK-LABEL: @concat_sparse_sparse_dynamic(
// 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: %[[TMP_23:.*]] = sparse_tensor.load %[[RET_3]] hasInserts
// CHECK: %[[TMP_22:.*]] = sparse_tensor.convert %[[TMP_23]] : tensor<?x?xf64, #sparse_tensor
// CHECK: return %[[TMP_22]] : tensor<?x?xf64, #sparse_tensor
func.func @concat_sparse_sparse_dynamic(%arg0: tensor<2x4xf64, #DCSR>,
%arg1: tensor<3x4xf64, #DCSR>,
%arg2: tensor<4x4xf64, #DCSR>)
-> tensor<?x?xf64, #DCSR> {
%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
: tensor<2x4xf64, #DCSR>,
tensor<3x4xf64, #DCSR>,
tensor<4x4xf64, #DCSR> to tensor<?x?xf64, #DCSR>
return %0 : tensor<?x?xf64, #DCSR>
}
No newline at end of file