intel/llvm
1
0
Fork 0
mirror of https://github.com/intel/llvm.git synced 2026-01-25 10:55:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

[mlir] Skip scalar operands when tiling to linalg.tiled_loop.

Browse Source 
We are interested only in tensors/memrefs when creating a TiledLoopOp.

Differential Revision: https://reviews.llvm.org/D105059
This commit is contained in:
Alexander Belyaev
2021-06-28 22:20:30 +02:00
parent 9d0bf7699c
commit 69046b4a79
2 changed files with 27 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
mlir/lib/Dialect/Linalg/Utils/Utils.cpp
View File

@@ -284,6 +284,19 @@ void GenerateLoopNest<TiledLoopOp>::doit(
SmallVector<Value, 4> lbs, ubs, steps;
unpackRanges(loopRanges, lbs, ubs, steps);
auto dropNonShapedValues =
[](ArrayRef<OpOperand *> operands) -> SmallVector<Value, 2> {
SmallVector<Value, 2> filteredOperands;
for (OpOperand *operand : operands) {
Type type = operand->get().getType();
if (type.isa<ShapedType>())
filteredOperands.push_back(operand->get());
}
return filteredOperands;
};
auto inputOperands = dropNonShapedValues(linalgOp.getInputOperands());
auto outputOperands = dropNonShapedValues(linalgOp.getOutputOperands());
auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange ivs, ValueRange inputs,
ValueRange outputs) {
@@ -292,9 +305,6 @@ void GenerateLoopNest<TiledLoopOp>::doit(
bodyBuilderFn(nestedBuilder, nestedLoc, ivs, outputTensors);
nestedBuilder.create<linalg::YieldOp>(nestedLoc, results);
};
SmallVector<Value> inputOperands = linalgOp.getInputOperands();
SmallVector<Value> outputOperands = linalgOp.getOutputOperands();
auto tiledLoop =
b.create<TiledLoopOp>(loc, lbs, ubs, steps, inputOperands, outputOperands,
b.getArrayAttr(iteratorTypes), wrappedBuilderFn);

14
mlir/test/Dialect/Linalg/tile-tensors.mlir
View File

@@ -130,3 +130,17 @@ func @generic_op_tensors(
// TLOOP-SAME: ins (%{{.*}} = %[[ARG_0]]: [[TY]], %{{.*}} = %[[ARG_1]]: [[TY]])
// TLOOP-SAME: outs (%{{.*}} = %[[INIT]]: [[TY]])
// TLOOP-SAME: distribution["block_x", "block_y", "none"] {
func @fill(%arg0 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
%c0 = constant 0.0 : f32
%0 = linalg.fill(%c0, %arg0) : f32, tensor<?x?x?xf32> -> tensor<?x?x?xf32>
return %0 : tensor<?x?x?xf32>
}
// CHECK-LABEL: func @fill
// TLOOP-LABEL: func @fill
// TLOOP-NOT: ins
// TLOOP: tensor.extract_slice
// TLOOP-NEXT: linalg.fill
// TLOOP-NEXT: tensor.insert_slice