llvm/test/CodeGen/NVPTX/bf16-instructions.ll

   1 ; RUN: llc < %s -march=nvptx64 -mcpu=sm_80 -mattr=+ptx70 | FileCheck --check-prefixes=CHECK,SM80 %s
   2 ; RUN: llc < %s -march=nvptx64 -mcpu=sm_90 -mattr=+ptx78 | FileCheck --check-prefixes=CHECK,SM90 %s
   3 ; RUN: %if ptxas-11.8 %{ llc < %s -march=nvptx64 -mcpu=sm_80 -mattr=+ptx71 | %ptxas-verify -arch=sm_80 %}
   4 ; RUN: %if ptxas-11.8 %{ llc < %s -march=nvptx64 -mcpu=sm_90 -mattr=+ptx78 | %ptxas-verify -arch=sm_90 %}
   5
   6 ; LDST: .b8 bfloat_array[8] = {1, 2, 3, 4, 5, 6, 7, 8};
   7 @"bfloat_array" = addrspace(1) constant [4 x bfloat]
   8                 [bfloat 0xR0201, bfloat 0xR0403, bfloat 0xR0605, bfloat 0xR0807]
   9
  10 ; CHECK-LABEL: test_fadd(
  11 ; CHECK-DAG:  ld.param.b16    [[A:%rs[0-9]+]], [test_fadd_param_0];
  12 ; CHECK-DAG:  ld.param.b16    [[B:%rs[0-9]+]], [test_fadd_param_1];
  13 ; SM90:       add.rn.bf16     [[R:%rs[0-9]+]], [[A]], [[B]];
  14 ;
  15 ; SM80-DAG:   cvt.f32.bf16    [[FA:%f[0-9]+]], [[A]];
  16 ; SM80-DAG:   cvt.f32.bf16    [[FB:%f[0-9]+]], [[B]];
  17 ; SM80:       add.rn.f32      [[FR:%f[0-9]+]], [[FA]], [[FB]];
  18 ; SM80:       cvt.rn.bf16.f32 [[R:%rs[0-9]+]], [[FR]];
  19 ; CHECK-NEXT: st.param.b16    [func_retval0+0], [[R]];
  20 ; CHECK-NEXT: ret;
  21
  22 define bfloat @test_fadd(bfloat %0, bfloat %1) {
  23   %3 = fadd bfloat %0, %1
  24   ret bfloat %3
  25 }
  26
  27 ; CHECK-LABEL: test_fsub(
  28 ; CHECK-DAG:  ld.param.b16    [[A:%rs[0-9]+]], [test_fsub_param_0];
  29 ; CHECK-DAG:  ld.param.b16    [[B:%rs[0-9]+]], [test_fsub_param_1];
  30 ; SM90:       sub.rn.bf16     [[R:%rs[0-9]+]], [[A]], [[B]];
  31 ;
  32 ; SM80-DAG:   cvt.f32.bf16    [[FA:%f[0-9]+]], [[A]];
  33 ; SM80-DAG:   cvt.f32.bf16    [[FB:%f[0-9]+]], [[B]];
  34 ; SM80:       sub.rn.f32      [[FR:%f[0-9]+]], [[FA]], [[FB]];
  35 ; SM80:       cvt.rn.bf16.f32 [[R:%rs[0-9]+]], [[FR]];
  36 ; CHECK-NEXT: st.param.b16    [func_retval0+0], [[R]];
  37 ; CHECK-NEXT: ret;
  38
  39 define bfloat @test_fsub(bfloat %0, bfloat %1) {
  40   %3 = fsub bfloat %0, %1
  41   ret bfloat %3
  42 }
  43
  44 ; CHECK-LABEL: test_faddx2(
  45 ; CHECK-DAG:  ld.param.b32    [[A:%r[0-9]+]], [test_faddx2_param_0];
  46 ; CHECK-DAG:  ld.param.b32    [[B:%r[0-9]+]], [test_faddx2_param_1];
  47 ; SM90:       add.rn.bf16x2   [[R:%r[0-9]+]], [[A]], [[B]];
  48
  49 ; SM80-DAG:   mov.b32         {[[A0:%rs[0-9]+]], [[A1:%rs[0-9]+]]}, [[A]];
  50 ; SM80-DAG:   mov.b32         {[[B0:%rs[0-9]+]], [[B1:%rs[0-9]+]]}, [[B]];
  51 ; SM80-DAG:   cvt.f32.bf16    [[FA1:%f[0-9]+]], [[A1]];
  52 ; SM80-DAG:   cvt.f32.bf16    [[FA0:%f[0-9]+]], [[A0]];
  53 ; SM80-DAG:   cvt.f32.bf16    [[FB0:%f[0-9]+]], [[B0]];
  54 ; SM80-DAG:   cvt.f32.bf16    [[FB1:%f[0-9]+]], [[B1]];
  55 ; SM80-DAG:   add.rn.f32      [[FR0:%f[0-9]+]], [[FA0]], [[FB0]];
  56 ; SM80-DAG:   add.rn.f32      [[FR1:%f[0-9]+]], [[FA1]], [[FB1]];
  57 ; SM80-DAG:   cvt.rn.bf16.f32 [[R0:%rs[0-9]+]], [[FR0]];
  58 ; SM80-DAG:   cvt.rn.bf16.f32 [[R1:%rs[0-9]+]], [[FR1]];
  59 ; SM80:       mov.b32         [[R:%r[0-9]+]], {[[R0]], [[R1]]};
  60 ; CHECK:      st.param.b32    [func_retval0+0], [[R]];
  61 ; CHECK:      ret;
  62
  63 define <2 x bfloat> @test_faddx2(<2 x bfloat> %a, <2 x bfloat> %b) #0 {
  64   %r = fadd <2 x bfloat> %a, %b
  65   ret <2 x bfloat> %r
  66 }
  67
  68 ; CHECK-LABEL: test_fsubx2(
  69 ; CHECK-DAG:  ld.param.b32    [[A:%r[0-9]+]], [test_fsubx2_param_0];
  70 ; CHECK-DAG:  ld.param.b32    [[B:%r[0-9]+]], [test_fsubx2_param_1];
  71 ; SM90:       sub.rn.bf16x2   [[R:%r[0-9]+]], [[A]], [[B]];
  72
  73 ; SM80-DAG:   mov.b32         {[[A0:%rs[0-9]+]], [[A1:%rs[0-9]+]]}, [[A]];
  74 ; SM80-DAG:   mov.b32         {[[B0:%rs[0-9]+]], [[B1:%rs[0-9]+]]}, [[B]];
  75 ; SM80-DAG:   cvt.f32.bf16    [[FA1:%f[0-9]+]], [[A1]];
  76 ; SM80-DAG:   cvt.f32.bf16    [[FA0:%f[0-9]+]], [[A0]];
  77 ; SM80-DAG:   cvt.f32.bf16    [[FB0:%f[0-9]+]], [[B0]];
  78 ; SM80-DAG:   cvt.f32.bf16    [[FB1:%f[0-9]+]], [[B1]];
  79 ; SM80-DAG:   sub.rn.f32      [[FR0:%f[0-9]+]], [[FA0]], [[FB0]];
  80 ; SM80-DAG:   sub.rn.f32      [[FR1:%f[0-9]+]], [[FA1]], [[FB1]];
  81 ; SM80-DAG:   cvt.rn.bf16.f32 [[R0:%rs[0-9]+]], [[FR0]];
  82 ; SM80-DAG:   cvt.rn.bf16.f32 [[R1:%rs[0-9]+]], [[FR1]];
  83 ; SM80:       mov.b32         [[R:%r[0-9]+]], {[[R0]], [[R1]]};
  84
  85 ; CHECK:      st.param.b32    [func_retval0+0], [[R]];
  86 ; CHECK:      ret;
  87
  88 define <2 x bfloat> @test_fsubx2(<2 x bfloat> %a, <2 x bfloat> %b) #0 {
  89   %r = fsub <2 x bfloat> %a, %b
  90   ret <2 x bfloat> %r
  91 }
  92
  93 ; CHECK-LABEL: test_fmulx2(
  94 ; CHECK-DAG:  ld.param.b32    [[A:%r[0-9]+]], [test_fmulx2_param_0];
  95 ; CHECK-DAG:  ld.param.b32    [[B:%r[0-9]+]], [test_fmulx2_param_1];
  96 ; SM90:       mul.rn.bf16x2   [[R:%r[0-9]+]], [[A]], [[B]];
  97
  98 ; SM80-DAG:   mov.b32         {[[A0:%rs[0-9]+]], [[A1:%rs[0-9]+]]}, [[A]];
  99 ; SM80-DAG:   mov.b32         {[[B0:%rs[0-9]+]], [[B1:%rs[0-9]+]]}, [[B]];
 100 ; SM80-DAG:   cvt.f32.bf16    [[FA1:%f[0-9]+]], [[A1]];
 101 ; SM80-DAG:   cvt.f32.bf16    [[FA0:%f[0-9]+]], [[A0]];
 102 ; SM80-DAG:   cvt.f32.bf16    [[FB0:%f[0-9]+]], [[B0]];
 103 ; SM80-DAG:   cvt.f32.bf16    [[FB1:%f[0-9]+]], [[B1]];
 104 ; SM80-DAG:   mul.rn.f32      [[FR0:%f[0-9]+]], [[FA0]], [[FB0]];
 105 ; SM80-DAG:   mul.rn.f32      [[FR1:%f[0-9]+]], [[FA1]], [[FB1]];
 106 ; SM80-DAG:   cvt.rn.bf16.f32 [[R0:%rs[0-9]+]], [[FR0]];
 107 ; SM80-DAG:   cvt.rn.bf16.f32 [[R1:%rs[0-9]+]], [[FR1]];
 108 ; SM80:       mov.b32         [[R:%r[0-9]+]], {[[R0]], [[R1]]};
 109
 110 ; CHECK:      st.param.b32    [func_retval0+0], [[R]];
 111 ; CHECK:      ret;
 112
 113 define <2 x bfloat> @test_fmulx2(<2 x bfloat> %a, <2 x bfloat> %b) #0 {
 114   %r = fmul <2 x bfloat> %a, %b
 115   ret <2 x bfloat> %r
 116 }
 117
 118 ; CHECK-LABEL: test_fdiv(
 119 ; CHECK-DAG:  ld.param.b32    [[A:%r[0-9]+]], [test_fdiv_param_0];
 120 ; CHECK-DAG:  ld.param.b32    [[B:%r[0-9]+]], [test_fdiv_param_1];
 121 ; CHECK-DAG:  mov.b32         {[[A0:%rs[0-9]+]], [[A1:%rs[0-9]+]]}, [[A]]
 122 ; CHECK-DAG:  mov.b32         {[[B0:%rs[0-9]+]], [[B1:%rs[0-9]+]]}, [[B]]
 123 ; CHECK-DAG:  cvt.f32.bf16     [[FA0:%f[0-9]+]], [[A0]];
 124 ; CHECK-DAG:  cvt.f32.bf16     [[FA1:%f[0-9]+]], [[A1]];
 125 ; CHECK-DAG:  cvt.f32.bf16     [[FB0:%f[0-9]+]], [[B0]];
 126 ; CHECK-DAG:  cvt.f32.bf16     [[FB1:%f[0-9]+]], [[B1]];
 127 ; CHECK-DAG:  div.rn.f32      [[FR0:%f[0-9]+]], [[FA0]], [[FB0]];
 128 ; CHECK-DAG:  div.rn.f32      [[FR1:%f[0-9]+]], [[FA1]], [[FB1]];
 129 ; CHECK-DAG:  cvt.rn.bf16.f32  [[R0:%rs[0-9]+]], [[FR0]];
 130 ; CHECK-DAG:  cvt.rn.bf16.f32  [[R1:%rs[0-9]+]], [[FR1]];
 131 ; CHECK-NEXT: mov.b32         [[R:%r[0-9]+]], {[[R0]], [[R1]]}
 132 ; CHECK-NEXT: st.param.b32    [func_retval0+0], [[R]];
 133 ; CHECK-NEXT: ret;
 134
 135 define <2 x bfloat> @test_fdiv(<2 x bfloat> %a, <2 x bfloat> %b) #0 {
 136   %r = fdiv <2 x bfloat> %a, %b
 137   ret <2 x bfloat> %r
 138 }
 139
 140 ; CHECK-LABEL: test_extract_0(
 141 ; CHECK:      ld.param.b16    [[A:%rs[0-9]+]], [test_extract_0_param_0];
 142 ; CHECK:      st.param.b16    [func_retval0+0], [[A]];
 143 ; CHECK:      ret;
 144
 145 define bfloat @test_extract_0(<2 x bfloat> %a) #0 {
 146   %e = extractelement <2 x bfloat> %a, i32 0
 147   ret bfloat %e
 148 }
 149
 150 ; CHECK-LABEL: test_extract_1(
 151 ; CHECK:      ld.param.b16    [[A:%rs[0-9]+]], [test_extract_1_param_0+2];
 152 ; CHECK:      st.param.b16    [func_retval0+0], [[A]];
 153 ; CHECK:      ret;
 154
 155 define bfloat @test_extract_1(<2 x bfloat> %a) #0 {
 156   %e = extractelement <2 x bfloat> %a, i32 1
 157   ret bfloat %e
 158 }
 159
 160 ; CHECK-LABEL: test_fpext_float(
 161 ; CHECK:      ld.param.b16    [[A:%rs[0-9]+]], [test_fpext_float_param_0];
 162 ; CHECK:      cvt.f32.bf16     [[R:%f[0-9]+]], [[A]];
 163 ; CHECK:      st.param.f32    [func_retval0+0], [[R]];
 164 ; CHECK:      ret;
 165 define float @test_fpext_float(bfloat %a) #0 {
 166   %r = fpext bfloat %a to float
 167   ret float %r
 168 }
 169
 170 ; CHECK-LABEL: test_fptrunc_float(
 171 ; CHECK:      ld.param.f32    [[A:%f[0-9]+]], [test_fptrunc_float_param_0];
 172 ; CHECK:      cvt.rn.bf16.f32  [[R:%rs[0-9]+]], [[A]];
 173 ; CHECK:      st.param.b16    [func_retval0+0], [[R]];
 174 ; CHECK:      ret;
 175 define bfloat @test_fptrunc_float(float %a) #0 {
 176   %r = fptrunc float %a to bfloat
 177   ret bfloat %r
 178 }
 179
 180 ; CHECK-LABEL: test_fadd_imm_1(
 181 ; CHECK:      ld.param.b16    [[A:%rs[0-9]+]], [test_fadd_imm_1_param_0];
 182 ; SM90:       mov.b16         [[B:%rs[0-9]+]], 0x3F80;
 183 ; SM90:       add.rn.bf16     [[R:%rs[0-9]+]], [[A]], [[B]];
 184
 185 ; SM80-DAG:   cvt.f32.bf16    [[FA:%f[0-9]+]], [[A]];
 186 ; SM80:       add.rn.f32      [[FR:%f[0-9]+]], [[FA]], 0f3F800000;
 187 ; SM80:       cvt.rn.bf16.f32 [[R:%rs[0-9]+]], [[FR]];
 188
 189 ; CHECK:      st.param.b16    [func_retval0+0], [[R]];
 190 ; CHECK-NEXT: ret;
 191 define bfloat @test_fadd_imm_1(bfloat %a) #0 {
 192   %r = fadd bfloat %a, 1.0
 193   ret bfloat %r
 194 }
 195
 196 ; CHECK-LABEL: test_select_cc_bf16_f64(
 197 ; CHECK-DAG:      ld.param.f64    [[A:%fd[0-9]+]], [test_select_cc_bf16_f64_param_0];
 198 ; CHECK-DAG:      ld.param.f64    [[B:%fd[0-9]+]], [test_select_cc_bf16_f64_param_1];
 199 ; CHECK:          setp.lt.f64     [[P:%p[0-9]+]], [[A]], [[B]];
 200 ; CHECK-DAG:      ld.param.b16    [[C:%rs[0-9]+]], [test_select_cc_bf16_f64_param_2];
 201 ; CHECK-DAG:      ld.param.b16    [[D:%rs[0-9]+]], [test_select_cc_bf16_f64_param_3];
 202 ; CHECK:          selp.b16        [[R:%rs[0-9]+]], [[C]], [[D]], [[P]];
 203 ; CHECK-NEXT:     st.param.b16    [func_retval0+0], [[R]];
 204 ; CHECK-NEXT:     ret;
 205 define bfloat @test_select_cc_bf16_f64(double %a, double %b, bfloat %c, bfloat %d) #0 {
 206   %cc = fcmp olt double %a, %b
 207   %r = select i1 %cc, bfloat %c, bfloat %d
 208   ret bfloat %r
 209 }