[GPU] Avoid fusing slices of already tiled ops (#19404)

This fixes a bug in FuseAndHoistParallelLoops that can cause an infinite
loop for imperfectly aligned unpack producer fusion. The
FuseTilableSliceProducers pattern does not check that the producer
operation is not already in the tiled loop, and the tiling
implementation of imperfectly aligned unpack ops generates an
extract_slice on the result of the tiled unpack. This would cause an
infinite recursive application of the producer fusion pattern. This
fixes the bug by enforcing that the producer operation is outside of the
slice's parent loop.

Signed-off-by: Max Dawkins <maxdawkins19@gmail.com>
Co-authored-by: Max Dawkins <maxdawkins19@gmail.com>

compiler/src/iree/compiler/Codegen/Common/GPU/GPUFuseAndHoistParallelLoops.cpp

diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUFuseAndHoistParallelLoops.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUFuseAndHoistParallelLoops.cpp
index a8db945..f974e51 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUFuseAndHoistParallelLoops.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUFuseAndHoistParallelLoops.cpp
@@ -256,6 +256,11 @@
       return failure();
     }
 
+    auto producerParent = tilableProducer->getParentOfType<scf::ForallOp>();
+    if (producerParent && producerParent == parentForall) {
+      return failure();
+    }
+
     SmallVector<LoopLikeOpInterface> loops = {parentForall};
     std::optional<scf::SCFFuseProducerOfSliceResult> fusionResult =
         mlir::scf::tileAndFuseProducerOfSlice(rewriter, sliceOp, loops);

compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_fuse_and_hoist_forall.mlir

diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_fuse_and_hoist_forall.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_fuse_and_hoist_forall.mlir
index fbf53eb..b7a805f 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_fuse_and_hoist_forall.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_fuse_and_hoist_forall.mlir
@@ -518,3 +518,31 @@
 //       CHECK:   scf.forall.in_parallel
 //       CHECK:   linalg.add
 //       CHECK:   return
+
+// -----
+
+#map = affine_map<(d0) -> (d0 * 64)>
+func.func @fuse_imperfectly_aligned_unpack(%arg0: tensor<5x31xf16>, %arg1: index) -> tensor<128xf16> {
+  %c4 = arith.constant 4 : index
+  %c128 = arith.constant 128 : index
+  %c0 = arith.constant 0 : index
+  %0 = tensor.empty() : tensor<128xf16>
+  %unpack = tensor.unpack %arg0 inner_dims_pos = [0] inner_tiles = [31] into %0 : tensor<5x31xf16> -> tensor<128xf16>
+  %1 = scf.forall (%arg2) in (2) shared_outs(%arg3 = %0) -> (tensor<128xf16>) {
+    %2 = affine.apply #map(%arg2)
+    %extracted_slice = tensor.extract_slice %unpack[%2] [64] [1] : tensor<128xf16> to tensor<64xf16>
+    %extracted_slice_0 = tensor.extract_slice %arg3[%2] [64] [1] : tensor<128xf16> to tensor<64xf16>
+    %3 = linalg.copy ins(%extracted_slice : tensor<64xf16>) outs(%extracted_slice_0 : tensor<64xf16>) -> tensor<64xf16>
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %3 into %arg3[%2] [64] [1] : tensor<64xf16> into tensor<128xf16>
+    }
+  } {mapping = [#gpu.thread<linear_dim_0>]}
+  return %1 : tensor<128xf16>
+}
+
+// CHECK-LABEL: func @fuse_imperfectly_aligned_unpack
+//       CHECK:   scf.forall
+//       CHECK:     tensor.unpack
+//       CHECK:     linalg.copy
+//       CHECK:   scf.forall.in_parallel
+//       CHECK:   return