[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

Author: pifon2a

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);

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