llvm/test/CodeGen/AArch64/partial-reduction-sub.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"

;
; Ensure base cases are supported.
;

define <vscale x 8 x i16> @umlslbt_i8_i16(<vscale x 8 x i16> %acc, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) #0 {
; CHECK-LABEL: umlslbt_i8_i16:
; CHECK:       // %bb.0:
; CHECK-NEXT:    umlslb z0.h, z1.b, z2.b
; CHECK-NEXT:    umlslt z0.h, z1.b, z2.b
; CHECK-NEXT:    ret
  %a.zext = zext <vscale x 16 x i8> %a to <vscale x 16 x i16>
  %b.zext = zext <vscale x 16 x i8> %b to <vscale x 16 x i16>
  %mul = mul <vscale x 16 x i16> %a.zext, %b.zext
  %mul.neg = sub <vscale x 16 x i16> zeroinitializer, %mul
  %res = call <vscale x 8 x i16> @llvm.vector.partial.reduce.add(<vscale x 8 x i16> %acc, <vscale x 16 x i16> %mul.neg)
  ret <vscale x 8 x i16> %res
}

define <vscale x 8 x i16> @smlslbt_i8_i16(<vscale x 8 x i16> %acc, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) #0 {
; CHECK-LABEL: smlslbt_i8_i16:
; CHECK:       // %bb.0:
; CHECK-NEXT:    smlslb z0.h, z1.b, z2.b
; CHECK-NEXT:    smlslt z0.h, z1.b, z2.b
; CHECK-NEXT:    ret
  %a.sext = sext <vscale x 16 x i8> %a to <vscale x 16 x i16>
  %b.sext = sext <vscale x 16 x i8> %b to <vscale x 16 x i16>
  %mul = mul <vscale x 16 x i16> %a.sext, %b.sext
  %mul.neg = sub <vscale x 16 x i16> zeroinitializer, %mul
  %res = call <vscale x 8 x i16> @llvm.vector.partial.reduce.add(<vscale x 8 x i16> %acc, <vscale x 16 x i16> %mul.neg)
  ret <vscale x 8 x i16> %res
}

define <vscale x 4 x i32> @umlslbt_i16_i32(<vscale x 4 x i32> %acc, <vscale x 8 x i16> %a, <vscale x 8 x i16> %b) #0 {
; CHECK-LABEL: umlslbt_i16_i32:
; CHECK:       // %bb.0:
; CHECK-NEXT:    umlslb z0.s, z1.h, z2.h
; CHECK-NEXT:    umlslt z0.s, z1.h, z2.h
; CHECK-NEXT:    ret
  %a.zext = zext <vscale x 8 x i16> %a to <vscale x 8 x i32>
  %b.zext = zext <vscale x 8 x i16> %b to <vscale x 8 x i32>
  %mul = mul <vscale x 8 x i32> %a.zext, %b.zext
  %mul.neg = sub <vscale x 8 x i32> zeroinitializer, %mul
  %res = call <vscale x 4 x i32> @llvm.vector.partial.reduce.add(<vscale x 4 x i32> %acc, <vscale x 8 x i32> %mul.neg)
  ret <vscale x 4 x i32> %res
}

define <vscale x 4 x i32> @smlslbt_i16_i32(<vscale x 4 x i32> %acc, <vscale x 8 x i16> %a, <vscale x 8 x i16> %b) #0 {
; CHECK-LABEL: smlslbt_i16_i32:
; CHECK:       // %bb.0:
; CHECK-NEXT:    smlslb z0.s, z1.h, z2.h
; CHECK-NEXT:    smlslt z0.s, z1.h, z2.h
; CHECK-NEXT:    ret
  %a.sext = sext <vscale x 8 x i16> %a to <vscale x 8 x i32>
  %b.sext = sext <vscale x 8 x i16> %b to <vscale x 8 x i32>
  %mul = mul <vscale x 8 x i32> %a.sext, %b.sext
  %mul.neg = sub <vscale x 8 x i32> zeroinitializer, %mul
  %res = call <vscale x 4 x i32> @llvm.vector.partial.reduce.add(<vscale x 4 x i32> %acc, <vscale x 8 x i32> %mul.neg)
  ret <vscale x 4 x i32> %res
}

define <vscale x 2 x i64> @umlslbt_i32_i64(<vscale x 2 x i64> %acc, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b) #0 {
; CHECK-LABEL: umlslbt_i32_i64:
; CHECK:       // %bb.0:
; CHECK-NEXT:    umlslb z0.d, z1.s, z2.s
; CHECK-NEXT:    umlslt z0.d, z1.s, z2.s
; CHECK-NEXT:    ret
  %a.zext = zext <vscale x 4 x i32> %a to <vscale x 4 x i64>
  %b.zext = zext <vscale x 4 x i32> %b to <vscale x 4 x i64>
  %mul = mul <vscale x 4 x i64> %a.zext, %b.zext
  %mul.neg = sub <vscale x 4 x i64> zeroinitializer, %mul
  %res = call <vscale x 2 x i64> @llvm.vector.partial.reduce.add(<vscale x 2 x i64> %acc, <vscale x 4 x i64> %mul.neg)
  ret <vscale x 2 x i64> %res
}

define <vscale x 2 x i64> @smlslbt_i32_i64(<vscale x 2 x i64> %acc, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b) #0 {
; CHECK-LABEL: smlslbt_i32_i64:
; CHECK:       // %bb.0:
; CHECK-NEXT:    smlslb z0.d, z1.s, z2.s
; CHECK-NEXT:    smlslt z0.d, z1.s, z2.s
; CHECK-NEXT:    ret
  %a.sext = sext <vscale x 4 x i32> %a to <vscale x 4 x i64>
  %b.sext = sext <vscale x 4 x i32> %b to <vscale x 4 x i64>
  %mul = mul <vscale x 4 x i64> %a.sext, %b.sext
  %mul.neg = sub <vscale x 4 x i64> zeroinitializer, %mul
  %res = call <vscale x 2 x i64> @llvm.vector.partial.reduce.add(<vscale x 2 x i64> %acc, <vscale x 4 x i64> %mul.neg)
  ret <vscale x 2 x i64> %res
}

;
; Ensure fixed-length codegen for streaming-compatible functions.
;

define <8 x i16> @fixed_umlslbt_i8_i16(<8 x i16> %acc, <16 x i8> %a, <16 x i8> %b) #0 "aarch64_pstate_sm_compatible" {
; CHECK-LABEL: fixed_umlslbt_i8_i16:
; CHECK:       // %bb.0:
; CHECK-NEXT:    // kill: def $q0 killed $q0 def $z0
; CHECK-NEXT:    // kill: def $q2 killed $q2 def $z2
; CHECK-NEXT:    // kill: def $q1 killed $q1 def $z1
; CHECK-NEXT:    umlslb z0.h, z1.b, z2.b
; CHECK-NEXT:    umlslt z0.h, z1.b, z2.b
; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
; CHECK-NEXT:    ret
  %a.zext = zext <16 x i8> %a to <16 x i16>
  %b.zext = zext <16 x i8> %b to <16 x i16>
  %mul = mul <16 x i16> %a.zext, %b.zext
  %mul.neg = sub <16 x i16> zeroinitializer, %mul
  %res = call <8 x i16> @llvm.vector.partial.reduce.add(<8 x i16> %acc, <16 x i16> %mul.neg)
  ret <8 x i16> %res
}

define <8 x i16> @fixed_smlslbt_i8_i16(<8 x i16> %acc, <16 x i8> %a, <16 x i8> %b) #0 "aarch64_pstate_sm_compatible" {
; CHECK-LABEL: fixed_smlslbt_i8_i16:
; CHECK:       // %bb.0:
; CHECK-NEXT:    // kill: def $q0 killed $q0 def $z0
; CHECK-NEXT:    // kill: def $q2 killed $q2 def $z2
; CHECK-NEXT:    // kill: def $q1 killed $q1 def $z1
; CHECK-NEXT:    smlslb z0.h, z1.b, z2.b
; CHECK-NEXT:    smlslt z0.h, z1.b, z2.b
; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
; CHECK-NEXT:    ret
  %a.sext = sext <16 x i8> %a to <16 x i16>
  %b.sext = sext <16 x i8> %b to <16 x i16>
  %mul = mul <16 x i16> %a.sext, %b.sext
  %mul.neg = sub <16 x i16> zeroinitializer, %mul
  %res = call <8 x i16> @llvm.vector.partial.reduce.add(<8 x i16> %acc, <16 x i16> %mul.neg)
  ret <8 x i16> %res
}

define <4 x i32> @fixed_umlslbt_i16_i32(<4 x i32> %acc, <8 x i16> %a, <8 x i16> %b) #0 "aarch64_pstate_sm_compatible" {
; CHECK-LABEL: fixed_umlslbt_i16_i32:
; CHECK:       // %bb.0:
; CHECK-NEXT:    // kill: def $q0 killed $q0 def $z0
; CHECK-NEXT:    // kill: def $q2 killed $q2 def $z2
; CHECK-NEXT:    // kill: def $q1 killed $q1 def $z1
; CHECK-NEXT:    umlslb z0.s, z1.h, z2.h
; CHECK-NEXT:    umlslt z0.s, z1.h, z2.h
; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
; CHECK-NEXT:    ret
  %a.zext = zext <8 x i16> %a to <8 x i32>
  %b.zext = zext <8 x i16> %b to <8 x i32>
  %mul = mul <8 x i32> %a.zext, %b.zext
  %mul.neg = sub <8 x i32> zeroinitializer, %mul
  %res = call <4 x i32> @llvm.vector.partial.reduce.add(<4 x i32> %acc, <8 x i32> %mul.neg)
  ret <4 x i32> %res
}

define <4 x i32> @fixed_smlslbt_i16_i32(<4 x i32> %acc, <8 x i16> %a, <8 x i16> %b) #0 "aarch64_pstate_sm_compatible" {
; CHECK-LABEL: fixed_smlslbt_i16_i32:
; CHECK:       // %bb.0:
; CHECK-NEXT:    // kill: def $q0 killed $q0 def $z0
; CHECK-NEXT:    // kill: def $q2 killed $q2 def $z2
; CHECK-NEXT:    // kill: def $q1 killed $q1 def $z1
; CHECK-NEXT:    smlslb z0.s, z1.h, z2.h
; CHECK-NEXT:    smlslt z0.s, z1.h, z2.h
; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
; CHECK-NEXT:    ret
  %a.sext = sext <8 x i16> %a to <8 x i32>
  %b.sext = sext <8 x i16> %b to <8 x i32>
  %mul = mul <8 x i32> %a.sext, %b.sext
  %mul.neg = sub <8 x i32> zeroinitializer, %mul
  %res = call <4 x i32> @llvm.vector.partial.reduce.add(<4 x i32> %acc, <8 x i32> %mul.neg)
  ret <4 x i32> %res
}

define <2 x i64> @fixed_umlslbt_i32_i64(<2 x i64> %acc, <4 x i32> %a, <4 x i32> %b) #0 "aarch64_pstate_sm_compatible" {
; CHECK-LABEL: fixed_umlslbt_i32_i64:
; CHECK:       // %bb.0:
; CHECK-NEXT:    // kill: def $q0 killed $q0 def $z0
; CHECK-NEXT:    // kill: def $q2 killed $q2 def $z2
; CHECK-NEXT:    // kill: def $q1 killed $q1 def $z1
; CHECK-NEXT:    umlslb z0.d, z1.s, z2.s
; CHECK-NEXT:    umlslt z0.d, z1.s, z2.s
; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
; CHECK-NEXT:    ret
  %a.zext = zext <4 x i32> %a to <4 x i64>
  %b.zext = zext <4 x i32> %b to <4 x i64>
  %mul = mul <4 x i64> %a.zext, %b.zext
  %mul.neg = sub <4 x i64> zeroinitializer, %mul
  %res = call <2 x i64> @llvm.vector.partial.reduce.add(<2 x i64> %acc, <4 x i64> %mul.neg)
  ret <2 x i64> %res
}

define <2 x i64> @fixed_smlslbt_i32_i64(<2 x i64> %acc, <4 x i32> %a, <4 x i32> %b) #0 "aarch64_pstate_sm_compatible" {
; CHECK-LABEL: fixed_smlslbt_i32_i64:
; CHECK:       // %bb.0:
; CHECK-NEXT:    // kill: def $q0 killed $q0 def $z0
; CHECK-NEXT:    // kill: def $q2 killed $q2 def $z2
; CHECK-NEXT:    // kill: def $q1 killed $q1 def $z1
; CHECK-NEXT:    smlslb z0.d, z1.s, z2.s
; CHECK-NEXT:    smlslt z0.d, z1.s, z2.s
; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
; CHECK-NEXT:    ret
  %a.sext = sext <4 x i32> %a to <4 x i64>
  %b.sext = sext <4 x i32> %b to <4 x i64>
  %mul = mul <4 x i64> %a.sext, %b.sext
  %mul.neg = sub <4 x i64> zeroinitializer, %mul
  %res = call <2 x i64> @llvm.vector.partial.reduce.add(<2 x i64> %acc, <4 x i64> %mul.neg)
  ret <2 x i64> %res
}
;
; Test type legalisation for sub-reductions.
;

; FIXME: The subr's could be removed with a DAG combine.
define <vscale x 8 x i16> @legalization_split_i8_i16(<vscale x 8 x i16> %acc, <vscale x 32 x i8> %a, <vscale x 32 x i8> %b) #0 {
; CHECK-LABEL: legalization_split_i8_i16:
; CHECK:       // %bb.0:
; CHECK-NEXT:    subr z0.h, z0.h, #0 // =0x0
; CHECK-NEXT:    umlalb z0.h, z1.b, z3.b
; CHECK-NEXT:    umlalt z0.h, z1.b, z3.b
; CHECK-NEXT:    umlalb z0.h, z2.b, z4.b
; CHECK-NEXT:    umlalt z0.h, z2.b, z4.b
; CHECK-NEXT:    subr z0.h, z0.h, #0 // =0x0
; CHECK-NEXT:    ret
  %a.zext = zext <vscale x 32 x i8> %a to <vscale x 32 x i16>
  %b.zext = zext <vscale x 32 x i8> %b to <vscale x 32 x i16>
  %mul = mul <vscale x 32 x i16> %a.zext, %b.zext
  %mul.neg = sub <vscale x 32 x i16> zeroinitializer, %mul
  %res = call <vscale x 8 x i16> @llvm.vector.partial.reduce.add(<vscale x 8 x i16> %acc, <vscale x 32 x i16> %mul.neg)
  ret <vscale x 8 x i16> %res
}

define <vscale x 4 x i16> @legalization_promote_acc_i8_i16(<vscale x 4 x i16> %acc, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) #0 {
; CHECK-LABEL: legalization_promote_acc_i8_i16:
; CHECK:       // %bb.0:
; CHECK-NEXT:    subr z0.s, z0.s, #0 // =0x0
; CHECK-NEXT:    udot z0.s, z1.b, z2.b
; CHECK-NEXT:    subr z0.s, z0.s, #0 // =0x0
; CHECK-NEXT:    ret
  %a.zext = zext <vscale x 16 x i8> %a to <vscale x 16 x i16>
  %b.zext = zext <vscale x 16 x i8> %b to <vscale x 16 x i16>
  %mul = mul <vscale x 16 x i16> %a.zext, %b.zext
  %mul.neg = sub <vscale x 16 x i16> zeroinitializer, %mul
  %res = call <vscale x 4 x i16> @llvm.vector.partial.reduce.add(<vscale x 4 x i16> %acc, <vscale x 16 x i16> %mul.neg)
  ret <vscale x 4 x i16> %res
}

; Test that MLSB/T are still generated when there is no 'mul'.
define <vscale x 2 x i64> @extended_sub_i32_i64(<vscale x 2 x i64> %acc, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b) #0 {
; CHECK-LABEL: extended_sub_i32_i64:
; CHECK:       // %bb.0:
; CHECK-NEXT:    mov z2.s, #1 // =0x1
; CHECK-NEXT:    smlslb z0.d, z1.s, z2.s
; CHECK-NEXT:    smlslt z0.d, z1.s, z2.s
; CHECK-NEXT:    ret
  %a.sext = sext <vscale x 4 x i32> %a to <vscale x 4 x i64>
  %a.sext.neg = sub <vscale x 4 x i64> zeroinitializer, %a.sext
  %res = call <vscale x 2 x i64> @llvm.vector.partial.reduce.add(<vscale x 2 x i64> %acc, <vscale x 4 x i64> %a.sext.neg)
  ret <vscale x 2 x i64> %res
}

; Test that 'predication' is still supported.
define <vscale x 2 x i64> @predicated_smlslbt_i32_i64(<vscale x 4 x i1> %pred, <vscale x 2 x i64> %acc, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b) #0 {
; CHECK-LABEL: predicated_smlslbt_i32_i64:
; CHECK:       // %bb.0:
; CHECK-NEXT:    movi v3.2d, #0000000000000000
; CHECK-NEXT:    sel z2.s, p0, z2.s, z3.s
; CHECK-NEXT:    smlslb z0.d, z1.s, z2.s
; CHECK-NEXT:    smlslt z0.d, z1.s, z2.s
; CHECK-NEXT:    ret
  %a.sext = sext <vscale x 4 x i32> %a to <vscale x 4 x i64>
  %b.sext = sext <vscale x 4 x i32> %b to <vscale x 4 x i64>
  %mul = mul <vscale x 4 x i64> %a.sext, %b.sext
  %mul.neg = sub <vscale x 4 x i64> zeroinitializer, %mul
  %mul.neg.sel = select <vscale x 4 x i1> %pred, <vscale x 4 x i64> %mul.neg, <vscale x 4 x i64> zeroinitializer
  %res = call <vscale x 2 x i64> @llvm.vector.partial.reduce.add(<vscale x 2 x i64> %acc, <vscale x 4 x i64> %mul.neg.sel)
  ret <vscale x 2 x i64> %res
}

; There is no sub dot-reduction, so we need to introduce explicit instructions for negation.
define <vscale x 4 x i32> @negative_test_no_sub_dot_inst(<vscale x 4 x i32> %acc, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) #0 {
; CHECK-LABEL: negative_test_no_sub_dot_inst:
; CHECK:       // %bb.0:
; CHECK-NEXT:    subr z0.s, z0.s, #0 // =0x0
; CHECK-NEXT:    udot z0.s, z1.b, z2.b
; CHECK-NEXT:    subr z0.s, z0.s, #0 // =0x0
; CHECK-NEXT:    ret
  %a.zext = zext <vscale x 16 x i8> %a to <vscale x 16 x i32>
  %b.zext = zext <vscale x 16 x i8> %b to <vscale x 16 x i32>
  %mul = mul <vscale x 16 x i32> %a.zext, %b.zext
  %mul.neg = sub <vscale x 16 x i32> zeroinitializer, %mul
  %res = call <vscale x 4 x i32> @llvm.vector.partial.reduce.add(<vscale x 4 x i32> %acc, <vscale x 16 x i32> %mul.neg)
  ret <vscale x 4 x i32> %res
}

; Make sure wider types are supported when vscale_range supports it
define void @wide_fixed_umlslbt_i8_i16(ptr %acc.ptr, ptr %a.ptr, ptr %b.ptr, ptr %dest.ptr) #0 vscale_range(2,0) {
; CHECK-LABEL: wide_fixed_umlslbt_i8_i16:
; CHECK:       // %bb.0:
; CHECK-NEXT:    ptrue p0.b, vl32
; CHECK-NEXT:    ptrue p1.h, vl16
; CHECK-NEXT:    ld1b { z0.b }, p0/z, [x1]
; CHECK-NEXT:    ld1b { z1.b }, p0/z, [x2]
; CHECK-NEXT:    ld1h { z2.h }, p1/z, [x0]
; CHECK-NEXT:    umlslb z2.h, z0.b, z1.b
; CHECK-NEXT:    umlslt z2.h, z0.b, z1.b
; CHECK-NEXT:    st1h { z2.h }, p1, [x3]
; CHECK-NEXT:    ret
  %a = load <32 x i8>, ptr %a.ptr
  %b = load <32 x i8>, ptr %b.ptr
  %acc = load <16 x i16>, ptr %acc.ptr
  %a.zext = zext <32 x i8> %a to <32 x i16>
  %b.zext = zext <32 x i8> %b to <32 x i16>
  %mul = mul <32 x i16> %a.zext, %b.zext
  %mul.neg = sub <32 x i16> zeroinitializer, %mul
  %res = call <16 x i16> @llvm.vector.partial.reduce.add(<16 x i16> %acc, <32 x i16> %mul.neg)
  store <16 x i16> %res, ptr %dest.ptr
  ret void
}

attributes #0 = { "target-features"="+sve2" }