TypeAnalysis: accept nusw as well as inbounds on GEP indices

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp

@@ -2115,8 +2115,17 @@ void TypeAnalyzer::visitGEPOperator(GEPOperator &gep) {
       }
     }
 
+    // An index is an offset rather than a pointer if the GEP promises not to
+    // wrap. `nusw` suffices; `inbounds` also guarantees the result stays within
+    // the object, which this rule does not need. flang emits `nusw nuw`, never
+    // `inbounds`, for array element addresses.
+    bool indexCannotWrap = gep.isInBounds();
+#if LLVM_VERSION_MAJOR >= 19
+    indexCannotWrap |= gep.hasNoUnsignedSignedWrap();
+#endif
+
     if (has_non_const_idx &&
-        (gep.isInBounds() ||
+        (indexCannotWrap ||
          (!EnzymeStrictAliasing &&
           pointerAnalysis.Inner0() == BaseType::Pointer &&
           getAnalysis(&gep).Inner0() == BaseType::Pointer))) {

enzyme/test/Fortran/CMakeLists.txt

@@ -2,6 +2,7 @@ message("Building Fortran tests")
 
 add_subdirectory(ForwardMode)
 add_subdirectory(ReverseMode)
+add_subdirectory(TypeAnalysis)
 
 # Run regression and unit tests
 add_lit_testsuite(check-enzyme-fortran "Running enzyme fortran integration tests"

enzyme/test/Fortran/TypeAnalysis/CMakeLists.txt

@@ -0,0 +1,8 @@
+# Run regression and unit tests
+add_lit_testsuite(check-enzyme-fortran-typeanalysis "Running enzyme fortran type analysis tests"
+    ${CMAKE_CURRENT_BINARY_DIR}
+    DEPENDS ${ENZYME_TEST_DEPS} LLVMEnzyme-${LLVM_VERSION_MAJOR}
+    ARGS -v
+)
+
+set_target_properties(check-enzyme-fortran-typeanalysis PROPERTIES FOLDER "Tests")

enzyme/test/Fortran/TypeAnalysis/gep_nusw.f90

@@ -0,0 +1,23 @@
+! REQUIRES: flangenzyme
+! RUN: %fc -S -emit-llvm -O1 %s -o %t.ll
+! RUN: FileCheck %s --check-prefix=IR < %t.ll
+! RUN: %opt < %t.ll %newLoadEnzyme -passes="print-type-analysis" -type-analysis-func=stride_index_ -S -o /dev/null | FileCheck %s --check-prefix=TA
+
+! Companion to test/TypeAnalysis/gepnusw.ll: shows that the `nusw nuw` GEP the
+! rule is about is what flang actually emits for an array element address.
+
+subroutine stride_index(a, n, i, res)
+  implicit none
+  integer, intent(in) :: n, i
+  double precision, intent(in) :: a(n, *)
+  double precision, intent(out) :: res
+  res = a(i, 2)
+end subroutine stride_index
+
+! IR: getelementptr nusw nuw
+
+! The leading dimension `n`, and the stride computation indexing with it:
+! TA: ptr %{{[0-9]+}}: {[-1]:Pointer, [-1,0]:Integer, [-1,1]:Integer, [-1,2]:Integer, [-1,3]:Integer}
+! TA: %{{[0-9]+}} = load i32, ptr %{{[0-9]+}}, align 4{{.*}}: {[-1]:Integer}
+! TA-NEXT: %{{[0-9]+}} = tail call i32 @llvm.smax.i32({{.*}}): {[-1]:Integer}
+! TA-NEXT: %{{[0-9]+}} = zext nneg i32 %{{[0-9]+}} to i64: {[-1]:Integer}

enzyme/test/TypeAnalysis/gepnusw.ll

@@ -0,0 +1,51 @@
+; RUN: if [ %llvmver -ge 19 ]; then %opt < %s %newLoadEnzyme -passes="print-type-analysis" -type-analysis-func=f_nusw -S -o /dev/null | FileCheck %s --check-prefix=NUSW; fi
+; RUN: if [ %llvmver -ge 19 ]; then %opt < %s %newLoadEnzyme -passes="print-type-analysis" -type-analysis-func=f_inbounds -S -o /dev/null | FileCheck %s --check-prefix=INBOUNDS; fi
+
+; A `nusw` GEP must type its indices exactly as an `inbounds` one does. The two
+; functions below are identical apart from the no-wrap flag.
+
+declare i32 @llvm.smax.i32(i32, i32)
+
+define double @f_nusw(ptr %yh, ptr %ldyh, i64 %i) {
+entry:
+  %n = load i32, ptr %ldyh, align 4
+  %m = call i32 @llvm.smax.i32(i32 %n, i32 0)
+  %z = zext i32 %m to i64
+  %idx = mul i64 %i, %z
+  %p = getelementptr nusw nuw double, ptr %yh, i64 %idx
+  %v = load double, ptr %p, align 8
+  ret double %v
+}
+
+define double @f_inbounds(ptr %yh, ptr %ldyh, i64 %i) {
+entry:
+  %n = load i32, ptr %ldyh, align 4
+  %m = call i32 @llvm.smax.i32(i32 %n, i32 0)
+  %z = zext i32 %m to i64
+  %idx = mul i64 %i, %z
+  %p = getelementptr inbounds double, ptr %yh, i64 %idx
+  %v = load double, ptr %p, align 8
+  ret double %v
+}
+
+; NUSW: ptr %yh: {[-1]:Pointer}
+; NUSW-NEXT: ptr %ldyh: {[-1]:Pointer, [-1,0]:Integer, [-1,1]:Integer, [-1,2]:Integer, [-1,3]:Integer}
+; NUSW-NEXT: i64 %i: {[-1]:Integer}
+; NUSW-NEXT: entry
+; NUSW-NEXT:   %n = load i32, ptr %ldyh, align 4: {[-1]:Integer}
+; NUSW-NEXT:   %m = call i32 @llvm.smax.i32(i32 %n, i32 0): {[-1]:Integer}
+; NUSW-NEXT:   %z = zext i32 %m to i64: {[-1]:Integer}
+; NUSW-NEXT:   %idx = mul i64 %i, %z: {[-1]:Integer}
+; NUSW-NEXT:   %p = getelementptr nusw nuw double, ptr %yh, i64 %idx: {[-1]:Pointer, [-1,0]:Float@double}
+; NUSW-NEXT:   %v = load double, ptr %p, align 8: {[-1]:Float@double}
+
+; INBOUNDS: ptr %yh: {[-1]:Pointer}
+; INBOUNDS-NEXT: ptr %ldyh: {[-1]:Pointer, [-1,0]:Integer, [-1,1]:Integer, [-1,2]:Integer, [-1,3]:Integer}
+; INBOUNDS-NEXT: i64 %i: {[-1]:Integer}
+; INBOUNDS-NEXT: entry
+; INBOUNDS-NEXT:   %n = load i32, ptr %ldyh, align 4: {[-1]:Integer}
+; INBOUNDS-NEXT:   %m = call i32 @llvm.smax.i32(i32 %n, i32 0): {[-1]:Integer}
+; INBOUNDS-NEXT:   %z = zext i32 %m to i64: {[-1]:Integer}
+; INBOUNDS-NEXT:   %idx = mul i64 %i, %z: {[-1]:Integer}
+; INBOUNDS-NEXT:   %p = getelementptr inbounds double, ptr %yh, i64 %idx: {[-1]:Pointer, [-1,0]:Float@double}
+; INBOUNDS-NEXT:   %v = load double, ptr %p, align 8: {[-1]:Float@double}