; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
; RUN: llc < %s -march=nvptx64 -mcpu=sm_90 -mattr=+ptx88 | FileCheck %s -check-prefixes=CHECK,SM90
; RUN: %if ptxas-sm_90 && ptxas-isa-8.8 %{ llc < %s -march=nvptx64 -mcpu=sm_90 -mattr=+ptx88 | %ptxas-verify -arch=sm_90 %}
; RUN: llc < %s -march=nvptx64 -mcpu=sm_100 -mattr=+ptx88 | FileCheck %s -check-prefixes=CHECK,SM100
; RUN: %if ptxas-sm_100 && ptxas-isa-8.8 %{ llc < %s -march=nvptx64 -mcpu=sm_100 -mattr=+ptx88 | %ptxas-verify -arch=sm_100 %}

; This test is based on load-store-vectors.ll,
; and contains testing for lowering 256-bit masked vector stores

; Types we are checking: i32, i64, f32, f64

; Address spaces we are checking: generic, global
; - Global is the only address space that currently supports masked stores.
; - The generic stores will get legalized before the backend via scalarization,
;   this file tests that even though we won't be generating them in the LSV.

; 256-bit vector loads/stores are only legal for blackwell+, so on sm_90, the vectors will be split

; generic address space

define void @generic_8xi32(ptr %a, ptr %b) {
; CHECK-LABEL: generic_8xi32(
; CHECK:       {
; CHECK-NEXT:    .reg .b32 %r<9>;
; CHECK-NEXT:    .reg .b64 %rd<3>;
; CHECK-EMPTY:
; CHECK-NEXT:  // %bb.0:
; CHECK-NEXT:    ld.param.b64 %rd1, [generic_8xi32_param_0];
; CHECK-NEXT:    ld.v4.b32 {%r1, %r2, %r3, %r4}, [%rd1+16];
; CHECK-NEXT:    ld.v4.b32 {%r5, %r6, %r7, %r8}, [%rd1];
; CHECK-NEXT:    ld.param.b64 %rd2, [generic_8xi32_param_1];
; CHECK-NEXT:    st.b32 [%rd2], %r5;
; CHECK-NEXT:    st.b32 [%rd2+8], %r7;
; CHECK-NEXT:    st.b32 [%rd2+28], %r4;
; CHECK-NEXT:    ret;
  %a.load = load <8 x i32>, ptr %a
  tail call void @llvm.masked.store.v8i32.p0(<8 x i32> %a.load, ptr align 32 %b, <8 x i1> <i1 true, i1 false, i1 true, i1 false, i1 false, i1 false, i1 false, i1 true>)
  ret void
}

define void @generic_4xi64(ptr %a, ptr %b) {
; CHECK-LABEL: generic_4xi64(
; CHECK:       {
; CHECK-NEXT:    .reg .b64 %rd<7>;
; CHECK-EMPTY:
; CHECK-NEXT:  // %bb.0:
; CHECK-NEXT:    ld.param.b64 %rd1, [generic_4xi64_param_0];
; CHECK-NEXT:    ld.v2.b64 {%rd2, %rd3}, [%rd1+16];
; CHECK-NEXT:    ld.v2.b64 {%rd4, %rd5}, [%rd1];
; CHECK-NEXT:    ld.param.b64 %rd6, [generic_4xi64_param_1];
; CHECK-NEXT:    st.b64 [%rd6], %rd4;
; CHECK-NEXT:    st.b64 [%rd6+16], %rd2;
; CHECK-NEXT:    ret;
  %a.load = load <4 x i64>, ptr %a
  tail call void @llvm.masked.store.v4i64.p0(<4 x i64> %a.load, ptr align 32 %b, <4 x i1> <i1 true, i1 false, i1 true, i1 false>)
  ret void
}

define void @generic_8xfloat(ptr %a, ptr %b) {
; CHECK-LABEL: generic_8xfloat(
; CHECK:       {
; CHECK-NEXT:    .reg .b32 %r<9>;
; CHECK-NEXT:    .reg .b64 %rd<3>;
; CHECK-EMPTY:
; CHECK-NEXT:  // %bb.0:
; CHECK-NEXT:    ld.param.b64 %rd1, [generic_8xfloat_param_0];
; CHECK-NEXT:    ld.v4.b32 {%r1, %r2, %r3, %r4}, [%rd1+16];
; CHECK-NEXT:    ld.v4.b32 {%r5, %r6, %r7, %r8}, [%rd1];
; CHECK-NEXT:    ld.param.b64 %rd2, [generic_8xfloat_param_1];
; CHECK-NEXT:    st.b32 [%rd2], %r5;
; CHECK-NEXT:    st.b32 [%rd2+8], %r7;
; CHECK-NEXT:    st.b32 [%rd2+28], %r4;
; CHECK-NEXT:    ret;
  %a.load = load <8 x float>, ptr %a
  tail call void @llvm.masked.store.v8f32.p0(<8 x float> %a.load, ptr align 32 %b, <8 x i1> <i1 true, i1 false, i1 true, i1 false, i1 false, i1 false, i1 false, i1 true>)
  ret void
}

define void @generic_4xdouble(ptr %a, ptr %b) {
; CHECK-LABEL: generic_4xdouble(
; CHECK:       {
; CHECK-NEXT:    .reg .b64 %rd<7>;
; CHECK-EMPTY:
; CHECK-NEXT:  // %bb.0:
; CHECK-NEXT:    ld.param.b64 %rd1, [generic_4xdouble_param_0];
; CHECK-NEXT:    ld.v2.b64 {%rd2, %rd3}, [%rd1+16];
; CHECK-NEXT:    ld.v2.b64 {%rd4, %rd5}, [%rd1];
; CHECK-NEXT:    ld.param.b64 %rd6, [generic_4xdouble_param_1];
; CHECK-NEXT:    st.b64 [%rd6], %rd4;
; CHECK-NEXT:    st.b64 [%rd6+16], %rd2;
; CHECK-NEXT:    ret;
  %a.load = load <4 x double>, ptr %a
  tail call void @llvm.masked.store.v4f64.p0(<4 x double> %a.load, ptr align 32 %b, <4 x i1> <i1 true, i1 false, i1 true, i1 false>)
  ret void
}

; global address space

define void @global_8xi32(ptr addrspace(1) %a, ptr addrspace(1) %b) {
; SM90-LABEL: global_8xi32(
; SM90:       {
; SM90-NEXT:    .reg .b32 %r<9>;
; SM90-NEXT:    .reg .b64 %rd<3>;
; SM90-EMPTY:
; SM90-NEXT:  // %bb.0:
; SM90-NEXT:    ld.param.b64 %rd1, [global_8xi32_param_0];
; SM90-NEXT:    ld.global.v4.b32 {%r1, %r2, %r3, %r4}, [%rd1+16];
; SM90-NEXT:    ld.global.v4.b32 {%r5, %r6, %r7, %r8}, [%rd1];
; SM90-NEXT:    ld.param.b64 %rd2, [global_8xi32_param_1];
; SM90-NEXT:    st.global.b32 [%rd2], %r5;
; SM90-NEXT:    st.global.b32 [%rd2+8], %r7;
; SM90-NEXT:    st.global.b32 [%rd2+28], %r4;
; SM90-NEXT:    ret;
;
; SM100-LABEL: global_8xi32(
; SM100:       {
; SM100-NEXT:    .reg .b32 %r<9>;
; SM100-NEXT:    .reg .b64 %rd<3>;
; SM100-EMPTY:
; SM100-NEXT:  // %bb.0:
; SM100-NEXT:    ld.param.b64 %rd1, [global_8xi32_param_0];
; SM100-NEXT:    ld.global.v8.b32 {%r1, %r2, %r3, %r4, %r5, %r6, %r7, %r8}, [%rd1];
; SM100-NEXT:    ld.param.b64 %rd2, [global_8xi32_param_1];
; SM100-NEXT:    st.global.v8.b32 [%rd2], {%r1, _, %r3, _, _, _, _, %r8};
; SM100-NEXT:    ret;
  %a.load = load <8 x i32>, ptr addrspace(1) %a
  tail call void @llvm.masked.store.v8i32.p1(<8 x i32> %a.load, ptr addrspace(1) align 32 %b, <8 x i1> <i1 true, i1 false, i1 true, i1 false, i1 false, i1 false, i1 false, i1 true>)
  ret void
}

define void @global_4xi64(ptr addrspace(1) %a, ptr addrspace(1) %b) {
; SM90-LABEL: global_4xi64(
; SM90:       {
; SM90-NEXT:    .reg .b64 %rd<7>;
; SM90-EMPTY:
; SM90-NEXT:  // %bb.0:
; SM90-NEXT:    ld.param.b64 %rd1, [global_4xi64_param_0];
; SM90-NEXT:    ld.global.v2.b64 {%rd2, %rd3}, [%rd1+16];
; SM90-NEXT:    ld.global.v2.b64 {%rd4, %rd5}, [%rd1];
; SM90-NEXT:    ld.param.b64 %rd6, [global_4xi64_param_1];
; SM90-NEXT:    st.global.b64 [%rd6], %rd4;
; SM90-NEXT:    st.global.b64 [%rd6+16], %rd2;
; SM90-NEXT:    ret;
;
; SM100-LABEL: global_4xi64(
; SM100:       {
; SM100-NEXT:    .reg .b64 %rd<7>;
; SM100-EMPTY:
; SM100-NEXT:  // %bb.0:
; SM100-NEXT:    ld.param.b64 %rd1, [global_4xi64_param_0];
; SM100-NEXT:    ld.global.v4.b64 {%rd2, %rd3, %rd4, %rd5}, [%rd1];
; SM100-NEXT:    ld.param.b64 %rd6, [global_4xi64_param_1];
; SM100-NEXT:    st.global.v4.b64 [%rd6], {%rd2, _, %rd4, _};
; SM100-NEXT:    ret;
  %a.load = load <4 x i64>, ptr addrspace(1) %a
  tail call void @llvm.masked.store.v4i64.p1(<4 x i64> %a.load, ptr addrspace(1) align 32 %b, <4 x i1> <i1 true, i1 false, i1 true, i1 false>)
  ret void
}

define void @global_8xfloat(ptr addrspace(1) %a, ptr addrspace(1) %b) {
; SM90-LABEL: global_8xfloat(
; SM90:       {
; SM90-NEXT:    .reg .b32 %r<9>;
; SM90-NEXT:    .reg .b64 %rd<3>;
; SM90-EMPTY:
; SM90-NEXT:  // %bb.0:
; SM90-NEXT:    ld.param.b64 %rd1, [global_8xfloat_param_0];
; SM90-NEXT:    ld.global.v4.b32 {%r1, %r2, %r3, %r4}, [%rd1+16];
; SM90-NEXT:    ld.global.v4.b32 {%r5, %r6, %r7, %r8}, [%rd1];
; SM90-NEXT:    ld.param.b64 %rd2, [global_8xfloat_param_1];
; SM90-NEXT:    st.global.b32 [%rd2], %r5;
; SM90-NEXT:    st.global.b32 [%rd2+8], %r7;
; SM90-NEXT:    st.global.b32 [%rd2+28], %r4;
; SM90-NEXT:    ret;
;
; SM100-LABEL: global_8xfloat(
; SM100:       {
; SM100-NEXT:    .reg .b32 %r<9>;
; SM100-NEXT:    .reg .b64 %rd<3>;
; SM100-EMPTY:
; SM100-NEXT:  // %bb.0:
; SM100-NEXT:    ld.param.b64 %rd1, [global_8xfloat_param_0];
; SM100-NEXT:    ld.global.v8.b32 {%r1, %r2, %r3, %r4, %r5, %r6, %r7, %r8}, [%rd1];
; SM100-NEXT:    ld.param.b64 %rd2, [global_8xfloat_param_1];
; SM100-NEXT:    st.global.v8.b32 [%rd2], {%r1, _, %r3, _, _, _, _, %r8};
; SM100-NEXT:    ret;
  %a.load = load <8 x float>, ptr addrspace(1) %a
  tail call void @llvm.masked.store.v8f32.p1(<8 x float> %a.load, ptr addrspace(1) align 32 %b, <8 x i1> <i1 true, i1 false, i1 true, i1 false, i1 false, i1 false, i1 false, i1 true>)
  ret void
}

define void @global_4xdouble(ptr addrspace(1) %a, ptr addrspace(1) %b) {
; SM90-LABEL: global_4xdouble(
; SM90:       {
; SM90-NEXT:    .reg .b64 %rd<7>;
; SM90-EMPTY:
; SM90-NEXT:  // %bb.0:
; SM90-NEXT:    ld.param.b64 %rd1, [global_4xdouble_param_0];
; SM90-NEXT:    ld.global.v2.b64 {%rd2, %rd3}, [%rd1+16];
; SM90-NEXT:    ld.global.v2.b64 {%rd4, %rd5}, [%rd1];
; SM90-NEXT:    ld.param.b64 %rd6, [global_4xdouble_param_1];
; SM90-NEXT:    st.global.b64 [%rd6], %rd4;
; SM90-NEXT:    st.global.b64 [%rd6+16], %rd2;
; SM90-NEXT:    ret;
;
; SM100-LABEL: global_4xdouble(
; SM100:       {
; SM100-NEXT:    .reg .b64 %rd<7>;
; SM100-EMPTY:
; SM100-NEXT:  // %bb.0:
; SM100-NEXT:    ld.param.b64 %rd1, [global_4xdouble_param_0];
; SM100-NEXT:    ld.global.v4.b64 {%rd2, %rd3, %rd4, %rd5}, [%rd1];
; SM100-NEXT:    ld.param.b64 %rd6, [global_4xdouble_param_1];
; SM100-NEXT:    st.global.v4.b64 [%rd6], {%rd2, _, %rd4, _};
; SM100-NEXT:    ret;
  %a.load = load <4 x double>, ptr addrspace(1) %a
  tail call void @llvm.masked.store.v4f64.p1(<4 x double> %a.load, ptr addrspace(1) align 32 %b, <4 x i1> <i1 true, i1 false, i1 true, i1 false>)
  ret void
}

; edge cases
define void @global_8xi32_all_mask_on(ptr addrspace(1) %a, ptr addrspace(1) %b) {
; SM90-LABEL: global_8xi32_all_mask_on(
; SM90:       {
; SM90-NEXT:    .reg .b32 %r<9>;
; SM90-NEXT:    .reg .b64 %rd<3>;
; SM90-EMPTY:
; SM90-NEXT:  // %bb.0:
; SM90-NEXT:    ld.param.b64 %rd1, [global_8xi32_all_mask_on_param_0];
; SM90-NEXT:    ld.global.v4.b32 {%r1, %r2, %r3, %r4}, [%rd1];
; SM90-NEXT:    ld.global.v4.b32 {%r5, %r6, %r7, %r8}, [%rd1+16];
; SM90-NEXT:    ld.param.b64 %rd2, [global_8xi32_all_mask_on_param_1];
; SM90-NEXT:    st.global.v4.b32 [%rd2+16], {%r5, %r6, %r7, %r8};
; SM90-NEXT:    st.global.v4.b32 [%rd2], {%r1, %r2, %r3, %r4};
; SM90-NEXT:    ret;
;
; SM100-LABEL: global_8xi32_all_mask_on(
; SM100:       {
; SM100-NEXT:    .reg .b64 %rd<7>;
; SM100-EMPTY:
; SM100-NEXT:  // %bb.0:
; SM100-NEXT:    ld.param.b64 %rd1, [global_8xi32_all_mask_on_param_0];
; SM100-NEXT:    ld.global.v4.b64 {%rd2, %rd3, %rd4, %rd5}, [%rd1];
; SM100-NEXT:    ld.param.b64 %rd6, [global_8xi32_all_mask_on_param_1];
; SM100-NEXT:    st.global.v4.b64 [%rd6], {%rd2, %rd3, %rd4, %rd5};
; SM100-NEXT:    ret;
  %a.load = load <8 x i32>, ptr addrspace(1) %a
  tail call void @llvm.masked.store.v8i32.p1(<8 x i32> %a.load, ptr addrspace(1) align 32 %b, <8 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>)
  ret void
}

define void @global_8xi32_all_mask_off(ptr addrspace(1) %a, ptr addrspace(1) %b) {
; CHECK-LABEL: global_8xi32_all_mask_off(
; CHECK:       {
; CHECK-EMPTY:
; CHECK-EMPTY:
; CHECK-NEXT:  // %bb.0:
; CHECK-NEXT:    ret;
  %a.load = load <8 x i32>, ptr addrspace(1) %a
  tail call void @llvm.masked.store.v8i32.p1(<8 x i32> %a.load, ptr addrspace(1) align 32 %b, <8 x i1> <i1 false, i1 false, i1 false, i1 false, i1 false, i1 false, i1 false, i1 false>)
  ret void
}

; This is an example pattern for the LSV's output of these masked stores
define void @vectorizerOutput(ptr addrspace(1) %in, ptr addrspace(1) %out) {
; SM90-LABEL: vectorizerOutput(
; SM90:       {
; SM90-NEXT:    .reg .b32 %r<9>;
; SM90-NEXT:    .reg .b64 %rd<3>;
; SM90-EMPTY:
; SM90-NEXT:  // %bb.0:
; SM90-NEXT:    ld.param.b64 %rd1, [vectorizerOutput_param_0];
; SM90-NEXT:    ld.global.v4.b32 {%r1, %r2, %r3, %r4}, [%rd1+16];
; SM90-NEXT:    ld.global.v4.b32 {%r5, %r6, %r7, %r8}, [%rd1];
; SM90-NEXT:    ld.param.b64 %rd2, [vectorizerOutput_param_1];
; SM90-NEXT:    st.global.b32 [%rd2], %r5;
; SM90-NEXT:    st.global.b32 [%rd2+4], %r6;
; SM90-NEXT:    st.global.b32 [%rd2+12], %r8;
; SM90-NEXT:    st.global.b32 [%rd2+16], %r1;
; SM90-NEXT:    ret;
;
; SM100-LABEL: vectorizerOutput(
; SM100:       {
; SM100-NEXT:    .reg .b32 %r<9>;
; SM100-NEXT:    .reg .b64 %rd<3>;
; SM100-EMPTY:
; SM100-NEXT:  // %bb.0:
; SM100-NEXT:    ld.param.b64 %rd1, [vectorizerOutput_param_0];
; SM100-NEXT:    ld.global.v8.b32 {%r1, %r2, %r3, %r4, %r5, %r6, %r7, %r8}, [%rd1];
; SM100-NEXT:    ld.param.b64 %rd2, [vectorizerOutput_param_1];
; SM100-NEXT:    st.global.v8.b32 [%rd2], {%r1, %r2, _, %r4, %r5, _, _, _};
; SM100-NEXT:    ret;
  %1 = load <8 x i32>, ptr addrspace(1) %in, align 32
  %load05 = extractelement <8 x i32> %1, i32 0
  %load16 = extractelement <8 x i32> %1, i32 1
  %load38 = extractelement <8 x i32> %1, i32 3
  %load49 = extractelement <8 x i32> %1, i32 4
  %2 = insertelement <8 x i32> poison, i32 %load05, i32 0
  %3 = insertelement <8 x i32> %2, i32 %load16, i32 1
  %4 = insertelement <8 x i32> %3, i32 poison, i32 2
  %5 = insertelement <8 x i32> %4, i32 %load38, i32 3
  %6 = insertelement <8 x i32> %5, i32 %load49, i32 4
  %7 = insertelement <8 x i32> %6, i32 poison, i32 5
  %8 = insertelement <8 x i32> %7, i32 poison, i32 6
  %9 = insertelement <8 x i32> %8, i32 poison, i32 7
  call void @llvm.masked.store.v8i32.p1(<8 x i32> %9, ptr addrspace(1) align 32 %out, <8 x i1> <i1 true, i1 true, i1 false, i1 true, i1 true, i1 false, i1 false, i1 false>)
  ret void
}

declare void @llvm.masked.store.v8i32.p0(<8 x i32>, ptr, <8 x i1>)
declare void @llvm.masked.store.v4i64.p0(<4 x i64>, ptr, <4 x i1>)
declare void @llvm.masked.store.v8f32.p0(<8 x float>, ptr, <8 x i1>)
declare void @llvm.masked.store.v4f64.p0(<4 x double>, ptr, <4 x i1>)

declare void @llvm.masked.store.v8i32.p1(<8 x i32>, ptr addrspace(1), <8 x i1>)
declare void @llvm.masked.store.v4i64.p1(<4 x i64>, ptr addrspace(1), <4 x i1>)
declare void @llvm.masked.store.v8f32.p1(<8 x float>, ptr addrspace(1), <8 x i1>)
declare void @llvm.masked.store.v4f64.p1(<4 x double>, ptr addrspace(1), <4 x i1>)