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

[Uday Bondhugula]

https://github.com/bondhugula created https://github.com/llvm/llvm-project/pull/130526

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


>From 9070b971dee8d510a2718bbe7aea522c2bc5ddce Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday at polymagelabs.com>
Date: Sun, 9 Mar 2025 16:11:09 +0530
Subject: [PATCH] [MLIR][Affine] Fix affine-loop-tile zero cache size corner
 case crash

Fixes: https://github.com/llvm/llvm-project/issues/64979
---
 mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp | 9 ++++++++-
 mlir/test/Dialect/Affine/loop-tiling.mlir         | 5 +++--
 2 files changed, 11 insertions(+), 3 deletions(-)

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 {