Revert "[mlir] Skip scalar operands when tiling to linalg.tiled_loop."
authorAlexander Belyaev <pifon@google.com>
Tue, 29 Jun 2021 06:55:25 +0000 (08:55 +0200)
committerAlexander Belyaev <pifon@google.com>
Tue, 29 Jun 2021 06:55:25 +0000 (08:55 +0200)
This reverts commit 69046b4a79e2670053362112aa467f89faf9e53e. It did not
really break anything, but it was decided to allow scalars and other
non-shaped operands for tiled_loop.

mlir/lib/Dialect/Linalg/Utils/Utils.cpp
mlir/test/Dialect/Linalg/tile-tensors.mlir

index c82e7eb..9d7286c 100644 (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);
index 63bddb5..f446d9d 100644 (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