Skip to content

Commit 53d5d34

Browse files
author
Peiming Liu
committed
[mlir][sparse] extend foreach operation to accept reduction arguments.
Reviewed By: aartbik Differential Revision: https://reviews.llvm.org/D137463
1 parent 026ddce commit 53d5d34

File tree

5 files changed

+172
-54
lines changed

5 files changed

+172
-54
lines changed

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

Lines changed: 42 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -857,32 +857,55 @@ def SparseTensor_YieldOp : SparseTensor_Op<"yield", [Pure, Terminator]>,
857857

858858
def SparseTensor_ForeachOp : SparseTensor_Op<"foreach",
859859
[SingleBlockImplicitTerminator<"YieldOp">]>,
860-
Arguments<(ins AnyTensor:$tensor)>{
860+
Arguments<(ins AnyTensor:$tensor,
861+
Variadic<AnyType>:$initArgs)>,
862+
Results<(outs Variadic<AnyType>:$results)> {
861863
let summary = "Iterates over elements in a tensor";
862864
let description = [{
863865
Iterates over stored elements in a tensor (which are typically, but not always,
864866
non-zero for sparse tensors) and executes the block.
865867

866-
For an input tensor with rank n, the block must take n + 1 arguments. The
867-
first n arguments must be Index type, together indicating the current coordinates
868-
of the element being visited. The last argument must have the same type as the
868+
For an input tensor with rank n, the block must take n + 1 (and additional loop
869+
carried variables as described below) arguments. The first n arguments must be
870+
Index type, together indicating the current coordinates of the element being visited.
871+
The last argument must have the same type as the
869872
tensor's element type, representing the actual value loaded from the input
870873
tensor at the given coordinates.
871874

872-
Note that foreach generated loop iterates over the stored elements in the storage
873-
order. However, no matter what storage order is used, the indices passed to the block
874-
always obey the original dimension order.
875+
`sparse_tensor.foreach` can also operate on loop-carried variables and returns
876+
the final values after loop termination. The initial values of the variables are
877+
passed as additional SSA operands to the "sparse_tensor.foreach" following the n + 1
878+
SSA values mentioned above (n coordinate and 1 value).
879+
880+
The region must terminate with a "sparse_tensor.yield" that passes the current
881+
values of all loop-carried variables to the next iteration, or to the
882+
result, if at the last iteration. The number and static types of loop-carried
883+
variables may not change with iterations.
884+
885+
For example:
886+
```mlir
887+
%c0 = arith.constant 0 : i32
888+
%ret = sparse_tensor.foreach in %0 init(%c0): tensor<?x?xi32, #DCSR>, i32 -> i32 do {
889+
^bb0(%arg1: index, %arg2: index, %arg3: i32, %iter: i32):
890+
%sum = arith.add %iter, %arg3
891+
sparse_tensor.yield %sum
892+
}
893+
```
894+
895+
It is important to note that foreach generated loop iterates over the stored elements
896+
in the storage order. However, no matter what storage order is used, the indices passed
897+
to the block always obey the original dimension order.
875898

876899
For example:
877900
```mlir
878901
#COL_MAJOR = #sparse_tensor.encoding<{
879902
dimLevelType = [ "compressed", "compressed" ],
880903
dimOrdering = affine_map<(i,j) -> (j,i)>
881904
}>
882-
905+
883906
// foreach on a column-major sparse tensor
884907
sparse_tensor.foreach in %0 : tensor<2x3xf64, #COL_MAJOR> do {
885-
^bb0(%row: index, %col: index, %arg3: f64):
908+
^bb0(%row: index, %col: index, %arg3: f64):
886909
// [%row, %col] -> [0, 0], [1, 0], [2, 0], [0, 1], [1, 1], [2, 1]
887910
}
888911

@@ -892,30 +915,25 @@ def SparseTensor_ForeachOp : SparseTensor_Op<"foreach",
892915

893916
// foreach on a row-major sparse tensor
894917
sparse_tensor.foreach in %0 : tensor<2x3xf64, #ROW_MAJOR> do {
895-
^bb0(%row: index, %col: index, %arg3: f64):
918+
^bb0(%row: index, %col: index, %arg3: f64):
896919
// [%row, %col] -> [0, 0], [0, 1], [1, 0], [1, 1], [2, 0], [2, 1]
897920
}
898921

899922
```
900-
901-
Example:
902-
903-
```mlir
904-
sparse_tensor.foreach in %0 : tensor<?x?xf64, #DCSR> do {
905-
^bb0(%arg1: index, %arg2: index, %arg3: f64):
906-
do something...
907-
}
908-
```
909923
}];
910924

911925
let builders = [
912-
OpBuilder<(
913-
ins "Value":$tensor,
914-
"function_ref<void(OpBuilder &, Location, ValueRange)>")>
926+
OpBuilder<(ins "Value":$tensor,
927+
"function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>")>,
928+
OpBuilder<(ins "Value":$tensor,
929+
"ValueRange":$iterArgs,
930+
"function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>")>
915931
];
916932

917-
let regions = (region AnyRegion:$region);
918-
let assemblyFormat = "`in` $tensor attr-dict `:` type($tensor) `do` $region";
933+
let regions = (region SizedRegion<1>:$region);
934+
let assemblyFormat = "`in` $tensor (`init``(`$initArgs^`)`)? attr-dict"
935+
" `:` type($tensor) (`,` type($initArgs)^)?"
936+
" (`->` type($results)^)? `do` $region";
919937
let hasVerifier = 1;
920938
}
921939

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

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -581,11 +581,20 @@ LogicalResult CompressOp::verify() {
581581

582582
void ForeachOp::build(
583583
OpBuilder &builder, OperationState &result, Value tensor,
584-
function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder) {
585-
build(builder, result, tensor);
584+
function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>
585+
bodyBuilder) {
586+
build(builder, result, tensor, llvm::None, bodyBuilder);
587+
}
588+
589+
void ForeachOp::build(
590+
OpBuilder &builder, OperationState &result, Value tensor,
591+
ValueRange initArgs,
592+
function_ref<void(OpBuilder &, Location, ValueRange, Value, ValueRange)>
593+
bodyBuilder) {
594+
build(builder, result, initArgs.getTypes(), tensor, initArgs);
595+
// Builds foreach body.
586596
if (!bodyBuilder)
587597
return;
588-
589598
auto rtp = tensor.getType().cast<RankedTensorType>();
590599
int64_t rank = rtp.getRank();
591600

@@ -602,23 +611,38 @@ void ForeachOp::build(
602611
auto &region = *result.regions.front();
603612
Block *bodyBlock =
604613
builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
605-
bodyBuilder(builder, result.location, bodyBlock->getArguments());
614+
bodyBuilder(builder, result.location,
615+
bodyBlock->getArguments().slice(0, rank),
616+
bodyBlock->getArguments()[rank],
617+
bodyBlock->getArguments().drop_front(rank + 1));
606618
}
607619

608620
LogicalResult ForeachOp::verify() {
609621
auto t = getTensor().getType().cast<RankedTensorType>();
610622
auto args = getBody()->getArguments();
611623

612-
if (static_cast<size_t>(t.getRank()) + 1 != args.size())
624+
if (static_cast<size_t>(t.getRank()) + 1 + getInitArgs().size() !=
625+
args.size())
613626
return emitError("Unmatched number of arguments in the block");
614627

628+
if (getNumResults() != getInitArgs().size())
629+
return emitError("Mismatch in number of init arguments and results");
630+
631+
if (getResultTypes() != getInitArgs().getTypes())
632+
return emitError("Mismatch in types of init arguments and results");
633+
634+
auto yield = cast<YieldOp>(getBody()->getTerminator());
635+
if (yield.getNumOperands() != getNumResults() ||
636+
yield.getOperands().getTypes() != getResultTypes())
637+
return emitError("Mismatch in types of yield values and results");
638+
615639
for (int64_t i = 0, e = t.getRank(); i < e; i++)
616640
if (args[i].getType() != IndexType::get(getContext()))
617641
emitError(
618642
llvm::formatv("Expecting Index type for argument at index {0}", i));
619643

620644
auto elemTp = t.getElementType();
621-
auto valueTp = args.back().getType();
645+
auto valueTp = args[t.getRank()].getType();
622646
if (elemTp != valueTp)
623647
emitError(llvm::formatv("Unmatched element type between input tensor and "
624648
"block argument, expected:{0}, got: {1}",

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

Lines changed: 35 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -357,7 +357,9 @@ struct Sparse2SparseReshapeRewriter : public OpRewritePattern<ReshapeOp> {
357357
auto cooBuffer =
358358
rewriter.create<AllocTensorOp>(loc, cooTp, dstDynSizes).getResult();
359359
rewriter.create<ForeachOp>(
360-
loc, srcTensor, [&](OpBuilder &builder, Location loc, ValueRange args) {
360+
loc, srcTensor, llvm::None,
361+
[&](OpBuilder &builder, Location loc, ValueRange args, Value v,
362+
ValueRange reduc) {
361363
SmallVector<Value, 4> srcIndices;
362364
SmallVector<Value, 4> dstIndices;
363365
for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
@@ -366,7 +368,7 @@ struct Sparse2SparseReshapeRewriter : public OpRewritePattern<ReshapeOp> {
366368
}
367369
translateIndicesArray(builder, loc, op.getReassociationIndices(),
368370
srcIndices, srcSizes, dstSizes, dstIndices);
369-
builder.create<InsertOp>(loc, args.back(), cooBuffer, dstIndices);
371+
builder.create<InsertOp>(loc, v, cooBuffer, dstIndices);
370372
builder.create<sparse_tensor::YieldOp>(loc);
371373
});
372374

@@ -446,7 +448,9 @@ struct ConcatenateRewriter : public OpRewritePattern<ConcatenateOp> {
446448
// Build a for op for each input tensor to append new values into the
447449
// output tensor.
448450
rewriter.create<ForeachOp>(
449-
loc, input, [&](OpBuilder &builder, Location loc, ValueRange args) {
451+
loc, input, llvm::None,
452+
[&](OpBuilder &builder, Location loc, ValueRange args, Value v,
453+
ValueRange reduc) {
450454
SmallVector<Value, 4> indices;
451455
for (int64_t i = 0; i < rank; i++) {
452456
uint64_t dim =
@@ -457,7 +461,7 @@ struct ConcatenateRewriter : public OpRewritePattern<ConcatenateOp> {
457461
idx = builder.create<arith::AddIOp>(loc, idx, offset);
458462
indices.push_back(idx);
459463
}
460-
builder.create<InsertOp>(loc, args.back(), cooBuffer, indices);
464+
builder.create<InsertOp>(loc, v, cooBuffer, indices);
461465
builder.create<sparse_tensor::YieldOp>(loc);
462466
});
463467
// Accumulates the offset. Note that only static-shaped inputs are allowed
@@ -558,12 +562,13 @@ struct ConvertRewriter : public OpRewritePattern<ConvertOp> {
558562
sizesForTensor(rewriter, sizes, loc, srcTp, src);
559563
Value dst = allocDenseTensor(rewriter, loc, dstTp, sizes);
560564

561-
rewriter.create<ForeachOp>(
562-
loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
563-
builder.create<memref::StoreOp>(loc, args.back(), dst,
564-
args.drop_back());
565-
builder.create<sparse_tensor::YieldOp>(loc);
566-
});
565+
rewriter.create<ForeachOp>(loc, src, llvm::None,
566+
[&](OpBuilder &builder, Location loc,
567+
ValueRange args, Value v, ValueRange reduc) {
568+
builder.create<memref::StoreOp>(loc, v, dst,
569+
args);
570+
builder.create<sparse_tensor::YieldOp>(loc);
571+
});
567572

568573
rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstTp, dst);
569574
return success();
@@ -598,13 +603,15 @@ struct ConvertRewriter : public OpRewritePattern<ConvertOp> {
598603
tmpCoo =
599604
rewriter.create<AllocTensorOp>(loc, srcTp, dynSrcSizes).getResult();
600605
rewriter.create<ForeachOp>(
601-
loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
606+
loc, src, llvm::None,
607+
[&](OpBuilder &builder, Location loc, ValueRange args, Value v,
608+
ValueRange reduc) {
602609
SmallVector<Value, 4> indices;
603610
for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
604611
uint64_t dim = toStoredDim(encSrc, i);
605612
indices.push_back(args[dim]);
606613
}
607-
builder.create<InsertOp>(loc, args.back(), tmpCoo, indices);
614+
builder.create<InsertOp>(loc, v, tmpCoo, indices);
608615
builder.create<sparse_tensor::YieldOp>(loc);
609616
});
610617
src = tmpCoo;
@@ -646,16 +653,18 @@ struct ConvertRewriter : public OpRewritePattern<ConvertOp> {
646653
getDynamicSizes(dstTp, srcSizes, dynDstSizes);
647654
Value dst =
648655
rewriter.create<AllocTensorOp>(loc, dstTp, dynDstSizes).getResult();
649-
rewriter.create<ForeachOp>(
650-
loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
651-
SmallVector<Value, 4> indices;
652-
for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
653-
uint64_t dim = toStoredDim(encDst, i);
654-
indices.push_back(args[dim]);
655-
}
656-
builder.create<InsertOp>(loc, args.back(), dst, indices);
657-
builder.create<sparse_tensor::YieldOp>(loc);
658-
});
656+
rewriter.create<ForeachOp>(loc, src, llvm::None,
657+
[&](OpBuilder &builder, Location loc,
658+
ValueRange args, Value v, ValueRange reduc) {
659+
SmallVector<Value, 4> indices;
660+
for (int64_t i = 0, e = srcTp.getRank(); i < e;
661+
i++) {
662+
uint64_t dim = toStoredDim(encDst, i);
663+
indices.push_back(args[dim]);
664+
}
665+
builder.create<InsertOp>(loc, v, dst, indices);
666+
builder.create<sparse_tensor::YieldOp>(loc);
667+
});
659668

660669
// Release the temporary COO if it is created.
661670
if (tmpCoo)
@@ -866,12 +875,14 @@ struct OutRewriter : public OpRewritePattern<OutOp> {
866875
ModuleOp module = op->getParentOfType<ModuleOp>();
867876
// For each element in the source tensor, output the element.
868877
rewriter.create<ForeachOp>(
869-
loc, src, [&](OpBuilder &builder, Location loc, ValueRange args) {
878+
loc, src, llvm::None,
879+
[&](OpBuilder &builder, Location loc, ValueRange args, Value v,
880+
ValueRange reduc) {
870881
for (uint64_t i = 0; i < rank; i++) {
871882
rewriter.create<memref::StoreOp>(loc, args[i], indices,
872883
constantIndex(builder, loc, i));
873884
}
874-
rewriter.create<memref::StoreOp>(loc, args.back(), value);
885+
rewriter.create<memref::StoreOp>(loc, v, value);
875886
SmallVector<Value, 4> operands{writer, rankValue, indices, value};
876887
FlatSymbolRefAttr fn = getFunc(module, outNextFuncName, {}, operands,
877888
EmitCInterface::On);

‎mlir/test/Dialect/SparseTensor/invalid.mlir

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -551,6 +551,51 @@ func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
551551

552552
// -----
553553

554+
#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
555+
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
556+
// expected-error@+1 {{Unmatched element type between input tensor and block argument}}
557+
sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
558+
^bb0(%1: index, %2: index, %v: f32) :
559+
}
560+
return
561+
}
562+
563+
// -----
564+
565+
#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
566+
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
567+
// expected-error@+1 {{Mismatch in number of init arguments and results}}
568+
sparse_tensor.foreach in %arg0 init(%arg1) : tensor<2x4xf64, #DCSR>, f32 do {
569+
^bb0(%1: index, %2: index, %v: f32, %r1 : i32) :
570+
}
571+
return
572+
}
573+
574+
// -----
575+
576+
#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
577+
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
578+
// expected-error@+1 {{Mismatch in types of init arguments and results}}
579+
%1 = sparse_tensor.foreach in %arg0 init(%arg1) : tensor<2x4xf64, #DCSR>, f32 -> i32 do {
580+
^bb0(%1: index, %2: index, %v: f32, %r0 : f32) :
581+
}
582+
return
583+
}
584+
585+
// -----
586+
587+
#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
588+
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
589+
// expected-error@+1 {{Mismatch in types of yield values and results}}
590+
%1 = sparse_tensor.foreach in %arg0 init(%arg1) : tensor<2x4xf64, #DCSR>, f32 -> f32 do {
591+
^bb0(%1: index, %2: index, %v: f32, %r0 : f32) :
592+
sparse_tensor.yield %1 : index
593+
}
594+
return
595+
}
596+
597+
// -----
598+
554599
// TODO: a test case with empty xs doesn't work due to some parser issues.
555600

556601
func.func @sparse_sort_x_type( %arg0: index, %arg1: memref<?xf32>) {

‎mlir/test/Dialect/SparseTensor/roundtrip.mlir

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -411,6 +411,26 @@ func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
411411
return
412412
}
413413

414+
// -----
415+
416+
#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
417+
418+
// CHECK-LABEL: func @sparse_tensor_foreach(
419+
// CHECK-SAME: %[[A0:.*]]: tensor<2x4xf64, #sparse_tensor.encoding<{{{.*}}}>>,
420+
// CHECK-SAME: %[[A1:.*]]: f32
421+
// CHECK-NEXT: %[[RET:.*]] = sparse_tensor.foreach in %[[A0]] init(%[[A1]])
422+
// CHECK-NEXT: ^bb0(%[[TMP_1:.*]]: index, %[[TMP_2:.*]]: index, %[[TMP_v:.*]]: f64, %[[TMP_r:.*]]: f32)
423+
// CHECK: sparse_tensor.yield %[[TMP_r]] : f32
424+
// CHECK: }
425+
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
426+
%ret = sparse_tensor.foreach in %arg0 init(%arg1): tensor<2x4xf64, #DCSR>, f32 -> f32
427+
do {
428+
^bb0(%1: index, %2: index, %v: f64, %r: f32) :
429+
sparse_tensor.yield %r : f32
430+
}
431+
return
432+
}
433+
414434
// ----
415435

416436
// CHECK-LABEL: func @sparse_sort_1d0v(

0 commit comments

Comments
 (0)