[CPU] Fix check for scalable tile sizes (#15264)

This was meant to check `tileScalableFlags` not `tileSizes`. Tests happened to pass as the lowering configs contained 1s.
iree-org · Oct 23, 2023 · 7d650ae · 7d650ae
1 parent aa5602d
commit 7d650ae
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Utils.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/Utils.cpp
@@ -115,7 +115,7 @@ void setSCFTileSizes(scf::SCFTilingOptions &options, TilingInterface consumerOp,
   int numLoops = consumerOp.getLoopIteratorTypes().size();
   tileSizes.resize(numLoops, /*default=*/0);
   tileScalableFlags.resize(numLoops, /*default=*/false);
-  if (!llvm::is_contained(tileSizes, 1)) {
+  if (!llvm::is_contained(tileScalableFlags, true)) {
     // Non-scalable case: All constant tile sizes.
     options.setTileSizes(
         getAsIndexOpFoldResult(consumerOp.getContext(), tileSizes));

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile.mlir
@@ -59,3 +59,13 @@ func.func @scalable_matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<
 //  CHECK-SAME:       step %[[SCALABLE_TILE_SIZE]]
 //       CHECK:     scf.for
 //  CHECK-SAME:         step %[[C1]]
+
+// -----
+
+// CHECK-LABEL: scalable_lowering_config_with_no_1s
+// CHECK: vector.vscale
+func.func @scalable_lowering_config_with_no_1s(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %1 = linalg.matmul {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[8, [32], 0]]>} ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)
+            outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %1 : tensor<?x?xf32>
+}