From ad4c93f4d4f5dc454bb987e2da250f1deb0fdde5 Mon Sep 17 00:00:00 2001
From: Anna Gringauze <agringauze@nvidia.com>
Date: Sat, 1 Feb 2025 17:35:08 -0800
Subject: [PATCH] Normalize loops with negative upper bounds

Signed-off-by: Anna Gringauze <agringauze@nvidia.com>
---
 lib/Optimizer/Transforms/LoopAnalysis.cpp | 10 ++-
 test/Quake/loop_normalize.qke             | 95 +++++++++++++++++++++++
 2 files changed, 104 insertions(+), 1 deletion(-)
 create mode 100644 test/Quake/loop_normalize.qke

diff --git a/lib/Optimizer/Transforms/LoopAnalysis.cpp b/lib/Optimizer/Transforms/LoopAnalysis.cpp
index f002d5ee8a..c40e2b7e30 100644
--- a/lib/Optimizer/Transforms/LoopAnalysis.cpp
+++ b/lib/Optimizer/Transforms/LoopAnalysis.cpp
@@ -73,6 +73,14 @@ static bool isaConstantOf(Value v, std::int64_t hasVal) {
   return false;
 }
 
+static bool isNegativeConstant(Value v) {
+  v = peelCastOps(v);
+  if (auto c = v.getDefiningOp<arith::ConstantOp>())
+    if (auto ia = dyn_cast<IntegerAttr>(c.getValue()))
+      return ia.getInt() < 0;
+  return false;
+}
+
 static bool isClosedIntervalForm(arith::CmpIPredicate p) {
   return p == arith::CmpIPredicate::ule || p == arith::CmpIPredicate::sle;
 }
@@ -274,7 +282,7 @@ bool opt::isaMonotonicLoop(Operation *op, bool allowEarlyExit,
 
 bool opt::isaInvariantLoop(const LoopComponents &c, bool allowClosedInterval) {
   if (isaConstantOf(c.initialValue, 0) && isaConstantOf(c.stepValue, 1) &&
-      isa<arith::AddIOp>(c.stepOp)) {
+      isa<arith::AddIOp>(c.stepOp) && !isNegativeConstant(c.compareValue)) {
     auto cmp = cast<arith::CmpIOp>(c.compareOp);
     return validCountedLoopIntervalForm(cmp, allowClosedInterval);
   }
diff --git a/test/Quake/loop_normalize.qke b/test/Quake/loop_normalize.qke
new file mode 100644
index 0000000000..f040f87843
--- /dev/null
+++ b/test/Quake/loop_normalize.qke
@@ -0,0 +1,95 @@
+// ========================================================================== //
+// Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates.                 //
+// All rights reserved.                                                       //
+//                                                                            //
+// This source code and the accompanying materials are made available under   //
+// the terms of the Apache License 2.0 which accompanies this distribution.   //
+// ========================================================================== //
+
+// RUN: cudaq-opt -cc-loop-normalize %s | FileCheck %s
+
+module {
+  func.func @test_positive_boundaries() {
+    %c0_i64 = arith.constant 0 : i64
+    %c1_i64 = arith.constant 1 : i64
+    %0 = quake.alloca !quake.veq<0>
+    %1 = cc.loop while ((%arg0 = %c1_i64) -> (i64)) {
+      %2 = arith.cmpi ult, %arg0, %c0_i64 : i64
+      cc.condition %2(%arg0 : i64)
+    } do {
+    ^bb0(%arg0: i64):
+      %2 = arith.subi %arg0, %c1_i64 : i64
+      %3 = quake.extract_ref %0[%2] : (!quake.veq<0>, i64) -> !quake.ref
+      quake.x %3 : (!quake.ref) -> ()
+      cc.continue %arg0 : i64
+    } step {
+    ^bb0(%arg0: i64):
+      %2 = arith.addi %arg0, %c1_i64 : i64
+      cc.continue %2 : i64
+    }
+    return
+  }
+
+// CHECK-LABEL:   func.func @test_positive_boundaries() {
+// CHECK:     %[[VAL_0:.*]] = arith.constant 0 : i64
+// CHECK:     %[[VAL_1:.*]] = arith.constant 1 : i64
+// CHECK:     %[[VAL_2:.*]] = quake.alloca !quake.veq<0>
+// CHECK:     %[[VAL_3:.*]] = cc.loop while ((%arg0 = %[[VAL_0]]) -> (i64)) {
+// CHECK:       %[[VAL_4:.*]] = arith.cmpi ne, %arg0, %[[VAL_0]] : i64
+// CHECK:       cc.condition %[[VAL_4]](%arg0 : i64)
+// CHECK:     } do {
+// CHECK:     ^bb0(%arg0: i64):
+// CHECK:       %[[VAL_4:.*]] = quake.extract_ref %[[VAL_2]][%arg0] : (!quake.veq<0>, i64) -> !quake.ref
+// CHECK:       quake.x %[[VAL_4]] : (!quake.ref) -> ()
+// CHECK:       cc.continue %arg0 : i64
+// CHECK:     } step {
+// CHECK:     ^bb0(%arg0: i64):
+// CHECK:       %[[VAL_4:.*]] = arith.addi %arg0, %[[VAL_1]] : i64
+// CHECK:       cc.continue %[[VAL_4]] : i64
+// CHECK:     } {normalized}
+// CHECK:     return
+// CHECK:   }
+
+  func.func @test_negative_boundaries() {
+    %c-1_i32 = arith.constant -1 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %c0_i32 = arith.constant 0 : i32
+    %0 = quake.alloca !quake.veq<0>
+    %1 = cc.loop while ((%arg0 = %c0_i32) -> (i32)) {
+      %2 = arith.cmpi slt, %arg0, %c-1_i32 : i32
+      cc.condition %2(%arg0 : i32)
+    } do {
+    ^bb0(%arg0: i32):
+      %2 = cc.cast signed %arg0 : (i32) -> i64
+      %3 = quake.extract_ref %0[%2] : (!quake.veq<0>, i64) -> !quake.ref
+      quake.x %3 : (!quake.ref) -> ()
+      cc.continue %arg0 : i32
+    } step {
+    ^bb0(%arg0: i32):
+      %2 = arith.addi %arg0, %c1_i32 : i32
+      cc.continue %2 : i32
+    }
+    return
+  }
+
+// CHECK-LABEL:   func.func @test_negative_boundaries() {
+// CHECK:           %[[VAL_0:.*]] = arith.constant 0 : i32
+// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : i32
+// CHECK:           %[[VAL_2:.*]] = quake.alloca !quake.veq<0>
+// CHECK:           %[[VAL_3:.*]] = cc.loop while ((%arg0 = %[[VAL_0]]) -> (i32)) {
+// CHECK:             %[[VAL_4:.*]] = arith.cmpi ne, %arg0, %[[VAL_0]] : i32
+// CHECK:             cc.condition %[[VAL_4]](%arg0 : i32)
+// CHECK:           } do {
+// CHECK:           ^bb0(%arg0: i32):
+// CHECK:             %[[VAL_4:.*]] = cc.cast signed %arg0 : (i32) -> i64
+// CHECK:             %[[VAL_5:.*]] = quake.extract_ref %[[VAL_2]][%[[VAL_4]]] : (!quake.veq<0>, i64) -> !quake.ref
+// CHECK:             quake.x %[[VAL_5]] : (!quake.ref) -> ()
+// CHECK:             cc.continue %arg0 : i32
+// CHECK:           } step {
+// CHECK:           ^bb0(%arg0: i32):
+// CHECK:             %[[VAL_4:.*]] = arith.addi %arg0, %[[VAL_1]] : i32
+// CHECK:             cc.continue %[[VAL_4]] : i32
+// CHECK:           } {normalized}
+// CHECK:           return
+// CHECK:         }
+}