[Mlir-commits] [mlir] 2146dab - [mlir][affine] Add trip count check for affine-pipeline-data-transfer pass (#206755)

[llvmlistbot at llvm.org]

Author: lonely eagle
Date: 2026-07-01T20:55:07+08:00
New Revision: 2146dab17c298b296e16c22cac64bf9ccd9f4e2f

URL: https://github.com/llvm/llvm-project/commit/2146dab17c298b296e16c22cac64bf9ccd9f4e2f
DIFF: https://github.com/llvm/llvm-project/commit/2146dab17c298b296e16c22cac64bf9ccd9f4e2f.diff

LOG: [mlir][affine] Add trip count check for affine-pipeline-data-transfer pass (#206755)

Fix `affine-pipeline-data-transfer` to skip loops with a constant trip
count ≤0, avoiding invalid DMA generation that leads to wrong code.

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
    mlir/test/Dialect/Affine/pipeline-data-transfer.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
index 575b529658127..9a1c731b5a97e 100644
--- a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
@@ -245,10 +245,16 @@ static void findMatchingStartFinishInsts(
 /// 'forOp' is deleted, and a prologue, a new pipelined loop, and epilogue are
 /// inserted right before where it was.
 void PipelineDataTransfer::runOnAffineForOp(AffineForOp forOp) {
-  if (!forOp.getStaticTripCount()) {
+  std::optional<llvm::APInt> tripCount = forOp.getStaticTripCount();
+  if (!tripCount) {
     LLVM_DEBUG(forOp.emitRemark("won't pipeline due to unknown trip count"));
     return;
   }
+  if (tripCount->getSExtValue() <= 0) {
+    LLVM_DEBUG(
+        forOp.emitRemark("won't pipeline due to trip count is non-positive"));
+    return;
+  }
 
   SmallVector<std::pair<Operation *, Operation *>, 4> startWaitPairs;
   findMatchingStartFinishInsts(forOp, startWaitPairs);

diff  --git a/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir b/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir
index 35507c37be79b..2abf28a7a9249 100644
--- a/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir
+++ b/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir
@@ -396,3 +396,30 @@ func.func @same_memref_source_and_tag(%arg0: index, %arg1: index) {
   return
 }
 // CHECK: affine.for
+
+// -----
+
+// CHECK-LABEL: func @loop_trip_count_non_positive
+func.func @loop_trip_count_non_positive() {
+  %A = memref.alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 0>
+  %Ah = memref.alloc() : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
+  %tag = memref.alloc() : memref<1 x f32>
+  %zero = arith.constant 0 : index
+  %num_elts = arith.constant 32 : index
+  affine.for %i = 0 to -1 {
+    affine.dma_start %A[%i], %Ah[%i], %tag[%zero], %num_elts : memref<256 x f32>, memref<32 x f32, 1>, memref<1 x f32>
+    affine.dma_wait %tag[%zero], %num_elts : memref<1 x f32>
+    %v = affine.load %Ah[%i] : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
+    %r = "compute"(%v) : (f32) -> (f32)
+    affine.store %r, %Ah[%i] : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
+  }
+  memref.dealloc %tag : memref<1 x f32>
+  memref.dealloc %Ah : memref<32 x f32, affine_map<(d0) -> (d0)>, 1>
+  return
+}
+
+//  CHECK-NOT: affine.dma_start
+//  CHECK-NOT: affine.dma_wait
+//      CHECK: affine.for %{{.*}} = 0 to -1
+// CHECK-NEXT:   affine.dma_start
+// CHECK-NEXT:   affine.dma_wait