From d98e60069975e003ab5216b7ac7f64be0fd4ea01 Mon Sep 17 00:00:00 2001
From: Vimal Patel <vimal@polymagelabs.com>
Date: Thu, 6 Apr 2023 10:44:43 +0530
Subject: [PATCH] Emit error message instead of asserting

test-affine-parametric-tile pass was crashing on an incorrect input.
Instead it should emit an error message and exit.

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

Reviewed By: mehdi_amini, bondhugula

Differential Revision: https://reviews.llvm.org/D147595
---
 .../Dialect/Affine/loop-tiling-parametric.mlir     | 30 +++++++++++++++++++++-
 .../Affine/TestAffineLoopParametricTiling.cpp      | 25 ++++++++++++------
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/mlir/test/Dialect/Affine/loop-tiling-parametric.mlir b/mlir/test/Dialect/Affine/loop-tiling-parametric.mlir
index 11ad83a..ba045a87 100644
--- a/mlir/test/Dialect/Affine/loop-tiling-parametric.mlir
+++ b/mlir/test/Dialect/Affine/loop-tiling-parametric.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -split-input-file -test-affine-parametric-tile | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -test-affine-parametric-tile -verify-diagnostics | FileCheck %s
 // Test cases to test the utility introduced to tile affine for loops using
 // SSA values as tiling parameters(tile sizes). The tile sizes are expected
 // to be passed as input arguments(before any other argument) to the function
@@ -271,3 +271,31 @@ func.func @tile_with_upper_bounds_in_dimensions_and_symbols_non_unit_steps(%t12
   }
   return
 }
+
+// -----
+
+func.func @too_few_tile_size_params() {
+  // expected-error@+1 {{too few tile sizes provided in the argument list of the function which contains the current band}}
+  affine.for %i = 0 to 256 {
+    affine.for %j = 0 to 512 {
+      affine.for %k = 0 to 1024 {
+        "test.foo"(%i, %j, %k) : (index, index, index) -> ()
+      }
+    }
+  }
+  return
+}
+
+// -----
+
+func.func @invalid_type_for_tile_size_params(%arg0: f32, %arg1: f32, %arg2: f32) {
+  // expected-error@+1 {{expected tiling parameters to be of index type}}
+  affine.for %i = 0 to 256 {
+    affine.for %j = 0 to 512 {
+      affine.for %k = 0 to 1024 {
+        "test.foo"(%i, %j, %k) : (index, index, index) -> ()
+      }
+    }
+  }
+  return
+}
diff --git a/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp b/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp
index ef72792..5629a8d 100644
--- a/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp
@@ -37,20 +37,25 @@ struct TestAffineLoopParametricTiling
 /// Checks if the function enclosing the loop nest has any arguments passed to
 /// it, which can be used as tiling parameters. Assumes that atleast 'n'
 /// arguments are passed, where 'n' is the number of loops in the loop nest.
-static void checkIfTilingParametersExist(ArrayRef<AffineForOp> band) {
+static LogicalResult checkIfTilingParametersExist(ArrayRef<AffineForOp> band) {
   assert(!band.empty() && "no loops in input band");
   AffineForOp topLoop = band[0];
 
   if (func::FuncOp funcOp = dyn_cast<func::FuncOp>(topLoop->getParentOp()))
-    assert(funcOp.getNumArguments() >= band.size() && "Too few tile sizes");
+    if (funcOp.getNumArguments() < band.size())
+      return topLoop->emitError(
+          "too few tile sizes provided in the argument list of the function "
+          "which contains the current band");
+  return success();
 }
 
 /// Captures tiling parameters, which are expected to be passed as arguments
 /// to the function enclosing the loop nest. Also checks if the required
 /// parameters are of index type. This approach is temporary for testing
 /// purposes.
-static void getTilingParameters(ArrayRef<AffineForOp> band,
-                                SmallVectorImpl<Value> &tilingParameters) {
+static LogicalResult
+getTilingParameters(ArrayRef<AffineForOp> band,
+                    SmallVectorImpl<Value> &tilingParameters) {
   AffineForOp topLoop = band[0];
   Region *funcOpRegion = topLoop->getParentRegion();
   unsigned nestDepth = band.size();
@@ -58,11 +63,13 @@ static void getTilingParameters(ArrayRef<AffineForOp> band,
   for (BlockArgument blockArgument :
        funcOpRegion->getArguments().take_front(nestDepth)) {
     if (blockArgument.getArgNumber() < nestDepth) {
-      assert(blockArgument.getType().isIndex() &&
-             "expected tiling parameters to be of index type.");
+      if (!blockArgument.getType().isIndex())
+        return topLoop->emitError(
+            "expected tiling parameters to be of index type");
       tilingParameters.push_back(blockArgument);
     }
   }
+  return success();
 }
 
 void TestAffineLoopParametricTiling::runOnOperation() {
@@ -77,10 +84,12 @@ void TestAffineLoopParametricTiling::runOnOperation() {
     SmallVector<AffineForOp, 6> tiledNest;
     SmallVector<Value, 6> tilingParameters;
     // Check if tiling parameters are present.
-    checkIfTilingParametersExist(band);
+    if (checkIfTilingParametersExist(band).failed())
+      return;
 
     // Get function arguments as tiling parameters.
-    getTilingParameters(band, tilingParameters);
+    if (getTilingParameters(band, tilingParameters).failed())
+      return;
 
     (void)tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest);
   }
-- 
2.7.4