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

Fixes: #64979

Author: bondhugula

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp

@@ -110,14 +110,21 @@ 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);
     return;
   }
   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;

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 {