Skip to content

Commit f6af51e

Browse files
tgymnichtgymnich
authored
Implement vector mode: intrinsics (rust-lang#457)
1 parent 7bd7b47 commit f6af51e

File tree

18 files changed

+765
-79
lines changed

18 files changed

+765
-79
lines changed

‎enzyme/Enzyme/AdjointGenerator.h

Lines changed: 162 additions & 79 deletions
Large diffs are not rendered by default.

‎enzyme/test/Enzyme/ForwardMode/maxnum-inactive.ll

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -early-cse -simplifycfg -S | FileCheck %s
2+
3+
; Function Attrs: noinline nounwind readnone uwtable
4+
define double @tester(double %x, double %y) {
5+
entry:
6+
%0 = tail call double @llvm.maxnum.f64(double %x, double %y)
7+
ret double %0
8+
}
9+
10+
define double @test_derivative(double %x, double %y) {
11+
entry:
12+
%0 = tail call double (double (double, double)*, ...) @__enzyme_fwddiff(double (double, double)* nonnull @tester, double %x, double 1.0, metadata !"enzyme_const", double %y)
13+
ret double %0
14+
}
15+
16+
; Function Attrs: nounwind readnone speculatable
17+
declare double @llvm.maxnum.f64(double, double)
18+
19+
; Function Attrs: nounwind
20+
declare double @__enzyme_fwddiff(double (double, double)*, ...)
21+
22+
23+
; CHECK: define internal double @fwddiffetester(double %x, double %"x'", double %y)
24+
; CHECK-NEXT: entry:
25+
; CHECK-NEXT: %0 = fcmp fast olt double %x, %y
26+
; CHECK-NEXT: %1 = select {{(fast )?}}i1 %0, double %"x'", double 0.000000e+00
27+
; CHECK-NEXT: ret double %1
28+
; CHECK-NEXT: }

‎enzyme/test/Enzyme/ForwardModeVector/cos.ll

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -O3 -S | FileCheck %s
2+
3+
%struct.Gradients = type { double, double }
4+
5+
; Function Attrs: nounwind
6+
declare %struct.Gradients @__enzyme_fwddiff(double (double)*, ...)
7+
8+
; Function Attrs: nounwind readnone uwtable
9+
define double @tester(double %x) {
10+
entry:
11+
%0 = tail call fast double @llvm.cos.f64(double %x)
12+
ret double %0
13+
}
14+
15+
define %struct.Gradients @test_derivative(double %x) {
16+
entry:
17+
%0 = tail call %struct.Gradients (double (double)*, ...) @__enzyme_fwddiff(double (double)* nonnull @tester, metadata !"enzyme_width", i64 2, double %x, double 0.000000e+00, double 1.000000e+00)
18+
ret %struct.Gradients %0
19+
}
20+
21+
; Function Attrs: nounwind readnone speculatable
22+
declare double @llvm.cos.f64(double)
23+
24+
; CHECK: define %struct.Gradients @test_derivative(double %x)
25+
; CHECK-NEXT: entry:
26+
; CHECK-NEXT: %0 = tail call fast double @llvm.sin.f64(double %x)
27+
; CHECK-NEXT: %1 = {{(fsub fast double -0.000000e\+00,|fneg fast double)}} %0
28+
; CHECK-NEXT: %2 = insertvalue %struct.Gradients zeroinitializer, double %1, 1
29+
; CHECK-NEXT: ret %struct.Gradients %2
30+
; CHECK-NEXT: }

‎enzyme/test/Enzyme/ForwardModeVector/exp.ll

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -O3 -S | FileCheck %s
2+
3+
%struct.Gradients = type { double, double, double }
4+
5+
; Function Attrs: nounwind
6+
declare %struct.Gradients @__enzyme_fwddiff(double (double)*, ...)
7+
8+
; Function Attrs: nounwind readnone uwtable
9+
define double @tester(double %x) {
10+
entry:
11+
%0 = tail call fast double @llvm.exp.f64(double %x)
12+
ret double %0
13+
}
14+
15+
define %struct.Gradients @test_derivative(double %x) {
16+
entry:
17+
%0 = tail call %struct.Gradients (double (double)*, ...) @__enzyme_fwddiff(double (double)* nonnull @tester, metadata !"enzyme_width", i64 3, double %x, double 1.0, double 2.0, double 3.0)
18+
ret %struct.Gradients %0
19+
}
20+
21+
; Function Attrs: nounwind readnone speculatable
22+
declare double @llvm.exp.f64(double)
23+
24+
25+
; CHECK: define %struct.Gradients @test_derivative(double %x)
26+
; CHECK-NEXT: entry:
27+
; CHECK-NEXT: %0 = tail call fast double @llvm.exp.f64(double %x)
28+
; CHECK-NEXT: %1 = fmul fast double %0, 2.000000e+00
29+
; CHECK-NEXT: %2 = fmul fast double %0, 3.000000e+00
30+
; CHECK-NEXT: %3 = insertvalue %struct.Gradients zeroinitializer, double %0, 0
31+
; CHECK-NEXT: %4 = insertvalue %struct.Gradients %3, double %1, 1
32+
; CHECK-NEXT: %5 = insertvalue %struct.Gradients %4, double %2, 2
33+
; CHECK-NEXT: ret %struct.Gradients %5
34+
; CHECK-NEXT: }

‎enzyme/test/Enzyme/ForwardModeVector/exp2.ll

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -O3 -S | FileCheck %s
2+
3+
%struct.Gradients = type { double, double }
4+
5+
; Function Attrs: nounwind
6+
declare %struct.Gradients @__enzyme_fwddiff(double (double)*, ...)
7+
8+
; Function Attrs: nounwind readnone uwtable
9+
define double @tester(double %x) {
10+
entry:
11+
%0 = tail call fast double @llvm.exp2.f64(double %x)
12+
ret double %0
13+
}
14+
15+
define %struct.Gradients @test_derivative(double %x) {
16+
entry:
17+
%0 = tail call %struct.Gradients (double (double)*, ...) @__enzyme_fwddiff(double (double)* nonnull @tester, metadata !"enzyme_width", i64 2, double %x, double 1.0, double 2.5)
18+
ret %struct.Gradients %0
19+
}
20+
21+
; Function Attrs: nounwind readnone speculatable
22+
declare double @llvm.exp2.f64(double)
23+
24+
25+
; CHECK: define %struct.Gradients @test_derivative(double %x)
26+
; CHECK-NEXT: entry:
27+
; CHECK-NEXT: %0 = tail call fast double @llvm.exp2.f64(double %x)
28+
; CHECK-NEXT: %1 = fmul fast double %0, 0x3FE62E42FEFA39EF
29+
; CHECK-NEXT: %2 = fmul fast double %0, 0x3FFBB9D3BEB8C86B
30+
; CHECK-NEXT: %3 = insertvalue %struct.Gradients zeroinitializer, double %1, 0
31+
; CHECK-NEXT: %4 = insertvalue %struct.Gradients %3, double %2, 1
32+
; CHECK-NEXT: ret %struct.Gradients %4
33+
; CHECK-NEXT: }

‎enzyme/test/Enzyme/ForwardModeVector/experimental_vector_reduce_v2_fadd.ll

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
; RUN: if [ %llvmver -ge 9 ] && [ %llvmver -le 11 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi
2+
3+
%struct.Gradients = type { float, float }
4+
5+
; Function Attrs: nounwind
6+
declare %struct.Gradients @__enzyme_fwddiff(float (float, <4 x float>)*, ...)
7+
8+
define float @tester(float %start_value, <4 x float> %input) {
9+
entry:
10+
%ord = call float @llvm.experimental.vector.reduce.v2.fadd.f32.v4f32(float %start_value, <4 x float> %input)
11+
ret float %ord
12+
}
13+
14+
define %struct.Gradients @test_derivative(float %start_value, <4 x float> %input) {
15+
entry:
16+
%0 = tail call %struct.Gradients (float (float, <4 x float>)*, ...) @__enzyme_fwddiff(float (float, <4 x float>)* nonnull @tester, metadata !"enzyme_width", i64 2, float %start_value, float 1.0, float 2.0, <4 x float> %input, <4 x float> <float 1.0, float 1.0, float 1.0, float 1.0>, <4 x float> <float 1.0, float 1.0, float 1.0, float 1.0>)
17+
ret %struct.Gradients %0
18+
}
19+
20+
declare float @llvm.experimental.vector.reduce.v2.fadd.f32.v4f32(float, <4 x float>)
21+
22+
; CHECK: define internal [2 x float] @fwddiffe2tester(float %start_value, [2 x float] %"start_value'", <4 x float> %input, [2 x <4 x float>] %"input'")
23+
; CHECK-NEXT: entry:
24+
; CHECK-NEXT: %0 = extractvalue [2 x float] %"start_value'", 0
25+
; CHECK-NEXT: %1 = extractvalue [2 x <4 x float>] %"input'", 0
26+
; CHECK-NEXT: %2 = call {{(fast )?}}float @llvm.experimental.vector.reduce.v2.fadd.f32.v4f32(float %0, <4 x float> %1)
27+
; CHECK-NEXT: %3 = insertvalue [2 x float] undef, float %2, 0
28+
; CHECK-NEXT: %4 = extractvalue [2 x float] %"start_value'", 1
29+
; CHECK-NEXT: %5 = extractvalue [2 x <4 x float>] %"input'", 1
30+
; CHECK-NEXT: %6 = call {{(fast )?}}float @llvm.experimental.vector.reduce.v2.fadd.f32.v4f32(float %4, <4 x float> %5)
31+
; CHECK-NEXT: %7 = insertvalue [2 x float] %3, float %6, 1
32+
; CHECK-NEXT: ret [2 x float] %7
33+
; CHECK-NEXT: }
34+

‎enzyme/test/Enzyme/ForwardModeVector/fabs.ll

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -early-cse -simplifycfg -S | FileCheck %s
2+
3+
%struct.Gradients = type { double, double }
4+
5+
; Function Attrs: nounwind
6+
declare %struct.Gradients @__enzyme_fwddiff(double (double)*, ...)
7+
8+
; Function Attrs: nounwind readnone uwtable
9+
define double @tester(double %x) {
10+
entry:
11+
%0 = tail call fast double @llvm.fabs.f64(double %x)
12+
ret double %0
13+
}
14+
15+
define %struct.Gradients @test_derivative(double %x) {
16+
entry:
17+
%0 = tail call %struct.Gradients (double (double)*, ...) @__enzyme_fwddiff(double (double)* nonnull @tester, metadata !"enzyme_width", i64 2, double %x, double 1.0, double 2.0)
18+
ret %struct.Gradients %0
19+
}
20+
21+
; Function Attrs: nounwind readnone speculatable
22+
declare double @llvm.fabs.f64(double)
23+
24+
25+
; CHECK: define internal [2 x double] @fwddiffe2tester(double %x, [2 x double] %"x'")
26+
; CHECK-NEXT: entry:
27+
; CHECK-NEXT: %0 = extractvalue [2 x double] %"x'", 0
28+
; CHECK-NEXT: %1 = fcmp fast olt double %x, 0.000000e+00
29+
; CHECK-NEXT: %2 = select {{(fast )?}}i1 %1, double -1.000000e+00, double 1.000000e+00
30+
; CHECK-NEXT: %3 = fmul fast double %2, %0
31+
; CHECK-NEXT: %4 = insertvalue [2 x double] undef, double %3, 0
32+
; CHECK-NEXT: %5 = extractvalue [2 x double] %"x'", 1
33+
; CHECK-NEXT: %6 = fmul fast double %2, %5
34+
; CHECK-NEXT: %7 = insertvalue [2 x double] %4, double %6, 1
35+
; CHECK-NEXT: ret [2 x double] %7
36+
; CHECK-NEXT: }

‎enzyme/test/Enzyme/ForwardModeVector/log.ll

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -O3 -S | FileCheck %s
2+
3+
%struct.Gradients = type { double, double, double }
4+
5+
; Function Attrs: nounwind
6+
declare %struct.Gradients @__enzyme_fwddiff(double (double)*, ...)
7+
8+
; Function Attrs: nounwind readnone uwtable
9+
define double @tester(double %x) {
10+
entry:
11+
%0 = tail call fast double @llvm.log.f64(double %x)
12+
ret double %0
13+
}
14+
15+
define %struct.Gradients @test_derivative(double %x) {
16+
entry:
17+
%0 = tail call %struct.Gradients (double (double)*, ...) @__enzyme_fwddiff(double (double)* nonnull @tester, metadata !"enzyme_width", i64 3, double %x, double 1.0, double 2.0, double 3.0)
18+
ret %struct.Gradients %0
19+
}
20+
21+
; Function Attrs: nounwind readnone speculatable
22+
declare double @llvm.log.f64(double)
23+
24+
25+
; CHECK: define %struct.Gradients @test_derivative(double %x)
26+
; CHECK-NEXT: entry:
27+
; CHECK-NEXT: %0 = fdiv fast double 1.000000e+00, %x
28+
; CHECK-NEXT: %1 = fdiv fast double 2.000000e+00, %x
29+
; CHECK-NEXT: %2 = fdiv fast double 3.000000e+00, %x
30+
; CHECK-NEXT: %3 = insertvalue %struct.Gradients zeroinitializer, double %0, 0
31+
; CHECK-NEXT: %4 = insertvalue %struct.Gradients %3, double %1, 1
32+
; CHECK-NEXT: %5 = insertvalue %struct.Gradients %4, double %2, 2
33+
; CHECK-NEXT: ret %struct.Gradients %5
34+
; CHECK-NEXT: }

‎enzyme/test/Enzyme/ForwardModeVector/log10.ll

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -O3 -S | FileCheck %s
2+
3+
%struct.Gradients = type { double, double, double }
4+
5+
; Function Attrs: nounwind
6+
declare %struct.Gradients @__enzyme_fwddiff(double (double)*, ...)
7+
8+
; Function Attrs: nounwind readnone uwtable
9+
define double @tester(double %x) {
10+
entry:
11+
%0 = tail call fast double @llvm.log10.f64(double %x)
12+
ret double %0
13+
}
14+
15+
define %struct.Gradients @test_derivative(double %x) {
16+
entry:
17+
%0 = tail call %struct.Gradients (double (double)*, ...) @__enzyme_fwddiff(double (double)* nonnull @tester, metadata !"enzyme_width", i64 3, double %x, double 1.0, double 2.0, double 3.0)
18+
ret %struct.Gradients %0
19+
}
20+
21+
; Function Attrs: nounwind readnone speculatable
22+
declare double @llvm.log10.f64(double)
23+
24+
25+
; CHECK: define %struct.Gradients @test_derivative(double %x)
26+
; CHECK-NEXT: entry:
27+
; CHECK-NEXT: %0 = fdiv fast double 0x3FDBCB7B1526E50D, %x
28+
; CHECK-NEXT: %1 = fdiv fast double 0x3FEBCB7B1526E50D, %x
29+
; CHECK-NEXT: %2 = fdiv fast double 0x3FF4D89C4FDD2BCA, %x
30+
; CHECK-NEXT: %3 = insertvalue %struct.Gradients zeroinitializer, double %0, 0
31+
; CHECK-NEXT: %4 = insertvalue %struct.Gradients %3, double %1, 1
32+
; CHECK-NEXT: %5 = insertvalue %struct.Gradients %4, double %2, 2
33+
; CHECK-NEXT: ret %struct.Gradients %5
34+
; CHECK-NEXT: }

‎enzyme/test/Enzyme/ForwardModeVector/log2.ll

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -O3 -S | FileCheck %s
2+
3+
%struct.Gradients = type { double, double, double }
4+
5+
; Function Attrs: nounwind
6+
declare %struct.Gradients @__enzyme_fwddiff(double (double)*, ...)
7+
8+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -O3 -S | FileCheck %s
9+
10+
; Function Attrs: nounwind readnone uwtable
11+
define double @tester(double %x) {
12+
entry:
13+
%0 = tail call fast double @llvm.log2.f64(double %x)
14+
ret double %0
15+
}
16+
17+
define %struct.Gradients @test_derivative(double %x) {
18+
entry:
19+
%0 = tail call %struct.Gradients (double (double)*, ...) @__enzyme_fwddiff(double (double)* nonnull @tester, metadata !"enzyme_width", i64 3, double %x, double 1.0, double 2.0, double 3.0)
20+
ret %struct.Gradients %0
21+
}
22+
23+
; Function Attrs: nounwind readnone speculatable
24+
declare double @llvm.log2.f64(double)
25+
26+
27+
; CHECK: define %struct.Gradients @test_derivative(double %x)
28+
; CHECK-NEXT: entry:
29+
; CHECK-NEXT: %0 = fdiv fast double 0x3FF71547652B82FE, %x
30+
; CHECK-NEXT: %1 = fdiv fast double 0x40071547652B82FE, %x
31+
; CHECK-NEXT: %2 = fdiv fast double 0x40114FF58BE0A23F, %x
32+
; CHECK-NEXT: %3 = insertvalue %struct.Gradients zeroinitializer, double %0, 0
33+
; CHECK-NEXT: %4 = insertvalue %struct.Gradients %3, double %1, 1
34+
; CHECK-NEXT: %5 = insertvalue %struct.Gradients %4, double %2, 2
35+
; CHECK-NEXT: ret %struct.Gradients %5
36+
; CHECK-NEXT: }

‎enzyme/test/Enzyme/ForwardModeVector/maxnum-inactive.ll

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -early-cse -simplifycfg -S | FileCheck %s
2+
3+
%struct.Gradients = type { double, double }
4+
5+
; Function Attrs: nounwind
6+
declare %struct.Gradients @__enzyme_fwddiff(double (double, double)*, ...)
7+
8+
; Function Attrs: noinline nounwind readnone uwtable
9+
define double @tester(double %x, double %y) {
10+
entry:
11+
%0 = tail call double @llvm.maxnum.f64(double %x, double %y)
12+
ret double %0
13+
}
14+
15+
define %struct.Gradients @test_derivative(double %x, double %y) {
16+
entry:
17+
%0 = tail call %struct.Gradients (double (double, double)*, ...) @__enzyme_fwddiff(double (double, double)* nonnull @tester, metadata !"enzyme_width", i64 2, double %x, double 1.0, double 0.0, metadata !"enzyme_const", double %y)
18+
ret %struct.Gradients %0
19+
}
20+
21+
; Function Attrs: nounwind readnone speculatable
22+
declare double @llvm.maxnum.f64(double, double)
23+
24+
25+
; CHECK: define internal [2 x double] @fwddiffe2tester(double %x, [2 x double] %"x'", double %y)
26+
; CHECK-NEXT: entry:
27+
; CHECK-NEXT: %0 = fcmp fast olt double %x, %y
28+
; CHECK-NEXT: %1 = extractvalue [2 x double] %"x'", 0
29+
; CHECK-NEXT: %2 = select {{(fast )?}}i1 %0, double %1, double 0.000000e+00
30+
; CHECK-NEXT: %3 = insertvalue [2 x double] undef, double %2, 0
31+
; CHECK-NEXT: %4 = extractvalue [2 x double] %"x'", 1
32+
; CHECK-NEXT: %5 = select {{(fast )?}}i1 %0, double %4, double 0.000000e+00
33+
; CHECK-NEXT: %6 = insertvalue [2 x double] %3, double %5, 1
34+
; CHECK-NEXT: ret [2 x double] %6
35+
; CHECK-NEXT: }

‎enzyme/test/Enzyme/ForwardModeVector/maxnum.ll

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -early-cse -simplifycfg -S | FileCheck %s
2+
3+
%struct.Gradients = type { double, double }
4+
5+
; Function Attrs: nounwind
6+
declare %struct.Gradients @__enzyme_fwddiff(double (double, double)*, ...)
7+
8+
; Function Attrs: noinline nounwind readnone uwtable
9+
define double @tester(double %x, double %y) {
10+
entry:
11+
%0 = tail call double @llvm.maxnum.f64(double %x, double %y)
12+
ret double %0
13+
}
14+
15+
define %struct.Gradients @test_derivative(double %x, double %y) {
16+
entry:
17+
%0 = tail call %struct.Gradients (double (double, double)*, ...) @__enzyme_fwddiff(double (double, double)* nonnull @tester, metadata !"enzyme_width", i64 2, double %x, double 1.0, double 0.0, double %y, double 0.0, double 1.0)
18+
ret %struct.Gradients %0
19+
}
20+
21+
; Function Attrs: nounwind readnone speculatable
22+
declare double @llvm.maxnum.f64(double, double)
23+
24+
25+
; CHECK: define internal [2 x double] @fwddiffe2tester(double %x, [2 x double] %"x'", double %y, [2 x double] %"y'")
26+
; CHECK-NEXT: entry:
27+
; CHECK-NEXT: %0 = fcmp fast olt double %x, %y
28+
; CHECK-NEXT: %1 = extractvalue [2 x double] %"x'", 0
29+
; CHECK-NEXT: %2 = extractvalue [2 x double] %"y'", 0
30+
; CHECK-NEXT: %3 = select {{(fast )?}}i1 %0, double %1, double %2
31+
; CHECK-NEXT: %4 = insertvalue [2 x double] undef, double %3, 0
32+
; CHECK-NEXT: %5 = extractvalue [2 x double] %"x'", 1
33+
; CHECK-NEXT: %6 = extractvalue [2 x double] %"y'", 1
34+
; CHECK-NEXT: %7 = select {{(fast )?}}i1 %0, double %5, double %6
35+
; CHECK-NEXT: %8 = insertvalue [2 x double] %4, double %7, 1
36+
; CHECK-NEXT: ret [2 x double] %8
37+
; CHECK-NEXT: }

0 commit comments

Comments
 (0)