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

[pifon2a]

Summary

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

[ThomasRaoux]

Thanks!

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
                                   ValueRange)>
         bodyBuilderFn,
     Optional<LinalgLoopDistributionOptions> distributionOptions,
     ArrayRef<StringRef> distributionTypes) {
   SmallVector<ProcInfo, 2> procInfo;
   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) {
     SmallVector<Value> outputTensors = linalgOp.getOutputTensorOperands();
     scf::ValueVector results =
         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);
   if (!distributionTypes.empty())
     tiledLoop.setDistributionTypes(b, distributionTypes);
 
   // Replace inputs/outputs with the corresponding region args.
   auto isInsideTiledLoop = [&](OpOperand &operand) {

mlir/test/Dialect/Linalg/tile-tensors.mlir

 
 // TLOOP:     %{{.*}} = linalg.tiled_loop (%{{.*}}, %{{.*}}, %{{.*}}) =
 // TLOOP-SAME: (%[[C0]], %[[C0]], %[[C0]])
 // TLOOP-SAME: to (%[[ARG_0_X]], %[[ARG_0_Y]], %[[ARG_0_Z]])
 // TLOOP-SAME: step (%[[C2]], %[[C3]], %[[C4]])
 // 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