From d15663710c91204e5987279c517a10f653efa2d9 Mon Sep 17 00:00:00 2001
From: Alexander Belyaev <pifon@google.com>
Date: Tue, 29 Jun 2021 08:55:25 +0200
Subject: [PATCH] Revert "[mlir] Skip scalar operands when tiling to
 linalg.tiled_loop."

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    | 16 +++-------------
 mlir/test/Dialect/Linalg/tile-tensors.mlir | 14 --------------
 2 files changed, 3 insertions(+), 27 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
index c82e7eb10df2..9d7286c08ded 100644
--- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -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);
diff --git a/mlir/test/Dialect/Linalg/tile-tensors.mlir b/mlir/test/Dialect/Linalg/tile-tensors.mlir
index 63bddb5a1605..f446d9da9179 100644
--- a/mlir/test/Dialect/Linalg/tile-tensors.mlir
+++ b/mlir/test/Dialect/Linalg/tile-tensors.mlir
@@ -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