Skip to content

Commit b6bb28e

Browse files
stijnhstijnh
committed
Add script to verify codegen and add it to github CI workflow
1 parent 4211713 commit b6bb28e

8 files changed

Lines changed: 908 additions & 0 deletions

File tree

‎.github/workflows/cmake-run-cuda.yml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +49,7 @@ jobs:
4949
# See https://cmake.org/cmake/help/latest/manual/ctest.1.html for more detail
5050
run: ./tests/kernel_float_tests --durations=yes --success --verbosity=high ~[GPU]
5151

52+
- name: Codegen instruction checks
53+
# These check exact instruction counts in the generated PTX, which can legitimately
54+
# differ between nvcc versions even when the code itself is correct.
55+
run: python3 tests_codegen/verify_codegen.py tests_codegen/*.cu --nvcc nvcc

‎tests_codegen/approx_ops.cu‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
#include "kernel_float.h"
2+
namespace kf = kernel_float;
3+
4+
extern "C" {
5+
6+
// CHECK-NOT: approx.f32
7+
// CHECK-COUNT-4: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
8+
__global__ void approx_half_sin1(kf::half2 *x, kf::half2* y) {
9+
*y = kf::approx_sin<1>(*x);
10+
}
11+
12+
// CHECK-NOT: approx.f32
13+
// CHECK-COUNT-6: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
14+
__global__ void approx_half_sin2(kf::half2 *x, kf::half2* y) {
15+
*y = kf::approx_sin<2>(*x);
16+
}
17+
18+
// CHECK-NOT: approx.f32
19+
// CHECK-COUNT-7: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
20+
__global__ void approx_half_sin3(kf::half2 *x, kf::half2* y) {
21+
*y = kf::approx_sin<3>(*x);
22+
}
23+
24+
// CHECK-NOT: approx.f32
25+
// CHECK-COUNT-8: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
26+
__global__ void approx_half_sin4(kf::half2 *x, kf::half2* y) {
27+
*y = kf::approx_sin<4>(*x);
28+
}
29+
30+
// CHECK-NOT: approx.f32
31+
// CHECK-NOT: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
32+
__global__ void approx_half_rcp0(kf::half2 *x, kf::half2* y) {
33+
*y = kf::approx_rcp<0>(*x);
34+
}
35+
36+
// CHECK-NOT: approx.f32
37+
// CHECK-COUNT-2: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
38+
__global__ void approx_half_rcp1(kf::half2 *x, kf::half2* y) {
39+
*y = kf::approx_rcp<1>(*x);
40+
}
41+
42+
// CHECK-NOT: approx.f32
43+
// CHECK-COUNT-4: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
44+
__global__ void approx_half_rcp2(kf::half2 *x, kf::half2* y) {
45+
*y = kf::approx_rcp<2>(*x);
46+
}
47+
48+
// CHECK-NOT: approx.f32
49+
// CHECK-COUNT-6: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
50+
__global__ void approx_half_rcp3(kf::half2 *x, kf::half2* y) {
51+
*y = kf::approx_rcp<3>(*x);
52+
}
53+
54+
// CHECK-NOT: approx.f32
55+
// CHECK-COUNT-8: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
56+
__global__ void approx_half_rcp4(kf::half2 *x, kf::half2* y) {
57+
*y = kf::approx_rcp<4>(*x);
58+
}
59+
60+
// CHECK-NOT: approx.f32
61+
// CHECK-COUNT-1: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
62+
__global__ void approx_half_sqrt0(kf::half2 *x, kf::half2* y) {
63+
*y = kf::approx_sqrt<0>(*x);
64+
}
65+
66+
// CHECK-NOT: approx.f32
67+
// CHECK-COUNT-4: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
68+
__global__ void approx_half_sqrt1(kf::half2 *x, kf::half2* y) {
69+
*y = kf::approx_sqrt<1>(*x);
70+
}
71+
72+
// CHECK-NOT: approx.f32
73+
// CHECK-COUNT-9: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
74+
__global__ void approx_half_sqrt2(kf::half2 *x, kf::half2* y) {
75+
*y = kf::approx_sqrt<2>(*x);
76+
}
77+
78+
// CHECK-NOT: approx.f32
79+
// CHECK-COUNT-13: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
80+
__global__ void approx_half_sqrt3(kf::half2 *x, kf::half2* y) {
81+
*y = kf::approx_sqrt<3>(*x);
82+
}
83+
84+
// CHECK-NOT: approx.f32
85+
// CHECK-COUNT-17: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
86+
__global__ void approx_half_sqrt4(kf::half2 *x, kf::half2* y) {
87+
*y = kf::approx_sqrt<4>(*x);
88+
}
89+
90+
// CHECK-NOT: approx.f32
91+
// CHECK-NOT: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
92+
__global__ void approx_half_rsqrt0(kf::half2 *x, kf::half2* y) {
93+
*y = kf::approx_rsqrt<0>(*x);
94+
}
95+
96+
// CHECK-NOT: approx.f32
97+
// CHECK-COUNT-4: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
98+
__global__ void approx_half_rsqrt1(kf::half2 *x, kf::half2* y) {
99+
*y = kf::approx_rsqrt<1>(*x);
100+
}
101+
102+
// CHECK-NOT: approx.f32
103+
// CHECK-COUNT-8: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
104+
__global__ void approx_half_rsqrt2(kf::half2 *x, kf::half2* y) {
105+
*y = kf::approx_rsqrt<2>(*x);
106+
}
107+
108+
// CHECK-NOT: approx.f32
109+
// CHECK-COUNT-12: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
110+
__global__ void approx_half_rsqrt3(kf::half2 *x, kf::half2* y) {
111+
*y = kf::approx_rsqrt<3>(*x);
112+
}
113+
114+
// CHECK-NOT: approx.f32
115+
// CHECK-COUNT-16: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
116+
__global__ void approx_half_rsqrt4(kf::half2 *x, kf::half2* y) {
117+
*y = kf::approx_rsqrt<4>(*x);
118+
}
119+
120+
121+
// CHECK-NOT: approx.f32
122+
// CHECK-COUNT-1: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
123+
__global__ void approx_half_exp(kf::half2 *x, kf::half2* y) {
124+
*y = kf::approx_exp(*x);
125+
}
126+
127+
// CHECK-NOT: approx.f32
128+
// CHECK-COUNT-1: (fma|mul|add)(.rn|.ftz|.sat)*.f16x2
129+
__global__ void approx_half_log(kf::half2 *x, kf::half2* y) {
130+
*y = kf::approx_log(*x);
131+
}
132+
133+
}

‎tests_codegen/bf16_ops.cu‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
#include "kernel_float.h"
2+
namespace kf = kernel_float;
3+
4+
extern "C" {
5+
6+
// CHECK: ld.global.u16
7+
// CHECK: st.global.u16
8+
__global__ void bfloat16x1_copy(const kf::bfloat16x1 *input, kf::bfloat16x1 *output) {
9+
*output = *input;
10+
}
11+
12+
// CHECK: ld.global.u32
13+
// CHECK: st.global.u32
14+
__global__ void bfloat16x2_copy(const kf::bfloat16x2 *input, kf::bfloat16x2 *output) {
15+
*output = *input;
16+
}
17+
18+
// CHECK: ld.global.v4.u16
19+
// CHECK: st.global.v4.u16
20+
__global__ void bfloat16x4_copy(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
21+
*output = *input;
22+
}
23+
24+
// CHECK: ld.global.v4.u32
25+
// CHECK: st.global.v4.u32
26+
__global__ void bfloat16x8_copy(const kf::bfloat16x8 *input, kf::bfloat16x8 *output) {
27+
*output = *input;
28+
}
29+
30+
// bfloat16x4 holds 2 packed bf16x2 pairs, so each op below lowers to the instruction twice
31+
// (once per pair). Unlike half2, bf16x2 has no native packed `add` instruction on sm_80: it
32+
// is emulated as `fma(x, 1.0, x)`, so this lowers to `fma.rn.bf16x2`, not `add.bf16x2`.
33+
// CHECK-COUNT-2: fma.rn.bf16x2
34+
// CHECK-NOT: add.bf16x2
35+
__global__ void bfloat16_add(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
36+
*output = *input + *input;
37+
}
38+
39+
// Same story as `add`: emulated as `fma(x, -1.0, x)`, lowering to `fma.rn.bf16x2` rather than
40+
// a `sub.bf16x2` instruction.
41+
// CHECK-COUNT-2: fma.rn.bf16x2
42+
// CHECK-NOT: sub.bf16x2
43+
__global__ void bfloat16_sub(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
44+
*output = *input - *input;
45+
}
46+
47+
// Same story again: emulated as `fma(x, x, -0.0)`, lowering to `fma.rn.bf16x2` rather than a
48+
// `mul.bf16x2` instruction.
49+
// CHECK-COUNT-2: fma.rn.bf16x2
50+
// CHECK-NOT: mul.bf16x2
51+
__global__ void bfloat16_mul(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
52+
*output = *input * *input;
53+
}
54+
55+
// This is a genuine, intentional packed FMA (not an add/sub/mul emulated via FMA like above).
56+
// CHECK-COUNT-2: fma.rn.bf16x2
57+
__global__ void bfloat16_fma(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
58+
*output = kf::fma(*input, *input, *input);
59+
}
60+
61+
// CHECK-COUNT-2: neg.bf16x2
62+
__global__ void bfloat16_neg(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
63+
*output = -*input;
64+
}
65+
66+
// CHECK-COUNT-2: min.bf16x2
67+
__global__ void bfloat16_min(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
68+
*output = kf::min(*input, *input);
69+
}
70+
71+
// CHECK-COUNT-2: max.bf16x2
72+
__global__ void bfloat16_max(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
73+
*output = kf::max(*input, *input);
74+
}
75+
76+
// Unlike half2's clean `set.eq.f16x2.f16x2`, bf16x2 equality has no native packed comparison:
77+
// it is emulated by splitting each of the 4 lanes out to `f32` and comparing with
78+
// `set.eq.f32.f32`, one per lane.
79+
// CHECK-COUNT-4: set.eq.f32.f32
80+
__global__ void bfloat16_eq(const kf::bfloat16x4 *input, bool *output) {
81+
output[0] = kf::all(*input == *input);
82+
}
83+
84+
// Division is not packed at all: each of the 4 lanes is extracted individually and divided
85+
// with a scalar `div.rn.f32`, then converted back with `cvt.rn.bf16.f32`.
86+
// CHECK-COUNT-4: div.rn.f32
87+
// CHECK-NOT: div.bf16x2
88+
__global__ void bfloat16_div(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
89+
*output = *input / *input;
90+
}
91+
92+
// Like half's `sqrt`, this is not native packed hardware: it expands to per-lane
93+
// `sqrt.approx.f32` (no `.ftz` here, unlike the half2 path) followed by `cvt.rn.bf16.f32`,
94+
// one per lane.
95+
// CHECK-COUNT-4: sqrt.approx.f32
96+
// CHECK-NOT: sqrt.bf16x2
97+
__global__ void bfloat16_sqrt(const kf::bfloat16x4 *input, kf::bfloat16x4 *output) {
98+
*output = kf::sqrt(*input);
99+
}
100+
101+
}

0 commit comments

Comments
 (0)