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

This reverts commit 69046b4a79. It did not
really break anything, but it was decided to allow scalars and other
non-shaped operands for tiled_loop.
This commit is contained in:
Alexander Belyaev 2021-06-29 08:55:25 +02:00
parent 487f74a6c4
commit d15663710c
2 changed files with 3 additions and 27 deletions

View File

@ -284,19 +284,6 @@ 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) {
@ -305,6 +292,9 @@ 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);

View File

@ -130,17 +130,3 @@ 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