Continued Type Analysis Updates (rust-lang#410)

* More non-strict type handling

* Merge global constants

Author: wsmoses

Diff for: enzyme/Enzyme/TypeAnalysis/ConcreteType.h

@@ -283,6 +283,13 @@ class ConcreteType {
   /// BaseType::Unknown
   bool operator&=(const ConcreteType CT) { return andIn(CT); }
 
+  /// Keep only mappings where the type is not an `Anything`
+  ConcreteType PurgeAnything() const {
+    if (SubTypeEnum == BaseType::Anything)
+      return BaseType::Unknown;
+    return *this;
+  }
+
   /// Set this to the logical `binop` of itself and RHS, using the Binop Op,
   /// returning true if this was changed.
   /// This function will error on an invalid type combination

Diff for: enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp

@@ -299,26 +299,30 @@ void getConstantAnalysis(Constant *Val, TypeAnalyzer &TA,
       int Off = (int)ai.getLimitedValue();
 
       getConstantAnalysis(Op, TA, analysis);
-      auto mid = analysis[Op].ShiftIndices(DL, /*init offset*/ 0,
-                                           /*maxSize*/ ObjSize,
-                                           /*addOffset*/ Off);
-
+      auto mid = analysis[Op];
       if (TA.fntypeinfo.Function->getParent()
               ->getDataLayout()
               .getTypeSizeInBits(CA->getType()) >= 16) {
         mid.ReplaceIntWithAnything();
       }
 
-      Result |= mid;
+      Result |= mid.ShiftIndices(DL, /*init offset*/ 0,
+                                 /*maxSize*/ ObjSize,
+                                 /*addOffset*/ Off);
     }
+    Result = Result.CanonicalizeValue(
+        (TA.fntypeinfo.Function->getParent()->getDataLayout().getTypeSizeInBits(
+             CA->getType()) +
+         7) /
+            8,
+        DL);
     return;
   }
 
   // Type of an sequence is the aggregation of
   // the subtypes
   if (auto CD = dyn_cast<ConstantDataSequential>(Val)) {
     TypeTree &Result = analysis[Val];
-
     for (unsigned i = 0, size = CD->getNumElements(); i < size; ++i) {
       assert(TA.fntypeinfo.Function);
       auto Op = CD->getElementAsConstant(i);
@@ -349,18 +353,24 @@ void getConstantAnalysis(Constant *Val, TypeAnalyzer &TA,
       int Off = (int)ai.getLimitedValue();
 
       getConstantAnalysis(Op, TA, analysis);
-      auto mid = analysis[Op].ShiftIndices(DL, /*init offset*/ 0,
-                                           /*maxSize*/ ObjSize,
-                                           /*addOffset*/ Off);
-
+      auto mid = analysis[Op];
       if (TA.fntypeinfo.Function->getParent()
               ->getDataLayout()
               .getTypeSizeInBits(CD->getType()) >= 16) {
         mid.ReplaceIntWithAnything();
       }
+      Result |= mid.ShiftIndices(DL, /*init offset*/ 0,
+                                 /*maxSize*/ ObjSize,
+                                 /*addOffset*/ Off);
 
       Result |= mid;
     }
+    Result = Result.CanonicalizeValue(
+        (TA.fntypeinfo.Function->getParent()->getDataLayout().getTypeSizeInBits(
+             CD->getType()) +
+         7) /
+            8,
+        DL);
     return;
   }
 
@@ -849,12 +859,12 @@ void TypeAnalyzer::considerTBAA() {
             (DL.getTypeSizeInBits(SI->getValueOperand()->getType()) + 7) / 8;
         updateAnalysis(SI->getPointerOperand(),
                        vdptr
+                           // Don't propagate "Anything" into ptr
+                           .PurgeAnything()
                            // Cut off any values outside of store
                            .ShiftIndices(DL, /*init offset*/ 0,
                                          /*max size*/ StoreSize,
                                          /*new offset*/ 0)
-                           // Don't propagate "Anything" into ptr
-                           .PurgeAnything()
                            .Only(-1),
                        SI);
         TypeTree req = vdptr.Only(-1);
@@ -863,12 +873,12 @@ void TypeAnalyzer::considerTBAA() {
         auto LoadSize = (DL.getTypeSizeInBits(LI->getType()) + 7) / 8;
         updateAnalysis(LI->getPointerOperand(),
                        vdptr
+                           // Don't propagate "Anything" into ptr
+                           .PurgeAnything()
                            // Cut off any values outside of load
                            .ShiftIndices(DL, /*init offset*/ 0,
                                          /*max size*/ LoadSize,
                                          /*new offset*/ 0)
-                           // Don't propagate "Anything" into ptr
-                           .PurgeAnything()
                            .Only(-1),
                        LI);
         TypeTree req = vdptr.Only(-1);
@@ -1125,12 +1135,12 @@ void TypeAnalyzer::visitCmpInst(CmpInst &cmp) {
   if (direction & UP) {
     updateAnalysis(
         cmp.getOperand(0),
-        TypeTree(getAnalysis(cmp.getOperand(1)).Data0().PurgeAnything()[{}])
+        TypeTree(getAnalysis(cmp.getOperand(1)).Inner0().PurgeAnything())
             .Only(-1),
         &cmp);
     updateAnalysis(
         cmp.getOperand(1),
-        TypeTree(getAnalysis(cmp.getOperand(0)).Data0().PurgeAnything()[{}])
+        TypeTree(getAnalysis(cmp.getOperand(0)).Inner0().PurgeAnything())
             .Only(-1),
         &cmp);
   }
@@ -1194,8 +1204,9 @@ void TypeAnalyzer::visitStoreInst(StoreInst &I) {
   // Only propagate mappings in range that aren't "Anything" into the pointer
   auto ptr = TypeTree(BaseType::Pointer);
   auto purged = getAnalysis(I.getValueOperand())
+                    .PurgeAnything()
                     .ShiftIndices(DL, /*start*/ 0, StoreSize, /*addOffset*/ 0)
-                    .PurgeAnything();
+                    .ReplaceMinus();
   ptr |= purged;
 
   if (direction & UP) {
@@ -2460,13 +2471,13 @@ void TypeAnalyzer::visitMemTransferCommon(llvm::CallInst &MTI) {
 
   auto &dl = MTI.getParent()->getParent()->getParent()->getDataLayout();
   TypeTree res = getAnalysis(MTI.getArgOperand(0))
+                     .PurgeAnything()
                      .Data0()
-                     .ShiftIndices(dl, 0, sz, 0)
-                     .PurgeAnything();
+                     .ShiftIndices(dl, 0, sz, 0);
   TypeTree res2 = getAnalysis(MTI.getArgOperand(1))
+                      .PurgeAnything()
                       .Data0()
-                      .ShiftIndices(dl, 0, sz, 0)
-                      .PurgeAnything();
+                      .ShiftIndices(dl, 0, sz, 0);
 
   bool Legal = true;
   res.checkedOrIn(res2, /*PointerIntSame*/ false, Legal);
@@ -3772,11 +3783,11 @@ void TypeAnalyzer::visitCallInst(CallInst &call) {
 
       auto &dl = call.getParent()->getParent()->getParent()->getDataLayout();
       TypeTree res = getAnalysis(call.getArgOperand(0))
+                         .PurgeAnything()
                          .Data0()
-                         .ShiftIndices(dl, 0, sz, 0)
-                         .PurgeAnything();
+                         .ShiftIndices(dl, 0, sz, 0);
       TypeTree res2 =
-          getAnalysis(&call).Data0().ShiftIndices(dl, 0, sz, 0).PurgeAnything();
+          getAnalysis(&call).PurgeAnything().Data0().ShiftIndices(dl, 0, sz, 0);
 
       res.orIn(res2, /*PointerIntSame*/ false);
       res.insert({}, BaseType::Pointer);
@@ -4524,6 +4535,9 @@ void TypeAnalyzer::visitIPOCall(CallInst &call, Function &fn) {
 
   TypeResults STR = interprocedural.analyzeFunction(typeInfo);
 
+  if (EnzymePrintType)
+    llvm::errs() << " ending IPO of " << call << "\n";
+
   if (hasUp) {
     auto a = fn.arg_begin();
 #if LLVM_VERSION_MAJOR >= 14
@@ -4533,6 +4547,9 @@ void TypeAnalyzer::visitIPOCall(CallInst &call, Function &fn) {
 #endif
     {
       auto dt = STR.query(a);
+      if (EnzymePrintType)
+        llvm::errs() << " updating " << *arg << " = " << dt.str()
+                     << "  via IPO of " << call << " arg " << *a << "\n";
       updateAnalysis(arg, dt, &call);
       ++a;
     }

Diff for: enzyme/Enzyme/TypeAnalysis/TypeTree.h

@@ -826,6 +826,21 @@ class TypeTree : public std::enable_shared_from_this<TypeTree> {
     return dat;
   }
 
+  /// Replace -1 with 0
+  TypeTree ReplaceMinus() const {
+    TypeTree dat;
+    for (const auto pair : mapping) {
+      if (pair.second == ConcreteType(BaseType::Anything))
+        continue;
+      std::vector<int> nex = pair.first;
+      for (auto &v : nex)
+        if (v == -1)
+          v = 0;
+      dat.insert(nex, pair.second);
+    }
+    return dat;
+  }
+
   /// Replace all integer subtypes with anything
   void ReplaceIntWithAnything() {
     for (auto &pair : mapping) {

Diff for: enzyme/test/TypeAnalysis/loadptrptr.ll

@@ -23,11 +23,11 @@ entry:
 !5 = !{!"Simple C++ TBAA"}
 
 ; CHECK: caller - {} |{}:{}
-; CHECK-NEXT: i64* %q: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Integer}
+; CHECK-NEXT: i64* %q: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,0]:Integer, [-1,0,1]:Integer, [-1,0,2]:Integer, [-1,0,3]:Integer, [-1,0,4]:Integer, [-1,0,5]:Integer, [-1,0,6]:Integer, [-1,0,7]:Integer}
 ; CHECK-NEXT: entry
-; CHECK-NEXT:   %p = bitcast i64* %q to i64**: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Integer}
+; CHECK-NEXT:   %p = bitcast i64* %q to i64**: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,0]:Integer, [-1,0,1]:Integer, [-1,0,2]:Integer, [-1,0,3]:Integer, [-1,0,4]:Integer, [-1,0,5]:Integer, [-1,0,6]:Integer, [-1,0,7]:Integer}
 ; CHECK-NEXT:   %x = alloca i64, align 8: {[-1]:Pointer, [-1,-1]:Integer}
 ; CHECK-NEXT:   store i64 132, i64* %x, align 8: {}
 ; CHECK-NEXT:   store i64* %x, i64** %p, align 8: {}
-; CHECK-NEXT:   %ld = load i64*, i64** %p, align 8: {[-1]:Pointer, [-1,-1]:Integer}
+; CHECK-NEXT:   %ld = load i64*, i64** %p, align 8: {[-1]:Pointer, [-1,0]:Integer, [-1,1]:Integer, [-1,2]:Integer, [-1,3]:Integer, [-1,4]:Integer, [-1,5]:Integer, [-1,6]:Integer, [-1,7]:Integer}
 ; CHECK-NEXT:   ret void: {}

Diff for: enzyme/test/TypeAnalysis/mergeglob.ll

@@ -0,0 +1,14 @@
+; RUN: %opt < %s %loadEnzyme -print-type-analysis -type-analysis-func=callee -o /dev/null | FileCheck %s
+
+@ptr = internal constant [4 x double] [double 3.14000e+00, double 3.14000e+00, double 3.14000e+00, double 3.14000e+00], align 8
+
+define void @callee() {
+entry:
+  %self = getelementptr inbounds [4 x double], [4 x double]* @ptr, i32 0, i32 0
+  ret void
+}
+
+; CHECK: callee - {} |
+; CHECK-NEXT: entry
+; CHECK-NEXT:   %self = getelementptr inbounds [4 x double], [4 x double]* @ptr, i32 0, i32 0: {[-1]:Pointer, [-1,-1]:Float@double}
+; CHECK-NEXT:   ret void: {}

Diff for: enzyme/test/TypeAnalysis/transf.ll

@@ -0,0 +1,52 @@
+; RUN: %opt < %s %loadEnzyme -print-type-analysis -type-analysis-func=caller -o /dev/null | FileCheck %s
+
+define void @caller(double* %p) {
+entry:
+  %q = alloca [2 x double], align 8
+  %gep0 = getelementptr [2 x double], [2 x double]* %q, i32 0, i32 0
+  %gep = getelementptr [2 x double], [2 x double]* %q, i32 0, i32 1
+  store double 0.000000e+00, double* %gep, align 8, !tbaa !2
+  %o = call double* @max(double* %p, double* %gep0, i1 true)
+  ret void
+}
+
+define double* @max(double* %a, double* %b, i1 %cmp) {
+entry:
+  %retval = alloca double*, align 8
+  %av = load double, double* %a, align 8
+  %bv = load double, double* %b, align 8
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  store double* %a, double** %retval, align 8
+  br label %return
+
+if.end:
+  store double* %b, double** %retval, align 8
+  br label %return
+
+return:
+  %res = load double*, double** %retval, align 8
+  ret double* %res
+}
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 7.1.0 "}
+!2 = !{!3, !3, i64 0, i64 8}
+!3 = !{!4, i64 8, !"long"}
+!4 = !{!5, i64 1, !"omnipotent char"}
+!5 = !{!"Simple C++ TBAA"}
+
+
+; CHECK: caller - {} |{[-1]:Pointer, [-1,-1]:Float@double}:{}
+; CHECK-NEXT: double* %p: {[-1]:Pointer, [-1,-1]:Float@double}
+; CHECK-NEXT: entry
+; CHECK-NEXT:   %q = alloca [2 x double], align 8: {[-1]:Pointer, [-1,0]:Float@double, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer}
+; CHECK-NEXT:   %gep0 = getelementptr [2 x double], [2 x double]* %q, i32 0, i32 0: {[-1]:Pointer, [-1,0]:Float@double}
+; CHECK-NEXT:   %gep = getelementptr [2 x double], [2 x double]* %q, i32 0, i32 1: {[-1]:Pointer, [-1,0]:Integer, [-1,1]:Integer, [-1,2]:Integer, [-1,3]:Integer, [-1,4]:Integer, [-1,5]:Integer, [-1,6]:Integer, [-1,7]:Integer}
+; CHECK-NEXT:   store double 0.000000e+00, double* %gep, align 8, !tbaa !2: {}
+; CHECK-NEXT:   %o = call double* @max(double* %p, double* %gep0, i1 true): {[-1]:Pointer, [-1,0]:Float@double}
+; CHECK-NEXT:   ret void: {}