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

[llvmlistbot at llvm.org]

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

>From ac5dc66c74a36425bc2e0f49fe35e5e11bdc5e9e Mon Sep 17 00:00:00 2001
From: rishabhmadan19 <rishabhkec at gmail.com>
Date: Fri, 30 Jan 2026 00:24:49 +0530
Subject: [PATCH 1/2] 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

>From 68a70e3cffdb3c78fea5d162720a5aec66ff24bc Mon Sep 17 00:00:00 2001
From: rishabhmadan19 <rishabhkec at gmail.com>
Date: Fri, 30 Jan 2026 00:30:50 +0530
Subject: [PATCH 2/2] Fixes #115849 Added Utils.cpp

---
 mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 92d8846a4640d..75e98ca11d92f 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -1648,6 +1648,23 @@ mlir::affine::computeSliceUnion(ArrayRef<Operation *> opsA,
       dstAccess = MemRefAccess(b);
       if (srcAccess.memref != dstAccess.memref)
         continue;
+
+      // Bailout if producer and consumer access different-sized element types.
+      // This prevents illegal fusion when vector types have different shapes.
+      Type srcElementType =
+          cast<MemRefType>(srcAccess.memref.getType()).getElementType();
+      Type dstElementType =
+          cast<MemRefType>(dstAccess.memref.getType()).getElementType();
+      if (auto srcVecType = dyn_cast<VectorType>(srcElementType)) {
+        if (auto dstVecType = dyn_cast<VectorType>(dstElementType)) {
+          if (srcVecType.getShape() != dstVecType.getShape()) {
+            LDBG() << "Cannot fuse: different vector element sizes in "
+                      "producer/consumer";
+            return SliceComputationResult::GenericFailure;
+          }
+        }
+      }
+
       // Check if 'loopDepth' exceeds nesting depth of src/dst ops.
       if ((!isBackwardSlice && loopDepth > getNestingDepth(a)) ||
           (isBackwardSlice && loopDepth > getNestingDepth(b))) {