Fix insert choice (rust-lang#988)

* fix insert choice

* fix tests

Author: tgymnich

enzyme/Enzyme/TraceUtils.h

@@ -489,36 +489,35 @@ class TraceUtils {
 
   CallInst *InsertChoice(IRBuilder<> &Builder, Value *address, Value *score,
                          Value *choice) {
-    auto size = choice->getType()->getPrimitiveSizeInBits() / 8;
     Type *size_type = interface->getChoiceTy()->getParamType(3);
-
-    auto M = interface->getSampleFunction()->getParent();
-    auto &DL = M->getDataLayout();
-    auto pointersize = DL.getPointerSizeInBits();
+    auto choicesize = choice->getType()->getPrimitiveSizeInBits();
 
     Value *retval;
     if (choice->getType()->isPointerTy()) {
       retval = Builder.CreatePointerCast(choice, Builder.getInt8PtrTy());
     } else {
-      IRBuilder<> AllocaBuilder(
-          newFunc->getEntryBlock().getFirstNonPHIOrDbgOrLifetime());
-      auto alloca = AllocaBuilder.CreateAlloca(choice->getType(), nullptr,
-                                               choice->getName() + ".ptr");
-      Builder.CreateStore(choice, alloca);
-      bool fitsInPointer =
-          choice->getType()->getPrimitiveSizeInBits() == pointersize;
-      if (fitsInPointer) {
-        auto dblptr =
-            PointerType::get(Builder.getInt8PtrTy(), DL.getAllocaAddrSpace());
-        retval = Builder.CreateLoad(Builder.getInt8PtrTy(),
-                                    Builder.CreatePointerCast(alloca, dblptr));
+      auto M = interface->getSampleFunction()->getParent();
+      auto &DL = M->getDataLayout();
+      auto pointersize = DL.getPointerSizeInBits();
+      if (choicesize <= pointersize) {
+        auto cast = Builder.CreateBitCast(
+            choice, IntegerType::get(M->getContext(), choicesize));
+        cast = choicesize == pointersize
+                   ? cast
+                   : Builder.CreateZExt(cast, Builder.getIntPtrTy(DL));
+        retval = Builder.CreateIntToPtr(cast, Builder.getInt8PtrTy());
       } else {
+        IRBuilder<> AllocaBuilder(
+            newFunc->getEntryBlock().getFirstNonPHIOrDbgOrLifetime());
+        auto alloca = AllocaBuilder.CreateAlloca(choice->getType(), nullptr,
+                                                 choice->getName() + ".ptr");
+        Builder.CreateStore(choice, alloca);
         retval = alloca;
       }
     }
 
     Value *args[] = {trace, address, score, retval,
-                     ConstantInt::get(size_type, size)};
+                     ConstantInt::get(size_type, choicesize / 8)};
 
     auto call = Builder.CreateCall(interface->insertChoiceTy(),
                                    interface->insertChoice(), args);

enzyme/test/Enzyme/ProbProg/condition-dynamic.ll

@@ -86,12 +86,10 @@ entry:
 ; CHECK-NEXT:   %4 = getelementptr inbounds i8*, i8** %interface, i32 6
 ; CHECK-NEXT:   %5 = load i8*, i8** %4
 ; CHECK-NEXT:   %has_call = bitcast i8* %5 to i1 (i8*, i8*)*
-; CHECK-NEXT:   %call1.ptr3 = alloca double
 ; CHECK-NEXT:   %call1.ptr = alloca double
 ; CHECK-NEXT:   %6 = getelementptr inbounds i8*, i8** %interface, i32 3
 ; CHECK-NEXT:   %7 = load i8*, i8** %6
 ; CHECK-NEXT:   %insert_choice = bitcast i8* %7 to void (i8*, i8*, double, i8*, i64)*
-; CHECK-NEXT:   %call.ptr2 = alloca double
 ; CHECK-NEXT:   %8 = getelementptr inbounds i8*, i8** %interface, i32 1
 ; CHECK-NEXT:   %9 = load i8*, i8** %8
 ; CHECK-NEXT:   %get_choice = bitcast i8* %9 to i64 (i8*, i8*, i8*, i64)*
@@ -119,9 +117,8 @@ entry:
 ; CHECK: entry.cntd:                                       ; preds = %condition.call.without.trace, %condition.call.with.trace
 ; CHECK-NEXT:   %call = phi double [ %from.trace.call, %condition.call.with.trace ], [ %sample.call, %condition.call.without.trace ]
 ; CHECK-NEXT:   %likelihood.call = call double @normal_logpdf(double 0.000000e+00, double 1.000000e+00, double %call)
-; CHECK-NEXT:   store double %call, double* %call.ptr2
-; CHECK-NEXT:   %15 = bitcast double* %call.ptr2 to i8**
-; CHECK-NEXT:   %16 = load i8*, i8** %15
+; CHECK-NEXT:   %15 = bitcast double %call to i64
+; CHECK-NEXT:   %16 = inttoptr i64 %15 to i8*
 ; CHECK-NEXT:   call void %insert_choice(i8* %trace, i8* nocapture readonly getelementptr inbounds ([2 x i8], [2 x i8]* @.str.1, i64 0, i64 0), double %likelihood.call, i8* %16, i64 8)
 ; CHECK-NEXT:   %has.choice.call1 = call i1 %has_choice(i8* %observations, i8* nocapture readonly getelementptr inbounds ([2 x i8], [2 x i8]* @.str.2, i64 0, i64 0))
 ; CHECK-NEXT:   br i1 %has.choice.call1, label %condition.call1.with.trace, label %condition.call1.without.trace
@@ -139,9 +136,8 @@ entry:
 ; CHECK: entry.cntd.cntd:                                  ; preds = %condition.call1.without.trace, %condition.call1.with.trace
 ; CHECK-NEXT:   %call1 = phi double [ %from.trace.call1, %condition.call1.with.trace ], [ %sample.call1, %condition.call1.without.trace ]
 ; CHECK-NEXT:   %likelihood.call1 = call double @normal_logpdf(double 0.000000e+00, double 1.000000e+00, double %call1)
-; CHECK-NEXT:   store double %call1, double* %call1.ptr3
-; CHECK-NEXT:   %18 = bitcast double* %call1.ptr3 to i8**
-; CHECK-NEXT:   %19 = load i8*, i8** %18
+; CHECK-NEXT:   %18 = bitcast double %call1 to i64
+; CHECK-NEXT:   %19 = inttoptr i64 %18 to i8*
 ; CHECK-NEXT:   call void %insert_choice(i8* %trace, i8* nocapture readonly getelementptr inbounds ([2 x i8], [2 x i8]* @.str.2, i64 0, i64 0), double %likelihood.call1, i8* %19, i64 8)
 ; CHECK-NEXT:   %has.call.call2 = call i1 %has_call(i8* %observations, i8* nocapture readonly getelementptr inbounds ([21 x i8], [21 x i8]* @0, i32 0, i32 0))
 ; CHECK-NEXT:   br i1 %has.call.call2, label %condition.call2.with.trace, label %condition.call2.without.trace
@@ -156,9 +152,9 @@ entry:
 ; CHECK-NEXT:   br label %entry.cntd.cntd.cntd
 
 ; CHECK: entry.cntd.cntd.cntd:                             ; preds = %condition.call2.without.trace, %condition.call2.with.trace
-; CHECK-NEXT:   %call24 = phi { double, i8* } [ %condition.calculate_loss, %condition.call2.with.trace ], [ %trace.calculate_loss, %condition.call2.without.trace ]
-; CHECK-NEXT:   %call2 = extractvalue { double, i8* } %call24, 0
-; CHECK-NEXT:   %newtrace.calculate_loss = extractvalue { double, i8* } %call24, 1
+; CHECK-NEXT:   %call22 = phi { double, i8* } [ %condition.calculate_loss, %condition.call2.with.trace ], [ %trace.calculate_loss, %condition.call2.without.trace ]
+; CHECK-NEXT:   %call2 = extractvalue { double, i8* } %call22, 0
+; CHECK-NEXT:   %newtrace.calculate_loss = extractvalue { double, i8* } %call22, 1
 ; CHECK-NEXT:   call void %insert_call(i8* %trace, i8* nocapture readonly getelementptr inbounds ([21 x i8], [21 x i8]* @0, i32 0, i32 0), i8* %newtrace.calculate_loss)
 ; CHECK-NEXT:   %mrv = insertvalue { double, i8* } {{(undef|poison)}}, double %call2, 0
 ; CHECK-NEXT:   %mrv1 = insertvalue { double, i8* } %mrv, i8* %trace, 1
@@ -171,7 +167,6 @@ entry:
 ; CHECK-NEXT:   %0 = getelementptr inbounds i8*, i8** %interface, i32 3
 ; CHECK-NEXT:   %1 = load i8*, i8** %0
 ; CHECK-NEXT:   %insert_choice = bitcast i8* %1 to void (i8*, i8*, double, i8*, i64)*
-; CHECK-NEXT:   %call.ptr2 = alloca double
 ; CHECK-NEXT:   %2 = getelementptr inbounds i8*, i8** %interface, i32 1
 ; CHECK-NEXT:   %3 = load i8*, i8** %2
 ; CHECK-NEXT:   %get_choice = bitcast i8* %3 to i64 (i8*, i8*, i8*, i64)*
@@ -219,9 +214,8 @@ entry:
 ; CHECK: for.body.cntd:                                    ; preds = %condition.call.without.trace, %condition.call.with.trace
 ; CHECK-NEXT:   %call = phi double [ %from.trace.call, %condition.call.with.trace ], [ %sample.call, %condition.call.without.trace ]
 ; CHECK-NEXT:   %likelihood.call = call double @normal_logpdf(double %9, double 1.000000e+00, double %call)
-; CHECK-NEXT:   store double %call, double* %call.ptr2
-; CHECK-NEXT:   %11 = bitcast double* %call.ptr2 to i8**
-; CHECK-NEXT:   %12 = load i8*, i8** %11
+; CHECK-NEXT:   %11 = bitcast double %call to i64
+; CHECK-NEXT:   %12 = inttoptr i64 %11 to i8*
 ; CHECK-NEXT:   call void %insert_choice(i8* %trace, i8* nocapture readonly getelementptr inbounds ([11 x i8], [11 x i8]* @.str, i64 0, i64 0), double %likelihood.call, i8* %12, i64 8)
 ; CHECK-NEXT:   %arrayidx3 = getelementptr inbounds double, double* %data, i64 %indvars.iv
 ; CHECK-NEXT:   %13 = load double, double* %arrayidx3

enzyme/test/Enzyme/ProbProg/condition-static.ll

@@ -80,9 +80,7 @@ entry:
 
 ; CHECK: define internal { double, i8* } @condition_loss(double* %data, i32 %n, i8* %observations)
 ; CHECK-NEXT: entry:
-; CHECK-NEXT:   %call1.ptr3 = alloca double
 ; CHECK-NEXT:   %call1.ptr = alloca double
-; CHECK-NEXT:   %call.ptr2 = alloca double
 ; CHECK-NEXT:   %call.ptr = alloca double
 ; CHECK-NEXT:   %trace = call i8* @__enzyme_newtrace()
 ; CHECK-NEXT:   %has.choice.call = call i1 @__enzyme_has_choice(i8* %observations, i8* nocapture readonly getelementptr inbounds ([2 x i8], [2 x i8]* @.str.1, i64 0, i64 0))
@@ -101,9 +99,8 @@ entry:
 ; CHECK: entry.cntd:                                       ; preds = %condition.call.without.trace, %condition.call.with.trace
 ; CHECK-NEXT:   %call = phi double [ %from.trace.call, %condition.call.with.trace ], [ %sample.call, %condition.call.without.trace ]
 ; CHECK-NEXT:   %likelihood.call = call double @normal_logpdf(double 0.000000e+00, double 1.000000e+00, double %call)
-; CHECK-NEXT:   store double %call, double* %call.ptr2
-; CHECK-NEXT:   %1 = bitcast double* %call.ptr2 to i8**
-; CHECK-NEXT:   %2 = load i8*, i8** %1
+; CHECK-NEXT:   %1 = bitcast double %call to i64
+; CHECK-NEXT:   %2 = inttoptr i64 %1 to i8*
 ; CHECK-NEXT:   call void @__enzyme_insert_choice(i8* %trace, i8* nocapture readonly getelementptr inbounds ([2 x i8], [2 x i8]* @.str.1, i64 0, i64 0), double %likelihood.call, i8* %2, i64 8)
 ; CHECK-NEXT:   %has.choice.call1 = call i1 @__enzyme_has_choice(i8* %observations, i8* nocapture readonly getelementptr inbounds ([2 x i8], [2 x i8]* @.str.2, i64 0, i64 0))
 ; CHECK-NEXT:   br i1 %has.choice.call1, label %condition.call1.with.trace, label %condition.call1.without.trace
@@ -121,9 +118,8 @@ entry:
 ; CHECK: entry.cntd.cntd:                                  ; preds = %condition.call1.without.trace, %condition.call1.with.trace
 ; CHECK-NEXT:   %call1 = phi double [ %from.trace.call1, %condition.call1.with.trace ], [ %sample.call1, %condition.call1.without.trace ]
 ; CHECK-NEXT:   %likelihood.call1 = call double @normal_logpdf(double 0.000000e+00, double 1.000000e+00, double %call1)
-; CHECK-NEXT:   store double %call1, double* %call1.ptr3
-; CHECK-NEXT:   %4 = bitcast double* %call1.ptr3 to i8**
-; CHECK-NEXT:   %5 = load i8*, i8** %4
+; CHECK-NEXT:   %4 = bitcast double %call1 to i64
+; CHECK-NEXT:   %5 = inttoptr i64 %4 to i8*
 ; CHECK-NEXT:   call void @__enzyme_insert_choice(i8* %trace, i8* nocapture readonly getelementptr inbounds ([2 x i8], [2 x i8]* @.str.2, i64 0, i64 0), double %likelihood.call1, i8* %5, i64 8)
 ; CHECK-NEXT:   %has.call.call2 = call i1 @__enzyme_has_call(i8* %observations, i8* nocapture readonly getelementptr inbounds ([21 x i8], [21 x i8]* @0, i32 0, i32 0))
 ; CHECK-NEXT:   br i1 %has.call.call2, label %condition.call2.with.trace, label %condition.call2.without.trace
@@ -138,9 +134,9 @@ entry:
 ; CHECK-NEXT:   br label %entry.cntd.cntd.cntd
 
 ; CHECK: entry.cntd.cntd.cntd:                             ; preds = %condition.call2.without.trace, %condition.call2.with.trace
-; CHECK-NEXT:   %call24 = phi { double, i8* } [ %condition.calculate_loss, %condition.call2.with.trace ], [ %trace.calculate_loss, %condition.call2.without.trace ]
-; CHECK-NEXT:   %call2 = extractvalue { double, i8* } %call24, 0
-; CHECK-NEXT:   %newtrace.calculate_loss = extractvalue { double, i8* } %call24, 1
+; CHECK-NEXT:   %call22 = phi { double, i8* } [ %condition.calculate_loss, %condition.call2.with.trace ], [ %trace.calculate_loss, %condition.call2.without.trace ]
+; CHECK-NEXT:   %call2 = extractvalue { double, i8* } %call22, 0
+; CHECK-NEXT:   %newtrace.calculate_loss = extractvalue { double, i8* } %call22, 1
 ; CHECK-NEXT:   call void @__enzyme_insert_call(i8* %trace, i8* nocapture readonly getelementptr inbounds ([21 x i8], [21 x i8]* @0, i32 0, i32 0), i8* %newtrace.calculate_loss)
 ; CHECK-NEXT:   %mrv = insertvalue { double, i8* } {{(undef|poison)}}, double %call2, 0
 ; CHECK-NEXT:   %mrv1 = insertvalue { double, i8* } %mrv, i8* %trace, 1
@@ -150,7 +146,6 @@ entry:
 
 ; CHECK: define internal { double, i8* } @condition_calculate_loss(double %m, double %b, double* %data, i32 %n, i8* %observations)
 ; CHECK-NEXT: entry:
-; CHECK-NEXT:   %call.ptr2 = alloca double
 ; CHECK-NEXT:   %call.ptr = alloca double
 ; CHECK-NEXT:   %trace = call i8* @__enzyme_newtrace()
 ; CHECK-NEXT:   %cmp19 = icmp sgt i32 %n, 0
@@ -189,9 +184,8 @@ entry:
 ; CHECK: for.body.cntd:                                    ; preds = %condition.call.without.trace, %condition.call.with.trace
 ; CHECK-NEXT:   %call = phi double [ %from.trace.call, %condition.call.with.trace ], [ %sample.call, %condition.call.without.trace ]
 ; CHECK-NEXT:   %likelihood.call = call double @normal_logpdf(double %1, double 1.000000e+00, double %call)
-; CHECK-NEXT:   store double %call, double* %call.ptr2
-; CHECK-NEXT:   %3 = bitcast double* %call.ptr2 to i8**
-; CHECK-NEXT:   %4 = load i8*, i8** %3
+; CHECK-NEXT:   %3 = bitcast double %call to i64
+; CHECK-NEXT:   %4 = inttoptr i64 %3 to i8*
 ; CHECK-NEXT:   call void @__enzyme_insert_choice(i8* %trace, i8* nocapture readonly getelementptr inbounds ([11 x i8], [11 x i8]* @.str, i64 0, i64 0), double %likelihood.call, i8* %4, i64 8)
 ; CHECK-NEXT:   %arrayidx3 = getelementptr inbounds double, double* %data, i64 %indvars.iv
 ; CHECK-NEXT:   %5 = load double, double* %arrayidx3

enzyme/test/Enzyme/ProbProg/simple-condition.ll

@@ -51,9 +51,7 @@ entry:
 
 ; CHECK: define internal i8* @condition_test(i8* %observations)
 ; CHECK-NEXT: entry:
-; CHECK-NEXT:   %x.ptr2 = alloca double
 ; CHECK-NEXT:   %x.ptr = alloca double
-; CHECK-NEXT:   %mu.ptr1 = alloca double
 ; CHECK-NEXT:   %mu.ptr = alloca double
 ; CHECK-NEXT:   %trace = call i8* @__enzyme_newtrace()
 ; CHECK-NEXT:   %has.choice.mu = call i1 @__enzyme_has_choice(i8* %observations, i8* nocapture readonly getelementptr inbounds ([3 x i8], [3 x i8]* @.str, i64 0, i64 0))
@@ -72,9 +70,8 @@ entry:
 ; CHECK: entry.cntd:                                       ; preds = %condition.mu.without.trace, %condition.mu.with.trace
 ; CHECK-NEXT:   %mu = phi double [ %from.trace.mu, %condition.mu.with.trace ], [ %sample.mu, %condition.mu.without.trace ]
 ; CHECK-NEXT:   %likelihood.mu = call double @normal_logpdf(double 0.000000e+00, double 1.000000e+00, double %mu)
-; CHECK-NEXT:   store double %mu, double* %mu.ptr1
-; CHECK-NEXT:   %1 = bitcast double* %mu.ptr1 to i8**
-; CHECK-NEXT:   %2 = load i8*, i8** %1
+; CHECK-NEXT:   %1 = bitcast double %mu to i64
+; CHECK-NEXT:   %2 = inttoptr i64 %1 to i8*
 ; CHECK-NEXT:   call void @__enzyme_insert_choice(i8* %trace, i8* nocapture readonly getelementptr inbounds ([3 x i8], [3 x i8]* @.str, i64 0, i64 0), double %likelihood.mu, i8* %2, i64 8)
 ; CHECK-NEXT:   %has.choice.x = call i1 @__enzyme_has_choice(i8* %observations, i8* nocapture readonly getelementptr inbounds ([2 x i8], [2 x i8]* @.str.1, i64 0, i64 0))
 ; CHECK-NEXT:   br i1 %has.choice.x, label %condition.x.with.trace, label %condition.x.without.trace
@@ -92,9 +89,8 @@ entry:
 ; CHECK: entry.cntd.cntd:                                  ; preds = %condition.x.without.trace, %condition.x.with.trace
 ; CHECK-NEXT:   %x = phi double [ %from.trace.x, %condition.x.with.trace ], [ %sample.x, %condition.x.without.trace ]
 ; CHECK-NEXT:   %likelihood.x = call double @normal_logpdf(double %mu, double 1.000000e+00, double %x)
-; CHECK-NEXT:   store double %x, double* %x.ptr2
-; CHECK-NEXT:   %4 = bitcast double* %x.ptr2 to i8**
-; CHECK-NEXT:   %5 = load i8*, i8** %4
+; CHECK-NEXT:   %4 = bitcast double %x to i64
+; CHECK-NEXT:   %5 = inttoptr i64 %4 to i8*
 ; CHECK-NEXT:   call void @__enzyme_insert_choice(i8* %trace, i8* nocapture readonly getelementptr inbounds ([2 x i8], [2 x i8]* @.str.1, i64 0, i64 0), double %likelihood.x, i8* %5, i64 8)
 ; CHECK-NEXT:   ret i8* %trace
 ; CHECK-NEXT: }

enzyme/test/Enzyme/ProbProg/simple-trace.ll

@@ -40,20 +40,16 @@ entry:
 
 ; CHECK: define internal i8* @trace_test()
 ; CHECK-NEXT: entry:
-; CHECK-NEXT:   %x.ptr = alloca double
-; CHECK-NEXT:   %mu.ptr = alloca double
 ; CHECK-NEXT:   %trace = call i8* @__enzyme_newtrace()
 ; CHECK-NEXT:   %mu = call double @normal(double 0.000000e+00, double 1.000000e+00)
 ; CHECK-NEXT:   %likelihood.mu = call double @normal_logpdf(double 0.000000e+00, double 1.000000e+00, double %mu)
-; CHECK-NEXT:   store double %mu, double* %mu.ptr
-; CHECK-NEXT:   %0 = bitcast double* %mu.ptr to i8**
-; CHECK-NEXT:   %1 = load i8*, i8** %0
+; CHECK-NEXT:   %0 = bitcast double %mu to i64
+; CHECK-NEXT:   %1 = inttoptr i64 %0 to i8*
 ; CHECK-NEXT:   call void @__enzyme_insert_choice(i8* %trace, i8* nocapture readonly getelementptr inbounds ([3 x i8], [3 x i8]* @.str, i64 0, i64 0), double %likelihood.mu, i8* %1, i64 8)
 ; CHECK-NEXT:   %x = call double @normal(double %mu, double 1.000000e+00)
 ; CHECK-NEXT:   %likelihood.x = call double @normal_logpdf(double %mu, double 1.000000e+00, double %x)
-; CHECK-NEXT:   store double %x, double* %x.ptr
-; CHECK-NEXT:   %2 = bitcast double* %x.ptr to i8**
-; CHECK-NEXT:   %3 = load i8*, i8** %2
+; CHECK-NEXT:   %2 = bitcast double %x to i64
+; CHECK-NEXT:   %3 = inttoptr i64 %2 to i8*
 ; CHECK-NEXT:   call void @__enzyme_insert_choice(i8* %trace, i8* nocapture readonly getelementptr inbounds ([2 x i8], [2 x i8]* @.str.1, i64 0, i64 0), double %likelihood.x, i8* %3, i64 8)
 ; CHECK-NEXT:   ret i8* %trace
 ; CHECK-NEXT: }