diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp index 5f530a585ddb9..baa57f8920094 100644 --- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp +++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp @@ -3890,7 +3890,9 @@ DiagnosedSilenceableFailure transform::tileToForallOpImpl( tilingResult = *maybeTilingResult; - if (mixedNumThreads.empty()) { + // Rank-0 ops produce no loops; skip normalization when there is nothing + // to normalize. + if (mixedNumThreads.empty() && !tilingResult.loops.empty()) { auto generatedForallOp = cast(tilingResult.loops.front()); OpBuilder::InsertionGuard g(rewriter); rewriter.setInsertionPoint(generatedForallOp); @@ -3938,7 +3940,8 @@ DiagnosedSilenceableFailure transform::TileUsingForallOp::apply( getMapping(), tilingResult); if (!diag.succeeded()) return diag; - tileOps.push_back(tilingResult.loops.front()); + if (!tilingResult.loops.empty()) + tileOps.push_back(tilingResult.loops.front()); tiledOps.append(tilingResult.tiledOps); } diff --git a/mlir/test/Dialect/Linalg/tile-to-forall.mlir b/mlir/test/Dialect/Linalg/tile-to-forall.mlir index 1b0bade728b44..f0ba58c55e11f 100644 --- a/mlir/test/Dialect/Linalg/tile-to-forall.mlir +++ b/mlir/test/Dialect/Linalg/tile-to-forall.mlir @@ -723,3 +723,61 @@ module attributes {transform.with_named_sequence} { transform.yield } } + +// ----- + +// Tiling a rank-0 linalg op with tile_sizes should not crash even when +// tileUsingSCF produces no loops. (https://github.com/llvm/llvm-project/issues/187073) +#rankZeroMap = affine_map<() -> ()> +// CHECK-LABEL: @tile_rank_zero_op_no_forall +func.func @tile_rank_zero_op_no_forall(%arg0: tensor) -> tensor { + // CHECK-NOT: scf.forall + // CHECK: return %{{.*}} : tensor + %empty = tensor.empty() : tensor + %copy = linalg.generic {indexing_maps = [#rankZeroMap, #rankZeroMap], iterator_types = []} + ins(%arg0 : tensor) outs(%empty : tensor) { + ^bb0(%in: i64, %out: i64): + linalg.yield %in : i64 + } -> tensor + return %copy : tensor +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main( + %arg0: !transform.any_op {transform.readonly}) { + %ops = transform.structured.match ops{["linalg.generic"]} in %arg0 + : (!transform.any_op) -> !transform.any_op + %tiled, %forall = transform.structured.tile_using_forall %ops tile_sizes [32] + : (!transform.any_op) -> (!transform.any_op, !transform.any_op) + transform.yield + } +} + +// ----- + +// Tiling a rank-0 linalg op with num_threads should also not crash. +// (https://github.com/llvm/llvm-project/issues/187073) +#rankZeroMap = affine_map<() -> ()> +// CHECK-LABEL: @tile_rank_zero_op_no_forall_num_threads +func.func @tile_rank_zero_op_no_forall_num_threads(%arg0: tensor) -> tensor { + // CHECK-NOT: scf.forall + // CHECK: return %{{.*}} : tensor + %empty = tensor.empty() : tensor + %copy = linalg.generic {indexing_maps = [#rankZeroMap, #rankZeroMap], iterator_types = []} + ins(%arg0 : tensor) outs(%empty : tensor) { + ^bb0(%in: i64, %out: i64): + linalg.yield %in : i64 + } -> tensor + return %copy : tensor +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main( + %arg0: !transform.any_op {transform.readonly}) { + %ops = transform.structured.match ops{["linalg.generic"]} in %arg0 + : (!transform.any_op) -> !transform.any_op + %tiled, %forall = transform.structured.tile_using_forall %ops num_threads [4] + : (!transform.any_op) -> (!transform.any_op, !transform.any_op) + transform.yield + } +}