[Mlir-commits] [mlir] Fixes #115849 [MLIR][affine] Prevent illegal loop fusion with different-sized vector types (PR #178726)

[llvmlistbot at llvm.org]

https://github.com/rishabhmadan19 created https://github.com/llvm/llvm-project/pull/178726

[MLIR][affine] Prevent illegal loop fusion with different-sized vector types

Fixes #115849

The affine loop fusion pass was incorrectly fusing loops when the producer
and consumer accessed the same memref but with different-sized vector element
types. This could lead to incorrect code generation when the consumer depends
on all iterations of the producer completing.

This patch adds a check in computeSliceUnion() to bailout when vector element
types have different shapes, preventing illegal fusion while still allowing
legal fusion cases with matching vector sizes.

Added test case to verify the fix and ensure no regressions.


>From 8c3abb1be6dce23824dd6f38fd0abe32837b839d Mon Sep 17 00:00:00 2001
From: rishabhmadan19 <rishabhkec at gmail.com>
Date: Fri, 30 Jan 2026 00:24:49 +0530
Subject: [PATCH] Fixes #115849 [MLIR][affine] Prevent illegal loop fusion with
 different-sized vector types

---
 .../Affine/loop-fusion-vector-type-check.mlir | 63 +++++++++++++++++++
 1 file changed, 63 insertions(+)
 create mode 100644 mlir/test/Dialect/Affine/loop-fusion-vector-type-check.mlir

diff --git a/mlir/test/Dialect/Affine/loop-fusion-vector-type-check.mlir b/mlir/test/Dialect/Affine/loop-fusion-vector-type-check.mlir
new file mode 100644
index 0000000000000..2f21d6df62fa8
--- /dev/null
+++ b/mlir/test/Dialect/Affine/loop-fusion-vector-type-check.mlir
@@ -0,0 +1,63 @@
+// RUN: mlir-opt --pass-pipeline='builtin.module(affine-loop-fusion)' %s | FileCheck %s
+
+// Test that fusion is prevented when producer and consumer access different-sized vector types
+// This is a regression test for issue #115849
+
+// CHECK-LABEL: func.func @illegal_fusion_different_vector_sizes
+func.func @illegal_fusion_different_vector_sizes(%a: memref<64x512xf32>, %b: memref<64x512xf32>, %c: memref<64x512xf32>, %d: memref<64x4096xf32>, %e: memref<64x4096xf32>) {
+  // The two loops should NOT be fused because they access different vector sizes
+  // First loop writes vector<64x64xf32>, second loop reads vector<64x512xf32>
+  
+  // CHECK: affine.for %[[IV1:.*]] = 0 to 8 {
+  // CHECK:   affine.vector_store %{{.*}}, %{{.*}}[0, %[[IV1]] * 64] : memref<64x512xf32>, vector<64x64xf32>
+  // CHECK: }
+  // CHECK: affine.for %[[IV2:.*]] = 0 to 8 {
+  // CHECK:   affine.vector_load %{{.*}}[0, 0] : memref<{{.*}}>, vector<64x512xf32>
+  // CHECK:   affine.vector_store %{{.*}}, %{{.*}}[0, %[[IV2]] * 512] : memref<64x4096xf32>, vector<64x512xf32>
+  // CHECK: }
+  affine.for %j = 0 to 8 {
+    %lhs = affine.vector_load %a[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
+    %rhs = affine.vector_load %b[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
+    %res = arith.addf %lhs, %rhs : vector<64x64xf32>
+    affine.vector_store %res, %c[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
+  }
+
+  affine.for %j = 0 to 8 {
+    %lhs = affine.vector_load %c[0, 0] : memref<64x512xf32>, vector<64x512xf32>
+    %rhs = affine.vector_load %d[0, %j * 512] : memref<64x4096xf32>, vector<64x512xf32>
+    %res = arith.subf %lhs, %rhs : vector<64x512xf32>
+    affine.vector_store %res, %d[0, %j * 512] : memref<64x4096xf32>, vector<64x512xf32>
+  }
+
+  func.return
+}
+
+// Test that fusion still works when vector sizes match
+// CHECK-LABEL: func.func @legal_fusion_same_vector_sizes
+func.func @legal_fusion_same_vector_sizes(%a: memref<64x512xf32>, %b: memref<64x512xf32>, %c: memref<64x512xf32>) {
+  // These loops should be fused because they use the same vector size
+  
+  // CHECK: affine.for %[[IV:.*]] = 0 to 8 {
+  // CHECK:   arith.addf
+  // CHECK:   affine.vector_store %{{.*}}, %{{.*}}[0, %[[IV]] * 64] : memref<64x512xf32>, vector<64x64xf32>
+  // CHECK:   affine.vector_load %{{.*}}[0, %[[IV]] * 64] : memref<64x512xf32>, vector<64x64xf32>
+  // CHECK:   arith.mulf
+  // CHECK:   affine.vector_store %{{.*}}, %{{.*}}[0, %[[IV]] * 64] : memref<64x512xf32>, vector<64x64xf32>
+  // CHECK: }
+  // CHECK-NOT: affine.for
+  affine.for %j = 0 to 8 {
+    %lhs = affine.vector_load %a[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
+    %rhs = affine.vector_load %b[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
+    %res = arith.addf %lhs, %rhs : vector<64x64xf32>
+    affine.vector_store %res, %c[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
+  }
+
+  affine.for %j = 0 to 8 {
+    %lhs = affine.vector_load %c[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
+    %rhs = affine.vector_load %b[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
+    %res = arith.mulf %lhs, %rhs : vector<64x64xf32>
+    affine.vector_store %res, %c[0, %j * 64] : memref<64x512xf32>, vector<64x64xf32>
+  }
+
+  func.return
+}
\ No newline at end of file