Fix TypeAnalysis

enzyme/Enzyme/GradientUtils.cpp

@@ -56,6 +56,10 @@ llvm::cl::opt<bool>
     EnzymeNewCache("enzyme-new-cache", cl::init(true), cl::Hidden,
                    cl::desc("Use new cache decision algorithm"));
 
+llvm::cl::opt<bool> EnzymeLoopInvariantCache(
+    "enzyme-loop-invariant-cache", cl::init(true), cl::Hidden,
+    cl::desc("Attempt to hoist cache outside of loop"));
+
 bool isPotentialLastLoopValue(Value *val, const BasicBlock *loc,
                               const LoopInfo &LI) {
   if (Instruction *inst = dyn_cast<Instruction>(val)) {
@@ -2688,7 +2692,7 @@ Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM,
 
       // this is guarded because havent told cacheForReverse how to move
       if (mode == DerivativeMode::Both)
-        if (!li->isVolatile()) {
+        if (!li->isVolatile() && EnzymeLoopInvariantCache) {
           if (auto AI = dyn_cast<AllocaInst>(liobj)) {
             assert(isa<AllocaInst>(orig_liobj));
             if (auto AT = dyn_cast<ArrayType>(AI->getAllocatedType()))

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp

@@ -1246,10 +1246,22 @@ void TypeAnalyzer::visitPHINode(PHINode &phi) {
 }
 
 void TypeAnalyzer::visitTruncInst(TruncInst &I) {
+  auto &DL = fntypeinfo.Function->getParent()->getDataLayout();
+  size_t inSize = (DL.getTypeSizeInBits(I.getOperand(0)->getType()) + 7) / 8;
+  size_t outSize = (DL.getTypeSizeInBits(I.getType()) + 7) / 8;
   if (direction & DOWN)
-    updateAnalysis(&I, getAnalysis(I.getOperand(0)), &I);
+    updateAnalysis(&I,
+                   getAnalysis(I.getOperand(0))
+                       .ShiftIndices(DL, /*off*/ 0, inSize, /*addOffset*/ 0)
+                       .ShiftIndices(DL, /*off*/ 0, outSize, /*addOffset*/ 0)
+                       .CanonicalizeValue(outSize, DL),
+                   &I);
   if (direction & UP)
-    updateAnalysis(I.getOperand(0), getAnalysis(&I), &I);
+    updateAnalysis(I.getOperand(0),
+                   getAnalysis(&I)
+                       .ShiftIndices(DL, /*off*/ 0, outSize, /*addOffset*/ 0)
+                       .CanonicalizeValue(inSize, DL),
+                   &I);
 }
 
 void TypeAnalyzer::visitZExtInst(ZExtInst &I) {
@@ -1465,7 +1477,7 @@ void TypeAnalyzer::visitExtractElementInst(ExtractElementInst &I) {
 }
 
 void TypeAnalyzer::visitInsertElementInst(InsertElementInst &I) {
-  updateAnalysis(I.getOperand(2), BaseType::Integer, &I);
+  updateAnalysis(I.getOperand(2), TypeTree(BaseType::Integer).Only(-1), &I);
 
   auto &dl = fntypeinfo.Function->getParent()->getDataLayout();
   VectorType *vecType = cast<VectorType>(I.getOperand(0)->getType());

enzyme/test/Integration/eigentensorreal.cpp

@@ -5,9 +5,9 @@
 // TODO: %clang++ -I/usr/include/eigen3 -Xclang -new-struct-path-tbaa -fno-unroll-loops -fno-vectorize -fno-slp-vectorize -fno-exceptions -O1 -Xclang -disable-llvm-optzns %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S | %lli - 
 // RUN: %clang++ -I/usr/include/eigen3 -Xclang -new-struct-path-tbaa -fno-unroll-loops -fno-vectorize -fno-slp-vectorize -fno-exceptions -O3 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - 
 // RUN: %clang++ -I/usr/include/eigen3 -Xclang -new-struct-path-tbaa -fno-unroll-loops -fno-vectorize -fno-slp-vectorize -fno-exceptions -O2 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - 
-// RUN: %clang++ -I/usr/include/eigen3 -Xclang -new-struct-path-tbaa -fno-unroll-loops -fno-vectorize -fno-slp-vectorize -fno-exceptions -O1 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - 
+// RUN: if [ %llvmver -ge 8 ]; then %clang++ -I/usr/include/eigen3 -Xclang -new-struct-path-tbaa -fno-unroll-loops -fno-vectorize -fno-slp-vectorize -fno-exceptions -O1 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - ; fi
 //   note not doing O0 below as to ensure we get tbaa
-// RUN: %clang++ -I/usr/include/eigen3 -Xclang -new-struct-path-tbaa -fno-unroll-loops -fno-vectorize -fno-slp-vectorize -fno-exceptions -O1 -Xclang -disable-llvm-optzns %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - 
+// RUN: if [ %llvmver -ge 8 ]; then %clang++ -I/usr/include/eigen3 -Xclang -new-struct-path-tbaa -fno-unroll-loops -fno-vectorize -fno-slp-vectorize -fno-exceptions -O1 -Xclang -disable-llvm-optzns %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - ; fi
 
 #define EIGEN_NO_AUTOMATIC_RESIZING 1
 #define EIGEN_DONT_ALIGN 1

enzyme/test/Integration/multivecmax.cpp

@@ -2,9 +2,9 @@
 // RUN: %clang++ -ffast-math -O1 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S | %lli - 
 // RUN: %clang++ -ffast-math -O2 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S | %lli - 
 // RUN: %clang++ -ffast-math -O3 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S | %lli - 
-// RUN: %clang++ -ffast-math -O1 -Xclang -disable-llvm-optzns -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - 
-// RUN: %clang++ -ffast-math -O1 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - 
-// RUN: %clang++ -ffast-math -O2 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - 
+// RUN: if [ %llvmver -ge 8 ]; then %clang++ -ffast-math -O1 -Xclang -disable-llvm-optzns -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - ; fi
+// RUN: if [ %llvmver -ge 8 ]; then %clang++ -ffast-math -O1 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - ; fi
+// RUN: if [ %llvmver -ge 8 ]; then %clang++ -ffast-math -O2 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - ; fi
 // RUN: %clang++ -ffast-math -O3 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S | %lli - 
 
 #include <stdio.h>

enzyme/test/TypeAnalysis/sveci1.ll

@@ -11,7 +11,7 @@ entry:
 
 ; CHECK: callee - {} |{[-1]:Integer}:{}
 ; CHECK-NEXT: <4 x i1> %arg: {[-1]:Integer}
-; CHECK-NEXT: i64 %idx: {[]:Integer, [-1]:Integer}
+; CHECK-NEXT: i64 %idx: {[-1]:Integer}
 ; CHECK-NEXT: entry
 ; CHECK-NEXT:   %ai = alloca <4 x i1>, align 4: {[-1]:Pointer, [-1,0]:Integer}
 ; CHECK-NEXT:   %r = shufflevector <4 x i1> %arg, <4 x i1> undef, <4 x i32> zeroinitializer: {[-1]:Integer}

enzyme/test/TypeAnalysis/vecptr.ll

@@ -0,0 +1,26 @@
+; RUN: %opt < %s %loadEnzyme -print-type-analysis -type-analysis-func=callee -o /dev/null | FileCheck %s
+
+define void @callee(<2 x i64>* %argp) {
+entry:
+  %arg = load <2 x i64>, <2 x i64>* %argp, align 16
+  %bv = bitcast <2 x i64> %arg to i128
+  %tr = trunc i128 %bv to i64
+  %ptr = inttoptr i64 %tr to double*
+  %ld = load double, double* %ptr, align 8, !tbaa !8
+  ret void
+}
+
+!5 = !{!"omnipotent char", !6, i64 0}
+!6 = !{!"Simple C++ TBAA"}
+!7 = !{!"double", !5, i64 0}
+!8 = !{!7, !7, i64 0}
+
+; CHECK: callee - {} |{}:{} 
+; CHECK-NEXT: <2 x i64>* %argp: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,0]:Float@double}
+; CHECK-NEXT: entry
+; CHECK-NEXT:   %arg = load <2 x i64>, <2 x i64>* %argp, align 16: {[0]:Pointer, [0,0]:Float@double}
+; CHECK-NEXT:   %bv = bitcast <2 x i64> %arg to i128: {[0]:Pointer, [0,0]:Float@double}
+; CHECK-NEXT:   %tr = trunc i128 %bv to i64: {[-1]:Pointer, [-1,0]:Float@double}
+; CHECK-NEXT:   %ptr = inttoptr i64 %tr to double*: {[-1]:Pointer, [-1,0]:Float@double}
+; CHECK-NEXT:   %ld = load double, double* %ptr, align 8, !tbaa !0: {[-1]:Float@double}
+; CHECK-NEXT:   ret void: {}