added test for f32 and f64

Signed-off-by: Karan Janthe <karanjanthe@gmail.com>

Author: KMJ-007

tests/codegen/autodiff/f32.rs

@@ -1,10 +1,18 @@
-//@ compile-flags: -Zautodiff=Enable -Zautodiff=NoPostopt -C opt-level=3 -Clto=fat -g
+//@ compile-flags: -Zautodiff=Enable,PrintTAFn=callee -Zautodiff=NoPostopt -C opt-level=3 -Clto=fat -g
 //@ no-prefer-dynamic
 //@ needs-enzyme
 #![feature(autodiff)]
 
 use std::autodiff::autodiff_reverse;
 
+// CHECK-LABEL: callee
+// CHECK: callee - {[-1]:Float@float} |{[-1]:Pointer}:{}
+// CHECK: ptr %x: {[-1]:Pointer, [-1,0]:Float@float}
+// CHECK: start
+// CHECK: %_2 = load float, ptr %x, align 4, !noundef !3: {[-1]:Float@float}
+// CHECK: %_0 = fmul float %_2, %_2: {[-1]:Float@float}
+// CHECK: ret float %_0: {}
+
 #[autodiff_reverse(d_square, Duplicated, Active)]
 #[no_mangle]
 fn callee(x: &f32) -> f32 {
@@ -14,22 +22,9 @@ fn callee(x: &f32) -> f32 {
 fn main() {
     let x: f32 = 7.0;
     let mut df_dx: f32 = 0.0;
-    d_square(&x, &mut df_dx, 1.0);
-}
-
-// CHECK: define float @callee(ptr align 4 {{.*}}) {{.*}} !dbg {
-// CHECK-NEXT: start:
-// CHECK: #dbg_value
-// CHECK: load float
-// CHECK: fmul float
-// CHECK: ret float
-
-// CHECK: define void @_ZN9f32_debug8d_square{{.*}}(ptr align 4 {{.*}}, ptr align 4 {{.*}}, float {{.*}}) {{.*}} {
-// CHECK: call {{.*}} @diffecallee
-// CHECK: ret void
+    let out = callee(&x);
+    let out_ = d_square(&x, &mut df_dx, 1.0);
 
-// CHECK: define {{.*}} @diffecallee(ptr {{.*}} align 4 {{.*}}, ptr {{.*}} align 4 {{.*}}, float {{.*}}) {{.*}} {
-// CHECK: load float
-// CHECK: fmul float
-// CHECK: store float
-// CHECK: ret {{.*}}
+    assert_eq!(out, out_);
+    assert_eq!(14.0, df_dx);
+}

tests/codegen/autodiff/f64.rs

@@ -0,0 +1,31 @@
+//@ compile-flags: -Zautodiff=Enable,PrintTAFn=square -C opt-level=3  -Clto=fat 
+//@ no-prefer-dynamic
+//@ needs-enzyme
+#![feature(autodiff)]
+use std::autodiff::autodiff_reverse;
+
+// CHECK-LABEL: square
+// CHECK: square - {[-1]:Float@double} |{[-1]:Pointer}:{}
+// CHECK: ptr %x: {[-1]:Pointer, [-1,0]:Float@double}
+// CHECK: start
+// CHECK: %_2 = load double, ptr %x, align 8, !noundef !3: {[-1]:Float@double}
+// CHECK: %_0 = fmul double %_2, %_2: {[-1]:Float@double}
+// CHECK: ret double %_0: {}
+
+#[autodiff_reverse(d_square, Duplicated, Active)]
+#[no_mangle]
+fn square(x: &f64) -> f64 {
+    x * x
+}
+
+fn main() {
+    let x = std::hint::black_box(3.0);
+    
+    let output = square(&x);
+    assert_eq!(9.0, output);
+
+    let mut df_dx = 0.0;
+    let output_ = d_square(&x, &mut df_dx, 1.0);
+    assert_eq!(output, output_);
+    assert_eq!(6.0, df_dx);
+}