test/Transforms/SLPVectorizer/AMDGPU/packed-fp32-softmax.ll

; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
; RUN: opt -S -mtriple=amdgpu9.00-amd-amdhsa -passes=slp-vectorizer < %s | FileCheck -check-prefixes=GFX900 %s
; RUN: opt -S -mtriple=amdgpu9.50-amd-amdhsa -passes=slp-vectorizer < %s | FileCheck -check-prefixes=GFX950 %s
; RUN: opt -S -mtriple=amdgpu12.50-amd-amdhsa -passes=slp-vectorizer < %s | FileCheck -check-prefixes=GFX1250 %s

; Reduced from an MXFP softmax kernel on gfx1250. Packed-f32 pair formation
; is priced too high for these shapes, so both stay scalar. gfx900 has no
; packed f32 ops; gfx950 does. The three targets produce the same IR, and
; each has its own checks.

; The softmax scale is multiplied into two lanes of a loaded pair, and each
; product is then broadcast over a <4 x float> operand. The two fmuls stay
; scalar, and each product is splatted with its own insertelement.
define amdgpu_kernel void @softmax_scale_splat(ptr addrspace(3) %in, ptr addrspace(3) %out, float %scale, <4 x float> %x0, <4 x float> %x1) {
; GFX900-LABEL: define amdgpu_kernel void @softmax_scale_splat(
; GFX900-SAME: ptr addrspace(3) [[IN:%.*]], ptr addrspace(3) [[OUT:%.*]], float [[SCALE:%.*]], <4 x float> [[X0:%.*]], <4 x float> [[X1:%.*]]) {
; GFX900-NEXT:    [[M:%.*]] = load <2 x float>, ptr addrspace(3) [[IN]], align 16
; GFX900-NEXT:    [[M0:%.*]] = extractelement <2 x float> [[M]], i64 0
; GFX900-NEXT:    [[M1:%.*]] = extractelement <2 x float> [[M]], i64 1
; GFX900-NEXT:    [[B0:%.*]] = fmul float [[SCALE]], [[M0]]
; GFX900-NEXT:    [[B1:%.*]] = fmul float [[SCALE]], [[M1]]
; GFX900-NEXT:    [[I0:%.*]] = insertelement <4 x float> poison, float [[B0]], i64 0
; GFX900-NEXT:    [[S0:%.*]] = shufflevector <4 x float> [[I0]], <4 x float> poison, <4 x i32> zeroinitializer
; GFX900-NEXT:    [[I1:%.*]] = insertelement <4 x float> poison, float [[B1]], i64 0
; GFX900-NEXT:    [[S1:%.*]] = shufflevector <4 x float> [[I1]], <4 x float> poison, <4 x i32> zeroinitializer
; GFX900-NEXT:    [[R0:%.*]] = fsub <4 x float> [[X0]], [[S0]]
; GFX900-NEXT:    [[R1:%.*]] = fsub <4 x float> [[X1]], [[S1]]
; GFX900-NEXT:    store <4 x float> [[R0]], ptr addrspace(3) [[OUT]], align 16
; GFX900-NEXT:    [[O1:%.*]] = getelementptr inbounds <4 x float>, ptr addrspace(3) [[OUT]], i32 1
; GFX900-NEXT:    store <4 x float> [[R1]], ptr addrspace(3) [[O1]], align 16
; GFX900-NEXT:    ret void
;
; GFX950-LABEL: define amdgpu_kernel void @softmax_scale_splat(
; GFX950-SAME: ptr addrspace(3) [[IN:%.*]], ptr addrspace(3) [[OUT:%.*]], float [[SCALE:%.*]], <4 x float> [[X0:%.*]], <4 x float> [[X1:%.*]]) {
; GFX950-NEXT:    [[M:%.*]] = load <2 x float>, ptr addrspace(3) [[IN]], align 16
; GFX950-NEXT:    [[M0:%.*]] = extractelement <2 x float> [[M]], i64 0
; GFX950-NEXT:    [[M1:%.*]] = extractelement <2 x float> [[M]], i64 1
; GFX950-NEXT:    [[B0:%.*]] = fmul float [[SCALE]], [[M0]]
; GFX950-NEXT:    [[B1:%.*]] = fmul float [[SCALE]], [[M1]]
; GFX950-NEXT:    [[I0:%.*]] = insertelement <4 x float> poison, float [[B0]], i64 0
; GFX950-NEXT:    [[S0:%.*]] = shufflevector <4 x float> [[I0]], <4 x float> poison, <4 x i32> zeroinitializer
; GFX950-NEXT:    [[I1:%.*]] = insertelement <4 x float> poison, float [[B1]], i64 0
; GFX950-NEXT:    [[S1:%.*]] = shufflevector <4 x float> [[I1]], <4 x float> poison, <4 x i32> zeroinitializer
; GFX950-NEXT:    [[R0:%.*]] = fsub <4 x float> [[X0]], [[S0]]
; GFX950-NEXT:    [[R1:%.*]] = fsub <4 x float> [[X1]], [[S1]]
; GFX950-NEXT:    store <4 x float> [[R0]], ptr addrspace(3) [[OUT]], align 16
; GFX950-NEXT:    [[O1:%.*]] = getelementptr inbounds <4 x float>, ptr addrspace(3) [[OUT]], i32 1
; GFX950-NEXT:    store <4 x float> [[R1]], ptr addrspace(3) [[O1]], align 16
; GFX950-NEXT:    ret void
;
; GFX1250-LABEL: define amdgpu_kernel void @softmax_scale_splat(
; GFX1250-SAME: ptr addrspace(3) [[IN:%.*]], ptr addrspace(3) [[OUT:%.*]], float [[SCALE:%.*]], <4 x float> [[X0:%.*]], <4 x float> [[X1:%.*]]) {
; GFX1250-NEXT:    [[M:%.*]] = load <2 x float>, ptr addrspace(3) [[IN]], align 16
; GFX1250-NEXT:    [[M0:%.*]] = extractelement <2 x float> [[M]], i64 0
; GFX1250-NEXT:    [[M1:%.*]] = extractelement <2 x float> [[M]], i64 1
; GFX1250-NEXT:    [[B0:%.*]] = fmul float [[SCALE]], [[M0]]
; GFX1250-NEXT:    [[B1:%.*]] = fmul float [[SCALE]], [[M1]]
; GFX1250-NEXT:    [[I0:%.*]] = insertelement <4 x float> poison, float [[B0]], i64 0
; GFX1250-NEXT:    [[S0:%.*]] = shufflevector <4 x float> [[I0]], <4 x float> poison, <4 x i32> zeroinitializer
; GFX1250-NEXT:    [[I1:%.*]] = insertelement <4 x float> poison, float [[B1]], i64 0
; GFX1250-NEXT:    [[S1:%.*]] = shufflevector <4 x float> [[I1]], <4 x float> poison, <4 x i32> zeroinitializer
; GFX1250-NEXT:    [[R0:%.*]] = fsub <4 x float> [[X0]], [[S0]]
; GFX1250-NEXT:    [[R1:%.*]] = fsub <4 x float> [[X1]], [[S1]]
; GFX1250-NEXT:    store <4 x float> [[R0]], ptr addrspace(3) [[OUT]], align 16
; GFX1250-NEXT:    [[O1:%.*]] = getelementptr inbounds <4 x float>, ptr addrspace(3) [[OUT]], i32 1
; GFX1250-NEXT:    store <4 x float> [[R1]], ptr addrspace(3) [[O1]], align 16
; GFX1250-NEXT:    ret void
;
  %m = load <2 x float>, ptr addrspace(3) %in, align 16
  %m0 = extractelement <2 x float> %m, i64 0
  %m1 = extractelement <2 x float> %m, i64 1
  %b0 = fmul float %scale, %m0
  %b1 = fmul float %scale, %m1
  %i0 = insertelement <4 x float> poison, float %b0, i64 0
  %s0 = shufflevector <4 x float> %i0, <4 x float> poison, <4 x i32> zeroinitializer
  %i1 = insertelement <4 x float> poison, float %b1, i64 0
  %s1 = shufflevector <4 x float> %i1, <4 x float> poison, <4 x i32> zeroinitializer
  %r0 = fsub <4 x float> %x0, %s0
  %r1 = fsub <4 x float> %x1, %s1
  store <4 x float> %r0, ptr addrspace(3) %out, align 16
  %o1 = getelementptr inbounds <4 x float>, ptr addrspace(3) %out, i32 1
  store <4 x float> %r1, ptr addrspace(3) %o1, align 16
  ret void
}

; The row sums of two <4 x float> accumulators are reduced lane by lane and
; the two sums are stored as a pair. The reduction stays six scalar fadds, and
; the stored pair is built with insertelement.
define amdgpu_kernel void @softmax_row_sum(ptr addrspace(3) %out, <4 x float> %p, <4 x float> %q) {
; GFX900-LABEL: define amdgpu_kernel void @softmax_row_sum(
; GFX900-SAME: ptr addrspace(3) [[OUT:%.*]], <4 x float> [[P:%.*]], <4 x float> [[Q:%.*]]) {
; GFX900-NEXT:    [[P0:%.*]] = extractelement <4 x float> [[P]], i64 0
; GFX900-NEXT:    [[P1:%.*]] = extractelement <4 x float> [[P]], i64 1
; GFX900-NEXT:    [[P2:%.*]] = extractelement <4 x float> [[P]], i64 2
; GFX900-NEXT:    [[P3:%.*]] = extractelement <4 x float> [[P]], i64 3
; GFX900-NEXT:    [[Q0:%.*]] = extractelement <4 x float> [[Q]], i64 0
; GFX900-NEXT:    [[Q1:%.*]] = extractelement <4 x float> [[Q]], i64 1
; GFX900-NEXT:    [[Q2:%.*]] = extractelement <4 x float> [[Q]], i64 2
; GFX900-NEXT:    [[Q3:%.*]] = extractelement <4 x float> [[Q]], i64 3
; GFX900-NEXT:    [[PA:%.*]] = fadd float [[P0]], [[P1]]
; GFX900-NEXT:    [[PB:%.*]] = fadd float [[P2]], [[P3]]
; GFX900-NEXT:    [[PS:%.*]] = fadd float [[PA]], [[PB]]
; GFX900-NEXT:    [[QA:%.*]] = fadd float [[Q0]], [[Q1]]
; GFX900-NEXT:    [[QB:%.*]] = fadd float [[Q2]], [[Q3]]
; GFX900-NEXT:    [[QS:%.*]] = fadd float [[QA]], [[QB]]
; GFX900-NEXT:    [[V0:%.*]] = insertelement <2 x float> poison, float [[PS]], i64 0
; GFX900-NEXT:    [[V1:%.*]] = insertelement <2 x float> [[V0]], float [[QS]], i64 1
; GFX900-NEXT:    store <2 x float> [[V1]], ptr addrspace(3) [[OUT]], align 8
; GFX900-NEXT:    ret void
;
; GFX950-LABEL: define amdgpu_kernel void @softmax_row_sum(
; GFX950-SAME: ptr addrspace(3) [[OUT:%.*]], <4 x float> [[P:%.*]], <4 x float> [[Q:%.*]]) {
; GFX950-NEXT:    [[P0:%.*]] = extractelement <4 x float> [[P]], i64 0
; GFX950-NEXT:    [[P1:%.*]] = extractelement <4 x float> [[P]], i64 1
; GFX950-NEXT:    [[P2:%.*]] = extractelement <4 x float> [[P]], i64 2
; GFX950-NEXT:    [[P3:%.*]] = extractelement <4 x float> [[P]], i64 3
; GFX950-NEXT:    [[Q0:%.*]] = extractelement <4 x float> [[Q]], i64 0
; GFX950-NEXT:    [[Q1:%.*]] = extractelement <4 x float> [[Q]], i64 1
; GFX950-NEXT:    [[Q2:%.*]] = extractelement <4 x float> [[Q]], i64 2
; GFX950-NEXT:    [[Q3:%.*]] = extractelement <4 x float> [[Q]], i64 3
; GFX950-NEXT:    [[PA:%.*]] = fadd float [[P0]], [[P1]]
; GFX950-NEXT:    [[PB:%.*]] = fadd float [[P2]], [[P3]]
; GFX950-NEXT:    [[PS:%.*]] = fadd float [[PA]], [[PB]]
; GFX950-NEXT:    [[QA:%.*]] = fadd float [[Q0]], [[Q1]]
; GFX950-NEXT:    [[QB:%.*]] = fadd float [[Q2]], [[Q3]]
; GFX950-NEXT:    [[QS:%.*]] = fadd float [[QA]], [[QB]]
; GFX950-NEXT:    [[V0:%.*]] = insertelement <2 x float> poison, float [[PS]], i64 0
; GFX950-NEXT:    [[V1:%.*]] = insertelement <2 x float> [[V0]], float [[QS]], i64 1
; GFX950-NEXT:    store <2 x float> [[V1]], ptr addrspace(3) [[OUT]], align 8
; GFX950-NEXT:    ret void
;
; GFX1250-LABEL: define amdgpu_kernel void @softmax_row_sum(
; GFX1250-SAME: ptr addrspace(3) [[OUT:%.*]], <4 x float> [[P:%.*]], <4 x float> [[Q:%.*]]) {
; GFX1250-NEXT:    [[P0:%.*]] = extractelement <4 x float> [[P]], i64 0
; GFX1250-NEXT:    [[P1:%.*]] = extractelement <4 x float> [[P]], i64 1
; GFX1250-NEXT:    [[P2:%.*]] = extractelement <4 x float> [[P]], i64 2
; GFX1250-NEXT:    [[P3:%.*]] = extractelement <4 x float> [[P]], i64 3
; GFX1250-NEXT:    [[Q0:%.*]] = extractelement <4 x float> [[Q]], i64 0
; GFX1250-NEXT:    [[Q1:%.*]] = extractelement <4 x float> [[Q]], i64 1
; GFX1250-NEXT:    [[Q2:%.*]] = extractelement <4 x float> [[Q]], i64 2
; GFX1250-NEXT:    [[Q3:%.*]] = extractelement <4 x float> [[Q]], i64 3
; GFX1250-NEXT:    [[PA:%.*]] = fadd float [[P0]], [[P1]]
; GFX1250-NEXT:    [[PB:%.*]] = fadd float [[P2]], [[P3]]
; GFX1250-NEXT:    [[PS:%.*]] = fadd float [[PA]], [[PB]]
; GFX1250-NEXT:    [[QA:%.*]] = fadd float [[Q0]], [[Q1]]
; GFX1250-NEXT:    [[QB:%.*]] = fadd float [[Q2]], [[Q3]]
; GFX1250-NEXT:    [[QS:%.*]] = fadd float [[QA]], [[QB]]
; GFX1250-NEXT:    [[V0:%.*]] = insertelement <2 x float> poison, float [[PS]], i64 0
; GFX1250-NEXT:    [[V1:%.*]] = insertelement <2 x float> [[V0]], float [[QS]], i64 1
; GFX1250-NEXT:    store <2 x float> [[V1]], ptr addrspace(3) [[OUT]], align 8
; GFX1250-NEXT:    ret void
;
  %p0 = extractelement <4 x float> %p, i64 0
  %p1 = extractelement <4 x float> %p, i64 1
  %p2 = extractelement <4 x float> %p, i64 2
  %p3 = extractelement <4 x float> %p, i64 3
  %q0 = extractelement <4 x float> %q, i64 0
  %q1 = extractelement <4 x float> %q, i64 1
  %q2 = extractelement <4 x float> %q, i64 2
  %q3 = extractelement <4 x float> %q, i64 3
  %pa = fadd float %p0, %p1
  %pb = fadd float %p2, %p3
  %ps = fadd float %pa, %pb
  %qa = fadd float %q0, %q1
  %qb = fadd float %q2, %q3
  %qs = fadd float %qa, %qb
  %v0 = insertelement <2 x float> poison, float %ps, i64 0
  %v1 = insertelement <2 x float> %v0, float %qs, i64 1
  store <2 x float> %v1, ptr addrspace(3) %out, align 8
  ret void
}