[Mlir-commits] [mlir] 69046b4 - [mlir] Skip scalar operands when tiling to linalg.tiled_loop.

[Alexander Belyaev]

Author: Alexander Belyaev
Date: 2021-06-28T23:01:17+02:00
New Revision: 69046b4a79e2670053362112aa467f89faf9e53e

URL: https://github.com/llvm/llvm-project/commit/69046b4a79e2670053362112aa467f89faf9e53e
DIFF: https://github.com/llvm/llvm-project/commit/69046b4a79e2670053362112aa467f89faf9e53e.diff

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

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

Differential Revision: https://reviews.llvm.org/D105059

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
index 9d7286c08dedf..c82e7eb10df21 100644
--- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -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);

diff  --git a/mlir/test/Dialect/Linalg/tile-tensors.mlir b/mlir/test/Dialect/Linalg/tile-tensors.mlir
index f446d9da9179d..63bddb5a16055 100644
--- a/mlir/test/Dialect/Linalg/tile-tensors.mlir
+++ b/mlir/test/Dialect/Linalg/tile-tensors.mlir
@@ -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