[Mlir-commits] [mlir] [mlir][scf] Fix `for-loop-peeling` crash (PR #77697)

[Felix Schneider]

https://github.com/ubfx updated https://github.com/llvm/llvm-project/pull/77697

>From c5fe22d67ea7e98189b2babe834e5decd38cf709 Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Wed, 10 Jan 2024 23:01:15 +0100
Subject: [PATCH 1/3] [mlir][scf] Fix crash in `for-loop-peeling`

Before applying the peeling patterns, it can happen that the `ForOp`
gets a step of zero during folding. This leads to a division-by-zero
down the line.

This patch adds an additional check for a constant-zero step and a
test.

Fix https://github.com/llvm/llvm-project/issues/75758
Depends on https://github.com/llvm/llvm-project/pull/77691
---
 .../SCF/Transforms/LoopSpecialization.cpp     |  4 ++--
 mlir/test/Dialect/SCF/for-loop-peeling.mlir   | 19 +++++++++++++++++--
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 9fda4861d40a3b..7aa1f768b303f7 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -123,8 +123,8 @@ static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp,
   auto ubInt = getConstantIntValue(forOp.getUpperBound());
   auto stepInt = getConstantIntValue(forOp.getStep());
 
-  // No specialization necessary if step size is 1.
-  if (getConstantIntValue(forOp.getStep()) == static_cast<int64_t>(1))
+  // No specialization necessary if step size is 1. Also bail out in case of a zero step which might have happened during folding.
+  if (stepInt == static_cast<int64_t>(1) || stepInt == static_cast<int64_t>(0))
     return failure();
 
   // No specialization necessary if step already divides upper bound evenly.
diff --git a/mlir/test/Dialect/SCF/for-loop-peeling.mlir b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
index 9a6d1c8c0a14cd..aa8fd7ec7ef84d 100644
--- a/mlir/test/Dialect/SCF/for-loop-peeling.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-peeling.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt %s -scf-for-loop-peeling -canonicalize -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -scf-for-loop-peeling=skip-partial=false -canonicalize -split-input-file | FileCheck %s -check-prefix=CHECK-NO-SKIP
+// RUN: mlir-opt %s -scf-for-loop-peeling -canonicalize -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: mlir-opt %s -scf-for-loop-peeling=skip-partial=false -canonicalize -split-input-file -verify-diagnostics | FileCheck %s -check-prefix=CHECK-NO-SKIP
 
 //  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (-s0 + s1) mod s2)>
 //  CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>
@@ -289,3 +289,18 @@ func.func @regression(%arg0: memref<i64>, %arg1: index) {
   }
   return
 }
+
+// -----
+
+// Check that this doesn't crash but trigger the verifier.
+func.func @zero_step(%arg0: memref<i64>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %foldto0 = arith.subi %c1, %c1 : index
+  // expected-error @+1 {{'scf.for' op constant step operand must be positive}}
+  scf.for %arg2 = %c0 to %c1 step %foldto0 {
+    %2 = arith.index_cast %arg2 : index to i64
+    memref.store %2, %arg0[] : memref<i64>
+  }
+  return
+}

>From 7e5c24c211e3322a8511089874140f613bdb1aef Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Thu, 11 Jan 2024 06:54:00 +0100
Subject: [PATCH 2/3] Bail on negative steps, fix formatting

---
 mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 7aa1f768b303f7..342213507486af 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -123,8 +123,9 @@ static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp,
   auto ubInt = getConstantIntValue(forOp.getUpperBound());
   auto stepInt = getConstantIntValue(forOp.getStep());
 
-  // No specialization necessary if step size is 1. Also bail out in case of a zero step which might have happened during folding.
-  if (stepInt == static_cast<int64_t>(1) || stepInt == static_cast<int64_t>(0))
+  // No specialization necessary if step size is 1. Also bail out in case of an
+  // invalid zero or negative step which might have happened during folding.
+  if (stepInt && *stepInt <= 1)
     return failure();
 
   // No specialization necessary if step already divides upper bound evenly.