This is an archive of the discontinued LLVM Phabricator instance.

Differential D132500

[mlir][sparse] Folding operations that try to insert zero into an all-zero sparse tensor
ClosedPublic

Authored by Peiming on Aug 23 2022, 2:16 PM.

Details

Reviewers
aartbik
bixia
wrengr
anlunx
Commits
rGec495b53f8b2: [mlir][sparse] Folding operations that try to insert zero into an all-zero…
Summary

The operations to fill zero into newly allocated sparse tensor are redundant, plus it failed
to lowering the test cases provided in the patch as well.

Diff Detail

Event Timeline

Peiming created this revision.Aug 23 2022, 2:16 PM
Herald added a reviewer: aartbik. · View Herald TranscriptAug 23 2022, 2:16 PM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: anlunx, bzcheeseman, sdasgup3 and 19 others. · View Herald Transcript
Peiming requested review of this revision.Aug 23 2022, 2:16 PM
Herald added a project: Restricted Project. · View Herald TranscriptAug 23 2022, 2:16 PM
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
Peiming added reviewers: bixia, wrengr, anlunx.Aug 23 2022, 2:17 PM
Peiming updated this revision to Diff 454959.Aug 23 2022, 2:18 PM

format...

Peiming retitled this revision from [mlir][sparse] Folding operations that tries to insert zero into an all-zero sparse tensor to [mlir][sparse] Folding operations that try to insert zero into an all-zero sparse tensor.Aug 23 2022, 2:44 PM
Harbormaster completed remote builds in B182929: Diff 454959.Aug 23 2022, 4:16 PM
bixia added inline comments.Aug 24 2022, 8:23 AM
mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
1803 ↗(On Diff #454959)

Not needed.

mlir/test/Dialect/SparseTensor/sparse_fill_zero.mlir
5

The function doesn't need to be enclosed inside a module.

115

The function doesn't need to be public.

Peiming updated this revision to Diff 455259.Aug 24 2022, 9:55 AM

address comments

Peiming marked 2 inline comments as done.Aug 24 2022, 9:56 AM
Peiming updated this revision to Diff 455262.Aug 24 2022, 10:02 AM

address comment

Peiming marked an inline comment as done.Aug 24 2022, 10:02 AM
Harbormaster completed remote builds in B183145: Diff 455262.Aug 24 2022, 11:00 AM
Peiming updated this revision to Diff 455617.Aug 25 2022, 9:23 AM

Move folding to pre rewrite before sparsification takes place

aartbik added inline comments.Aug 25 2022, 9:27 AM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
132

Yes, I like this much better. Move it even before static shape test.

Peiming updated this revision to Diff 455621.Aug 25 2022, 9:30 AM

Allow folding zero-yielding generic op on dynamic sized sparse tensor.

Peiming marked an inline comment as done.Aug 25 2022, 9:31 AM
aartbik added inline comments.Aug 25 2022, 9:47 AM
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
139

Maybe move this comment above L135 now and change it into

// Incorporate zero value into a dense, statically allocation copy.

so that the sparse/dense cases are more clear, and the test on L135/136 is not dangling

mlir/test/Dialect/SparseTensor/sparse_fill_zero.mlir
4

unrelated to your change, but this test shows that perhaps we should consider making this pre-rewriting its own pass, so we can test the input/output in isolation :-) :-)

But this is a good test to have regardless.

4

Nit: perhaps use DCSR instead of MAT

Peiming updated this revision to Diff 455623.Aug 25 2022, 9:47 AM

rebase

Peiming updated this revision to Diff 455624.Aug 25 2022, 9:52 AM
Peiming marked 3 inline comments as done.

Address comments

aartbik accepted this revision.Aug 25 2022, 9:58 AM

ship it!

This revision is now accepted and ready to land.Aug 25 2022, 9:58 AM
Peiming added inline comments.Aug 25 2022, 9:59 AM
mlir/test/Dialect/SparseTensor/sparse_fill_zero.mlir
4

Also, changed the function name to be more meaningful (instead of main).

This revision was landed with ongoing or failed builds.Aug 25 2022, 10:00 AM
Closed by commit rGec495b53f8b2: [mlir][sparse] Folding operations that try to insert zero into an all-zero… (authored by Peiming). · Explain Why
This revision was automatically updated to reflect the committed changes.
Peiming added a commit: rGec495b53f8b2: [mlir][sparse] Folding operations that try to insert zero into an all-zero….
Harbormaster completed remote builds in B183400: Diff 455624.Aug 25 2022, 10:07 AM

Revision Contents

PathSize
mlir/
lib/
Dialect/
SparseTensor/
Transforms/
10 lines
test/
Dialect/
SparseTensor/
122 lines

Diff 455617

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

Show First 20 Lines • Show All 120 Lines • ▼ Show 20 Linespublic:
using OpRewritePattern<GenericOp>::OpRewritePattern; using OpRewritePattern<GenericOp>::OpRewritePattern;
LogicalResult matchAndRewrite(GenericOp op, LogicalResult matchAndRewrite(GenericOp op,
PatternRewriter &rewriter) const override { PatternRewriter &rewriter) const override {
if (!op.hasTensorSemantics() || op.getNumResults() != 1 || if (!op.hasTensorSemantics() || op.getNumResults() != 1 ||
!isAlloc(op.getOutputOperand(0), /*isZero=*/false) || !isZeroYield(op)) !isAlloc(op.getOutputOperand(0), /*isZero=*/false) || !isZeroYield(op))
return failure(); return failure();
auto outputType = op.getResult(0).getType().cast<RankedTensorType>(); auto outputType = op.getResult(0).getType().cast<RankedTensorType>();
if (!outputType.hasStaticShape() || getSparseTensorEncoding(outputType)) if (!outputType.hasStaticShape())
return failure(); return failure();
// Yielding zero on newly allocated (all-zero) sparse tensor can be
aartbikUnsubmitted
Done
ReplyInline Actions

Yes, I like this much better. Move it even before static shape test.

aartbik: Yes, I like this much better. Move it even before static shape test.
// optimized out directly.
if (getSparseTensorEncoding(outputType)) {
rewriter.replaceOp(op, op.getOutputOperand(0)->get());
return success();
}
// Incorporate zero value into allocation copy. // Incorporate zero value into allocation copy.
aartbikUnsubmitted
Done
ReplyInline Actions

Maybe move this comment above L135 now and change it into

// Incorporate zero value into a dense, statically allocation copy.

so that the sparse/dense cases are more clear, and the test on L135/136 is not dangling

aartbik: Maybe move this comment above L135 now and change it into // Incorporate zero value into a…
Value zero = constantZero(rewriter, op.getLoc(), op.getResult(0).getType()); Value zero = constantZero(rewriter, op.getLoc(), op.getResult(0).getType());
AllocTensorOp a = AllocTensorOp a =
op.getOutputOperand(0)->get().getDefiningOp<AllocTensorOp>(); op.getOutputOperand(0)->get().getDefiningOp<AllocTensorOp>();
rewriter.updateRootInPlace(a, [&]() { a.getCopyMutable().assign(zero); }); rewriter.updateRootInPlace(a, [&]() { a.getCopyMutable().assign(zero); });
rewriter.replaceOp(op, op.getOutputOperand(0)->get()); rewriter.replaceOp(op, op.getOutputOperand(0)->get());
return success(); return success();
} }
}; };
▲ Show 20 Lines • Show All 160 Lines • Show Last 20 Lines

mlir/test/Dialect/SparseTensor/sparse_fill_zero.mlir

  • This file was added.
// RUN: mlir-opt %s --linalg-generalize-named-ops --sparsification --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
#MAT = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
// CHECK-LABEL: func.func @main(
aartbikUnsubmitted
Done
ReplyInline Actions

unrelated to your change, but this test shows that perhaps we should consider making this pre-rewriting its own pass, so we can test the input/output in isolation :-) :-)

But this is a good test to have regardless.

aartbik: unrelated to your change, but this test shows that perhaps we should consider making this pre…
aartbikUnsubmitted
Done
ReplyInline Actions

Nit: perhaps use DCSR instead of MAT

aartbik: Nit: perhaps use DCSR instead of MAT
PeimingAuthorUnsubmitted
Done
ReplyInline Actions

Also, changed the function name to be more meaningful (instead of main).

Peiming: Also, changed the function name to be more meaningful (instead of main).
// CHECK-SAME: %[[TMP_arg0:.*]]: !llvm.ptr<i8>,
bixiaUnsubmitted
Done
ReplyInline Actions

The function doesn't need to be enclosed inside a module.

bixia: The function doesn't need to be enclosed inside a module.
// CHECK-SAME: %[[TMP_arg1:.*]]: !llvm.ptr<i8>
// CHECK: %[[TMP_cst:.*]] = arith.constant 0.000000e+00 : f64
// CHECK: %[[TMP_c1_i32:.*]] = arith.constant 1 : i32
// CHECK: %[[TMP_c0_i32:.*]] = arith.constant 0 : i32
// CHECK: %[[TMP_c0:.*]] = arith.constant 0 : index
// CHECK: %[[TMP_c1:.*]] = arith.constant 1 : index
// CHECK: %[[TMP_false:.*]] = arith.constant false
// CHECK: %[[TMP_true:.*]] = arith.constant true
// CHECK: %[[TMP_c100:.*]] = arith.constant 100 : index
// CHECK: %[[TMP_c1_i8:.*]] = arith.constant 1 : i8
// CHECK: %[[TMP_0:.*]] = memref.alloca() : memref<2xi8>
// CHECK: %[[TMP_1:.*]] = memref.cast %[[TMP_0]] : memref<2xi8> to memref<?xi8>
// CHECK: memref.store %[[TMP_c1_i8]], %[[TMP_0]][%[[TMP_c0]]] : memref<2xi8>
// CHECK: memref.store %[[TMP_c1_i8]], %[[TMP_0]][%[[TMP_c1]]] : memref<2xi8>
// CHECK: %[[TMP_2:.*]] = memref.alloca() : memref<2xindex>
// CHECK: %[[TMP_3:.*]] = memref.cast %[[TMP_2]] : memref<2xindex> to memref<?xindex>
// CHECK: memref.store %[[TMP_c100]], %[[TMP_2]][%[[TMP_c0]]] : memref<2xindex>
// CHECK: memref.store %[[TMP_c100]], %[[TMP_2]][%[[TMP_c1]]] : memref<2xindex>
// CHECK: %[[TMP_4:.*]] = memref.alloca() : memref<2xindex>
// CHECK: %[[TMP_5:.*]] = memref.cast %[[TMP_4]] : memref<2xindex> to memref<?xindex>
// CHECK: memref.store %[[TMP_c0]], %[[TMP_4]][%[[TMP_c0]]] : memref<2xindex>
// CHECK: memref.store %[[TMP_c1]], %[[TMP_4]][%[[TMP_c1]]] : memref<2xindex>
// CHECK: %[[TMP_6:.*]] = llvm.mlir.null : !llvm.ptr<i8>
// CHECK: %[[TMP_7:.*]] = call @newSparseTensor(%[[TMP_1]], %[[TMP_3]], %[[TMP_5]], %[[TMP_c0_i32]], %[[TMP_c0_i32]], %[[TMP_c1_i32]], %[[TMP_c0_i32]], %[[TMP_6]])
// CHECK: %[[TMP_8:.*]] = call @sparseDimSize(%[[TMP_7]], %[[TMP_c1]])
// CHECK: %[[TMP_9:.*]] = memref.alloc(%[[TMP_8]]) : memref<?xf64>
// CHECK: %[[TMP_10:.*]] = memref.alloc(%[[TMP_8]]) : memref<?xi1>
// CHECK: %[[TMP_11:.*]] = memref.alloc(%[[TMP_8]]) : memref<?xindex>
// CHECK: linalg.fill ins(%[[TMP_cst]] : f64) outs(%[[TMP_9]] : memref<?xf64>)
// CHECK: linalg.fill ins(%[[TMP_false]] : i1) outs(%[[TMP_10]] : memref<?xi1>)
// CHECK: %[[TMP_12:.*]] = call @sparsePointers0(%[[TMP_arg0]], %[[TMP_c0]])
// CHECK: %[[TMP_13:.*]] = call @sparseIndices0(%[[TMP_arg0]], %[[TMP_c0]])
// CHECK: %[[TMP_14:.*]] = call @sparsePointers0(%[[TMP_arg0]], %[[TMP_c1]])
// CHECK: %[[TMP_15:.*]] = call @sparseIndices0(%[[TMP_arg0]], %[[TMP_c1]])
// CHECK: %[[TMP_16:.*]] = call @sparseValuesF64(%[[TMP_arg0]])
// CHECK: %[[TMP_17:.*]] = call @sparsePointers0(%[[TMP_arg1]], %[[TMP_c0]])
// CHECK: %[[TMP_18:.*]] = call @sparseIndices0(%[[TMP_arg1]], %[[TMP_c0]])
// CHECK: %[[TMP_19:.*]] = call @sparsePointers0(%[[TMP_arg1]], %[[TMP_c1]])
// CHECK: %[[TMP_20:.*]] = call @sparseIndices0(%[[TMP_arg1]], %[[TMP_c1]])
// CHECK: %[[TMP_21:.*]] = call @sparseValuesF64(%[[TMP_arg1]])
// CHECK: %[[TMP_22:.*]] = memref.alloca() : memref<2xindex>
// CHECK: %[[TMP_23:.*]] = memref.cast %[[TMP_22]] : memref<2xindex> to memref<?xindex>
// CHECK: %[[TMP_24:.*]] = memref.load %[[TMP_12]][%[[TMP_c0]]] : memref<?xindex>
// CHECK: %[[TMP_25:.*]] = memref.load %[[TMP_12]][%[[TMP_c1]]] : memref<?xindex>
// CHECK: scf.for %[[TMP_arg2:.*]] = %[[TMP_24]] to %[[TMP_25]] step %[[TMP_c1]] {
// CHECK: %[[TMP_26:.*]] = memref.load %[[TMP_13]][%[[TMP_arg2]]] : memref<?xindex>
// CHECK: memref.store %[[TMP_26]], %[[TMP_22]][%[[TMP_c0]]] : memref<2xindex>
// CHECK: %[[TMP_27:.*]] = memref.load %[[TMP_14]][%[[TMP_arg2]]] : memref<?xindex>
// CHECK: %[[TMP_28:.*]] = arith.addi %[[TMP_arg2]], %[[TMP_c1]] : index
// CHECK: %[[TMP_29:.*]] = memref.load %[[TMP_14]][%[[TMP_28]]] : memref<?xindex>
// CHECK: %[[TMP_30:.*]] = memref.load %[[TMP_17]][%[[TMP_c0]]] : memref<?xindex>
// CHECK: %[[TMP_31:.*]] = memref.load %[[TMP_17]][%[[TMP_c1]]] : memref<?xindex>
// CHECK: %[[TMP_32:.*]]:3 = scf.while (%[[TMP_arg3:.*]] = %[[TMP_27]], %[[TMP_arg4:.*]] = %[[TMP_30]], %[[TMP_arg5:.*]] = %[[TMP_c0]]) : (index, index, index) -> (index, index, index) {
// CHECK: %[[TMP_33:.*]] = arith.cmpi ult, %[[TMP_arg3]], %[[TMP_29]] : index
// CHECK: %[[TMP_34:.*]] = arith.cmpi ult, %[[TMP_arg4]], %[[TMP_31]] : index
// CHECK: %[[TMP_35:.*]] = arith.andi %[[TMP_33]], %[[TMP_34]] : i1
// CHECK: scf.condition(%[[TMP_35]]) %[[TMP_arg3]], %[[TMP_arg4]], %[[TMP_arg5]] : index, index, index
// CHECK: } do {
// CHECK: ^bb0(%[[TMP_arg3:.*]]: index, %[[TMP_arg4:.*]]: index, %[[TMP_arg5:.*]]: index):
// CHECK: %[[TMP_33:.*]] = memref.load %[[TMP_15]][%[[TMP_arg3]]] : memref<?xindex>
// CHECK: %[[TMP_34:.*]] = memref.load %[[TMP_18]][%[[TMP_arg4]]] : memref<?xindex>
// CHECK: %[[TMP_35:.*]] = arith.cmpi ult, %[[TMP_34]], %[[TMP_33]] : index
// CHECK: %[[TMP_36:.*]] = arith.select %[[TMP_35]], %[[TMP_34]], %[[TMP_33]] : index
// CHECK: %[[TMP_37:.*]] = arith.cmpi eq, %[[TMP_33]], %[[TMP_36]] : index
// CHECK: %[[TMP_38:.*]] = arith.cmpi eq, %[[TMP_34]], %[[TMP_36]] : index
// CHECK: %[[TMP_39:.*]] = arith.andi %[[TMP_37]], %[[TMP_38]] : i1
// CHECK: %[[TMP_40:.*]] = scf.if %[[TMP_39]] -> (index) {
// CHECK: %[[TMP_45:.*]] = memref.load %[[TMP_16]][%[[TMP_arg3]]] : memref<?xf64>
// CHECK: %[[TMP_46:.*]] = memref.load %[[TMP_19]][%[[TMP_arg4]]] : memref<?xindex>
// CHECK: %[[TMP_47:.*]] = arith.addi %[[TMP_arg4]], %[[TMP_c1]] : index
// CHECK: %[[TMP_48:.*]] = memref.load %[[TMP_19]][%[[TMP_47]]] : memref<?xindex>
// CHECK: %[[TMP_49:.*]] = scf.for %[[TMP_arg6:.*]] = %[[TMP_46]] to %[[TMP_48]] step %[[TMP_c1]] iter_args(%[[TMP_arg7:.*]] = %[[TMP_arg5]]) -> (index) {
// CHECK: %[[TMP_50:.*]] = memref.load %[[TMP_20]][%[[TMP_arg6]]] : memref<?xindex>
// CHECK: %[[TMP_51:.*]] = memref.load %[[TMP_9]][%[[TMP_50]]] : memref<?xf64>
// CHECK: %[[TMP_52:.*]] = memref.load %[[TMP_21]][%[[TMP_arg6]]] : memref<?xf64>
// CHECK: %[[TMP_53:.*]] = arith.mulf %[[TMP_45]], %[[TMP_52]] : f64
// CHECK: %[[TMP_54:.*]] = arith.addf %[[TMP_51]], %[[TMP_53]] : f64
// CHECK: %[[TMP_55:.*]] = memref.load %[[TMP_10]][%[[TMP_50]]] : memref<?xi1>
// CHECK: %[[TMP_56:.*]] = arith.cmpi eq, %[[TMP_55]], %[[TMP_false]] : i1
// CHECK: %[[TMP_57:.*]] = scf.if %[[TMP_56]] -> (index) {
// CHECK: memref.store %[[TMP_true]], %[[TMP_10]][%[[TMP_50]]] : memref<?xi1>
// CHECK: memref.store %[[TMP_50]], %[[TMP_11]][%[[TMP_arg7]]] : memref<?xindex>
// CHECK: %[[TMP_58:.*]] = arith.addi %[[TMP_arg7]], %[[TMP_c1]] : index
// CHECK: scf.yield %[[TMP_58]] : index
// CHECK: } else {
// CHECK: scf.yield %[[TMP_arg7]] : index
// CHECK: }
// CHECK: memref.store %[[TMP_54]], %[[TMP_9]][%[[TMP_50]]] : memref<?xf64>
// CHECK: scf.yield %[[TMP_57]] : index
// CHECK: }
// CHECK: scf.yield %[[TMP_49]] : index
// CHECK: } else {
// CHECK: scf.yield %[[TMP_arg5]] : index
// CHECK: }
// CHECK: %[[TMP_41:.*]] = arith.addi %[[TMP_arg3]], %[[TMP_c1]] : index
// CHECK: %[[TMP_42:.*]] = arith.select %[[TMP_37]], %[[TMP_41]], %[[TMP_arg3]] : index
// CHECK: %[[TMP_43:.*]] = arith.addi %[[TMP_arg4]], %[[TMP_c1]] : index
// CHECK: %[[TMP_44:.*]] = arith.select %[[TMP_38]], %[[TMP_43]], %[[TMP_arg4]] : index
// CHECK: scf.yield %[[TMP_42]], %[[TMP_44]], %[[TMP_40]] : index, index, index
// CHECK: }
// CHECK: func.call @expInsertF64(%[[TMP_7]], %[[TMP_23]], %[[TMP_9]], %[[TMP_10]], %[[TMP_11]], %[[TMP_32]]#2)
// CHECK: }
// CHECK: memref.dealloc %[[TMP_9]] : memref<?xf64>
// CHECK: memref.dealloc %[[TMP_10]] : memref<?xi1>
// CHECK: memref.dealloc %[[TMP_11]] : memref<?xindex>
// CHECK: call @endInsert(%[[TMP_7]]) : (!llvm.ptr<i8>) -> ()
// CHECK: return %[[TMP_7]] : !llvm.ptr<i8>
func.func @main(%arg0: tensor<100x100xf64, #MAT>,
%arg1: tensor<100x100xf64, #MAT>) -> tensor<100x100xf64, #MAT> {
%0 = bufferization.alloc_tensor() : tensor<100x100xf64, #MAT>
bixiaUnsubmitted
Done
ReplyInline Actions

The function doesn't need to be public.

bixia: The function doesn't need to be public.
%cst = arith.constant 0.000000e+00 : f64
%1 = linalg.fill ins(%cst : f64)
outs(%0 : tensor<100x100xf64, #MAT>) -> tensor<100x100xf64, #MAT>
%2 = linalg.matmul ins(%arg0, %arg1 : tensor<100x100xf64, #MAT>, tensor<100x100xf64, #MAT>)
outs(%1 : tensor<100x100xf64, #MAT>) -> tensor<100x100xf64, #MAT>
return %2 : tensor<100x100xf64, #MAT>
}