[Mlir-commits] [mlir] 4d17ae7 - [MLIR][Affine] Fix affine-loop-tile zero cache size corner case crash (#130526)

[llvmlistbot at llvm.org]

Author: Uday Bondhugula
Date: 2025-03-11T16:28:22+05:30
New Revision: 4d17ae7776ef6ffe2dd04c146632282ac173cae6

URL: https://github.com/llvm/llvm-project/commit/4d17ae7776ef6ffe2dd04c146632282ac173cae6
DIFF: https://github.com/llvm/llvm-project/commit/4d17ae7776ef6ffe2dd04c146632282ac173cae6.diff

LOG: [MLIR][Affine] Fix affine-loop-tile zero cache size corner case crash (#130526)

Fixes: https://github.com/llvm/llvm-project/issues/64979

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
    mlir/test/Dialect/Affine/loop-tiling.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
index c8400dfe8cd5c..45bda68bb8639 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
@@ -110,7 +110,7 @@ void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band,
     return;
   }
 
-  // Use tileSizes and fill them with default tile size if it's short.
+  // Use supplied tile sizes and fill them with default tile size if it's short.
   if (!this->tileSizes.empty()) {
     tileSizes->assign(this->tileSizes.begin(), this->tileSizes.end());
     tileSizes->resize(band.size(), kDefaultTileSize);
@@ -118,6 +118,13 @@ void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band,
   }
   tileSizes->resize(band.size());
 
+  // If the cache size is zero, set the minimum valid tile size. No good reason
+  // to pick another specific size over this.
+  if (cacheSizeInKiB == 0) {
+    std::fill(tileSizes->begin(), tileSizes->end(), 1);
+    return;
+  }
+
   // The first loop in the band.
   AffineForOp rootForOp = band[0];
   (void)rootForOp;

diff  --git a/mlir/test/Dialect/Affine/loop-tiling.mlir b/mlir/test/Dialect/Affine/loop-tiling.mlir
index c01e3e910d6a0..d2aca48e615ae 100644
--- a/mlir/test/Dialect/Affine/loop-tiling.mlir
+++ b/mlir/test/Dialect/Affine/loop-tiling.mlir
@@ -1,13 +1,14 @@
 // RUN: mlir-opt %s -split-input-file -affine-loop-tile="tile-size=32" | FileCheck %s
 // RUN: mlir-opt %s -split-input-file -affine-loop-tile="cache-size=512" | FileCheck %s --check-prefix=MODEL
+// RUN: mlir-opt %s -split-input-file -affine-loop-tile="cache-size=0" | FileCheck %s --check-prefix=ZERO-CACHE
 // RUN: mlir-opt %s -split-input-file -affine-loop-tile="tile-size=32 separate" | FileCheck %s --check-prefix=SEPARATE
 
-// -----
-
 // CHECK-DAG: [[$UB:#map[0-9]*]] = affine_map<(d0) -> (d0 + 32)>
 // CHECK-DAG: [[$UB_MIN:#map[0-9]*]] = affine_map<(d0) -> (d0 + 32, 50)>
 // CHECK-DAG: [[$ID:#map[0-9]*]] = affine_map<(d0) -> (d0)>
 // CHECK-DAG: [[$ID_PLUS_21:#map[0-9]*]] = affine_map<(d0) -> (d0 + 21)>
+// ZERO-CACHE-DAG: affine_map<(d0) -> (d0)>
+// ZERO-CACHE-DAG: affine_map<(d0) -> (d0 + 1)>
 
 // CHECK-LABEL: func @loop_tiling()
 // CHECK-NEXT:   affine.for %{{.*}} = 0 to 256 step 32 {