Normalize loops with negative upper bounds

Signed-off-by: Anna Gringauze <agringauze@nvidia.com>
NVIDIA · Feb 2, 2025 · ad4c93f · ad4c93f
1 parent 3b0f04c
commit ad4c93f
Show file tree

Hide file tree

Showing 2 changed files with 104 additions and 1 deletion.
diff --git 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
@@ -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:         }
+}