Fix gradient struct return on ARM (rust-lang#378)

* Fix gradient struct return of ARM where the type of CI is a struct

* add integration test

* add unit test

Author: tgymnich

enzyme/Enzyme/Enzyme.cpp

@@ -862,13 +862,23 @@ class Enzyme : public ModulePass {
         }
       }
     }
+    StructType *CIsty = dyn_cast<StructType>(CI->getType());
+    StructType *diffretsty = dyn_cast<StructType>(diffret->getType());
 
     if (!diffret->getType()->isEmptyTy() && !diffret->getType()->isVoidTy() &&
         !CI->getType()->isEmptyTy() &&
         (!CI->getType()->isVoidTy() ||
          CI->paramHasAttr(0, Attribute::StructRet))) {
       if (diffret->getType() == CI->getType()) {
         CI->replaceAllUsesWith(diffret);
+      } else if (CIsty && diffretsty && CIsty->isLayoutIdentical(diffretsty)) {
+        IRBuilder<> Builder(CI);
+        Value *newStruct = UndefValue::get(CIsty);
+        for (unsigned int i = 0; i < CIsty->getStructNumElements(); i++) {
+          Value *elem = Builder.CreateExtractValue(diffret, {i});
+          newStruct = Builder.CreateInsertValue(newStruct, elem, {i});
+        }
+        CI->replaceAllUsesWith(newStruct);
       } else if (mode == DerivativeMode::ReverseModePrimal) {
         auto &DL = cast<Function>(fn)->getParent()->getDataLayout();
         if (DL.getTypeSizeInBits(CI->getType()) >=

enzyme/test/Enzyme/ReverseMode/gradient-struct-ret.ll

@@ -0,0 +1,30 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -early-cse -instsimplify -simplifycfg -S | FileCheck %s
+
+%struct.Gradients = type { double, double }
+
+define dso_local double @muldd(double %x, double %y) #0 {
+entry:
+  %mul = fmul double %x, %y
+  ret double %mul
+}
+
+define dso_local %struct.Gradients @dmuldd(double %x, double %y) local_unnamed_addr #1 {
+entry:
+  %call = call %struct.Gradients (i8*, ...) @_Z17__enzyme_autodiffPvz(i8* bitcast (double (double, double)* @muldd to i8*), double %x, double %y)
+  ret %struct.Gradients %call
+}
+
+declare dso_local %struct.Gradients @_Z17__enzyme_autodiffPvz(i8*, ...) local_unnamed_addr #2
+
+attributes #0 = { norecurse nounwind readnone "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="non-leaf" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+neon" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #1 = { "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="non-leaf" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+neon" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #2 = { "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="non-leaf" "less-precise-fpmad"="false" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+neon" "unsafe-fp-math"="false" "use-soft-float"="false" }
+
+; CHECK: define internal {{(dso_local )?}}{ double, double } @diffemuldd(double %x, double %y, double %differeturn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %m0diffex = fmul fast double %differeturn, %y
+; CHECK-NEXT:   %m1diffey = fmul fast double %differeturn, %x
+; CHECK-NEXT:   %0 = insertvalue { double, double } undef, double %m0diffex, 0
+; CHECK-NEXT:   %1 = insertvalue { double, double } %0, double %m1diffey, 1
+; CHECK-NEXT:   ret { double, double } %1
+; CHECK-NEXT: }

enzyme/test/Integration/ReverseMode/gradient-struct-return.c

@@ -0,0 +1,33 @@
+// RUN: %clang -std=c11 -O0 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S | %lli - 
+// RUN: %clang -std=c11 -O1 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S | %lli - 
+// RUN: %clang -std=c11 -O2 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S | %lli - 
+// RUN: %clang -std=c11 -O3 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S | %lli - 
+// RUN: %clang -std=c11 -O0 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme --enzyme-inline=1 -S | %lli - 
+// RUN: %clang -std=c11 -O1 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme --enzyme-inline=1 -S | %lli - 
+// RUN: %clang -std=c11 -O2 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme --enzyme-inline=1 -S | %lli - 
+// RUN: %clang -std=c11 -O3 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme --enzyme-inline=1 -S | %lli - 
+
+#include <stdio.h>
+#include "test_utils.h"
+
+typedef struct {
+    double dx,dy;
+} Gradients;
+
+extern Gradients __enzyme_autodiff(void*, ...);
+
+double mul(double x, double y) {
+    return x * y;
+}
+Gradients dmul(double x, double y) {
+    return __enzyme_autodiff((void*)mul, x, y);
+}
+int main() {
+    double x = 1.0;
+    double y = 2.0;
+    printf("mul(x=%f,y%f)=%f\n", x, y, mul(x,y));
+    printf("ddx dmul(x=%f,y%f)=%f\n", x, y, dmul(x,y).dx);
+    printf("ddy dmul(x=%f,y%f)=%f\n", x, y, dmul(x,y).dy);
+    APPROX_EQ(dmul(x,y).dx, 2.0, 10e-10);
+    APPROX_EQ(dmul(x,y).dy, 1.0, 10e-10);
+}