llvm/test/CodeGen/AArch64/partial-reduction-add-predicated.ll

; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
; RUN: llc < %s | FileCheck %s

target triple = "aarch64"

define <4 x i32> @predicate_dot_fixed_length(<4 x i32> %acc, <16 x i1> %p, <16 x i8> %a, <16 x i8> %b) #0 {
; CHECK-LABEL: predicate_dot_fixed_length:
; CHECK:       // %bb.0:
; CHECK-NEXT:    shl v1.16b, v1.16b, #7
; CHECK-NEXT:    cmlt v1.16b, v1.16b, #0
; CHECK-NEXT:    and v1.16b, v1.16b, v3.16b
; CHECK-NEXT:    sdot v0.4s, v2.16b, v1.16b
; CHECK-NEXT:    ret
 %ext.1 = sext <16 x i8> %a to <16 x i32>
 %ext.2 = sext <16 x i8> %b to <16 x i32>
 %mul = mul nsw <16 x i32> %ext.1, %ext.2
 %sel = select <16 x i1> %p, <16 x i32> %mul, <16 x i32> zeroinitializer
 %red = call <4 x i32> @llvm.vector.partial.reduce.add(<4 x i32> %acc, <16 x i32> %sel)
 ret <4 x i32> %red
}

define <4 x i32> @predicate_dot_by_C_fixed_length(<4 x i32> %acc, <16 x i1> %p, <16 x i8> %a) #0 {
; CHECK-LABEL: predicate_dot_by_C_fixed_length:
; CHECK:       // %bb.0:
; CHECK-NEXT:    shl v1.16b, v1.16b, #7
; CHECK-NEXT:    movi v3.16b, #127
; CHECK-NEXT:    cmlt v1.16b, v1.16b, #0
; CHECK-NEXT:    and v1.16b, v1.16b, v3.16b
; CHECK-NEXT:    sdot v0.4s, v2.16b, v1.16b
; CHECK-NEXT:    ret
 %ext.1 = sext <16 x i8> %a to <16 x i32>
 %mul = mul nsw <16 x i32> %ext.1, splat(i32 127)
 %sel = select <16 x i1> %p, <16 x i32> %mul, <16 x i32> zeroinitializer
 %red = call <4 x i32> @llvm.vector.partial.reduce.add(<4 x i32> %acc, <16 x i32> %sel)
 ret <4 x i32> %red
}

define <vscale x 4 x i32> @predicate_dot_scalable(<vscale x 4 x i32> %acc, <vscale x 16 x i1> %p, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) #0 {
; CHECK-LABEL: predicate_dot_scalable:
; CHECK:       // %bb.0:
; CHECK-NEXT:    movi v3.2d, #0000000000000000
; CHECK-NEXT:    sel z2.b, p0, z2.b, z3.b
; CHECK-NEXT:    sdot z0.s, z1.b, z2.b
; CHECK-NEXT:    ret
 %ext.1 = sext <vscale x 16 x i8> %a to <vscale x 16 x i32>
 %ext.2 = sext <vscale x 16 x i8> %b to <vscale x 16 x i32>
 %mul = mul nsw <vscale x 16 x i32> %ext.1, %ext.2
 %sel = select <vscale x 16 x i1> %p, <vscale x 16 x i32> %mul, <vscale x 16 x i32> zeroinitializer
 %red = call <vscale x 4 x i32> @llvm.vector.partial.reduce.add(<vscale x 4 x i32> %acc, <vscale x 16 x i32> %sel)
 ret <vscale x 4 x i32> %red
}

define <vscale x 4 x i32> @predicate_dot_by_C_scalable(<vscale x 4 x i32> %acc, <vscale x 16 x i1> %p, <vscale x 16 x i8> %a) #0 {
; CHECK-LABEL: predicate_dot_by_C_scalable:
; CHECK:       // %bb.0:
; CHECK-NEXT:    mov z2.b, p0/z, #127 // =0x7f
; CHECK-NEXT:    sdot z0.s, z1.b, z2.b
; CHECK-NEXT:    ret
 %ext.1 = sext <vscale x 16 x i8> %a to <vscale x 16 x i32>
 %mul = mul nsw <vscale x 16 x i32> %ext.1, splat(i32 127)
 %sel = select <vscale x 16 x i1> %p, <vscale x 16 x i32> %mul, <vscale x 16 x i32> zeroinitializer
 %red = call <vscale x 4 x i32> @llvm.vector.partial.reduce.add(<vscale x 4 x i32> %acc, <vscale x 16 x i32> %sel)
 ret <vscale x 4 x i32> %red
}

define <4 x i32> @predicate_ext_mul_fixed_length(<4 x i32> %acc, <16 x i1> %p, <16 x i8> %a) #0 {
; CHECK-LABEL: predicate_ext_mul_fixed_length:
; CHECK:       // %bb.0:
; CHECK-NEXT:    movi v3.16b, #1
; CHECK-NEXT:    and v1.16b, v1.16b, v3.16b
; CHECK-NEXT:    sdot v0.4s, v2.16b, v1.16b
; CHECK-NEXT:    ret
 %ext = sext <16 x i8> %a to <16 x i32>
 %sel = select <16 x i1> %p, <16 x i32> %ext, <16 x i32> zeroinitializer
 %red = call <4 x i32> @llvm.vector.partial.reduce.add(<4 x i32> %acc, <16 x i32> %sel)
 ret <4 x i32> %red
}

define <vscale x 4 x i32> @predicate_ext_mul_scalable(<vscale x 4 x i32> %acc, <vscale x 16 x i1> %p, <vscale x 16 x i8> %a) #0 {
; CHECK-LABEL: predicate_ext_mul_scalable:
; CHECK:       // %bb.0:
; CHECK-NEXT:    mov z2.b, p0/z, #1 // =0x1
; CHECK-NEXT:    sdot z0.s, z1.b, z2.b
; CHECK-NEXT:    ret
 %ext = sext <vscale x 16 x i8> %a to <vscale x 16 x i32>
 %sel = select <vscale x 16 x i1> %p, <vscale x 16 x i32> %ext, <vscale x 16 x i32> zeroinitializer
 %red = call <vscale x 4 x i32> @llvm.vector.partial.reduce.add(<vscale x 4 x i32> %acc, <vscale x 16 x i32> %sel)
 ret <vscale x 4 x i32> %red
}

define <4 x float> @predicated_fdot_fixed_length(<4 x float> %acc, <8 x i1> %p, <8 x half> %a, <8 x half> %b) #1 {
; CHECK-LABEL: predicated_fdot_fixed_length:
; CHECK:       // %bb.0:
; CHECK-NEXT:    ushll v1.8h, v1.8b, #0
; CHECK-NEXT:    // kill: def $q0 killed $q0 def $z0
; CHECK-NEXT:    // kill: def $q2 killed $q2 def $z2
; CHECK-NEXT:    shl v1.8h, v1.8h, #15
; CHECK-NEXT:    cmlt v1.8h, v1.8h, #0
; CHECK-NEXT:    and v1.16b, v1.16b, v3.16b
; CHECK-NEXT:    fdot z0.s, z2.h, z1.h
; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
; CHECK-NEXT:    ret
 %ext.1 = fpext <8 x half> %a to <8 x float>
 %ext.2 = fpext <8 x half> %b to <8 x float>
 %mul = fmul <8 x float> %ext.1, %ext.2
 %sel = select <8 x i1> %p, <8 x float> %mul, <8 x float> zeroinitializer
 %red = call <4 x float> @llvm.vector.partial.reduce.fadd(<4 x float> %acc, <8 x float> %sel)
 ret <4 x float> %red
}

define <vscale x 4 x float> @predicated_fdot_scalable(<vscale x 4 x float> %acc, <vscale x 8 x i1> %p, <vscale x 8 x half> %a, <vscale x 8 x half> %b) #1 {
; CHECK-LABEL: predicated_fdot_scalable:
; CHECK:       // %bb.0:
; CHECK-NEXT:    movi v3.2d, #0000000000000000
; CHECK-NEXT:    sel z2.h, p0, z2.h, z3.h
; CHECK-NEXT:    fdot z0.s, z1.h, z2.h
; CHECK-NEXT:    ret
 %ext.1 = fpext <vscale x 8 x half> %a to <vscale x 8 x float>
 %ext.2 = fpext <vscale x 8 x half> %b to <vscale x 8 x float>
 %mul = fmul <vscale x 8 x float> %ext.1, %ext.2
 %sel = select <vscale x 8 x i1> %p, <vscale x 8 x float> %mul, <vscale x 8 x float> zeroinitializer
 %red = call <vscale x 4 x float> @llvm.vector.partial.reduce.fadd(<vscale x 4 x float> %acc, <vscale x 8 x float> %sel)
 ret <vscale x 4 x float> %red
}

define <4 x float> @predicated_fpext_fmul_fixed_length(<4 x float> %acc, <8 x i1> %p, <8 x half> %a) #1 {
; CHECK-LABEL: predicated_fpext_fmul_fixed_length:
; CHECK:       // %bb.0:
; CHECK-NEXT:    ushll v1.8h, v1.8b, #0
; CHECK-NEXT:    movi v3.8h, #60, lsl #8
; CHECK-NEXT:    // kill: def $q0 killed $q0 def $z0
; CHECK-NEXT:    // kill: def $q2 killed $q2 def $z2
; CHECK-NEXT:    shl v1.8h, v1.8h, #15
; CHECK-NEXT:    cmlt v1.8h, v1.8h, #0
; CHECK-NEXT:    and v1.16b, v1.16b, v3.16b
; CHECK-NEXT:    fdot z0.s, z2.h, z1.h
; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
; CHECK-NEXT:    ret
 %ext = fpext <8 x half> %a to <8 x float>
 %sel = select <8 x i1> %p, <8 x float> %ext, <8 x float> zeroinitializer
 %red = call <4 x float> @llvm.vector.partial.reduce.fadd(<4 x float> %acc, <8 x float> %sel)
 ret <4 x float> %red
}

define <vscale x 4 x float> @predicated_fpext_fmul_scalable(<vscale x 4 x float> %acc, <vscale x 8 x i1> %p, <vscale x 8 x half> %a) #1 {
; CHECK-LABEL: predicated_fpext_fmul_scalable:
; CHECK:       // %bb.0:
; CHECK-NEXT:    movi v2.2d, #0000000000000000
; CHECK-NEXT:    fmov z2.h, p0/m, #1.00000000
; CHECK-NEXT:    fdot z0.s, z1.h, z2.h
; CHECK-NEXT:    ret
 %ext = fpext <vscale x 8 x half> %a to <vscale x 8 x float>
 %sel = select <vscale x 8 x i1> %p, <vscale x 8 x float> %ext, <vscale x 8 x float> zeroinitializer
 %red = call <vscale x 4 x float> @llvm.vector.partial.reduce.fadd(<vscale x 4 x float> %acc, <vscale x 8 x float> %sel)
 ret <vscale x 4 x float> %red
}

attributes #0 = { nounwind "target-features"="+sve,+dotprod" }
attributes #1 = { nounwind "target-features"="+sve2p1,+dotprod" }