32-bit zero loop iteration (#2939)

* 32-bit zero loop iteration

* fix

* fix

* fix

* fix

Author: wsmoses

enzyme/Enzyme/CacheUtility.cpp

@@ -725,8 +725,22 @@ bool CacheUtility::getContext(BasicBlock *BB, LoopContext &loopContext,
     }
   endOMP:;
 
-    if (Limit->getType() != CanonicalIV->getType())
-      Limit = SE.getZeroExtendExpr(Limit, CanonicalIV->getType());
+    if (Limit->getType() != CanonicalIV->getType()) {
+      const SCEV *Zero = SE.getZero(Limit->getType());
+      const Loop *outermost = L;
+      while (outermost->getParentLoop()) {
+        outermost = outermost->getParentLoop();
+      }
+      BasicBlock *outermostPreheader = outermost->getLoopPreheader();
+      Instruction *context = outermostPreheader
+                                 ? outermostPreheader->getTerminator()
+                                 : loopContexts[L].preheader->getTerminator();
+      if (SE.isKnownPredicateAt(ICmpInst::ICMP_SGE, Limit, Zero, context)) {
+        Limit = SE.getZeroExtendExpr(Limit, CanonicalIV->getType());
+      } else {
+        Limit = SE.getSignExtendExpr(Limit, CanonicalIV->getType());
+      }
+    }
 
 #if LLVM_VERSION_MAJOR >= 22
     SCEVExpander Exp(SE, "enzyme");

enzyme/test/Enzyme/ReverseMode/metasinfor.ll

@@ -79,10 +79,9 @@ attributes #8 = { noreturn nounwind }
 ; CHECK-NEXT:   store double 0.000000e+00, double* %"a.addr'ipa", align 8
 ; CHECK-NEXT:   %a.addr = alloca double, align 8
 ; CHECK-NEXT:   store double %a, double* %a.addr, align 8, !tbaa !2
-; CHECK-NEXT:   %0 = zext i32 %n to i64
-; CHECK-NEXT:   %1 = add {{(nuw nsw )?}}i64 %0, 1
-; CHECK-NEXT:   br label 
-; %invertfor.cond
+; CHECK-NEXT:   %0 = add {{(nuw )?}}i32 %n, 1
+; CHECK-NEXT:   %1 = sext i32 %0 to i64
+; CHECK-NEXT:   br label %invertfor.cond
 
 ; CHECK: invertentry:                                      ; preds = %invertfor.cond
 ; CHECK-NEXT:   %[[padd:.+]] = load double, double* %"a.addr'ipa", align 8

enzyme/test/Enzyme/ReverseMode/nested_loop_sext.ll

@@ -0,0 +1,53 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -enzyme -mem2reg -instsimplify -adce -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="enzyme,function(mem2reg,instsimplify,adce)" -S | FileCheck %s
+
+define void @nested(float* %x, i32 %outer, i32 %inner) {
+entry:
+  %cmp.out = icmp slt i32 %outer, 1
+  br i1 %cmp.out, label %exit, label %outer.preheader
+
+outer.preheader:
+  br label %outer.loop
+
+outer.loop:
+  %iv.out = phi i32 [ 1, %outer.preheader ], [ %iv.out.next, %outer.loop.latch ]
+  
+  ; Inner loop guard
+  %cmp.in = icmp slt i32 %inner, 1
+  br i1 %cmp.in, label %outer.loop.latch, label %inner.preheader
+
+inner.preheader:
+  br label %inner.loop
+
+inner.loop:
+  %iv.in = phi i32 [ 1, %inner.preheader ], [ %iv.in.next, %inner.loop ]
+  
+  %idx = getelementptr inbounds float, float* %x, i32 %iv.in
+  %val = load float, float* %idx, align 4
+  %val2 = fmul float %val, %val
+  store float %val2, float* %idx, align 4
+  
+  %iv.in.next = add i32 %iv.in, 1
+  %cmp.in.loop = icmp eq i32 %iv.in, %inner
+  br i1 %cmp.in.loop, label %outer.loop.latch, label %inner.loop
+
+outer.loop.latch:
+  %iv.out.next = add i32 %iv.out, 1
+  %cmp.out.loop = icmp eq i32 %iv.out, %outer
+  br i1 %cmp.out.loop, label %exit, label %outer.loop
+
+exit:
+  ret void
+}
+
+declare void @__enzyme_autodiff(i8*, float*, float*, i32, i32)
+
+define void @diffe_nested(float* %x, float* %xp, i32 %outer, i32 %inner) {
+entry:
+  call void @__enzyme_autodiff(i8* bitcast (void (float*, i32, i32)* @nested to i8*), float* %x, float* %xp, i32 %outer, i32 %inner)
+  ret void
+}
+
+; CHECK: define internal void @diffenested
+; CHECK:   %[[sub:.*]] = add i32 %inner, -1
+; CHECK:   %{{.*}} = sext i32 %[[sub]] to i64

enzyme/test/Enzyme/ReverseMode/nulluw.ll

@@ -124,7 +124,7 @@ declare dso_local double @__enzyme_autodiff(i8*, ...)
 ; CHECK-NEXT:   %[[i13:.+]] = getelementptr inbounds i32, i32* %i5_malloccache, i64 %"iv'ac.0"
 ; CHECK-NEXT:   %[[i14:.+]] = load i32, i32* %[[i13]], align 4, !invariant.group ![[ig6]]
 ; CHECK-NEXT:   %[[_unwrap1:.+]] = add i32 %[[i14]], -1
-; CHECK-NEXT:   %[[_unwrap2:.+]] = zext i32 %[[_unwrap1]] to i64
+; CHECK-NEXT:   %[[_unwrap2:.+]] = sext i32 %[[_unwrap1]] to i64
 ; CHECK-NEXT:   br label %remat_enter
 
 ; CHECK: invertbb7:                                        ; preds = %bb7, %incinvertbb2

enzyme/test/Enzyme/ReverseMode/nulluw_l10.ll

@@ -124,7 +124,7 @@ declare dso_local double @__enzyme_autodiff(i8*, ...)
 ; CHECK-NEXT:   %[[i13:.+]] = getelementptr inbounds i32, i32* %i5_malloccache, i64 %"iv'ac.0"
 ; CHECK-NEXT:   %[[i14:.+]] = load i32, i32* %[[i13]], align 4, !invariant.group ![[ig6]]
 ; CHECK-NEXT:   %[[_unwrap1:.+]] = add i32 %[[i14]], -1
-; CHECK-NEXT:   %[[_unwrap2:.+]] = zext i32 %[[_unwrap1]] to i64
+; CHECK-NEXT:   %[[_unwrap2:.+]] = sext i32 %[[_unwrap1]] to i64
 ; CHECK-NEXT:   br label %remat_enter
 
 ; CHECK: invertbb7:                                        ; preds = %bb7, %incinvertbb2

enzyme/test/Enzyme/ReverseMode/scase.ll

@@ -135,6 +135,6 @@ attributes #8 = { noreturn nounwind }
 
 ; CHECK: staging:                                          ; preds = %entry
 ; CHECK-NEXT:   %[[_unwrap1:.+]] = add i32 %SINCOSN, -2
-; CHECK-NEXT:   %[[_unwrap2]] = zext i32 %[[_unwrap1]] to i64
+; CHECK-NEXT:   %[[_unwrap2]] = sext i32 %[[_unwrap1]] to i64
 ; CHECK-NEXT:   br label %invertfor.body
 ; CHECK-NEXT: }