Add native self-hosted instance connection to fluxer_desktop
Trimmed monorepo checkout (fluxer_desktop + packages/voice_engine_v2 + tools/ci) with a "Connect to a Different Server" menu item and popout that lets the desktop app switch to any self-hosted Fluxer instance, plus fixes for well-known discovery on single-domain self-hosted deployments and a false-positive ERR_ABORTED on same-origin client redirects during the switch. Defaults to chat.fluxr.chat and uses an isolated userData directory from the official build.
This commit is contained in:
+533
@@ -0,0 +1,533 @@
|
||||
// vim: ft=arm
|
||||
.text
|
||||
.align 4
|
||||
|
||||
/* Z: 32x1
|
||||
z0[0] .. z0[15] z1[0] .. z1[15]
|
||||
*/
|
||||
|
||||
|
||||
.global {{G}}apple_amx_mmm_f16_64x1_{{suffix}}
|
||||
{{G}}apple_amx_mmm_f16_64x1_{{suffix}}:
|
||||
|
||||
{{ AMX_SET }}
|
||||
|
||||
// set x1 to a 128 bytes aligned block for loads
|
||||
mov x1, sp
|
||||
lsr x1, x1, #7
|
||||
lsl x1, x1, #7
|
||||
sub x1, x1, 128
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.leaky_relu:
|
||||
.q_scale:
|
||||
.q_shl:
|
||||
.q_shr:
|
||||
b .unsupported
|
||||
|
||||
.add_mat_mul:
|
||||
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x4, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
orr x4, x4, {{ 0 | setting(62) }} // load a pair of A
|
||||
|
||||
mov x5, {{ 0 | setting(43) }} // f16
|
||||
orr x5, x5, {{ 0 | setting(38) }} // Broadcast Y
|
||||
|
||||
orr x6, x5, {{ 0 | setting(20) }} // z offset
|
||||
orr x6, x6, {{ 0 | setting(16) }} // x offset
|
||||
|
||||
cmp x3, #32
|
||||
blt .packed_packed_loop_1
|
||||
|
||||
mov x9, {{ 0 | setting(32) }} // Y broadcast offset += 1
|
||||
|
||||
.packed_packed_loop_32:
|
||||
mov x7, x5
|
||||
mov x8, x6
|
||||
{{ amx("ldy", 2) }}
|
||||
{% for k in range(0, 32) %}
|
||||
{{ amx("ldx", 4) }}
|
||||
add x4, x4, 128
|
||||
{{ amx("vecfp", 7) }}
|
||||
{{ amx("vecfp", 8) }}
|
||||
add x7, x7, x9
|
||||
add x8, x8, x9
|
||||
{% endfor %}
|
||||
add x2, x2, #64
|
||||
sub x3, x3, #32
|
||||
cmp x3, #32
|
||||
bge .packed_packed_loop_32
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.packed_packed_loop_1:
|
||||
ldr w7, [x2], #2
|
||||
str w7, [x1]
|
||||
{{ amx("ldx", 4) }}
|
||||
{{ amx("ldy", 1) }}
|
||||
{{ amx("vecfp", 5) }}
|
||||
{{ amx("vecfp", 6) }}
|
||||
add x4, x4, 128
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.clear:
|
||||
// top left
|
||||
eor x2, x2, x2
|
||||
orr x2, x2, {{ 0 | setting(27) }}
|
||||
orr x2, x2, {{ 0 | setting(28) }}
|
||||
orr x2, x2, {{ 0 | setting(29) }} // Z = 0
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// top right
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(21) }} // Z row = 3
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom left
|
||||
eor x2, x2, {{ 0 | setting(20) }} // Z row = 2
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_sub:
|
||||
|
||||
// performs a unary neg on Z
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
// extr[hxyz] is suport confusing
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms16", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("extrx", 2) }} // extr[hxyz] is confusing
|
||||
{{ amx("fms16", 4) }}
|
||||
|
||||
// continue
|
||||
|
||||
.per_col_add:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
// broadcast value to x0
|
||||
ld1 { v0.h }[0], [x2]
|
||||
dup v0.8h, v0.h[0]
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
sub x1, x1, #64
|
||||
|
||||
{{ amx("ldx", 1) }} // load into x0 by default
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // z += y
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
orr x2, x2, {{ 0 | setting(20) }} // target is now z1
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_sub_flipped:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
// broadcast value to x0
|
||||
ld1 { v0.h }[0], [x2]
|
||||
dup v0.8h, v0.h[0]
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
sub x1, x1, #64
|
||||
|
||||
{{ amx("ldx", 1) }} // load into x0 by default
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // z += y
|
||||
{{ amx("fms16", 2) }}
|
||||
|
||||
orr x2, x2, {{ 0 | setting(20) }} // target is now z1
|
||||
{{ amx("fms16", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_sub_flipped:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2], #64
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1], #64
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
mov x2, {{ 0 | setting(63) }} // vector mode
|
||||
orr x2, x2, {{ 0 | setting(29) }} // z -= y
|
||||
|
||||
// top left
|
||||
{{ amx("fms16", 2) }}
|
||||
|
||||
// bottom left
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
orr x2, x2, {{ 0 | setting(6) }} // Y offset
|
||||
{{ amx("fms16", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_sub:
|
||||
// performs a unary neg on Z
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms16", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms16", 4) }}
|
||||
|
||||
// continue
|
||||
|
||||
.per_row_add:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2], #64
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1], #64
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
mov x2, {{ 0 | setting(63) }} // vector mode
|
||||
orr x2, x2, {{ 0 | setting(29) }} // z += y
|
||||
|
||||
// top left
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
// bottom left
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
orr x2, x2, {{ 0 | setting(6) }} // Y offset
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_min:
|
||||
mov x2, 5
|
||||
b .per_row_min_max
|
||||
.per_row_max:
|
||||
mov x2, 7
|
||||
.per_row_min_max:
|
||||
ldr x5, [x0, #8]
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x5], #64
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1], #64
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x5]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x5, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldx", 5) }}
|
||||
|
||||
lsl x2, x2, 47 // max(x,z) (or min)
|
||||
orr x2, x2, {{ 0 | setting(44) }} // f32
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
orr x2, x2, {{ 0 | setting(16) }} // x1
|
||||
orr x2, x2, {{ 0 | setting(20) }} // z1
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_min:
|
||||
mov x2, 5
|
||||
b .per_col_min_max
|
||||
.per_col_max:
|
||||
mov x2, 7
|
||||
.per_col_min_max:
|
||||
ldr x4, [x0, #8]
|
||||
|
||||
// broadcast value to x0
|
||||
ld1 { v0.h }[0], [x4]
|
||||
dup v0.8h, v0.h[0]
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
sub x1, x1, #64
|
||||
|
||||
{{ amx("ldx", 1) }}
|
||||
|
||||
lsl x2, x2, 47 // max(x,z) (or min)
|
||||
orr x2, x2, {{ 0 | setting(43) }} // f32
|
||||
|
||||
{{ amx("vecfp", 2) }}
|
||||
orr x2, x2, {{ 0 | setting(20) }} // z offset
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_mul:
|
||||
ldr x4, [x0, #8]
|
||||
|
||||
// broadcast value to y0
|
||||
ld1 { v0.h }[0], [x4]
|
||||
dup v0.8h, v0.h[0]
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
st1 { v0.8h }, [x1], #16
|
||||
sub x1, x1, #64
|
||||
|
||||
{{ amx("ldy", 1) }}
|
||||
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
{{ amx("extrx", 2) }}
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=X*Y
|
||||
{{ amx("fma16", 4) }}
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z1
|
||||
{{ amx("extrx", 2) }}
|
||||
orr x4, x4, {{ 0 | setting(20) }} // Z1
|
||||
{{ amx("fma16", 4) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_mul:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2], #64
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1], #64
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
{{ amx("extrx", 2) }}
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=X*Y
|
||||
{{ amx("fma16", 4) }}
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z1
|
||||
{{ amx("extrx", 2) }}
|
||||
orr x4, x4, {{ 0 | setting(20) }} // Z1
|
||||
orr x4, x4, {{ 0 | setting(6) }} // Y1
|
||||
{{ amx("fma16", 4) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_sub:
|
||||
// performs a unary neg on Z, then go to scalar_add
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms16", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms16", 4) }}
|
||||
|
||||
// continue on purpose
|
||||
|
||||
.scalar_add:
|
||||
ldr w5, [x0, #8]
|
||||
|
||||
fmov h0, w5
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // Z+=X
|
||||
{{ amx("fma16", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("fma16", 2) }}
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_sub_flipped:
|
||||
ldr w5, [x0, #8]
|
||||
fmov s0, w5
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // Z-=X
|
||||
{{ amx("fms16", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("fms16", 2) }}
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_mul:
|
||||
ldr w5, [x0, #8]
|
||||
fmov h0, w5
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
{{ amx("ldy", 1) }}
|
||||
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=X*Y
|
||||
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fma16", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fma16", 4) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_min:
|
||||
mov x2, 5
|
||||
b .scalar_min_max
|
||||
.scalar_max:
|
||||
mov x2, 7
|
||||
.scalar_min_max:
|
||||
ldr w5, [x0, #8]
|
||||
fmov h0, w5
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
lsl x2, x2, 47
|
||||
orr x2, x2, {{ 0 | setting(43) }} // f16
|
||||
|
||||
{{ amx("vecfp", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
{% for neon in range(0, 8) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
ld1 { v{{neon}}.h }[{{lane}}], [x5], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
mov x8, x1
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x8], #64
|
||||
st1 { v4.8h, v5.8h, v6.8h, v7.8h }, [x8], #64
|
||||
|
||||
orr x8, x1, {{ 0 | setting(62) }} // pair
|
||||
{{ amx("ldy", 8) }}
|
||||
|
||||
eor x2, x2, x2
|
||||
orr x2, x2, {{ 0 | setting(63) }} // vector mode
|
||||
orr x2, x2, {{ 0 | setting(29) }} // perform Z0+=Y0
|
||||
{{ amx("fma16", 2) }}
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z1
|
||||
orr x2, x2, 64 // offset Y
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldp x5, x6, [x0, #8] // a base ptr, b base ptr
|
||||
|
||||
ld1 { v0.h }[0], [x6]
|
||||
st1 { v0.h }[0], [x1]
|
||||
{{ amx("ldy", 1) }}
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x5], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x5]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldx", 2) }}
|
||||
|
||||
mov x2, {{ 0 | setting(43) }} // f16
|
||||
orr x2, x2, {{ 0 | setting(38) }} // Broadcast Y
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
orr x2, x2, {{ 0 | setting(16) }} // X offset
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
ands x8, x5, 0x7f
|
||||
bne .store_generic
|
||||
cmp x6, 4
|
||||
bne .store_generic
|
||||
cmp x7, 4
|
||||
bne .store_generic
|
||||
|
||||
orr x5, x5, {{ 0 | setting(62) }} // pair
|
||||
{{ amx("stz", 5) }}
|
||||
b .non_linear_loop
|
||||
|
||||
.store_generic:
|
||||
|
||||
orr x8, x1, {{ 0 | setting(62) }} // pair
|
||||
{{ amx("stz", 8) }}
|
||||
|
||||
mov x8, x1
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x8], #64
|
||||
ld1 { v4.8h, v5.8h, v6.8h, v7.8h }, [x8], #64
|
||||
{% for neon in range(0, 8) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
st1 { v{{neon}}.h }[{{lane}}], [x5], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.load_tile:
|
||||
ldr x2, [x0, #16] // row major ptr
|
||||
orr x2, x2, {{ 0 | setting(62) }} // load pairs
|
||||
{{ amx("ldz", 2) }}
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
{{ AMX_CLR }}
|
||||
ret
|
||||
+658
@@ -0,0 +1,658 @@
|
||||
// vim: ft=arm
|
||||
.text
|
||||
.align 4
|
||||
|
||||
/* Z: 64x32 tile. each Z reg is f16x32
|
||||
Z0
|
||||
Z2
|
||||
...
|
||||
Z62
|
||||
|
||||
Z1
|
||||
Z3
|
||||
S63
|
||||
*/
|
||||
|
||||
|
||||
.global {{G}}apple_amx_mmm_f16_64x32_{{suffix}}
|
||||
{{G}}apple_amx_mmm_f16_64x32_{{suffix}}:
|
||||
|
||||
{{ AMX_SET }}
|
||||
|
||||
// set x1 to a 128 bytes aligned block for loads
|
||||
mov x1, sp
|
||||
lsr x1, x1, #7
|
||||
lsl x1, x1, #7
|
||||
sub x1, x1, 128
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.leaky_relu:
|
||||
.q_scale:
|
||||
.q_shl:
|
||||
.q_shr:
|
||||
b .unsupported
|
||||
|
||||
.add_mat_mul:
|
||||
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x4, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
orr x4, x4, {{ 0 | setting(62) }} // load pairs (A)
|
||||
|
||||
eor x5, x5, x5 // top left
|
||||
|
||||
orr x7, x5, {{ 0 | setting(20) }}
|
||||
orr x7, x7, {{ 0 | setting(6) }} // bottom left
|
||||
|
||||
.packed_packed_loop_1:
|
||||
{{ amx("ldx", 2) }}
|
||||
{{ amx("ldy", 4) }}
|
||||
add x2, x2, 64
|
||||
add x4, x4, 128
|
||||
|
||||
{{ amx("fma16", 5) }}
|
||||
{{ amx("fma16", 7) }}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.clear:
|
||||
// top left
|
||||
eor x2, x2, x2
|
||||
orr x2, x2, {{ 0 | setting(27) }}
|
||||
orr x2, x2, {{ 0 | setting(28) }}
|
||||
orr x2, x2, {{ 0 | setting(29) }} // Z = 0
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// top right
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(21) }} // Z row = 3
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom left
|
||||
eor x2, x2, {{ 0 | setting(20) }} // Z row = 2
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
mov x3, #16
|
||||
str x3, [x1]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_sub:
|
||||
|
||||
// performs a unary neg on Z
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
|
||||
mov x6, 64
|
||||
.per_col_sub_loop:
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms16", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
subs x6, x6, 1
|
||||
bne .per_col_sub_loop
|
||||
|
||||
// continue
|
||||
|
||||
.per_col_add:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{{ amx("ldx", 1) }}
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // z += y
|
||||
|
||||
// top left
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
// bottom left
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 2
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_sub_flipped:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
|
||||
{{ amx("ldx", 1) }}
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // z += y
|
||||
|
||||
{{ amx("fms16", 2) }}
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
{{ amx("fms16", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_sub_flipped:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2], #64
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1], #64
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
mov x2, {{ 0 | setting(29) }} // z += y
|
||||
|
||||
// top left
|
||||
{{ amx("fms16", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 3
|
||||
orr x2, x2, {{ 0 | setting(6) }} // Y offset
|
||||
{{ amx("fms16", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_sub:
|
||||
// performs a unary neg on Z
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
|
||||
mov x6, 64
|
||||
.per_row_sub_loop:
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms16", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
subs x6, x6, 1
|
||||
bne .per_row_sub_loop
|
||||
|
||||
// continue
|
||||
|
||||
.per_row_add:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2], #64
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1], #64
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
mov x2, {{ 0 | setting(29) }} // z += y
|
||||
|
||||
// top left
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
orr x2, x2, {{ 0 | setting(6) }} // Y offset
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_min:
|
||||
mov x2, 5
|
||||
b .per_row_min_max
|
||||
.per_row_max:
|
||||
mov x2, 7
|
||||
.per_row_min_max:
|
||||
ldr x5, [x0, #8]
|
||||
|
||||
add x6, x5, 64
|
||||
|
||||
lsl x2, x2, 47 // max(x,z) (or min)
|
||||
orr x2, x2, {{ 0 | setting(43) }} // f16
|
||||
|
||||
orr x8, x2, {{ 0 | setting(20) }} // bottom left
|
||||
|
||||
mov x4, 32
|
||||
.loop_per_row_max:
|
||||
// top half
|
||||
ld1 { v0.h }[0], [x5], #2
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
|
||||
{{ amx("ldx", 1) }}
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
add x2, x2, {{ 0 | setting(21) }}
|
||||
|
||||
// bottom half
|
||||
ld1 { v0.h }[0], [x6], #2
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
|
||||
{{ amx("ldx", 1) }}
|
||||
{{ amx("vecfp", 8) }}
|
||||
|
||||
add x8, x8, {{ 0 | setting(21) }}
|
||||
|
||||
subs x4, x4, 1
|
||||
bne .loop_per_row_max
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_min:
|
||||
mov x2, 5
|
||||
b .per_col_min_max
|
||||
.per_col_max:
|
||||
mov x2, 7
|
||||
.per_col_min_max:
|
||||
ldr x4, [x0, #8]
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x4]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
{{ amx("ldx", 1) }}
|
||||
|
||||
lsl x2, x2, 47 // max(x,z) (or min)
|
||||
orr x2, x2, {{ 0 | setting(43) }} // f16
|
||||
|
||||
mov x4, 64
|
||||
.loop_per_col_max:
|
||||
{{ amx("vecfp", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }}
|
||||
subs x4, x4, 1
|
||||
bne .loop_per_col_max
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_mul:
|
||||
ldr x4, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x4]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{{ amx("ldy", 1) }}
|
||||
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=X*Y
|
||||
|
||||
mov x6, 64
|
||||
.loop_per_col_mul:
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fma16", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }}
|
||||
add x4, x4, {{ 0 | setting(20) }}
|
||||
subs x6, x6, 1
|
||||
bne .loop_per_col_mul
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_mul:
|
||||
ldr x14, [x0, #8]
|
||||
add x15, x14, 64
|
||||
|
||||
// extrx
|
||||
eor x2, x2, x2 // X[0] = Z[0] (top left)
|
||||
|
||||
eor x4, x4, x4
|
||||
orr x4, x4, {{ 0 | setting(20) }} // X[0] = Z[1] (bottom left)
|
||||
|
||||
// fma16
|
||||
eor x6, x6, x6
|
||||
orr x6, x6, {{ 0 | setting(63) }} // vector mode
|
||||
orr x6, x6, {{ 0 | setting(27) }} // Z=X*Y Z[0]=X[0]*Y[0]
|
||||
|
||||
orr x8, x6, {{ 0 | setting(20) }} // Z[1]
|
||||
|
||||
mov x10, 32
|
||||
.loop_per_row_mul:
|
||||
// top
|
||||
ld1 { v0.h }[0], [x14], #2
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
|
||||
{{ amx("ldy", 1) }}
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fma16", 6) }}
|
||||
|
||||
add x2, x2, {{ 0 | setting(21) }}
|
||||
add x6, x6, {{ 0 | setting(21) }}
|
||||
|
||||
// bottom
|
||||
ld1 { v0.h }[0], [x15], #2
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
|
||||
{{ amx("ldy", 1) }}
|
||||
{{ amx("extrx", 4) }}
|
||||
{{ amx("fma16", 8) }}
|
||||
|
||||
add x4, x4, {{ 0 | setting(21) }}
|
||||
add x8, x8, {{ 0 | setting(21) }}
|
||||
|
||||
subs x10, x10, 1
|
||||
bne .loop_per_row_mul
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_sub:
|
||||
// performs a unary neg on Z, then go to scalar_add
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
|
||||
mov x6, 64
|
||||
.scalar_sub_loop:
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms16", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
subs x6, x6, 1
|
||||
bne .scalar_sub_loop
|
||||
|
||||
// continue on purpose
|
||||
|
||||
.scalar_add:
|
||||
ldr w5, [x0, #8]
|
||||
|
||||
fmov h0, w5
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // Z+=X
|
||||
{{ amx("fma16", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // Z1
|
||||
{{ amx("fma16", 2) }}
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_sub_flipped:
|
||||
ldr w5, [x0, #8]
|
||||
|
||||
fmov h0, w5
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 32 values
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // Z-=X
|
||||
{{ amx("fms16", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("fms16", 2) }}
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_mul:
|
||||
ldr w5, [x0, #8]
|
||||
|
||||
fmov h0, w5
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
{{ amx("ldy", 1) }} // load 32 values
|
||||
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=X*Y
|
||||
|
||||
mov x6, 64
|
||||
.scalar_mul_loop:
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fma16", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
subs x6, x6, 1
|
||||
bne .scalar_mul_loop
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_min:
|
||||
mov x2, 5
|
||||
b .scalar_min_max
|
||||
.scalar_max:
|
||||
mov x2, 7
|
||||
.scalar_min_max:
|
||||
ldr w5, [x0, #8]
|
||||
|
||||
fmov h0, w5
|
||||
dup v0.8h, v0.h[0]
|
||||
dup v1.8h, v0.h[0]
|
||||
dup v2.8h, v0.h[0]
|
||||
dup v3.8h, v0.h[0]
|
||||
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
lsl x2, x2, 47
|
||||
orr x2, x2, {{ 0 | setting(43) }} // f32
|
||||
|
||||
mov x3, 64
|
||||
.loop_scalar_max:
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z
|
||||
{{ amx("vecfp", 2) }}
|
||||
subs x3, x3, 1
|
||||
bne .loop_scalar_max
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
mov x3, 0 // x3 is the row
|
||||
.loop_load:
|
||||
// z reg is (row % 32) * 2 + (row / 32)
|
||||
and x9, x3, 0x1f
|
||||
lsl x9, x9, 1
|
||||
lsr x10, x3, 5
|
||||
add x9, x9, x10
|
||||
|
||||
mov x4, x5
|
||||
{% for neon in range(0, 4) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
ld1 { v{{neon}}.h }[{{lane}}], [x4], x7
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
st1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
{{ amx("ldy", 1) }}
|
||||
|
||||
lsl x2, x9, 20 // Z register to update
|
||||
orr x2, x2, {{ 0 | setting(63) }} // vector mode
|
||||
orr x2, x2, {{ 0 | setting(29) }} // perform Z+=Y
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
add x5, x5, x6
|
||||
add x3, x3, 1
|
||||
cmp x3, 64
|
||||
bne .loop_load
|
||||
|
||||
/*
|
||||
mov x3, 0 // x3 is the row
|
||||
.loop_load:
|
||||
and x9, x3, 0xf // x9 = row % 16
|
||||
lsl x9, x9, 2 // x9 = (row % 16) * 4
|
||||
lsr x10, x3, 4 // x10 = row / 16
|
||||
lsl x10, x10, 1 // x10 = (row / 16) * 2
|
||||
add x9, x9, x10 // x9 = x9 + x10
|
||||
|
||||
mov x4, x5
|
||||
{% for neon in range(0, 4) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 { v{{neon}}.s }[{{lane}}], [x4], x7
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{% for neon in range(0, 4) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 { v{{neon}}.s }[{{lane}}], [x4], x7
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x8]
|
||||
|
||||
mov x2, x1
|
||||
orr x2, x2, {{ 0 | setting(62) }} // load 32 values
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
lsl x2, x9, 20 // left Z register to update
|
||||
orr x2, x2, {{ 0 | setting(63) }} // vector mode
|
||||
orr x2, x2, {{ 0 | setting(29) }} // perform Z+=Y
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
add x2, x2, {{ 0 | setting(20) }}
|
||||
orr x2, x2, 64 // offset Y by 16 values
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
add x5, x5, x6
|
||||
add x3, x3, 1
|
||||
cmp x3, 32
|
||||
bne .loop_load
|
||||
|
||||
*/
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldp x5, x6, [x0, #8] // a base ptr, b base ptr
|
||||
|
||||
add x8, x1, 64
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x5], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x5]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x6]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
|
||||
{{ amx("ldx", 1) }}
|
||||
|
||||
// top
|
||||
eor x2, x2, x2
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
orr x2, x2, {{ 0 | setting(6) }} // Y offset
|
||||
{{ amx("fma16", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x7, 2
|
||||
bne .store_generic
|
||||
ands x8, x5, 0x7f
|
||||
bne .store_generic
|
||||
ands x8, x6, 0x7f
|
||||
bne .store_generic
|
||||
|
||||
lsl x8, x6, 5
|
||||
add x8, x8, x5 // x8 = 32*rsc
|
||||
orr x8, x8, {{ 0 | setting(56) }} // first to x8 is z1
|
||||
|
||||
mov x4, {{ 0 | setting(57) }} // Zreg += 2
|
||||
add x4, x4, x6 // +rsc
|
||||
|
||||
mov x3, 32
|
||||
.loop_store_direct:
|
||||
{{ amx("stz", 5) }}
|
||||
{{ amx("stz", 8) }}
|
||||
add x5, x5, x4
|
||||
add x8, x8, x4
|
||||
subs x3, x3, 1
|
||||
bne .loop_store_direct
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_generic:
|
||||
|
||||
mov x3, 0 // row id
|
||||
.loop_store:
|
||||
// z reg is (row % 32) * 2 + (row / 32)
|
||||
and x9, x3, 0x1f
|
||||
lsl x9, x9, 1
|
||||
lsr x10, x3, 5
|
||||
add x9, x9, x10
|
||||
|
||||
lsl x2, x9, 56
|
||||
orr x2, x2, x1
|
||||
{{ amx("stz", 2) }} // f16 x 32
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x1]
|
||||
|
||||
mov x4, x5
|
||||
{% for neon in range(0, 4) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
st1 { v{{neon}}.h }[{{lane}}], [x4], x7
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x6
|
||||
|
||||
add x3, x3, 1
|
||||
cmp x3, 64
|
||||
bne .loop_store
|
||||
b .non_linear_loop
|
||||
|
||||
.load_tile:
|
||||
ldr x2, [x0, #16] // row major ptr
|
||||
orr x3, x2, {{ 0 | setting(56) }}
|
||||
add x3, x3, #2048
|
||||
|
||||
mov x4, {{ 0 | setting(57) }} // z+=2
|
||||
add x4, x4, #64
|
||||
|
||||
mov x8, 32
|
||||
.loop_load_tile:
|
||||
{{ amx("ldz", 2) }}
|
||||
{{ amx("ldz", 3) }}
|
||||
add x2, x2, x4
|
||||
add x3, x3, x4
|
||||
subs x8, x8, 1
|
||||
bne .loop_load_tile
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
{{ AMX_CLR }}
|
||||
ret
|
||||
+533
@@ -0,0 +1,533 @@
|
||||
// vim: ft=arm
|
||||
.text
|
||||
.align 4
|
||||
|
||||
/* Z: 32x1
|
||||
z0[0] .. z0[15] z1[0] .. z1[15]
|
||||
*/
|
||||
|
||||
|
||||
.global {{G}}apple_amx_mmm_f32_32x1_{{suffix}}
|
||||
{{G}}apple_amx_mmm_f32_32x1_{{suffix}}:
|
||||
|
||||
{{ AMX_SET }}
|
||||
|
||||
// set x1 to a 128 bytes aligned block for loads
|
||||
mov x1, sp
|
||||
lsr x1, x1, #7
|
||||
lsl x1, x1, #7
|
||||
sub x1, x1, 128
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.leaky_relu:
|
||||
.q_scale:
|
||||
.q_shl:
|
||||
.q_shr:
|
||||
b .unsupported
|
||||
|
||||
.add_mat_mul:
|
||||
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x4, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
orr x4, x4, {{ 0 | setting(62) }} // load a pair of A
|
||||
|
||||
mov x5, {{ 0 | setting(44) }} // f32
|
||||
orr x5, x5, {{ 0 | setting(38) }} // Broadcast Y
|
||||
|
||||
orr x6, x5, {{ 0 | setting(20) }} // z offset
|
||||
orr x6, x6, {{ 0 | setting(16) }} // x offset
|
||||
|
||||
cmp x3, #16
|
||||
blt .packed_packed_loop_1
|
||||
|
||||
mov x9, {{ 0 | setting(32) }} // Y broadcast offset += 1
|
||||
|
||||
.packed_packed_loop_16:
|
||||
mov x7, x5
|
||||
mov x8, x6
|
||||
{{ amx("ldy", 2) }}
|
||||
{% for k in range(0, 16) %}
|
||||
{{ amx("ldx", 4) }}
|
||||
add x4, x4, 128
|
||||
{{ amx("vecfp", 7) }}
|
||||
{{ amx("vecfp", 8) }}
|
||||
add x7, x7, x9
|
||||
add x8, x8, x9
|
||||
{% endfor %}
|
||||
add x2, x2, #64
|
||||
sub x3, x3, #16
|
||||
cmp x3, #16
|
||||
bge .packed_packed_loop_16
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.packed_packed_loop_1:
|
||||
ldr w7, [x2], #4
|
||||
str w7, [x1]
|
||||
{{ amx("ldx", 4) }}
|
||||
{{ amx("ldy", 1) }}
|
||||
{{ amx("vecfp", 5) }}
|
||||
{{ amx("vecfp", 6) }}
|
||||
add x4, x4, 128
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.clear:
|
||||
// top left
|
||||
eor x2, x2, x2
|
||||
orr x2, x2, {{ 0 | setting(27) }}
|
||||
orr x2, x2, {{ 0 | setting(28) }}
|
||||
orr x2, x2, {{ 0 | setting(29) }} // Z = 0
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// top right
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(21) }} // Z row = 3
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom left
|
||||
eor x2, x2, {{ 0 | setting(20) }} // Z row = 2
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_sub:
|
||||
|
||||
// performs a unary neg on Z
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
// extr[hxyz] is suport confusing
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms32", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("extrx", 2) }} // extr[hxyz] is confusing
|
||||
{{ amx("fms32", 4) }}
|
||||
|
||||
// continue
|
||||
|
||||
.per_col_add:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
// broadcast value to x0
|
||||
ld1 { v0.s }[0], [x2]
|
||||
dup v0.4s, v0.s[0]
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
sub x1, x1, #64
|
||||
|
||||
{{ amx("ldx", 1) }} // load into x0 by default
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // z += y
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
orr x2, x2, {{ 0 | setting(20) }} // target is now z1
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_sub_flipped:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
// broadcast value to x0
|
||||
ld1 { v0.s }[0], [x2]
|
||||
dup v0.4s, v0.s[0]
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
sub x1, x1, #64
|
||||
|
||||
{{ amx("ldx", 1) }} // load into x0 by default
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // z += y
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
orr x2, x2, {{ 0 | setting(20) }} // target is now z1
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_sub_flipped:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
mov x2, {{ 0 | setting(63) }} // vector mode
|
||||
orr x2, x2, {{ 0 | setting(29) }} // z -= y
|
||||
|
||||
// top left
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
// bottom left
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
orr x2, x2, {{ 0 | setting(6) }} // Y offset
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_sub:
|
||||
// performs a unary neg on Z
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms32", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms32", 4) }}
|
||||
|
||||
// continue
|
||||
|
||||
.per_row_add:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
mov x2, {{ 0 | setting(63) }} // vector mode
|
||||
orr x2, x2, {{ 0 | setting(29) }} // z += y
|
||||
|
||||
// top left
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom left
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
orr x2, x2, {{ 0 | setting(6) }} // Y offset
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_min:
|
||||
mov x2, 5
|
||||
b .per_row_min_max
|
||||
.per_row_max:
|
||||
mov x2, 7
|
||||
.per_row_min_max:
|
||||
ldr x5, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x5], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x5]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x5, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldx", 5) }}
|
||||
|
||||
lsl x2, x2, 47 // max(x,z) (or min)
|
||||
orr x2, x2, {{ 0 | setting(44) }} // f32
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
orr x2, x2, {{ 0 | setting(16) }} // x1
|
||||
orr x2, x2, {{ 0 | setting(20) }} // z1
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_min:
|
||||
mov x2, 5
|
||||
b .per_col_min_max
|
||||
.per_col_max:
|
||||
mov x2, 7
|
||||
.per_col_min_max:
|
||||
ldr x4, [x0, #8]
|
||||
|
||||
// broadcast value to x0
|
||||
ld1 { v0.s }[0], [x4]
|
||||
dup v0.4s, v0.s[0]
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
sub x1, x1, #64
|
||||
|
||||
{{ amx("ldx", 1) }}
|
||||
|
||||
lsl x2, x2, 47 // max(x,z) (or min)
|
||||
orr x2, x2, {{ 0 | setting(44) }} // f32
|
||||
|
||||
{{ amx("vecfp", 2) }}
|
||||
orr x2, x2, {{ 0 | setting(20) }} // z offset
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_mul:
|
||||
ldr x4, [x0, #8]
|
||||
|
||||
// broadcast value to y0
|
||||
ld1 { v0.s }[0], [x4]
|
||||
dup v0.4s, v0.s[0]
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
st1 { v0.4s }, [x1], #16
|
||||
sub x1, x1, #64
|
||||
|
||||
{{ amx("ldy", 1) }}
|
||||
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
{{ amx("extrx", 2) }}
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=X*Y
|
||||
{{ amx("fma32", 4) }}
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z1
|
||||
{{ amx("extrx", 2) }}
|
||||
orr x4, x4, {{ 0 | setting(20) }} // Z1
|
||||
{{ amx("fma32", 4) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_mul:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
{{ amx("extrx", 2) }}
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=X*Y
|
||||
{{ amx("fma32", 4) }}
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z1
|
||||
{{ amx("extrx", 2) }}
|
||||
orr x4, x4, {{ 0 | setting(20) }} // Z1
|
||||
orr x4, x4, {{ 0 | setting(6) }} // Y1
|
||||
{{ amx("fma32", 4) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_sub:
|
||||
// performs a unary neg on Z, then go to scalar_add
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms32", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms32", 4) }}
|
||||
|
||||
// continue on purpose
|
||||
|
||||
.scalar_add:
|
||||
ldr w5, [x0, #8]
|
||||
|
||||
fmov s0, w5
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // Z+=X
|
||||
{{ amx("fma32", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("fma32", 2) }}
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_sub_flipped:
|
||||
ldr w5, [x0, #8]
|
||||
fmov s0, w5
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // Z-=X
|
||||
{{ amx("fms32", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("fms32", 2) }}
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_mul:
|
||||
ldr w5, [x0, #8]
|
||||
fmov s0, w5
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{{ amx("ldy", 1) }} // load 16 values
|
||||
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=X*Y
|
||||
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fma32", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fma32", 4) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_min:
|
||||
mov x2, 5
|
||||
b .scalar_min_max
|
||||
.scalar_max:
|
||||
mov x2, 7
|
||||
.scalar_min_max:
|
||||
ldr w5, [x0, #8]
|
||||
fmov s0, w5
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
lsl x2, x2, 47
|
||||
orr x2, x2, {{ 0 | setting(44) }} // f32
|
||||
|
||||
{{ amx("vecfp", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
{% for neon in range(0, 8) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 { v{{neon}}.s }[{{lane}}], [x5], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
mov x8, x1
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x8], #64
|
||||
st1 { v4.4s, v5.4s, v6.4s, v7.4s }, [x8], #64
|
||||
|
||||
orr x8, x1, {{ 0 | setting(62) }} // pair
|
||||
{{ amx("ldy", 8) }}
|
||||
|
||||
eor x2, x2, x2
|
||||
orr x2, x2, {{ 0 | setting(63) }} // vector mode
|
||||
orr x2, x2, {{ 0 | setting(29) }} // perform Z0+=Y0
|
||||
{{ amx("fma32", 2) }}
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z1
|
||||
orr x2, x2, 64 // offset Y by 16 values
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldp x5, x6, [x0, #8] // a base ptr, b base ptr
|
||||
|
||||
ld1 { v0.s }[0], [x6]
|
||||
st1 { v0.s }[0], [x1]
|
||||
{{ amx("ldy", 1) }}
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x5], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x5]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldx", 2) }}
|
||||
|
||||
mov x2, {{ 0 | setting(44) }} // f32
|
||||
orr x2, x2, {{ 0 | setting(38) }} // Broadcast Y
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
orr x2, x2, {{ 0 | setting(16) }} // X offset
|
||||
{{ amx("vecfp", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
ands x8, x5, 0x7f
|
||||
bne .store_generic
|
||||
cmp x6, 4
|
||||
bne .store_generic
|
||||
cmp x7, 4
|
||||
bne .store_generic
|
||||
|
||||
orr x5, x5, {{ 0 | setting(62) }} // pair
|
||||
{{ amx("stz", 5) }}
|
||||
b .non_linear_loop
|
||||
|
||||
.store_generic:
|
||||
|
||||
orr x8, x1, {{ 0 | setting(62) }} // pair
|
||||
{{ amx("stz", 8) }}
|
||||
|
||||
mov x8, x1
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x8], #64
|
||||
ld1 { v4.4s, v5.4s, v6.4s, v7.4s }, [x8], #64
|
||||
{% for neon in range(0, 8) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{neon}}.s }[{{lane}}], [x5], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.load_tile:
|
||||
ldr x2, [x0, #16] // row major ptr
|
||||
orr x2, x2, {{ 0 | setting(62) }} // load pairs
|
||||
{{ amx("ldz", 2) }}
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
{{ AMX_CLR }}
|
||||
ret
|
||||
+764
@@ -0,0 +1,764 @@
|
||||
// vim: ft=arm
|
||||
.text
|
||||
.align 4
|
||||
|
||||
/* Z: 32x32
|
||||
z0[0] .. z0[15] z1[0] .. z1[15]
|
||||
z4[0] .. z4[15] z5[0] .. z5[15]
|
||||
..
|
||||
z60[0] .. z60[15] z61[0] .. z61[15]
|
||||
|
||||
z2[0] .. z2[15] z3[0] .. z3[15]
|
||||
z5[0] .. z5[15] z6[0] .. z6[15]
|
||||
..
|
||||
z62[0] .. z62[15] z63[0] .. z63[15]
|
||||
*/
|
||||
|
||||
|
||||
.global {{G}}apple_amx_mmm_f32_32x32_{{suffix}}
|
||||
{{G}}apple_amx_mmm_f32_32x32_{{suffix}}:
|
||||
|
||||
{{ AMX_SET }}
|
||||
|
||||
// set x1 to a 128 bytes aligned block for loads
|
||||
mov x1, sp
|
||||
lsr x1, x1, #7
|
||||
lsl x1, x1, #7
|
||||
sub x1, x1, 128
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.leaky_relu:
|
||||
.q_scale:
|
||||
.q_shl:
|
||||
.q_shr:
|
||||
b .unsupported
|
||||
|
||||
.add_mat_mul:
|
||||
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x4, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
orr x4, x4, {{ 0 | setting(62) }} // load pairs (A)
|
||||
orr x2, x2, {{ 0 | setting(62) }} // load pairs (B)
|
||||
|
||||
eor x5, x5, x5 // top left
|
||||
|
||||
orr x6, x5, {{ 0 | setting(20) }} // Z row = 1
|
||||
orr x6, x6, {{ 0 | setting(16) }} // top right
|
||||
|
||||
orr x7, x5, {{ 0 | setting(21) }}
|
||||
orr x7, x7, {{ 0 | setting(6) }} // bottom left
|
||||
|
||||
orr x8, x7, x6 // bottom right
|
||||
|
||||
.packed_packed_loop_1:
|
||||
{{ amx("ldx", 2) }}
|
||||
{{ amx("ldy", 4) }}
|
||||
add x2, x2, 128
|
||||
add x4, x4, 128
|
||||
|
||||
{{ amx("fma32", 5) }}
|
||||
{{ amx("fma32", 6) }}
|
||||
{{ amx("fma32", 7) }}
|
||||
{{ amx("fma32", 8) }}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.clear:
|
||||
// top left
|
||||
eor x2, x2, x2
|
||||
orr x2, x2, {{ 0 | setting(27) }}
|
||||
orr x2, x2, {{ 0 | setting(28) }}
|
||||
orr x2, x2, {{ 0 | setting(29) }} // Z = 0
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// top right
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(21) }} // Z row = 3
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom left
|
||||
eor x2, x2, {{ 0 | setting(20) }} // Z row = 2
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_sub:
|
||||
|
||||
// performs a unary neg on Z
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
|
||||
mov x6, 64
|
||||
.per_col_sub_loop:
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms32", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
subs x6, x6, 1
|
||||
bne .per_col_sub_loop
|
||||
|
||||
// continue
|
||||
|
||||
.per_col_add:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x1, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldx", 1) }}
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // z += y
|
||||
|
||||
// top left
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom left
|
||||
orr x2, x2, {{ 0 | setting(21) }} // Z row = 2
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(16) }} // X offset
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 3
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// top right
|
||||
eor x2, x2, {{ 0 | setting(21) }} // Z row = 1
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_sub_flipped:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x1, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldx", 1) }}
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // z += y
|
||||
|
||||
// top left
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
// bottom left
|
||||
orr x2, x2, {{ 0 | setting(21) }} // Z row = 2
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(16) }} // X offset
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 3
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
// top right
|
||||
eor x2, x2, {{ 0 | setting(21) }} // Z row = 1
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_sub_flipped:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
mov x2, {{ 0 | setting(29) }} // z += y
|
||||
|
||||
// top left
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
// top right
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(21) }} // Z row = 3
|
||||
orr x2, x2, {{ 0 | setting(6) }} // Y offset
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
// bottom left
|
||||
eor x2, x2, {{ 0 | setting(20) }} // Z row = 2
|
||||
{{ amx("fms32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_sub:
|
||||
// performs a unary neg on Z
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
|
||||
mov x6, 64
|
||||
.per_row_sub_loop:
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms32", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
subs x6, x6, 1
|
||||
bne .per_row_sub_loop
|
||||
|
||||
// continue
|
||||
|
||||
.per_row_add:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
mov x2, {{ 0 | setting(29) }} // z += y
|
||||
|
||||
// top left
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// top right
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(21) }} // Z row = 3
|
||||
orr x2, x2, {{ 0 | setting(6) }} // Y offset
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom left
|
||||
eor x2, x2, {{ 0 | setting(20) }} // Z row = 2
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_min:
|
||||
mov x2, 5
|
||||
b .per_row_min_max
|
||||
.per_row_max:
|
||||
mov x2, 7
|
||||
.per_row_min_max:
|
||||
ldr x5, [x0, #8]
|
||||
|
||||
add x6, x5, 64
|
||||
|
||||
lsl x2, x2, 47 // max(x,z) (or min)
|
||||
orr x2, x2, {{ 0 | setting(44) }} // f32
|
||||
orr x3, x2, {{ 0 | setting(20) }} // right half: z offset
|
||||
|
||||
orr x8, x2, {{ 0 | setting(21) }} // bottom left
|
||||
orr x9, x3, {{ 0 | setting(21) }} // bottom right
|
||||
|
||||
mov x4, 16
|
||||
.loop_per_row_max:
|
||||
// top half
|
||||
ld1 { v0.s }[0], [x5], #4
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
|
||||
{{ amx("ldx", 1) }}
|
||||
{{ amx("vecfp", 2) }}
|
||||
{{ amx("vecfp", 3) }}
|
||||
|
||||
add x2, x2, {{ 0 | setting(22) }}
|
||||
add x3, x3, {{ 0 | setting(22) }}
|
||||
|
||||
// bottom half
|
||||
ld1 { v0.s }[0], [x6], #4
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
|
||||
{{ amx("ldx", 1) }}
|
||||
{{ amx("vecfp", 8) }}
|
||||
{{ amx("vecfp", 9) }}
|
||||
|
||||
add x8, x8, {{ 0 | setting(22) }}
|
||||
add x9, x9, {{ 0 | setting(22) }}
|
||||
|
||||
subs x4, x4, 1
|
||||
bne .loop_per_row_max
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_min:
|
||||
mov x2, 5
|
||||
b .per_col_min_max
|
||||
.per_col_max:
|
||||
mov x2, 7
|
||||
.per_col_min_max:
|
||||
ldr x4, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x4], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x4]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x3, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldx", 3) }}
|
||||
|
||||
lsl x2, x2, 47 // max(x,z) (or min)
|
||||
orr x2, x2, {{ 0 | setting(44) }} // f32
|
||||
|
||||
orr x3, x2, {{ 0 | setting(16) }} // right half: x offset
|
||||
orr x3, x3, {{ 0 | setting(20) }} // right half: z offset
|
||||
|
||||
mov x4, 32
|
||||
.loop_per_col_max:
|
||||
{{ amx("vecfp", 2) }}
|
||||
{{ amx("vecfp", 3) }}
|
||||
add x2, x2, {{ 0 | setting(21) }}
|
||||
add x3, x3, {{ 0 | setting(21) }}
|
||||
subs x4, x4, 1
|
||||
bne .loop_per_col_max
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_col_mul:
|
||||
ldr x4, [x0, #8]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x4], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x4]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
eor x3, x3, x3
|
||||
orr x3, x3, {{ 0 | setting(20) }} // Z[1]
|
||||
orr x3, x3, {{ 0 | setting(16) }} // X[1]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=X*Y
|
||||
|
||||
mov x5, {{ 0 | setting(63) }} // vector mode
|
||||
orr x5, x5, {{ 0 | setting(27) }} // Z=X*Y
|
||||
orr x5, x5, {{ 0 | setting(20) }} // Z right
|
||||
orr x5, x5, {{ 0 | setting(16) }} // X[1] (right)
|
||||
orr x5, x5, {{ 0 | setting(6) }} // Y[1] (right)
|
||||
|
||||
mov x6, 32
|
||||
.loop_per_col_mul:
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("extrx", 3) }}
|
||||
{{ amx("fma32", 4) }}
|
||||
{{ amx("fma32", 5) }}
|
||||
add x2, x2, {{ 0 | setting(21) }}
|
||||
add x3, x3, {{ 0 | setting(21) }}
|
||||
add x4, x4, {{ 0 | setting(21) }}
|
||||
add x5, x5, {{ 0 | setting(21) }}
|
||||
subs x6, x6, 1
|
||||
bne .loop_per_col_mul
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.per_row_mul:
|
||||
ldr x14, [x0, #8]
|
||||
add x15, x14, 64
|
||||
|
||||
// extrx
|
||||
eor x2, x2, x2 // X[0] = Z[0] (top left)
|
||||
|
||||
eor x3, x3, x3
|
||||
orr x3, x3, {{ 0 | setting(20) }} // Z[1]
|
||||
orr x3, x3, {{ 0 | setting(16) }} // X[1] = Z[1] (top right)
|
||||
|
||||
eor x4, x4, x4
|
||||
orr x4, x4, {{ 0 | setting(21) }} // X[0] = Z[2] (bottom left)
|
||||
|
||||
orr x5, x4, {{ 0 | setting(20) }}
|
||||
orr x5, x5, {{ 0 | setting(16) }} // X[1] = Z[3] (bottom right)
|
||||
|
||||
// fma32
|
||||
eor x6, x6, x6
|
||||
orr x6, x6, {{ 0 | setting(63) }} // vector mode
|
||||
orr x6, x6, {{ 0 | setting(27) }} // Z=X*Y Z[0]=X[0]*Y[0]
|
||||
|
||||
orr x7, x6, {{ 0 | setting(20) }} // Z[1]
|
||||
orr x7, x7, {{ 0 | setting(16) }} // X[1] Z[1] = X[1]*Y[0]
|
||||
|
||||
orr x8, x6, {{ 0 | setting(21) }} // Z[2]
|
||||
orr x8, x8, {{ 0 | setting(21) }} // Z[2]
|
||||
|
||||
orr x9, x8, {{ 0 | setting(20) }} // Z[3]
|
||||
orr x9, x9, {{ 0 | setting(16) }} // X[1]
|
||||
|
||||
mov x10, 16
|
||||
.loop_per_row_mul:
|
||||
// top
|
||||
ld1 { v0.s }[0], [x14], #4
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
|
||||
{{ amx("ldy", 1) }}
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("extrx", 3) }}
|
||||
{{ amx("fma32", 6) }}
|
||||
{{ amx("fma32", 7) }}
|
||||
|
||||
add x2, x2, {{ 0 | setting(22) }}
|
||||
add x3, x3, {{ 0 | setting(22) }}
|
||||
add x6, x6, {{ 0 | setting(22) }}
|
||||
add x7, x7, {{ 0 | setting(22) }}
|
||||
|
||||
// bottom
|
||||
ld1 { v0.s }[0], [x15], #4
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
|
||||
{{ amx("ldy", 1) }}
|
||||
{{ amx("extrx", 4) }}
|
||||
{{ amx("extrx", 5) }}
|
||||
{{ amx("fma32", 8) }}
|
||||
{{ amx("fma32", 9) }}
|
||||
|
||||
add x4, x4, {{ 0 | setting(22) }}
|
||||
add x5, x5, {{ 0 | setting(22) }}
|
||||
add x8, x8, {{ 0 | setting(22) }}
|
||||
add x9, x9, {{ 0 | setting(22) }}
|
||||
|
||||
subs x10, x10, 1
|
||||
bne .loop_per_row_mul
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_sub:
|
||||
// performs a unary neg on Z, then go to scalar_add
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(28) }}
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=-X
|
||||
|
||||
mov x6, 64
|
||||
.scalar_sub_loop:
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fms32", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
subs x6, x6, 1
|
||||
bne .scalar_sub_loop
|
||||
|
||||
// continue on purpose
|
||||
|
||||
.scalar_add:
|
||||
ldr w5, [x0, #8]
|
||||
|
||||
fmov s0, w5
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // Z+=X
|
||||
{% for chunk in range(0, 4) %}
|
||||
{{ amx("fma32", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_sub_flipped:
|
||||
ldr w5, [x0, #8]
|
||||
fmov s0, w5
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
mov x2, {{ 0 | setting(28) }} // Z-=X
|
||||
{% for chunk in range(0, 4) %}
|
||||
{{ amx("fms32", 2) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_mul:
|
||||
ldr w5, [x0, #8]
|
||||
fmov s0, w5
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{{ amx("ldy", 1) }} // load 16 values
|
||||
|
||||
eor x2, x2, x2 // X[0] = Z[0]
|
||||
|
||||
mov x4, {{ 0 | setting(63) }} // vector mode
|
||||
orr x4, x4, {{ 0 | setting(27) }} // Z=X*Y
|
||||
|
||||
mov x6, 64
|
||||
.scalar_mul_loop:
|
||||
{{ amx("extrx", 2) }}
|
||||
{{ amx("fma32", 4) }}
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z row
|
||||
add x4, x4, {{ 0 | setting(20) }} // next Z row
|
||||
subs x6, x6, 1
|
||||
bne .scalar_mul_loop
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.scalar_min:
|
||||
mov x2, 5
|
||||
b .scalar_min_max
|
||||
.scalar_max:
|
||||
mov x2, 7
|
||||
.scalar_min_max:
|
||||
ldr w5, [x0, #8]
|
||||
fmov s0, w5
|
||||
dup v0.4s, v0.s[0]
|
||||
dup v1.4s, v0.s[0]
|
||||
dup v2.4s, v0.s[0]
|
||||
dup v3.4s, v0.s[0]
|
||||
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{{ amx("ldx", 1) }} // load 16 values
|
||||
|
||||
lsl x2, x2, 47
|
||||
orr x2, x2, {{ 0 | setting(44) }} // f32
|
||||
|
||||
mov x3, 64
|
||||
.loop_scalar_max:
|
||||
add x2, x2, {{ 0 | setting(20) }} // next Z
|
||||
{{ amx("vecfp", 2) }}
|
||||
subs x3, x3, 1
|
||||
bne .loop_scalar_max
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
add x8, x1, 64
|
||||
|
||||
mov x3, 0 // x3 is the row
|
||||
.loop_load:
|
||||
and x9, x3, 0xf // x9 = row % 16
|
||||
lsl x9, x9, 2 // x9 = (row % 16) * 4
|
||||
lsr x10, x3, 4 // x10 = row / 16
|
||||
lsl x10, x10, 1 // x10 = (row / 16) * 2
|
||||
add x9, x9, x10 // x9 = x9 + x10
|
||||
|
||||
mov x4, x5
|
||||
{% for neon in range(0, 4) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 { v{{neon}}.s }[{{lane}}], [x4], x7
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
{% for neon in range(0, 4) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 { v{{neon}}.s }[{{lane}}], [x4], x7
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x8]
|
||||
|
||||
mov x2, x1
|
||||
orr x2, x2, {{ 0 | setting(62) }} // load 32 values
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
lsl x2, x9, 20 // left Z register to update
|
||||
orr x2, x2, {{ 0 | setting(63) }} // vector mode
|
||||
orr x2, x2, {{ 0 | setting(29) }} // perform Z+=Y
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
add x2, x2, {{ 0 | setting(20) }}
|
||||
orr x2, x2, 64 // offset Y by 16 values
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
add x5, x5, x6
|
||||
add x3, x3, 1
|
||||
cmp x3, 32
|
||||
bne .loop_load
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldp x5, x6, [x0, #8] // a base ptr, b base ptr
|
||||
|
||||
add x8, x1, 64
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x5], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x5]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldy", 2) }}
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x6], #64
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1], #64
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x6]
|
||||
st1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
sub x1, x1, #64
|
||||
|
||||
orr x2, x1, {{ 0 | setting(62) }} // load a pair
|
||||
{{ amx("ldx", 2) }}
|
||||
|
||||
// top left
|
||||
eor x2, x2, x2
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// top right
|
||||
orr x2, x2, {{ 0 | setting(20) }} // Z row = 1
|
||||
orr x2, x2, {{ 0 | setting(16) }} // X offset
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom right
|
||||
orr x2, x2, {{ 0 | setting(21) }} // Z row = 3
|
||||
orr x2, x2, {{ 0 | setting(6) }} // Y offset
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
// bottom left
|
||||
eor x2, x2, {{ 0 | setting(20) }} // Z row = 2
|
||||
eor x2, x2, {{ 0 | setting(16) }} // X offset <-
|
||||
{{ amx("fma32", 2) }}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x7, 4
|
||||
bne .store_generic
|
||||
ands x8, x5, 0x7f
|
||||
bne .store_generic
|
||||
ands x8, x6, 0x7f
|
||||
bne .store_generic
|
||||
|
||||
orr x5, x5, {{ 0 | setting(62) }} // pair
|
||||
lsl x8, x6, 4
|
||||
add x8, x8, x5 // x8 = 16*rsc
|
||||
orr x8, x8, {{ 0 | setting(57) }} // first to x8 is z2
|
||||
|
||||
mov x4, {{ 0 | setting(58) }} // Zreg += 4
|
||||
add x4, x4, x6 // +rsc
|
||||
|
||||
mov x3, 16
|
||||
.loop_store_direct:
|
||||
{{ amx("stz", 5) }}
|
||||
{{ amx("stz", 8) }}
|
||||
add x5, x5, x4
|
||||
add x8, x8, x4
|
||||
subs x3, x3, 1
|
||||
bne .loop_store_direct
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_generic:
|
||||
|
||||
add x8, x1, 64
|
||||
|
||||
mov x3, 0 // row id
|
||||
.loop_store:
|
||||
and x9, x3, 0xf // x9 = row % 16
|
||||
lsl x9, x9, 2 // x9 = (row % 16) * 4
|
||||
lsr x10, x3, 4 // x10 = row / 16
|
||||
lsl x10, x10, 1 // x10 = (row / 16) * 2
|
||||
add x9, x9, x10 // x9 = x9 + x10
|
||||
|
||||
lsl x2, x9, 56
|
||||
orr x2, x2, {{ 0 | setting(62) }}
|
||||
orr x2, x2, x1
|
||||
{{ amx("stz", 2) }}
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x1]
|
||||
|
||||
mov x4, x5
|
||||
{% for neon in range(0, 4) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{neon}}.s }[{{lane}}], [x4], x7
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x8]
|
||||
{% for neon in range(0, 4) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{neon}}.s }[{{lane}}], [x4], x7
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x6
|
||||
|
||||
add x3, x3, 1
|
||||
cmp x3, 32
|
||||
bne .loop_store
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.load_tile:
|
||||
ldr x2, [x0, #16] // row major ptr
|
||||
orr x2, x2, {{ 0 | setting(62) }} // load pairs
|
||||
mov x3, x2
|
||||
orr x3, x3, {{ 0 | setting(57) }}
|
||||
add x3, x3, #2048
|
||||
|
||||
mov x4, {{ 0 | setting(58) }} // z+=4
|
||||
add x4, x4, #128
|
||||
|
||||
mov x8, 16
|
||||
.loop_load_tile:
|
||||
{{ amx("ldz", 2) }}
|
||||
{{ amx("ldz", 3) }}
|
||||
add x2, x2, x4
|
||||
add x3, x3, x4
|
||||
subs x8, x8, 1
|
||||
bne .loop_load_tile
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
{{ AMX_CLR }}
|
||||
ret
|
||||
Vendored
+37
@@ -0,0 +1,37 @@
|
||||
// vim: ft=arm
|
||||
|
||||
.non_linear:
|
||||
sub x0, x0, 40
|
||||
|
||||
.non_linear_loop:
|
||||
add x0, x0, 40
|
||||
ldr x2, [x0]
|
||||
|
||||
mov x4, #{{ jump_table | length }}
|
||||
|
||||
cmp x2, #{{ jump_table | length }}
|
||||
csel x2, x2, x4, lt
|
||||
cmp x2, #0
|
||||
csel x2, x4, x2, lt
|
||||
|
||||
adr x3, .jmp_table
|
||||
add x3, x3, x2, LSL#2
|
||||
br x3
|
||||
|
||||
.jmp_table:
|
||||
{% for j in jump_table %}
|
||||
b .{{j}}
|
||||
{% endfor %}
|
||||
b .unsupported
|
||||
|
||||
add x0, x2, #4000
|
||||
b .return
|
||||
|
||||
.unsupported:
|
||||
mov x0, #1
|
||||
b .return
|
||||
|
||||
.done:
|
||||
mov x0, 0
|
||||
b .return
|
||||
|
||||
+71
@@ -0,0 +1,71 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// no preservation either for v0-v7 and v16-v31
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
{% if needs_pragma == true %}
|
||||
.cpu generic+fp+simd+fp16
|
||||
{% endif %}
|
||||
.global {{G}}arm64fp16_leaky_relu_f16_8n_{{suffix}}
|
||||
{{G}}arm64fp16_leaky_relu_f16_8n_{{suffix}}:
|
||||
|
||||
cmp x1, #0
|
||||
beq .return
|
||||
|
||||
mov v31.h[0], w2
|
||||
dup v31.8h, v31.h[0]
|
||||
mov x2, x0
|
||||
|
||||
cmp x1, #64
|
||||
blt .loop
|
||||
|
||||
ld1 { v16.8h, v17.8h, v18.8h, v19.8h }, [x2], #64
|
||||
.loop4:
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [x2], #64
|
||||
|
||||
fmul v20.8h, v16.8h, v31.8h
|
||||
fmul v21.8h, v17.8h, v31.8h
|
||||
fmul v22.8h, v18.8h, v31.8h
|
||||
fmul v23.8h, v19.8h, v31.8h
|
||||
|
||||
fcmge v24.8h, v16.8h, #0.0
|
||||
fcmge v25.8h, v17.8h, #0.0
|
||||
fcmge v26.8h, v18.8h, #0.0
|
||||
fcmge v27.8h, v19.8h, #0.0
|
||||
|
||||
bsl v24.16b, v16.16b, v20.16b
|
||||
bsl v25.16b, v17.16b, v21.16b
|
||||
bsl v26.16b, v18.16b, v22.16b
|
||||
bsl v27.16b, v19.16b, v23.16b
|
||||
|
||||
st1 { v24.8h, v25.8h, v26.8h, v27.8h }, [x0], #64
|
||||
|
||||
and v16.16b, v0.16b, v0.16b
|
||||
and v17.16b, v1.16b, v1.16b
|
||||
and v18.16b, v2.16b, v2.16b
|
||||
and v19.16b, v3.16b, v3.16b
|
||||
|
||||
subs x1, x1, #32
|
||||
cmp x1, #64
|
||||
bge .loop4
|
||||
|
||||
cmp x1, #0
|
||||
beq .return
|
||||
|
||||
.loop:
|
||||
ld1 { v16.8h }, [x0]
|
||||
|
||||
fmul v17.8h, v16.8h, v31.8h
|
||||
fcmge v18.8h, v16.8h, #0.0
|
||||
bsl v18.16b, v16.16b, v17.16b
|
||||
|
||||
st1 { v18.8h }, [x0], #16
|
||||
|
||||
subs x1, x1, #8
|
||||
bne .loop
|
||||
|
||||
.return:
|
||||
ret
|
||||
+86
@@ -0,0 +1,86 @@
|
||||
|
||||
{% macro scalar(label, op, from, to, flipped=false) %}
|
||||
.{{label}}:
|
||||
add x2, x0, #8
|
||||
ld1 {v0.h}[0], [ x2 ]
|
||||
dup v0.8h, v0.h[0]
|
||||
{% if flipped %}
|
||||
{% for reg in range(from, to + 1) %}
|
||||
{{op}} v{{reg}}.8h, v{{reg}}.8h, v0.8h
|
||||
{% endfor %}
|
||||
{% else %}
|
||||
{% for reg in range(from, to + 1) %}
|
||||
{{op}} v{{reg}}.8h, v0.8h, v{{reg}}.8h
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{% macro per_row(label, op, mr, from, to, flipped=false) %}
|
||||
.{{label}}:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
{% set mr_over_8 = mr // 8 %}
|
||||
{% set mr_over_8_min_1 = mr // 8 - 1 %}
|
||||
|
||||
{% for reg in range(0, mr_over_8_min_1 + 1) %}
|
||||
ldr q{{reg}}, [ x2 ], #16
|
||||
{% endfor %}
|
||||
|
||||
{% if flipped %}
|
||||
{% for acc in range(from, to + 1) %}
|
||||
{% set other = (acc - from) % mr_over_8 %}
|
||||
{{op}} v{{acc}}.8h, v{{acc}}.8h, v{{other}}.8h
|
||||
{% endfor %}
|
||||
{% else %}
|
||||
{% for acc in range(from, to + 1) %}
|
||||
{% set other = (acc - from) % mr_over_8 %}
|
||||
{{op}} v{{acc}}.8h, v{{other}}.8h, v{{acc}}.8h
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{% macro per_col(label, op, mr, from, to, flipped=false) %}
|
||||
.{{label}}:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
{% set mr_over_8 = mr // 8 %}
|
||||
{% set cols = (to + 1 - from) // mr_over_8 %}
|
||||
|
||||
{% set loads = cols // 8 %}
|
||||
|
||||
{%if cols == 1 %}
|
||||
ld1 {v0.h}[0], [ x2 ]
|
||||
{% elif cols == 3 %}
|
||||
ld1 {v0.s}[0], [ x2 ], #4
|
||||
ld1 {v0.h}[2], [ x2 ]
|
||||
{% elif cols == 4 %}
|
||||
ldr d0, [ x2 ]
|
||||
{% elif cols == 6 %}
|
||||
ldr d0, [ x2 ], #8
|
||||
ld1 {v0.s}[2], [ x2 ]
|
||||
{% else %}
|
||||
{% for reg in range(1, loads + 1) %}
|
||||
ldr q{{ reg - 1 }}, [ x2 ], #16
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
// mr:{{mr}} {{ loads }} {{cols}}
|
||||
|
||||
{% for col in range(1, cols + 1) %}
|
||||
dup v3.8h, v{{ (col - 1) // 8 }}.h[{{ (col - 1) % 8 }}]
|
||||
{% for row in range(1, mr_over_8 + 1) %}
|
||||
{% set acc = (col - 1) * mr_over_8 + row - 1 + from %}
|
||||
{% if flipped %}
|
||||
{{op}} v{{acc}}.8h, v{{acc}}.8h, v3.8h
|
||||
{% else %}
|
||||
{{op}} v{{acc}}.8h, v3.8h, v{{acc}}.8h
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
+65
@@ -0,0 +1,65 @@
|
||||
fmla v16.4s, v0.4s, v8.s[0]
|
||||
ldr x5, [x1, #128]
|
||||
fmla v17.4s, v1.4s, v8.s[0]
|
||||
ldr x6, [x1, #136]
|
||||
fmla v18.4s, v2.4s, v8.s[0]
|
||||
ldr x7, [x1, #144]
|
||||
fmla v19.4s, v3.4s, v8.s[0]
|
||||
ldr x9, [x1, #152]
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [ x1 ], #64
|
||||
|
||||
fmla v20.4s, v4.4s, v8.s[0]
|
||||
ldr x10, [x1, #96]
|
||||
fmla v21.4s, v5.4s, v8.s[0]
|
||||
ldr x11, [x1, #104]
|
||||
fmla v22.4s, v6.4s, v8.s[0]
|
||||
ldr x12, [x1, #112]
|
||||
fmla v23.4s, v7.4s, v8.s[0]
|
||||
ldr x13, [x1, #120]
|
||||
|
||||
ld1 {{ v4.4s, v5.4s, v6.4s, v7.4s }}, [ x1 ]
|
||||
|
||||
fmla v24.4s, v0.4s, v8.s[0]
|
||||
ldr x14, [x1, #128]
|
||||
fmla v25.4s, v1.4s, v8.s[0]
|
||||
ldr x15, [x1, #136]
|
||||
fmla v26.4s, v2.4s, v8.s[0]
|
||||
ldr x20, [x1, #144]
|
||||
fmla v27.4s, v3.4s, v8.s[0]
|
||||
ldr x21, [x1, #152]
|
||||
fmla v28.4s, v4.4s, v8.s[0]
|
||||
ldr x22, [x1, #160]
|
||||
fmla v29.4s, v5.4s, v8.s[0]
|
||||
ldr x23, [x1, #168]
|
||||
fmla v30.4s, v6.4s, v8.s[0]
|
||||
ldr x24, [x1, #176]
|
||||
fmla v31.4s, v7.4s, v8.s[0]
|
||||
ldr x25, [x1, #184]
|
||||
|
||||
ld1 {{ v8.s }}[0], [ x2 ], #4
|
||||
|
||||
prfm pldl1keep, [x1, #1024]
|
||||
prfm pldl1keep, [x1, #1088]
|
||||
prfm pldl1keep, [x1, #1152]
|
||||
prfm pldl1keep, [x1, #1216]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
|
||||
ins v0.d[0], x5
|
||||
ins v1.d[0], x7
|
||||
ins v2.d[0], x10
|
||||
ins v3.d[0], x12
|
||||
ins v4.d[0], x14
|
||||
ins v5.d[0], x20
|
||||
ins v6.d[0], x22
|
||||
ins v7.d[0], x24
|
||||
|
||||
ins v0.d[1], x6
|
||||
ins v1.d[1], x9
|
||||
ins v2.d[1], x11
|
||||
ins v3.d[1], x13
|
||||
ins v4.d[1], x15
|
||||
ins v5.d[1], x21
|
||||
ins v6.d[1], x23
|
||||
ins v7.d[1], x25
|
||||
|
||||
add x1, x1, #192
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
ld1 {{ v9.8h, v10.8h, v11.8h, v12.8h }}, [x1], #64
|
||||
ld1 {{ v13.8h, v14.8h, v15.8h }}, [x1], #48
|
||||
|
||||
fmla v16.8h, v0.8h, v8.h[0]
|
||||
fmla v17.8h, v1.8h, v8.h[0]
|
||||
fmla v18.8h, v2.8h, v8.h[0]
|
||||
fmla v19.8h, v3.8h, v8.h[0]
|
||||
fmla v20.8h, v4.8h, v8.h[0]
|
||||
fmla v21.8h, v5.8h, v8.h[0]
|
||||
fmla v22.8h, v6.8h, v8.h[0]
|
||||
fmla v23.8h, v7.8h, v8.h[0]
|
||||
fmla v24.8h, v9.8h, v8.h[0]
|
||||
ld1 {{ v9.8h }}, [ x1 ], #16
|
||||
ld1 {{ v0.8h, v1.8h, v2.8h, v3.8h }}, [x1], #64
|
||||
ld1 {{ v4.8h, v5.8h, v6.8h, v7.8h }}, [x1], #64
|
||||
fmla v25.8h, v10.8h, v8.h[0]
|
||||
fmla v26.8h, v11.8h, v8.h[0]
|
||||
fmla v27.8h, v12.8h, v8.h[0]
|
||||
fmla v28.8h, v13.8h, v8.h[0]
|
||||
fmla v29.8h, v14.8h, v8.h[0]
|
||||
fmla v30.8h, v15.8h, v8.h[0]
|
||||
|
||||
fmla v31.8h, v9.8h, v8.h[0]
|
||||
|
||||
ld1 {{ v8.h }}[0], [ x2 ], #2
|
||||
|
||||
prfm pldl1keep, [x1, #1024]
|
||||
prfm pldl1keep, [x1, #1088]
|
||||
prfm pldl1keep, [x1, #1152]
|
||||
prfm pldl1keep, [x1, #1216]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
|
||||
+85
@@ -0,0 +1,85 @@
|
||||
ld1 {{ v9.4s, v10.4s, v11.4s }}, [x1], #48
|
||||
|
||||
fmla v16.8h, v0.8h, v8.h[0]
|
||||
ldr w8, [x2], #4
|
||||
fmla v17.8h, v1.8h, v8.h[0]
|
||||
ldr d12, [x1], #8
|
||||
fmla v18.8h, v2.8h, v8.h[0]
|
||||
ldr x12, [x1], #8
|
||||
fmla v19.8h, v3.8h, v8.h[0]
|
||||
ldr d13, [x1], #8
|
||||
fmla v20.8h, v4.8h, v8.h[0]
|
||||
ldr x13, [x1], #8
|
||||
fmla v21.8h, v5.8h, v8.h[0]
|
||||
ldr d14, [x1], #8
|
||||
fmla v22.8h, v6.8h, v8.h[0]
|
||||
ldr x14, [x1], #8
|
||||
fmla v23.8h, v7.8h, v8.h[0]
|
||||
ldr d15, [x1], #8
|
||||
fmla v24.8h, v9.8h, v8.h[0]
|
||||
ldr x15, [x1], #8
|
||||
|
||||
ld1 {{ v0.8h, v1.8h, v2.8h, v3.8h }}, [x1], #64
|
||||
ins v8.s[1], w8
|
||||
ld1 {{ v4.8h, v5.8h, v6.8h, v7.8h }}, [x1], #64
|
||||
|
||||
fmla v25.8h, v10.8h, v8.h[0]
|
||||
ins v12.d[1], x12
|
||||
fmla v26.8h, v11.8h, v8.h[0]
|
||||
ins v13.d[1], x13
|
||||
fmla v27.8h, v12.8h, v8.h[0]
|
||||
ins v14.d[1], x14
|
||||
fmla v28.8h, v13.8h, v8.h[0]
|
||||
ins v15.d[1], x15
|
||||
|
||||
ld1 {{ v9.8h, v10.8h, v11.8h, v12.8h }}, [x1], #64
|
||||
|
||||
fmla v29.8h, v14.8h, v8.h[0]
|
||||
ldr d13, [x1], #8
|
||||
fmla v30.8h, v15.8h, v8.h[0]
|
||||
ldr x13, [x1], #8
|
||||
fmla v31.8h, v0.8h, v8.h[0]
|
||||
ldr d14, [x1], #8
|
||||
|
||||
fmla v16.8h, v1.8h, v8.h[2]
|
||||
ldr x14, [x1], #8
|
||||
fmla v17.8h, v2.8h, v8.h[2]
|
||||
ldr d15, [x1], #8
|
||||
fmla v18.8h, v3.8h, v8.h[2]
|
||||
ldr x15, [x1], #8
|
||||
fmla v19.8h, v4.8h, v8.h[2]
|
||||
|
||||
ld1 {{ v0.8h }}, [x1], #16
|
||||
|
||||
fmla v20.8h, v5.8h, v8.h[2]
|
||||
ldr d1, [x1], #8
|
||||
fmla v21.8h, v6.8h, v8.h[2]
|
||||
ldr x10, [x1], #8
|
||||
|
||||
fmla v22.8h, v7.8h, v8.h[2]
|
||||
|
||||
fmla v23.8h, v9.8h, v8.h[2]
|
||||
ins v13.d[1], x13
|
||||
fmla v24.8h, v10.8h, v8.h[2]
|
||||
ins v14.d[1], x14
|
||||
fmla v25.8h, v11.8h, v8.h[2]
|
||||
ins v15.d[1], x15
|
||||
|
||||
fmla v26.8h, v12.8h, v8.h[2]
|
||||
prfm pldl1keep, [x1, #1024]
|
||||
fmla v27.8h, v13.8h, v8.h[2]
|
||||
ins v1.d[1], x10
|
||||
fmla v28.8h, v14.8h, v8.h[2]
|
||||
prfm pldl1keep, [x1, #1088]
|
||||
fmla v29.8h, v15.8h, v8.h[2]
|
||||
prfm pldl1keep, [x1, #1152]
|
||||
fmla v30.8h, v0.8h, v8.h[2]
|
||||
prfm pldl1keep, [x1, #1216]
|
||||
fmla v31.8h, v1.8h, v8.h[2]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [x1], #64
|
||||
ins v8.h[0], v8.h[3]
|
||||
ld1 {{ v4.4s, v5.4s, v6.4s, v7.4s }}, [x1], #64
|
||||
|
||||
|
||||
+203
@@ -0,0 +1,203 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs: v16 to v31, no need to preserve
|
||||
|
||||
// no preservation either for v0-v7...
|
||||
// v8..v15 are callee-preserved
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
{% if needs_pragma == true %}
|
||||
.cpu generic+fp+simd+fp16
|
||||
{% endif %}
|
||||
.global {{G}}arm64fp16_mmm_f16_128x1_{{core}}_{{suffix}}
|
||||
{{G}}arm64fp16_mmm_f16_128x1_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
sub x3, x3, #1
|
||||
|
||||
|
||||
ld1 { v8.h }[0], [ x2 ], #2
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
|
||||
ld1 { v4.4s, v5.4s, v6.4s, v7.4s }, [ x1 ], #64
|
||||
|
||||
cmp x3, #0
|
||||
beq .packed_packed_loop_1_last
|
||||
|
||||
cmp x3, #4
|
||||
blt .packed_packed_loop_1
|
||||
|
||||
{% set packed_packed_loop1 %}
|
||||
{% include "arm64fp16_mmm_f16_128x1/loop1/naive.S.raw" %}
|
||||
{% endset %}
|
||||
|
||||
{% set packed_packed_loop2 %}
|
||||
{% include "arm64fp16_mmm_f16_128x1/loop2/cortex_a55.S.raw" %}
|
||||
{% endset %}
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_4:
|
||||
{{ packed_packed_loop2 }}
|
||||
{{ packed_packed_loop2 }}
|
||||
|
||||
sub x3, x3, #4
|
||||
cmp x3, #4
|
||||
bge .packed_packed_loop_4
|
||||
|
||||
cmp x3, #0
|
||||
beq .packed_packed_loop_1_last
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
{{ packed_packed_loop1 }}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
// last loop can't read beyond actual input as it's likely not packed and padded
|
||||
.packed_packed_loop_1_last:
|
||||
ld1 { v9.8h, v10.8h, v11.8h, v12.8h }, [x1], #64
|
||||
ld1 { v13.8h, v14.8h, v15.8h }, [x1], #48
|
||||
|
||||
fmla v16.8h, v0.8h, v8.h[0]
|
||||
fmla v17.8h, v1.8h, v8.h[0]
|
||||
ld1 { v0.8h }, [ x1 ]
|
||||
fmla v18.8h, v2.8h, v8.h[0]
|
||||
fmla v19.8h, v3.8h, v8.h[0]
|
||||
fmla v20.8h, v4.8h, v8.h[0]
|
||||
fmla v21.8h, v5.8h, v8.h[0]
|
||||
fmla v22.8h, v6.8h, v8.h[0]
|
||||
fmla v23.8h, v7.8h, v8.h[0]
|
||||
|
||||
fmla v24.8h, v9.8h, v8.h[0]
|
||||
fmla v25.8h, v10.8h, v8.h[0]
|
||||
fmla v26.8h, v11.8h, v8.h[0]
|
||||
fmla v27.8h, v12.8h, v8.h[0]
|
||||
fmla v28.8h, v13.8h, v8.h[0]
|
||||
fmla v29.8h, v14.8h, v8.h[0]
|
||||
fmla v30.8h, v15.8h, v8.h[0]
|
||||
fmla v31.8h, v0.8h, v8.h[0]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_scalars.j2" %}
|
||||
{% set mr = 128 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_rows.j2" %}
|
||||
{% set mr = 128 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_cols.j2" %}
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
cmp x6, #2
|
||||
beq .do_per_row_add
|
||||
|
||||
{% for reg in range(16, 32) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
ld1 {v0.h}[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{reg}}.8h, v{{reg}}.8h, v0.8h
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.do_per_row_add:
|
||||
ld1 {v0.8h-v3.8h}, [x5], #64
|
||||
ld1 {v4.8h-v7.8h}, [x5], #64
|
||||
ld1 {v8.8h-v11.8h}, [x5], #64
|
||||
ld1 {v12.8h-v15.8h}, [x5], #64
|
||||
|
||||
{% for r in range(0, 16) %}
|
||||
fadd v{{ r + 16 }}.8h, v{{ r + 16 }}.8h, v{{r}}.8h
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x3, [x0, #16]
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 {v8.h}[0], [ x3 ]
|
||||
|
||||
{% for r in range(0, 8) %}
|
||||
ldr q{{r}}, [x2], #16
|
||||
{% endfor %}
|
||||
|
||||
fmla v16.8h, v0.8h, v8.h[0]
|
||||
ldr q0, [x2], #16
|
||||
fmla v17.8h, v1.8h, v8.h[0]
|
||||
ldr q1, [x2], #16
|
||||
fmla v18.8h, v2.8h, v8.h[0]
|
||||
ldr q2, [x2], #16
|
||||
fmla v19.8h, v3.8h, v8.h[0]
|
||||
ldr q3, [x2], #16
|
||||
fmla v20.8h, v4.8h, v8.h[0]
|
||||
ldr q4, [x2], #16
|
||||
fmla v21.8h, v5.8h, v8.h[0]
|
||||
ldr q5, [x2], #16
|
||||
fmla v22.8h, v6.8h, v8.h[0]
|
||||
ldr q6, [x2], #16
|
||||
fmla v23.8h, v7.8h, v8.h[0]
|
||||
ldr q7, [x2], #16
|
||||
|
||||
fmla v24.8h, v0.8h, v8.h[0]
|
||||
fmla v25.8h, v1.8h, v8.h[0]
|
||||
fmla v26.8h, v2.8h, v8.h[0]
|
||||
fmla v27.8h, v3.8h, v8.h[0]
|
||||
fmla v28.8h, v4.8h, v8.h[0]
|
||||
fmla v29.8h, v5.8h, v8.h[0]
|
||||
fmla v30.8h, v6.8h, v8.h[0]
|
||||
fmla v31.8h, v7.8h, v8.h[0]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc$
|
||||
|
||||
cmp x6, #2
|
||||
beq .store_strides_contig
|
||||
|
||||
{% for reg in range(16, 32) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
st1 { v{{reg}}.h }[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_contig:
|
||||
|
||||
{% for reg in range(16, 32) %}
|
||||
st1 { v{{reg}}.8h }, [ x5 ], #16
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
|
||||
fmla v16.8h, v0.8h, v4.h[0]
|
||||
fmla v17.8h, v1.8h, v4.h[0]
|
||||
fmla v18.8h, v0.8h, v4.h[1]
|
||||
fmla v19.8h, v1.8h, v4.h[1]
|
||||
fmla v20.8h, v0.8h, v4.h[2]
|
||||
fmla v21.8h, v1.8h, v4.h[2]
|
||||
fmla v22.8h, v0.8h, v4.h[3]
|
||||
fmla v23.8h, v1.8h, v4.h[3]
|
||||
|
||||
fmla v24.8h, v0.8h, v4.h[4]
|
||||
fmla v25.8h, v1.8h, v4.h[4]
|
||||
fmla v26.8h, v0.8h, v4.h[5]
|
||||
fmla v27.8h, v1.8h, v4.h[5]
|
||||
fmla v28.8h, v0.8h, v4.h[6]
|
||||
fmla v29.8h, v1.8h, v4.h[6]
|
||||
fmla v30.8h, v0.8h, v4.h[7]
|
||||
fmla v31.8h, v1.8h, v4.h[7]
|
||||
|
||||
ld1 {{ v0.8h, v1.8h }}, [x1], #32
|
||||
ld1 {{ v4.8h }}, [x2], #16
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
fmla v16.8h, v0.8h, v4.h[0]
|
||||
ldr d2, [x1], #8
|
||||
fmla v17.8h, v1.8h, v4.h[0]
|
||||
ldr d6, [x2], #8
|
||||
fmla v18.8h, v0.8h, v4.h[1]
|
||||
ldr x5, [x1], #8
|
||||
fmla v19.8h, v1.8h, v4.h[1]
|
||||
ldr x7, [x2], #8
|
||||
fmla v20.8h, v0.8h, v4.h[2]
|
||||
ldr d3, [x1], #8
|
||||
fmla v21.8h, v1.8h, v4.h[2]
|
||||
fmla v22.8h, v0.8h, v4.h[3]
|
||||
ldr x6, [x1], #8
|
||||
fmla v23.8h, v1.8h, v4.h[3]
|
||||
|
||||
fmla v24.8h, v0.8h, v4.h[4]
|
||||
fmla v25.8h, v1.8h, v4.h[4]
|
||||
fmla v26.8h, v0.8h, v4.h[5]
|
||||
fmla v27.8h, v1.8h, v4.h[5]
|
||||
fmla v28.8h, v0.8h, v4.h[6]
|
||||
ins v2.d[1], x5
|
||||
fmla v29.8h, v1.8h, v4.h[6]
|
||||
ins v6.d[1], x7
|
||||
fmla v30.8h, v0.8h, v4.h[7]
|
||||
ins v3.d[1], x6
|
||||
fmla v31.8h, v1.8h, v4.h[7]
|
||||
|
||||
fmla v16.8h, v2.8h, v6.h[0]
|
||||
ldr d0, [x1], #8
|
||||
fmla v17.8h, v3.8h, v6.h[0]
|
||||
ldr d4, [x2], #8
|
||||
fmla v18.8h, v2.8h, v6.h[1]
|
||||
ldr x5, [x1], #8
|
||||
fmla v19.8h, v3.8h, v6.h[1]
|
||||
ldr x7, [x2], #8
|
||||
fmla v20.8h, v2.8h, v6.h[2]
|
||||
ldr d1, [x1], #8
|
||||
fmla v21.8h, v3.8h, v6.h[2]
|
||||
fmla v22.8h, v2.8h, v6.h[3]
|
||||
ldr x6, [x1], #8
|
||||
fmla v23.8h, v3.8h, v6.h[3]
|
||||
|
||||
fmla v24.8h, v2.8h, v6.h[4]
|
||||
fmla v25.8h, v3.8h, v6.h[4]
|
||||
fmla v26.8h, v2.8h, v6.h[5]
|
||||
fmla v27.8h, v3.8h, v6.h[5]
|
||||
fmla v28.8h, v2.8h, v6.h[6]
|
||||
ins v0.d[1], x5
|
||||
fmla v29.8h, v3.8h, v6.h[6]
|
||||
ins v4.d[1], x7
|
||||
fmla v30.8h, v2.8h, v6.h[7]
|
||||
ins v1.d[1], x6
|
||||
fmla v31.8h, v3.8h, v6.h[7]
|
||||
|
||||
+174
@@ -0,0 +1,174 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// x20..x27 are used, callee-preserved
|
||||
|
||||
// C tile regs: v16 to v31, (scratch)
|
||||
//
|
||||
// v16[0] v18[0] v20[0] v22[0] v24[0] v26[0] v28[0] v30[0]
|
||||
// v16[1] v18[1]
|
||||
// v16[2] v18[2]
|
||||
// v16[3] v18[3]
|
||||
//
|
||||
// v17[0] v19[0] v21[0] v23[0] v25[0] v27[0] v29[0] v31[0]
|
||||
// v17[1] v19[1]
|
||||
// v17[2] v19[2]
|
||||
// v17[3] v19[3]
|
||||
|
||||
// v8 is used, d8 (lower half) must preserved
|
||||
// v0-v7 (scratch registers)
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
{% if needs_pragma == true %}
|
||||
.cpu generic+fp+simd+fp16
|
||||
{% endif %}
|
||||
.global {{G}}arm64fp16_mmm_f16_16x8_{{core}}_{{suffix}}
|
||||
{{G}}arm64fp16_mmm_f16_16x8_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
str q8, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.packed_packed:
|
||||
ld1 { v0.4s, v1.4s }, [ x1 ], #32
|
||||
ld1 { v4.4s }, [ x2 ], #16
|
||||
|
||||
{% set packed_packed_loop1 %}
|
||||
{% include "arm64fp16_mmm_f16_16x8/loop1/naive.S.raw" %}
|
||||
{% endset %}
|
||||
|
||||
{% set packed_packed_loop2 %}
|
||||
{% if core == "a55" %}
|
||||
{% include "arm64fp16_mmm_f16_16x8/loop2/cortex_a55.S.raw" %}
|
||||
{% else %}
|
||||
{{ packed_packed_loop1 }}
|
||||
{{ packed_packed_loop1 }}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
cmp x3, #4
|
||||
blt .packed_packed_loop_1
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_4:
|
||||
{{ packed_packed_loop2 }}
|
||||
{{ packed_packed_loop2 }}
|
||||
|
||||
sub x3, x3, #4
|
||||
cmp x3, #4
|
||||
bge .packed_packed_loop_4
|
||||
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
{{ packed_packed_loop1 }}
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_scalars.j2" %}
|
||||
{% set mr = 16 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_rows.j2" %}
|
||||
{% set mr = 16 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_cols.j2" %}
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
ld1 {v0.h}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{ col * 2 + reg }}.8h, v{{ col * 2 + reg }}.8h, v0.8h
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x2, [x0, #8]
|
||||
ldr x3, [x0, #16]
|
||||
|
||||
ld1 { v0.4s, v1.4s }, [ x2 ], #32
|
||||
ld1 { v4.4s }, [ x3 ], #16
|
||||
|
||||
fmla v16.8h, v0.8h, v4.h[0]
|
||||
fmla v17.8h, v1.8h, v4.h[0]
|
||||
fmla v18.8h, v0.8h, v4.h[1]
|
||||
fmla v19.8h, v1.8h, v4.h[1]
|
||||
fmla v20.8h, v0.8h, v4.h[2]
|
||||
fmla v21.8h, v1.8h, v4.h[2]
|
||||
fmla v22.8h, v0.8h, v4.h[3]
|
||||
fmla v23.8h, v1.8h, v4.h[3]
|
||||
|
||||
fmla v24.8h, v0.8h, v4.h[4]
|
||||
fmla v25.8h, v1.8h, v4.h[4]
|
||||
fmla v26.8h, v0.8h, v4.h[5]
|
||||
fmla v27.8h, v1.8h, v4.h[5]
|
||||
fmla v28.8h, v0.8h, v4.h[6]
|
||||
fmla v29.8h, v1.8h, v4.h[6]
|
||||
fmla v30.8h, v0.8h, v4.h[7]
|
||||
fmla v31.8h, v1.8h, v4.h[7]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x6, #2
|
||||
bne .store_strides_generic
|
||||
|
||||
{% for col in range(8, 16) %}
|
||||
str q{{ col * 2 }}, [ x5 ]
|
||||
str q{{ col * 2 + 1 }}, [ x5, #16 ]
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_generic:
|
||||
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
st1 { v{{ col * 2 + reg }}.h }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
ldr q8, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
|
||||
fmla v16.8h, v0.8h, v4.h[0]
|
||||
fmla v17.8h, v1.8h, v4.h[0]
|
||||
fmla v18.8h, v2.8h, v4.h[0]
|
||||
fmla v19.8h, v3.8h, v4.h[0]
|
||||
fmla v20.8h, v0.8h, v4.h[1]
|
||||
fmla v21.8h, v1.8h, v4.h[1]
|
||||
fmla v22.8h, v2.8h, v4.h[1]
|
||||
fmla v23.8h, v3.8h, v4.h[1]
|
||||
|
||||
fmla v24.8h, v0.8h, v4.h[2]
|
||||
fmla v25.8h, v1.8h, v4.h[2]
|
||||
fmla v26.8h, v2.8h, v4.h[2]
|
||||
fmla v27.8h, v3.8h, v4.h[2]
|
||||
fmla v28.8h, v0.8h, v4.h[3]
|
||||
fmla v29.8h, v1.8h, v4.h[3]
|
||||
fmla v30.8h, v2.8h, v4.h[3]
|
||||
fmla v31.8h, v3.8h, v4.h[3]
|
||||
|
||||
ld1 {{ v0.8h, v1.8h, v2.8h, v3.8h }}, [ x1 ], #64
|
||||
ldr d4, [x2], #8
|
||||
+71
@@ -0,0 +1,71 @@
|
||||
// mul a: v0, v1, v2, v3 b: v4
|
||||
// load a: v5(d5/x5), v6(d6,x6), v7(d7,x7), v8(d8, x8)
|
||||
// load b: v9 as d9
|
||||
|
||||
fmla v16.8h, v0.8h, v4.h[0]
|
||||
ldr d5, [x1], #8
|
||||
fmla v17.8h, v1.8h, v4.h[0]
|
||||
ldr d9, [x2], #8
|
||||
fmla v18.8h, v2.8h, v4.h[0]
|
||||
ldr x5, [x1], #8
|
||||
fmla v19.8h, v3.8h, v4.h[0]
|
||||
fmla v20.8h, v0.8h, v4.h[1]
|
||||
ldr d6, [x1], #8
|
||||
fmla v21.8h, v1.8h, v4.h[1]
|
||||
ldr x6, [x1], #8
|
||||
fmla v22.8h, v2.8h, v4.h[1]
|
||||
ldr d7, [x1], #8
|
||||
fmla v23.8h, v3.8h, v4.h[1]
|
||||
ldr x7, [x1], #8
|
||||
|
||||
fmla v24.8h, v0.8h, v4.h[2]
|
||||
ldr d8, [x1], #8
|
||||
fmla v25.8h, v1.8h, v4.h[2]
|
||||
ldr x8, [x1], #8
|
||||
fmla v26.8h, v2.8h, v4.h[2]
|
||||
ins v5.d[1], x5
|
||||
fmla v27.8h, v3.8h, v4.h[2]
|
||||
ins v6.d[1], x6
|
||||
fmla v28.8h, v0.8h, v4.h[3]
|
||||
ins v7.d[1], x7
|
||||
fmla v29.8h, v1.8h, v4.h[3]
|
||||
ins v8.d[1], x8
|
||||
fmla v30.8h, v2.8h, v4.h[3]
|
||||
ins v9.d[1], x9
|
||||
fmla v31.8h, v3.8h, v4.h[3]
|
||||
|
||||
// mul a: v5, v6, v7, v8 b: v9
|
||||
// load a: v0(d0/x5), v1(d1,x6), v2(d2,x7), v3(d3, x8)
|
||||
// load b: v4 as d4
|
||||
|
||||
fmla v16.8h, v5.8h, v9.h[0]
|
||||
ldr d0, [x1], #8
|
||||
fmla v17.8h, v6.8h, v9.h[0]
|
||||
ldr d4, [x2], #8
|
||||
fmla v18.8h, v7.8h, v9.h[0]
|
||||
ldr x5, [x1], #8
|
||||
fmla v19.8h, v8.8h, v9.h[0]
|
||||
fmla v20.8h, v5.8h, v9.h[1]
|
||||
ldr d1, [x1], #8
|
||||
fmla v21.8h, v6.8h, v9.h[1]
|
||||
ldr x6, [x1], #8
|
||||
fmla v22.8h, v7.8h, v9.h[1]
|
||||
ldr d2, [x1], #8
|
||||
fmla v23.8h, v8.8h, v9.h[1]
|
||||
ldr x7, [x1], #8
|
||||
|
||||
fmla v24.8h, v5.8h, v9.h[2]
|
||||
ldr d3, [x1], #8
|
||||
fmla v25.8h, v6.8h, v9.h[2]
|
||||
ldr x8, [x1], #8
|
||||
fmla v26.8h, v7.8h, v9.h[2]
|
||||
ins v0.d[1], x5
|
||||
fmla v27.8h, v8.8h, v9.h[2]
|
||||
ins v1.d[1], x6
|
||||
fmla v28.8h, v5.8h, v9.h[3]
|
||||
ins v2.d[1], x7
|
||||
fmla v29.8h, v6.8h, v9.h[3]
|
||||
ins v3.d[1], x8
|
||||
fmla v30.8h, v7.8h, v9.h[3]
|
||||
ins v4.d[1], x9
|
||||
fmla v31.8h, v8.8h, v9.h[3]
|
||||
+165
@@ -0,0 +1,165 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// x20..x27 are used, callee-preserved
|
||||
|
||||
// C tile regs: v16 to v31, (scratch)
|
||||
|
||||
// v8 is used, d8 (lower half) must preserved
|
||||
// v0-v7 (scratch registers)
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
{% if needs_pragma == true %}
|
||||
.cpu generic+fp+simd+fp16
|
||||
{% endif %}
|
||||
.global {{G}}arm64fp16_mmm_f16_32x4_{{core}}_{{suffix}}
|
||||
{{G}}arm64fp16_mmm_f16_32x4_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
|
||||
ldr d4, [x2], #8
|
||||
|
||||
{% set packed_packed_loop1 %}
|
||||
{% include "arm64fp16_mmm_f16_32x4/loop1/naive.S.raw" %}
|
||||
{% endset %}
|
||||
|
||||
{% set packed_packed_loop2 %}
|
||||
{% if core == "a55" %}
|
||||
{% include "arm64fp16_mmm_f16_32x4/loop2/cortex_a55.S.raw" %}
|
||||
{% else %}
|
||||
{{ packed_packed_loop1 }}
|
||||
{{ packed_packed_loop1 }}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
cmp x3, #4
|
||||
blt .packed_packed_loop_1
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_4:
|
||||
{{ packed_packed_loop2 }}
|
||||
{{ packed_packed_loop2 }}
|
||||
|
||||
sub x3, x3, #4
|
||||
cmp x3, #4
|
||||
bge .packed_packed_loop_4
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
{{ packed_packed_loop1 }}
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_scalars.j2" %}
|
||||
{% set mr = 32 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_rows.j2" %}
|
||||
{% set mr = 32 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_cols.j2" %}
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64fp16_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
{% for col in range(0, 4) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 4) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
ld1 {v0.h}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{ col * 4 + 16 + reg }}.8h, v{{ col * 4 + 16 + reg }}.8h, v0.8h
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x2, [x0, #8]
|
||||
ldr x3, [x0, #16]
|
||||
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [ x2 ]
|
||||
ldr d4, [x3]
|
||||
|
||||
fmla v16.8h, v0.8h, v4.h[0]
|
||||
fmla v17.8h, v1.8h, v4.h[0]
|
||||
fmla v18.8h, v2.8h, v4.h[0]
|
||||
fmla v19.8h, v3.8h, v4.h[0]
|
||||
fmla v20.8h, v0.8h, v4.h[1]
|
||||
fmla v21.8h, v1.8h, v4.h[1]
|
||||
fmla v22.8h, v2.8h, v4.h[1]
|
||||
fmla v23.8h, v3.8h, v4.h[1]
|
||||
|
||||
fmla v24.8h, v0.8h, v4.h[2]
|
||||
fmla v25.8h, v1.8h, v4.h[2]
|
||||
fmla v26.8h, v2.8h, v4.h[2]
|
||||
fmla v27.8h, v3.8h, v4.h[2]
|
||||
fmla v28.8h, v0.8h, v4.h[3]
|
||||
fmla v29.8h, v1.8h, v4.h[3]
|
||||
fmla v30.8h, v2.8h, v4.h[3]
|
||||
fmla v31.8h, v3.8h, v4.h[3]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x6, #2
|
||||
bne .store_strides_generic
|
||||
|
||||
{% for col in range(0, 4) %}
|
||||
str q{{ col * 4 + 16 + 0 }}, [ x5 ]
|
||||
str q{{ col * 4 + 16 + 1 }}, [ x5, #16 ]
|
||||
str q{{ col * 4 + 16 + 2 }}, [ x5, #32 ]
|
||||
str q{{ col * 4 + 16 + 3 }}, [ x5, #48 ]
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_generic:
|
||||
|
||||
{% for col in range(0, 4) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 4) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
st1 { v{{ col * 4 + 16 + reg }}.h }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+148
@@ -0,0 +1,148 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs: v16 to v31, no need to preserve
|
||||
|
||||
// no preservation either for v0-v7...
|
||||
// v8..v15 are callee-preserved
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
{% if needs_pragma == true %}
|
||||
.cpu generic+fp+simd+fp16
|
||||
{% endif %}
|
||||
.global {{G}}arm64fp16_mmm_f16_32x6_{{core}}_{{suffix}}
|
||||
{{G}}arm64fp16_mmm_f16_32x6_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldp x2, x4, [x0, #24] // b, packing
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
ld1 { v7.8h }, [ x2 ]
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [ x1 ], #64
|
||||
add x2, x2, 12
|
||||
|
||||
{% for row in range(0, 4) %}
|
||||
{% for col in range(0, 6) %}
|
||||
fmla v{{ col * 4 + 8 + row }}.8h, v{{row}}.8h, v7.h[{{col}}]
|
||||
{% endfor %}
|
||||
/*
|
||||
{% for col in range(0, 2) %}
|
||||
fmla v{{ (col + 4) * 4 + 8 + row }}.8h, v{{row}}.8h, v6.h[{{col}}]
|
||||
{% endfor %}
|
||||
*/
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 8 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_scalars.j2" %}
|
||||
{% set mr = 32 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_rows.j2" %}
|
||||
{% set mr = 32 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_cols.j2" %}
|
||||
{% set from = 8 %}{% set to = 31 %}{% include "arm64fp16_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
{% for col in range(0, 6) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 4) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
ld1 {v0.h}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{ col * 4 + 8 + reg }}.8h, v{{ col * 4 + 8 + reg }}.8h, v0.8h
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.do_per_row_add:
|
||||
ld1 {v0.8h-v3.8h}, [x5], #64
|
||||
ld1 {v4.8h-v7.8h}, [x5], #64
|
||||
|
||||
{% for r in range(0, 8) %}
|
||||
fadd v{{ r + 24 }}.8h, v{{ r + 24 }}.8h, v{{r}}.8h
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldp x2, x3, [x0, #8]
|
||||
|
||||
ld1 { v7.d }[0], [ x3 ], #8
|
||||
ld1 { v7.s }[2], [ x3 ], #4
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [ x2 ], #64
|
||||
|
||||
{% for row in range(0, 4) %}
|
||||
{% for col in range(0, 6) %}
|
||||
fmla v{{ col * 4 + 8 + row }}.8h, v{{row}}.8h, v7.h[{{col}}]
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x6, #2
|
||||
beq .store_strides_contig
|
||||
|
||||
{% for col in range(0, 6) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 4) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
st1 { v{{ col * 4 + 8 + reg }}.h }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_contig:
|
||||
|
||||
{% for col in range(0, 6) %}
|
||||
mov x4, x5
|
||||
{% for r in range(0, 4) %}
|
||||
st1 { v{{ col * 4 + 8 + r }}.8h }, [ x4 ], 16
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+264
@@ -0,0 +1,264 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs: v16 to v31, no need to preserve
|
||||
|
||||
// no preservation either for v0-v7...
|
||||
// v8..v15 are callee-preserved
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
{% if needs_pragma == true %}
|
||||
.cpu generic+fp+simd+fp16
|
||||
{% endif %}
|
||||
.global {{G}}arm64fp16_mmm_f16_64x1_{{core}}_{{suffix}}
|
||||
{{G}}arm64fp16_mmm_f16_64x1_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldp x2, x4, [x0, #24] // b, packing
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
cmp x4, #1
|
||||
beq .q4f16se
|
||||
|
||||
cmp x4, #2
|
||||
beq .q4f16
|
||||
|
||||
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
ld1 { v8.h }[0], [ x2 ], #2
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [ x1 ], #64
|
||||
ld1 { v4.8h, v5.8h, v6.8h, v7.8h }, [ x1 ], #64
|
||||
|
||||
fmla v24.8h, v0.8h, v8.h[0]
|
||||
fmla v25.8h, v1.8h, v8.h[0]
|
||||
fmla v26.8h, v2.8h, v8.h[0]
|
||||
fmla v27.8h, v3.8h, v8.h[0]
|
||||
fmla v28.8h, v4.8h, v8.h[0]
|
||||
fmla v29.8h, v5.8h, v8.h[0]
|
||||
fmla v30.8h, v6.8h, v8.h[0]
|
||||
fmla v31.8h, v7.8h, v8.h[0]
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.p2align 8
|
||||
.q40f16_const:
|
||||
.byte 0xc8, 0xc7, 0xc6, 0xc5, 0xc4, 0xc2, 0xc0, 0xbc
|
||||
.byte 0x00, 0x3c, 0x40, 0x42, 0x44, 0x45, 0x46, 0x47
|
||||
|
||||
.q4f16se:
|
||||
adr x4, .q40f16_const
|
||||
movi v15.16b, 15
|
||||
ld1 {v13.16b}, [ x4 ]
|
||||
eor v12.16b, v12.16b, v12.16b
|
||||
|
||||
.q4f16se_outerloop:
|
||||
{% for i in range(0, 8) %}
|
||||
eor v{{ i + 16 }}.16b, v{{ i + 16 }}.16b, v{{ i + 16 }}.16b
|
||||
{% endfor %}
|
||||
mov x4, #32
|
||||
|
||||
.p2align 4
|
||||
.q4f16se_innerloop:
|
||||
ld1 { v9.16b-v10.16b }, [x1], #32
|
||||
ld1 { v8.h }[0], [ x2 ], #2
|
||||
|
||||
and v0.16b, v9.16b, v15.16b
|
||||
ushr v2.16b, v9.16b, 4
|
||||
|
||||
and v4.16b, v10.16b, v15.16b
|
||||
ushr v6.16b, v10.16b, 4
|
||||
|
||||
tbl v0.16b, { v13.16b }, v0.16b
|
||||
tbl v2.16b, { v13.16b }, v2.16b
|
||||
tbl v4.16b, { v13.16b }, v4.16b
|
||||
tbl v6.16b, { v13.16b }, v6.16b
|
||||
|
||||
zip2 v1.16b, v12.16b, v0.16b
|
||||
zip2 v3.16b, v12.16b, v2.16b
|
||||
zip2 v5.16b, v12.16b, v4.16b
|
||||
zip2 v7.16b, v12.16b, v6.16b
|
||||
|
||||
zip1 v0.16b, v12.16b, v0.16b
|
||||
zip1 v2.16b, v12.16b, v2.16b
|
||||
zip1 v4.16b, v12.16b, v4.16b
|
||||
zip1 v6.16b, v12.16b, v6.16b
|
||||
|
||||
{% for i in range(0, 8) %}
|
||||
fmla v{{ i + 16 }}.8h, v{{i}}.8h, v8.h[0]
|
||||
{% endfor %}
|
||||
|
||||
subs x4, x4, #1
|
||||
bne .q4f16se_innerloop
|
||||
|
||||
// scales
|
||||
ld1 { v0.8h-v3.8h }, [ x1 ], #64
|
||||
ld1 { v4.8h-v7.8h }, [ x1 ], #64
|
||||
|
||||
{% for i in range(0, 8) %}
|
||||
fmla v{{ i + 24 }}.8h, v{{i}}.8h, v{{ i + 16 }}.8h
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #32
|
||||
bne .q4f16se_outerloop
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.q4f16:
|
||||
adr x4, .q40f16_const
|
||||
movi v15.16b, 15
|
||||
ld1 {v13.16b}, [ x4 ]
|
||||
eor v12.16b, v12.16b, v12.16b
|
||||
|
||||
.q4f16_outerloop:
|
||||
// scales
|
||||
ld1 { v16.8h-v19.8h }, [ x1 ], #64
|
||||
ld1 { v20.8h-v23.8h }, [ x1 ], #64
|
||||
mov x4, #32
|
||||
|
||||
.p2align 4
|
||||
.q4f16_innerloop:
|
||||
ld1 { v9.16b-v10.16b }, [x1], #32
|
||||
ld1 { v8.h }[0], [ x2 ], #2
|
||||
|
||||
and v0.16b, v9.16b, v15.16b
|
||||
ushr v2.16b, v9.16b, 4
|
||||
|
||||
and v4.16b, v10.16b, v15.16b
|
||||
ushr v6.16b, v10.16b, 4
|
||||
|
||||
tbl v0.16b, { v13.16b }, v0.16b
|
||||
tbl v2.16b, { v13.16b }, v2.16b
|
||||
tbl v4.16b, { v13.16b }, v4.16b
|
||||
tbl v6.16b, { v13.16b }, v6.16b
|
||||
|
||||
zip2 v1.16b, v12.16b, v0.16b
|
||||
zip2 v3.16b, v12.16b, v2.16b
|
||||
zip2 v5.16b, v12.16b, v4.16b
|
||||
zip2 v7.16b, v12.16b, v6.16b
|
||||
|
||||
zip1 v0.16b, v12.16b, v0.16b
|
||||
zip1 v2.16b, v12.16b, v2.16b
|
||||
zip1 v4.16b, v12.16b, v4.16b
|
||||
zip1 v6.16b, v12.16b, v6.16b
|
||||
|
||||
{% for i in range(0, 8) %}
|
||||
fmul v{{i}}.8h, v{{i}}.8h, v{{ i + 16 }}.8h
|
||||
{% endfor %}
|
||||
|
||||
{% for i in range(0, 8) %}
|
||||
fmla v{{ i + 24 }}.8h, v{{i}}.8h, v8.h[0]
|
||||
{% endfor %}
|
||||
|
||||
subs x4, x4, #1
|
||||
bne .q4f16_innerloop
|
||||
|
||||
subs x3, x3, #32
|
||||
bne .q4f16_outerloop
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 24 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_scalars.j2" %}
|
||||
{% set mr = 64 %}{% set from = 24 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_rows.j2" %}
|
||||
{% set mr = 64 %}{% set from = 24 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_cols.j2" %}
|
||||
{% set from = 24 %}{% set to = 31 %}{% include "arm64fp16_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
cmp x6, #2
|
||||
beq .do_per_row_add
|
||||
|
||||
{% for reg in range(24, 32) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
ld1 {v0.h}[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{reg}}.8h, v{{reg}}.8h, v0.8h
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.do_per_row_add:
|
||||
ld1 {v0.8h-v3.8h}, [x5], #64
|
||||
ld1 {v4.8h-v7.8h}, [x5], #64
|
||||
|
||||
{% for r in range(0, 8) %}
|
||||
fadd v{{ r + 24 }}.8h, v{{ r + 24 }}.8h, v{{r}}.8h
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x3, [x0, #16]
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 {v8.h}[0], [ x3 ]
|
||||
|
||||
{% for r in range(0, 8) %}
|
||||
ldr q{{r}}, [x2], #16
|
||||
{% endfor %}
|
||||
|
||||
fmla v24.8h, v0.8h, v8.h[0]
|
||||
fmla v25.8h, v1.8h, v8.h[0]
|
||||
fmla v26.8h, v2.8h, v8.h[0]
|
||||
fmla v27.8h, v3.8h, v8.h[0]
|
||||
fmla v28.8h, v4.8h, v8.h[0]
|
||||
fmla v29.8h, v5.8h, v8.h[0]
|
||||
fmla v30.8h, v6.8h, v8.h[0]
|
||||
fmla v31.8h, v7.8h, v8.h[0]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc$
|
||||
|
||||
cmp x6, #2
|
||||
beq .store_strides_contig
|
||||
|
||||
{% for reg in range(24, 32) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
st1 { v{{reg}}.h }[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_contig:
|
||||
|
||||
{% for reg in range(24, 32) %}
|
||||
st1 { v{{reg}}.8h }, [ x5 ], #16
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+165
@@ -0,0 +1,165 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs: v16 to v31, no need to preserve
|
||||
|
||||
// no preservation either for v0-v7...
|
||||
// v8..v15 are callee-preserved
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
{% if needs_pragma == true %}
|
||||
.cpu generic+fp+simd+fp16
|
||||
{% endif %}
|
||||
.global {{G}}arm64fp16_mmm_f16_64x3_{{core}}_{{suffix}}
|
||||
{{G}}arm64fp16_mmm_f16_64x3_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldp x2, x4, [x0, #24] // b, packing
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
ld1 { v7.4s }, [ x2 ]
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [ x1 ], #64
|
||||
ld1 { v4.8h, v5.8h, v6.8h }, [ x1 ], #48
|
||||
add x2, x2, #6
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 }}.8h, v0.8h, v7.h[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
ld1 { v0.8h }, [ x1 ], #16
|
||||
|
||||
{% for row in range(1, 7) %}
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 + row }}.8h, v{{row}}.8h, v7.h[{{ col }}]
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 15 }}.8h, v0.8h, v7.h[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 8 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_scalars.j2" %}
|
||||
{% set mr = 64 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_rows.j2" %}
|
||||
{% set mr = 64 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64fp16_mmm_f16_per_cols.j2" %}
|
||||
{% set from = 8 %}{% set to = 31 %}{% include "arm64fp16_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 8) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
ld1 {v0.h}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{ col * 8 + 8 + reg }}.8h, v{{ col * 8 + 8 + reg }}.8h, v0.8h
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.do_per_row_add:
|
||||
ld1 {v0.8h-v3.8h}, [x5], #64
|
||||
ld1 {v4.8h-v7.8h}, [x5], #64
|
||||
|
||||
{% for r in range(0, 8) %}
|
||||
fadd v{{ r + 24 }}.8h, v{{ r + 24 }}.8h, v{{r}}.8h
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldp x2, x3, [x0, #8]
|
||||
|
||||
ld1 { v7.s }[0], [ x3 ], #4
|
||||
ld1 { v7.h }[2], [ x3 ], #2
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [ x2 ], #64
|
||||
ld1 { v4.8h, v5.8h, v6.8h }, [ x2 ], #48
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 }}.8h, v0.8h, v7.h[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
ld1 { v0.8h }, [ x2 ], #16
|
||||
|
||||
{% for row in range(1, 7) %}
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 + row }}.8h, v{{row}}.8h, v7.h[{{ col }}]
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 15 }}.8h, v0.8h, v7.h[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x6, #2
|
||||
beq .store_strides_contig
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 8) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
st1 { v{{ col * 8 + 8 + reg }}.h }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_contig:
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
mov x4, x5
|
||||
{% for r in range(0, 8) %}
|
||||
st1 { v{{ col * 8 + 8 + r }}.8h }, [ x4 ], 16
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
// vim: ft=arm
|
||||
|
||||
{% from "arm64fp16_mmm_8h_ops.j2" import per_col %}
|
||||
|
||||
{{ per_col(label="per_col_min", op="fmin", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_max", op="fmax", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_mul", op="fmul", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_add", op="fadd", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_sub", op="fsub", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_sub_flipped", op="fsub", mr=mr, from=from, to=to, flipped=true) }}
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
// vim: ft=arm
|
||||
|
||||
{% from "arm64fp16_mmm_8h_ops.j2" import per_row %}
|
||||
|
||||
{{ per_row(label="per_row_min", op="fmin", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_max", op="fmax", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_mul", op="fmul", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_add", op="fadd", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_sub", op="fsub", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_sub_flipped", op="fsub", mr=mr, from=from, to=to, flipped=true) }}
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
// vim: ft=arm
|
||||
|
||||
{% from "arm64fp16_mmm_8h_ops.j2" import scalar %}
|
||||
|
||||
{{ scalar(label="scalar_min", op="fmin", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_max", op="fmax", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_mul", op="fmul", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_add", op="fadd", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_sub", op="fsub", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_sub_flipped", op="fsub", from=from, to=to, flipped=true) }}
|
||||
|
||||
.clear:
|
||||
{% for r in range(from, to + 1) %}
|
||||
eor v{{r}}.8b, v{{r}}.8b, v{{r}}.8b
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.leaky_relu:
|
||||
add x2, x0, #8
|
||||
ld1 {v4.s}[0], [ x2 ]
|
||||
dup v4.8h, v4.h[0]
|
||||
|
||||
// bsl cond/dst, then, else
|
||||
// fcmge dst, src, #0.0
|
||||
{% for r in range(from, to + 1) %}
|
||||
fmul v0.8h, v{{r}}.8h, v4.8h
|
||||
fcmge v1.8h, v{{r}}.8h, #0.0
|
||||
bsl v1.16b, v{{r}}.16b, v0.16b
|
||||
and v{{r}}.16b, v1.16b, v1.16b
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
|
||||
.q_scale:
|
||||
.q_shl:
|
||||
.q_shr:
|
||||
b .unsupported
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
// vim: ft=arm
|
||||
|
||||
.load_tile:
|
||||
ldr x2, [ x0, #8 ]
|
||||
{% for reg in range(from, to + 1) %}
|
||||
ld1 { v{{reg}}.4s }, [ x2 ], #16
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
+131
@@ -0,0 +1,131 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// no preservation either for v0-v7 and v16-v31
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
{% if needs_pragma == true %}
|
||||
.cpu generic+fp+simd+fp16
|
||||
{% endif %}
|
||||
.global {{G}}arm64fp16_sigmoid_f16_8n_{{suffix}}
|
||||
{{G}}arm64fp16_sigmoid_f16_8n_{{suffix}}:
|
||||
|
||||
cmp x1, #0
|
||||
beq .return
|
||||
|
||||
adr x2, .coeffs_num
|
||||
ld1 { v0.8h }, [x2]
|
||||
dup v5.8h, v0.h[0] // v5 <- low, broadcasted
|
||||
dup v6.8h, v0.h[1] // v6 <- high, broadcasted
|
||||
dup v7.8h, v0.h[7] // v7 <- half, broadcasted
|
||||
|
||||
cmp x1, #32
|
||||
blt .loop
|
||||
|
||||
.loop4:
|
||||
ld1 { v16.8h, v17.8h, v18.8h, v19.8h }, [x0]
|
||||
|
||||
fmax v16.8h, v16.8h, v5.8h
|
||||
fmax v17.8h, v17.8h, v5.8h
|
||||
fmax v18.8h, v18.8h, v5.8h
|
||||
fmax v19.8h, v19.8h, v5.8h
|
||||
|
||||
fmin v16.8h, v16.8h, v6.8h
|
||||
fmin v17.8h, v17.8h, v6.8h
|
||||
fmin v18.8h, v18.8h, v6.8h
|
||||
fmin v19.8h, v19.8h, v6.8h // v16 <- x
|
||||
|
||||
fmul v20.8h, v16.8h, v16.8h
|
||||
fmul v21.8h, v17.8h, v17.8h
|
||||
fmul v22.8h, v18.8h, v18.8h
|
||||
fmul v23.8h, v19.8h, v19.8h // v20 <- x2
|
||||
|
||||
dup v28.8h, v0.h[3]
|
||||
fmla v28.8h, v20.8h, v0.h[2]
|
||||
dup v29.8h, v0.h[3]
|
||||
fmla v29.8h, v21.8h, v0.h[2]
|
||||
dup v30.8h, v0.h[3]
|
||||
fmla v30.8h, v22.8h, v0.h[2]
|
||||
dup v31.8h, v0.h[3]
|
||||
fmla v31.8h, v23.8h, v0.h[2]
|
||||
|
||||
dup v24.8h, v0.h[4]
|
||||
fmla v24.8h, v20.8h, v28.8h
|
||||
dup v25.8h, v0.h[4]
|
||||
fmla v25.8h, v21.8h, v29.8h
|
||||
dup v26.8h, v0.h[4]
|
||||
fmla v26.8h, v22.8h, v30.8h
|
||||
dup v27.8h, v0.h[4]
|
||||
fmla v27.8h, v23.8h, v31.8h
|
||||
|
||||
fmul v16.8h, v16.8h, v24.8h
|
||||
fmul v17.8h, v17.8h, v25.8h
|
||||
fmul v18.8h, v18.8h, v26.8h
|
||||
fmul v19.8h, v19.8h, v27.8h // v16 <- numerator
|
||||
|
||||
dup v24.8h, v0.h[6]
|
||||
dup v25.8h, v0.h[6]
|
||||
dup v26.8h, v0.h[6]
|
||||
dup v27.8h, v0.h[6]
|
||||
fmla v24.8h, v20.8h, v0.h[5]
|
||||
fmla v25.8h, v21.8h, v0.h[5]
|
||||
fmla v26.8h, v22.8h, v0.h[5]
|
||||
fmla v27.8h, v23.8h, v0.h[5] // v24 <- denum
|
||||
|
||||
fdiv v16.8h, v16.8h, v24.8h
|
||||
fdiv v17.8h, v17.8h, v25.8h
|
||||
fdiv v18.8h, v18.8h, v26.8h
|
||||
fdiv v19.8h, v19.8h, v27.8h
|
||||
|
||||
fadd v16.8h, v16.8h, v7.8h
|
||||
fadd v17.8h, v17.8h, v7.8h
|
||||
fadd v18.8h, v18.8h, v7.8h
|
||||
fadd v19.8h, v19.8h, v7.8h
|
||||
|
||||
st1 { v16.8h, v17.8h, v18.8h, v19.8h }, [x0], #64
|
||||
|
||||
subs x1, x1, #32
|
||||
cmp x1, #32
|
||||
bge .loop4
|
||||
|
||||
cmp x1, #0
|
||||
beq .return
|
||||
|
||||
.loop:
|
||||
ld1 { v16.8h }, [x0]
|
||||
|
||||
fmax v16.8h, v16.8h, v5.8h
|
||||
fmin v16.8h, v16.8h, v6.8h // v16 <- x
|
||||
fmul v20.8h, v16.8h, v16.8h // v20 <- x2
|
||||
|
||||
dup v28.8h, v0.h[3]
|
||||
fmla v28.8h, v20.8h, v0.h[2]
|
||||
dup v24.8h, v0.h[4]
|
||||
fmla v24.8h, v20.8h, v28.8h
|
||||
fmul v16.8h, v16.8h, v24.8h // v16 <- numerator
|
||||
|
||||
dup v24.8h, v0.h[6]
|
||||
fmla v24.8h, v20.8h, v0.h[5] // v24 <- denum
|
||||
|
||||
fdiv v16.8h, v16.8h, v24.8h
|
||||
fadd v16.8h, v16.8h, v7.8h
|
||||
|
||||
st1 { v16.8h }, [x0], #16
|
||||
|
||||
subs x1, x1, #8
|
||||
bne .loop
|
||||
|
||||
.return:
|
||||
ret
|
||||
|
||||
.coeffs_num:
|
||||
{{ -6.92 | float16 }}
|
||||
{{ 6.92 | float16 }}
|
||||
{{ -0.0000124702 | float16 }}
|
||||
{{ 0.00400222 | float16 }}
|
||||
|
||||
{{ 0.249895 | float16 }}
|
||||
{{ 0.098734 | float16 }}
|
||||
{{ 1.0 | float16 }}
|
||||
{{ 0.5 | float16 }}
|
||||
+124
@@ -0,0 +1,124 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// no preservation either for v0-v7 and v16-v31
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
{% if needs_pragma == true %}
|
||||
.cpu generic+fp+simd+fp16
|
||||
{% endif %}
|
||||
.global {{G}}arm64fp16_tanh_f16_8n_{{suffix}}
|
||||
{{G}}arm64fp16_tanh_f16_8n_{{suffix}}:
|
||||
|
||||
cmp x1, #0
|
||||
beq .return
|
||||
|
||||
adr x2, .coeffs_num
|
||||
ld1 { v0.8h }, [x2]
|
||||
dup v5.8h, v0.h[0] // v5 <- low, broadcasted
|
||||
dup v6.8h, v0.h[1] // v6 <- high, broadcasted
|
||||
|
||||
cmp x1, #32
|
||||
blt .loop
|
||||
|
||||
.loop4:
|
||||
ld1 { v16.8h, v17.8h, v18.8h, v19.8h }, [x0]
|
||||
|
||||
fmax v16.8h, v16.8h, v5.8h
|
||||
fmax v17.8h, v17.8h, v5.8h
|
||||
fmax v18.8h, v18.8h, v5.8h
|
||||
fmax v19.8h, v19.8h, v5.8h
|
||||
|
||||
fmin v16.8h, v16.8h, v6.8h
|
||||
fmin v17.8h, v17.8h, v6.8h
|
||||
fmin v18.8h, v18.8h, v6.8h
|
||||
fmin v19.8h, v19.8h, v6.8h // v16 <- x
|
||||
|
||||
fmul v20.8h, v16.8h, v16.8h
|
||||
fmul v21.8h, v17.8h, v17.8h
|
||||
fmul v22.8h, v18.8h, v18.8h
|
||||
fmul v23.8h, v19.8h, v19.8h // v20 <- x2
|
||||
|
||||
dup v24.8h, v0.h[3]
|
||||
fmla v24.8h, v20.8h, v0.h[2]
|
||||
dup v25.8h, v0.h[3]
|
||||
fmla v25.8h, v21.8h, v0.h[2]
|
||||
dup v26.8h, v0.h[3]
|
||||
fmla v26.8h, v22.8h, v0.h[2]
|
||||
dup v27.8h, v0.h[3]
|
||||
fmla v27.8h, v23.8h, v0.h[2]
|
||||
|
||||
fmul v16.8h, v16.8h, v24.8h
|
||||
fmul v17.8h, v17.8h, v25.8h
|
||||
fmul v18.8h, v18.8h, v26.8h
|
||||
fmul v19.8h, v19.8h, v27.8h // v16 <- numerator
|
||||
|
||||
dup v28.8h, v0.h[5]
|
||||
fmla v28.8h, v20.8h, v0.h[4]
|
||||
dup v29.8h, v0.h[5]
|
||||
fmla v29.8h, v21.8h, v0.h[4]
|
||||
dup v30.8h, v0.h[5]
|
||||
fmla v30.8h, v22.8h, v0.h[4]
|
||||
dup v31.8h, v0.h[5]
|
||||
fmla v31.8h, v23.8h, v0.h[4]
|
||||
|
||||
dup v24.8h, v0.h[6]
|
||||
fmla v24.8h, v20.8h, v28.8h
|
||||
dup v25.8h, v0.h[6]
|
||||
fmla v25.8h, v21.8h, v29.8h
|
||||
dup v26.8h, v0.h[6]
|
||||
fmla v26.8h, v22.8h, v30.8h
|
||||
dup v27.8h, v0.h[6]
|
||||
fmla v27.8h, v23.8h, v31.8h // v24 <- denum
|
||||
|
||||
fdiv v16.8h, v16.8h, v24.8h
|
||||
fdiv v17.8h, v17.8h, v25.8h
|
||||
fdiv v18.8h, v18.8h, v26.8h
|
||||
fdiv v19.8h, v19.8h, v27.8h
|
||||
|
||||
st1 { v16.8h, v17.8h, v18.8h, v19.8h }, [x0], #64
|
||||
|
||||
subs x1, x1, #32
|
||||
cmp x1, #32
|
||||
bge .loop4
|
||||
|
||||
cmp x1, #0
|
||||
beq .return
|
||||
|
||||
.loop:
|
||||
ld1 { v16.8h }, [x0]
|
||||
|
||||
fmax v16.8h, v16.8h, v5.8h
|
||||
fmin v16.8h, v16.8h, v6.8h // v16 <- x
|
||||
fmul v20.8h, v16.8h, v16.8h // v20 <- x2
|
||||
|
||||
dup v24.8h, v0.h[3]
|
||||
fmla v24.8h, v20.8h, v0.h[2]
|
||||
fmul v16.8h, v16.8h, v24.8h // v16 <- numerator
|
||||
|
||||
dup v28.8h, v0.h[5]
|
||||
fmla v28.8h, v20.8h, v0.h[4]
|
||||
dup v24.8h, v0.h[6]
|
||||
fmla v24.8h, v20.8h, v28.8h // v24 <- denum
|
||||
|
||||
fdiv v16.8h, v16.8h, v24.8h
|
||||
|
||||
st1 { v16.8h }, [x0], #16
|
||||
|
||||
subs x1, x1, #8
|
||||
bne .loop
|
||||
|
||||
.return:
|
||||
ret
|
||||
|
||||
.coeffs_num:
|
||||
{{ -3.84 | float16 }}
|
||||
{{ 3.84 | float16 }}
|
||||
{{ 0.082654955 | float16 }} // alpha
|
||||
{{ 0.99963124 | float16 }}
|
||||
|
||||
{{ 0.0065383179 | float16 }} // beta
|
||||
{{ 0.41401828 | float16 }}
|
||||
{{ 1.0 | float16 }}
|
||||
{{ 0 | float16 }} // padding
|
||||
Vendored
+37
@@ -0,0 +1,37 @@
|
||||
// vim: ft=arm
|
||||
|
||||
.non_linear:
|
||||
sub x0, x0, 40
|
||||
|
||||
.non_linear_loop:
|
||||
add x0, x0, 40
|
||||
ldr x2, [x0]
|
||||
|
||||
mov x4, #{{ jump_table | length }}
|
||||
|
||||
cmp x2, #{{ jump_table | length }}
|
||||
csel x2, x2, x4, lt
|
||||
cmp x2, #0
|
||||
csel x2, x4, x2, lt
|
||||
|
||||
adr x3, .jmp_table
|
||||
add x3, x3, x2, LSL#2
|
||||
br x3
|
||||
|
||||
.jmp_table:
|
||||
{% for j in jump_table %}
|
||||
b .{{j}}
|
||||
{% endfor %}
|
||||
b .unsupported
|
||||
|
||||
add x0, x2, #4000
|
||||
b .return
|
||||
|
||||
.unsupported:
|
||||
mov x0, #1
|
||||
b .return
|
||||
|
||||
.done:
|
||||
mov x0, 0
|
||||
b .return
|
||||
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// serves as a canary build file to figure out which flag combination will accept half precision fmla
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
// .cpu generic+fp+simd+fp16
|
||||
.global foo
|
||||
foo:
|
||||
fmla v16.8h, v0.8h, v8.h[0]
|
||||
ret
|
||||
|
||||
Vendored
+13
@@ -0,0 +1,13 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// serves as a canary build file to figure out which flag combination will accept half precision fmla
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd+fp16
|
||||
.global foo
|
||||
foo:
|
||||
fmla v16.8h, v0.8h, v8.h[0]
|
||||
ret
|
||||
|
||||
+81
@@ -0,0 +1,81 @@
|
||||
|
||||
{% macro scalar(label, op, from, to, flipped=false) %}
|
||||
.{{label}}:
|
||||
add x2, x0, #8
|
||||
ld1 {v0.s}[0], [ x2 ]
|
||||
dup v0.4s, v0.s[0]
|
||||
{% if flipped %}
|
||||
{% for reg in range(from, to + 1) %}
|
||||
{{op}} v{{reg}}.4s, v{{reg}}.4s, v0.4s
|
||||
{% endfor %}
|
||||
{% else %}
|
||||
{% for reg in range(from, to + 1) %}
|
||||
{{op}} v{{reg}}.4s, v0.4s, v{{reg}}.4s
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{% macro per_row(label, op, mr, from, to, flipped=false) %}
|
||||
.{{label}}:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
{% set mr_over_4 = mr // 4 %}
|
||||
{% set mr_over_4_min_1 = mr // 4 - 1 %}
|
||||
|
||||
{% for reg in range(0, mr_over_4_min_1 + 1) %}
|
||||
ldr q{{reg}}, [ x2 ], #16
|
||||
{% endfor %}
|
||||
|
||||
{% if flipped %}
|
||||
{% for acc in range(from, to + 1) %}
|
||||
{% set other = (acc - from) % mr_over_4 %}
|
||||
{{op}} v{{acc}}.4s, v{{acc}}.4s, v{{other}}.4s
|
||||
{% endfor %}
|
||||
{% else %}
|
||||
{% for acc in range(from, to + 1) %}
|
||||
{% set other = (acc - from) % mr_over_4 %}
|
||||
{{op}} v{{acc}}.4s, v{{other}}.4s, v{{acc}}.4s
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{% macro per_col(label, op, mr, from, to, flipped=false) %}
|
||||
.{{label}}:
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
{% set mr_over_4 = mr // 4 %}
|
||||
{% set cols = (to + 1 - from) // mr_over_4 %}
|
||||
|
||||
{% set loads = cols // 4 %}
|
||||
|
||||
{%if cols == 1 %}
|
||||
ld1 {v0.s}[0], [ x2 ]
|
||||
{% elif cols == 3 %}
|
||||
ldr d0, [ x2 ], #8
|
||||
ld1 {v0.s}[2], [ x2 ]
|
||||
{% else %}
|
||||
{% for reg in range(1, loads + 1) %}
|
||||
ldr q{{ reg - 1 }}, [ x2 ], #16
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
|
||||
// {{mr}} {{cols}}
|
||||
|
||||
{% for col in range(1, cols + 1) %}
|
||||
dup v3.4s, v{{ (col - 1) // 4 }}.s[{{ (col - 1) % 4 }}]
|
||||
{% for row in range(1, mr_over_4 + 1) %}
|
||||
{% set acc = (col - 1) * mr_over_4 + row - 1 + from %}
|
||||
{% if flipped %}
|
||||
{{op}} v{{acc}}.4s, v{{acc}}.4s, v3.4s
|
||||
{% else %}
|
||||
{{op}} v{{acc}}.4s, v3.4s, v{{acc}}.4s
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
+69
@@ -0,0 +1,69 @@
|
||||
fmla v8.4s, v0.4s, v4.s[0]
|
||||
ldr w4, [x1], #4
|
||||
fmla v9.4s, v1.4s, v4.s[0]
|
||||
ldr w20, [x2], #4
|
||||
fmla v10.4s, v2.4s, v4.s[0]
|
||||
ldr w5, [x1], #4
|
||||
|
||||
fmla v11.4s, v0.4s, v4.s[1]
|
||||
ldr w21, [x2], #4
|
||||
fmla v12.4s, v1.4s, v4.s[1]
|
||||
ldr w6, [x1], #4
|
||||
fmla v13.4s, v2.4s, v4.s[1]
|
||||
ldr w22, [x2], #4
|
||||
|
||||
fmla v14.4s, v0.4s, v4.s[2]
|
||||
ldr w7, [x1], #4
|
||||
fmla v15.4s, v1.4s, v4.s[2]
|
||||
ldr w23, [x2], #4
|
||||
fmla v16.4s, v2.4s, v4.s[2]
|
||||
ldr w8, [x1], #4
|
||||
fmla v17.4s, v0.4s, v4.s[3]
|
||||
ldr w24, [x2], #4
|
||||
fmla v18.4s, v1.4s, v4.s[3]
|
||||
ldr w9, [x1], #4
|
||||
fmla v19.4s, v2.4s, v4.s[3]
|
||||
ldr w25, [x2], #4
|
||||
|
||||
fmla v20.4s, v0.4s, v5.s[0]
|
||||
ldr w10, [x1], #4
|
||||
fmla v21.4s, v1.4s, v5.s[0]
|
||||
ldr w26, [x2], #4
|
||||
fmla v22.4s, v2.4s, v5.s[0]
|
||||
ldr w11, [x1], #4
|
||||
fmla v23.4s, v0.4s, v5.s[1]
|
||||
ldr w27, [x2], #4
|
||||
fmla v24.4s, v1.4s, v5.s[1]
|
||||
ldr w12, [x1], #4
|
||||
fmla v25.4s, v2.4s, v5.s[1]
|
||||
|
||||
fmla v26.4s, v0.4s, v5.s[2]
|
||||
ldr w13, [x1], #4
|
||||
fmla v27.4s, v1.4s, v5.s[2]
|
||||
fmla v28.4s, v2.4s, v5.s[2]
|
||||
ldr w14, [x1], #4
|
||||
fmla v29.4s, v0.4s, v5.s[3]
|
||||
fmla v30.4s, v1.4s, v5.s[3]
|
||||
ldr w15, [x1], #4
|
||||
fmla v31.4s, v2.4s, v5.s[3]
|
||||
|
||||
ins v0.s[0], w4
|
||||
ins v1.s[0], w8
|
||||
ins v2.s[0], w12
|
||||
ins v4.s[0], w20
|
||||
ins v5.s[0], w24
|
||||
ins v0.s[1], w5
|
||||
ins v1.s[1], w9
|
||||
ins v2.s[1], w13
|
||||
ins v4.s[1], w21
|
||||
ins v5.s[1], w25
|
||||
ins v0.s[2], w6
|
||||
ins v1.s[2], w10
|
||||
ins v2.s[2], w14
|
||||
ins v4.s[2], w22
|
||||
ins v5.s[2], w26
|
||||
ins v0.s[3], w7
|
||||
ins v1.s[3], w11
|
||||
ins v2.s[3], w15
|
||||
ins v4.s[3], w23
|
||||
ins v5.s[3], w27
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
fmla v8.4s, v0.4s, v4.s[0]
|
||||
ldr w4, [x1]
|
||||
fmla v9.4s, v1.4s, v4.s[0]
|
||||
ldr w20, [x2], #4
|
||||
fmla v10.4s, v2.4s, v4.s[0]
|
||||
ldr w5, [x1, #4]
|
||||
|
||||
fmla v11.4s, v0.4s, v4.s[1]
|
||||
ldr w21, [x2], #4
|
||||
fmla v12.4s, v1.4s, v4.s[1]
|
||||
ldr w6, [x1, #8]
|
||||
fmla v13.4s, v2.4s, v4.s[1]
|
||||
ldr w22, [x2], #4
|
||||
|
||||
fmla v14.4s, v0.4s, v4.s[2]
|
||||
ldr w7, [x1, #12]
|
||||
fmla v15.4s, v1.4s, v4.s[2]
|
||||
ldr w23, [x2], #4
|
||||
fmla v16.4s, v2.4s, v4.s[2]
|
||||
ldr w8, [x1, #16]
|
||||
fmla v17.4s, v0.4s, v4.s[3]
|
||||
ldr w24, [x2], #4
|
||||
fmla v18.4s, v1.4s, v4.s[3]
|
||||
ldr w9, [x1, #20]
|
||||
fmla v19.4s, v2.4s, v4.s[3]
|
||||
ldr w25, [x2], #4
|
||||
|
||||
fmla v20.4s, v0.4s, v5.s[0]
|
||||
ldr w10, [x1, #24]
|
||||
fmla v21.4s, v1.4s, v5.s[0]
|
||||
ldr w26, [x2], #4
|
||||
fmla v22.4s, v2.4s, v5.s[0]
|
||||
ldr w11, [x1, #28]
|
||||
fmla v23.4s, v0.4s, v5.s[1]
|
||||
ldr w27, [x2], #4
|
||||
fmla v24.4s, v1.4s, v5.s[1]
|
||||
ldr w12, [x1, #32]
|
||||
fmla v25.4s, v2.4s, v5.s[1]
|
||||
ldr w13, [x1, #36]
|
||||
|
||||
fmla v26.4s, v0.4s, v5.s[2]
|
||||
ldr w14, [x1, #40]
|
||||
fmla v27.4s, v1.4s, v5.s[2]
|
||||
ldr w15, [x1, #44]
|
||||
fmla v28.4s, v2.4s, v5.s[2]
|
||||
prfm pldl1keep, [x1, #512]
|
||||
fmla v29.4s, v0.4s, v5.s[3]
|
||||
add x1, x1, #48
|
||||
fmla v30.4s, v1.4s, v5.s[3]
|
||||
prfm pldl1keep, [x2, #384]
|
||||
fmla v31.4s, v2.4s, v5.s[3]
|
||||
|
||||
ins v0.s[0], w4
|
||||
|
||||
ins v1.s[0], w8
|
||||
ins v2.s[0], w12
|
||||
|
||||
ins v4.s[0], w20
|
||||
ins v5.s[0], w24
|
||||
|
||||
ins v0.s[1], w5
|
||||
ins v1.s[1], w9
|
||||
|
||||
ins v2.s[1], w13
|
||||
ins v4.s[1], w21
|
||||
|
||||
ins v5.s[1], w25
|
||||
ins v0.s[2], w6
|
||||
|
||||
ins v1.s[2], w10
|
||||
ins v2.s[2], w14
|
||||
|
||||
ins v4.s[2], w22
|
||||
ins v5.s[2], w26
|
||||
|
||||
ins v0.s[3], w7
|
||||
ins v1.s[3], w11
|
||||
|
||||
ins v2.s[3], w15
|
||||
ins v4.s[3], w23
|
||||
ins v5.s[3], w27
|
||||
|
||||
+66
@@ -0,0 +1,66 @@
|
||||
fmla v8.4s, v0.4s, v4.s[0]
|
||||
ldr x4, [x1]
|
||||
fmla v9.4s, v1.4s, v4.s[0]
|
||||
ldr x20, [x2]
|
||||
fmla v10.4s, v2.4s, v4.s[0]
|
||||
ldr x5, [x1, #8]
|
||||
|
||||
fmla v11.4s, v0.4s, v4.s[1]
|
||||
ldr x21, [x2, #8]
|
||||
fmla v12.4s, v1.4s, v4.s[1]
|
||||
ldr x6, [x1, #16]
|
||||
fmla v13.4s, v2.4s, v4.s[1]
|
||||
ldr x22, [x2, #16]
|
||||
|
||||
fmla v14.4s, v0.4s, v4.s[2]
|
||||
ldr x7, [x1, #24]
|
||||
fmla v15.4s, v1.4s, v4.s[2]
|
||||
ldr x23, [x2, #24]
|
||||
fmla v16.4s, v2.4s, v4.s[2]
|
||||
ldr x8, [x1, #32]
|
||||
fmla v17.4s, v0.4s, v4.s[3]
|
||||
fmla v18.4s, v1.4s, v4.s[3]
|
||||
ldr x9, [x1, #40]
|
||||
fmla v19.4s, v2.4s, v4.s[3]
|
||||
|
||||
fmla v20.4s, v0.4s, v5.s[0]
|
||||
prfm pldl1keep, [x1, #256]
|
||||
fmla v21.4s, v1.4s, v5.s[0]
|
||||
prfm pldl1keep, [x1, #320]
|
||||
fmla v22.4s, v2.4s, v5.s[0]
|
||||
prfm pldl1keep, [x1, #384]
|
||||
fmla v23.4s, v0.4s, v5.s[1]
|
||||
prfm pldl1keep, [x1, #448]
|
||||
fmla v24.4s, v1.4s, v5.s[1]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
fmla v25.4s, v2.4s, v5.s[1]
|
||||
prfm pldl1keep, [x2, #320]
|
||||
|
||||
fmla v26.4s, v0.4s, v5.s[2]
|
||||
fmla v27.4s, v1.4s, v5.s[2]
|
||||
fmla v28.4s, v2.4s, v5.s[2]
|
||||
prfm pldl1keep, [x1, #512]
|
||||
fmla v29.4s, v0.4s, v5.s[3]
|
||||
add x1, x1, #48
|
||||
fmla v30.4s, v1.4s, v5.s[3]
|
||||
prfm pldl1keep, [x2, #384]
|
||||
fmla v31.4s, v2.4s, v5.s[3]
|
||||
add x2, x2, #32
|
||||
|
||||
|
||||
ins v0.d[0], x4
|
||||
ins v2.d[0], x8
|
||||
|
||||
ins v4.d[0], x20
|
||||
ins v5.d[0], x22
|
||||
|
||||
ins v0.d[1], x5
|
||||
ins v2.d[1], x9
|
||||
|
||||
ins v4.d[1], x21
|
||||
ins v1.d[0], x6
|
||||
|
||||
ins v1.d[1], x7
|
||||
|
||||
ins v5.d[1], x23
|
||||
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
fmla v8.4s, v0.4s, v4.s[0]
|
||||
fmla v9.4s, v1.4s, v4.s[0]
|
||||
fmla v10.4s, v2.4s, v4.s[0]
|
||||
|
||||
fmla v11.4s, v0.4s, v4.s[1]
|
||||
fmla v12.4s, v1.4s, v4.s[1]
|
||||
fmla v13.4s, v2.4s, v4.s[1]
|
||||
|
||||
fmla v14.4s, v0.4s, v4.s[2]
|
||||
fmla v15.4s, v1.4s, v4.s[2]
|
||||
fmla v16.4s, v2.4s, v4.s[2]
|
||||
|
||||
fmla v17.4s, v0.4s, v4.s[3]
|
||||
fmla v18.4s, v1.4s, v4.s[3]
|
||||
fmla v19.4s, v2.4s, v4.s[3]
|
||||
|
||||
fmla v20.4s, v0.4s, v5.s[0]
|
||||
fmla v21.4s, v1.4s, v5.s[0]
|
||||
fmla v22.4s, v2.4s, v5.s[0]
|
||||
|
||||
fmla v23.4s, v0.4s, v5.s[1]
|
||||
fmla v24.4s, v1.4s, v5.s[1]
|
||||
fmla v25.4s, v2.4s, v5.s[1]
|
||||
|
||||
fmla v26.4s, v0.4s, v5.s[2]
|
||||
fmla v27.4s, v1.4s, v5.s[2]
|
||||
fmla v28.4s, v2.4s, v5.s[2]
|
||||
|
||||
fmla v29.4s, v0.4s, v5.s[3]
|
||||
fmla v30.4s, v1.4s, v5.s[3]
|
||||
fmla v31.4s, v2.4s, v5.s[3]
|
||||
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s }}, [x1], #48
|
||||
ld1 {{ v4.4s, v5.4s }}, [x2], #32
|
||||
+107
@@ -0,0 +1,107 @@
|
||||
// mul a: v0, v1, v2, b: v4, v5
|
||||
// load a: d3/x23, d6/x26, d7/x27
|
||||
// load b: x4, x5, x6, x7
|
||||
|
||||
fmla v8.4s, v0.4s, v4.s[0]
|
||||
ldr d3, [x1], #8
|
||||
fmla v9.4s, v1.4s, v4.s[0]
|
||||
ldr x4, [x2], #8
|
||||
fmla v10.4s, v2.4s, v4.s[0]
|
||||
ldr x23, [x1], #8
|
||||
fmla v11.4s, v0.4s, v4.s[1]
|
||||
ldr x5, [x2], #8
|
||||
fmla v12.4s, v1.4s, v4.s[1]
|
||||
ldr d6, [x1], #8
|
||||
fmla v13.4s, v2.4s, v4.s[1]
|
||||
ldr x6, [x2], #8
|
||||
fmla v14.4s, v0.4s, v4.s[2]
|
||||
ldr x26, [x1], #8
|
||||
fmla v15.4s, v1.4s, v4.s[2]
|
||||
ldr x7, [x2], #8
|
||||
fmla v16.4s, v2.4s, v4.s[2]
|
||||
ldr d7, [x1], #8
|
||||
fmla v17.4s, v0.4s, v4.s[3]
|
||||
ldr x27, [x1], #8
|
||||
fmla v18.4s, v1.4s, v4.s[3]
|
||||
|
||||
fmla v19.4s, v2.4s, v4.s[3]
|
||||
fmla v20.4s, v0.4s, v5.s[0]
|
||||
|
||||
// ins b: v4 <- x4/x5
|
||||
// ins a: d3/x23, d6/x26, d7/x27
|
||||
|
||||
ins v4.d[0], x4
|
||||
fmla v21.4s, v1.4s, v5.s[0]
|
||||
ins v4.d[1], x5
|
||||
fmla v22.4s, v2.4s, v5.s[0]
|
||||
fmla v23.4s, v0.4s, v5.s[1]
|
||||
|
||||
fmla v24.4s, v1.4s, v5.s[1]
|
||||
fmla v25.4s, v2.4s, v5.s[1]
|
||||
fmla v26.4s, v0.4s, v5.s[2]
|
||||
fmla v27.4s, v1.4s, v5.s[2]
|
||||
fmla v28.4s, v2.4s, v5.s[2]
|
||||
fmla v29.4s, v0.4s, v5.s[3]
|
||||
ins v3.d[1], x23
|
||||
fmla v30.4s, v1.4s, v5.s[3]
|
||||
ins v6.d[1], x26
|
||||
fmla v31.4s, v2.4s, v5.s[3]
|
||||
ins v7.d[1], x27
|
||||
|
||||
// mul a: v3, v6, v7, b: v4, v5
|
||||
// ins b, v5 <- x6, x7
|
||||
// load a: d0/x20, d1/x21, d2/x22
|
||||
// load b: x4, x5
|
||||
|
||||
fmla v8.4s, v3.4s, v4.s[0]
|
||||
ins v5.d[0], x6
|
||||
fmla v9.4s, v6.4s, v4.s[0]
|
||||
ins v5.d[1], x7
|
||||
fmla v10.4s, v7.4s, v4.s[0]
|
||||
ldr d0, [x1], #8
|
||||
fmla v11.4s, v3.4s, v4.s[1]
|
||||
ldr x4, [x2], #8
|
||||
fmla v12.4s, v6.4s, v4.s[1]
|
||||
ldr x20, [x1], #8
|
||||
fmla v13.4s, v7.4s, v4.s[1]
|
||||
ldr x5, [x2], #8
|
||||
fmla v14.4s, v3.4s, v4.s[2]
|
||||
ldr d1, [x1], #8
|
||||
fmla v15.4s, v6.4s, v4.s[2]
|
||||
ldr x6, [x2], #8
|
||||
fmla v16.4s, v7.4s, v4.s[2]
|
||||
ldr x21, [x1], #8
|
||||
fmla v17.4s, v3.4s, v4.s[3]
|
||||
ldr x7, [x2], #8
|
||||
|
||||
// load b: x6, x7
|
||||
fmla v18.4s, v6.4s, v4.s[3]
|
||||
ldr d2, [x1], #8
|
||||
fmla v19.4s, v7.4s, v4.s[3]
|
||||
ldr x22, [x1], #8
|
||||
fmla v20.4s, v3.4s, v5.s[0]
|
||||
fmla v21.4s, v6.4s, v5.s[0]
|
||||
fmla v22.4s, v7.4s, v5.s[0]
|
||||
fmla v23.4s, v3.4s, v5.s[1]
|
||||
fmla v24.4s, v6.4s, v5.s[1]
|
||||
fmla v25.4s, v7.4s, v5.s[1]
|
||||
|
||||
// ins a: d0/x20, d1/x21, d2/x22
|
||||
fmla v26.4s, v3.4s, v5.s[2]
|
||||
ins v0.d[1], x20
|
||||
fmla v27.4s, v6.4s, v5.s[2]
|
||||
ins v1.d[1], x21
|
||||
fmla v28.4s, v7.4s, v5.s[2]
|
||||
ins v2.d[1], x22
|
||||
|
||||
// ins b: v4 <- x4, x5
|
||||
fmla v29.4s, v3.4s, v5.s[3]
|
||||
ins v4.d[0], x4
|
||||
fmla v30.4s, v6.4s, v5.s[3]
|
||||
ins v4.d[1], x5
|
||||
fmla v31.4s, v7.4s, v5.s[3]
|
||||
|
||||
// ins b: v5 <- x6, x7
|
||||
ins v5.d[0], x6
|
||||
ins v5.d[1], x7
|
||||
|
||||
+163
@@ -0,0 +1,163 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs:
|
||||
// - x19-x29 to preserve (but x19, x28, x29 not used)
|
||||
// - d8..d15 to preserve
|
||||
// - v16 to v31, no need to preserve
|
||||
//
|
||||
// v8 v11 v14 v17 v20 v23 v26 v29
|
||||
// v9 v12 v15 v18 v21 v24 v27 v30
|
||||
// v10 v13 v16 v19 v22 v25 v28 v31
|
||||
|
||||
// no preservation for v0-v7:
|
||||
// packed A buffering (2x8 values): rotating over v0..v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd
|
||||
.global {{G}}arm64simd_mmm_f32_12x8_{{core}}_{{suffix}}
|
||||
{{G}}arm64simd_mmm_f32_12x8_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s }, [ x1 ], #48
|
||||
ld1 { v4.4s, v5.4s }, [ x2 ], #32
|
||||
|
||||
{% set packed_packed_loop1 %}
|
||||
{% if core == "a53" %}
|
||||
{% include "arm64simd_mmm_f32_12x8/packed_packed_loop1/ldr_x_preload.S.raw" %}
|
||||
{% else %}
|
||||
{% include "arm64simd_mmm_f32_12x8/packed_packed_loop1/naive.S.raw" %}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
{% set packed_packed_loop2 %}
|
||||
{% if core == "a55" %}
|
||||
{% include "arm64simd_mmm_f32_12x8/packed_packed_loop2/cortex_a55.S.raw" %}
|
||||
{% else %}
|
||||
{{ packed_packed_loop1 }}
|
||||
{{ packed_packed_loop1 }}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
cmp x3, #4
|
||||
blt .packed_packed_loop_1
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_4:
|
||||
{{ packed_packed_loop2 }}
|
||||
{{ packed_packed_loop2 }}
|
||||
|
||||
sub x3, x3, #4
|
||||
cmp x3, #4
|
||||
bge .packed_packed_loop_4
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
{{ packed_packed_loop1 }}
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_scalars.j2" %}
|
||||
{% set mr = 12 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_rows.j2" %}
|
||||
{% set mr = 12 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_cols.j2" %}
|
||||
{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8 ] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
{% for col in range(0, 8) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 3) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.s}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{ col * 3 + 8 + reg }}.4s, v{{ col * 3 + 8 + reg }}.4s, v0.4s
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x2, [x0, #8]
|
||||
ldr x3, [x0, #16]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s }, [ x2 ]
|
||||
ld1 { v4.4s, v5.4s }, [ x3 ]
|
||||
|
||||
{% for col in range(0, 8) %}
|
||||
{% for reg in range(0, 3) %}
|
||||
fmla v{{ col * 3 + 8 + reg }}.4s, v{{reg}}.4s, v{{ col // 4 + 4 }}.s[{{ col % 4 }}]
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x6, #4
|
||||
bne .store_strides_generic
|
||||
|
||||
{% for col in range(0, 8) %}
|
||||
str q{{ col * 3 + 8 }}, [ x5 ]
|
||||
str q{{ col * 3 + 9 }}, [ x5, #16 ]
|
||||
str q{{ col * 3 + 10 }}, [ x5, #32 ]
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_generic:
|
||||
{% for col in range(0, 8) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 3) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{ col * 3 + 8 + reg }}.s }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+46
@@ -0,0 +1,46 @@
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
ldr x5, [x1]
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
ldr x6, [x1, #8]
|
||||
fmla v18.4s, v2.4s, v4.s[0]
|
||||
ldr x7, [x1, #16]
|
||||
fmla v19.4s, v3.4s, v4.s[0]
|
||||
ldr x8, [x1, #24]
|
||||
fmla v20.4s, v0.4s, v4.s[1]
|
||||
ldr x9, [x1, #32]
|
||||
fmla v21.4s, v1.4s, v4.s[1]
|
||||
ldr x10, [x1, #40]
|
||||
fmla v22.4s, v2.4s, v4.s[1]
|
||||
ldr x11, [x1, #48]
|
||||
fmla v23.4s, v3.4s, v4.s[1]
|
||||
ldr x12, [x1, #56]
|
||||
|
||||
fmla v24.4s, v0.4s, v4.s[2]
|
||||
ldr x24, [x2]
|
||||
fmla v25.4s, v1.4s, v4.s[2]
|
||||
ldr x25, [x2, #8]
|
||||
fmla v26.4s, v2.4s, v4.s[2]
|
||||
add x1, x1, #64
|
||||
fmla v27.4s, v3.4s, v4.s[2]
|
||||
add x2, x2, #16
|
||||
fmla v28.4s, v0.4s, v4.s[3]
|
||||
prfm pldl1keep, [x1, #256]
|
||||
fmla v29.4s, v1.4s, v4.s[3]
|
||||
prfm pldl1keep, [x1, #320]
|
||||
fmla v30.4s, v2.4s, v4.s[3]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
fmla v31.4s, v3.4s, v4.s[3]
|
||||
prfm pldl1keep, [x2, #320]
|
||||
|
||||
ins v0.d[0], x5
|
||||
ins v2.d[0], x9
|
||||
ins v1.d[0], x7
|
||||
ins v3.d[0], x11
|
||||
ins v4.d[0], x24
|
||||
|
||||
ins v0.d[1], x6
|
||||
ins v2.d[1], x10
|
||||
ins v1.d[1], x8
|
||||
ins v3.d[1], x12
|
||||
ins v4.d[1], x25
|
||||
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
fmla v18.4s, v2.4s, v4.s[0]
|
||||
fmla v19.4s, v3.4s, v4.s[0]
|
||||
fmla v20.4s, v0.4s, v4.s[1]
|
||||
fmla v21.4s, v1.4s, v4.s[1]
|
||||
fmla v22.4s, v2.4s, v4.s[1]
|
||||
fmla v23.4s, v3.4s, v4.s[1]
|
||||
|
||||
fmla v24.4s, v0.4s, v4.s[2]
|
||||
fmla v25.4s, v1.4s, v4.s[2]
|
||||
fmla v26.4s, v2.4s, v4.s[2]
|
||||
fmla v27.4s, v3.4s, v4.s[2]
|
||||
fmla v28.4s, v0.4s, v4.s[3]
|
||||
fmla v29.4s, v1.4s, v4.s[3]
|
||||
fmla v30.4s, v2.4s, v4.s[3]
|
||||
fmla v31.4s, v3.4s, v4.s[3]
|
||||
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [ x1 ], #64
|
||||
ld1 {{ v4.4s }}, [ x2 ], #16
|
||||
+73
@@ -0,0 +1,73 @@
|
||||
// mul a: v0, v1, v2, v3 b: v4
|
||||
// load a: v5(d5/x5), v6(d6,x6), v7(d7,x7), v8(d8, x8)
|
||||
// load b: v9(d9/x9)
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
ldr d5, [x1], #8
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
ldr d9, [x2], #8
|
||||
fmla v18.4s, v2.4s, v4.s[0]
|
||||
ldr x5, [x1], #8
|
||||
fmla v19.4s, v3.4s, v4.s[0]
|
||||
ldr x9, [x2], #8
|
||||
fmla v20.4s, v0.4s, v4.s[1]
|
||||
ldr d6, [x1], #8
|
||||
fmla v21.4s, v1.4s, v4.s[1]
|
||||
ldr x6, [x1], #8
|
||||
fmla v22.4s, v2.4s, v4.s[1]
|
||||
ldr d7, [x1], #8
|
||||
fmla v23.4s, v3.4s, v4.s[1]
|
||||
ldr x7, [x1], #8
|
||||
|
||||
fmla v24.4s, v0.4s, v4.s[2]
|
||||
ldr d8, [x1], #8
|
||||
fmla v25.4s, v1.4s, v4.s[2]
|
||||
ldr x8, [x1], #8
|
||||
fmla v26.4s, v2.4s, v4.s[2]
|
||||
ins v5.d[1], x5
|
||||
fmla v27.4s, v3.4s, v4.s[2]
|
||||
ins v6.d[1], x6
|
||||
fmla v28.4s, v0.4s, v4.s[3]
|
||||
ins v7.d[1], x7
|
||||
fmla v29.4s, v1.4s, v4.s[3]
|
||||
ins v8.d[1], x8
|
||||
fmla v30.4s, v2.4s, v4.s[3]
|
||||
ins v9.d[1], x9
|
||||
fmla v31.4s, v3.4s, v4.s[3]
|
||||
|
||||
// mul a: v5, v6, v7, v8 b: v9
|
||||
// load a: v0(d0/x5), v1(d1,x6), v2(d2,x7), v3(d3, x8)
|
||||
// load b: v4(d4/x9)
|
||||
|
||||
fmla v16.4s, v5.4s, v9.s[0]
|
||||
ldr d0, [x1], #8
|
||||
fmla v17.4s, v6.4s, v9.s[0]
|
||||
ldr d4, [x2], #8
|
||||
fmla v18.4s, v7.4s, v9.s[0]
|
||||
ldr x5, [x1], #8
|
||||
fmla v19.4s, v8.4s, v9.s[0]
|
||||
ldr x9, [x2], #8
|
||||
fmla v20.4s, v5.4s, v9.s[1]
|
||||
ldr d1, [x1], #8
|
||||
fmla v21.4s, v6.4s, v9.s[1]
|
||||
ldr x6, [x1], #8
|
||||
fmla v22.4s, v7.4s, v9.s[1]
|
||||
ldr d2, [x1], #8
|
||||
fmla v23.4s, v8.4s, v9.s[1]
|
||||
ldr x7, [x1], #8
|
||||
|
||||
fmla v24.4s, v5.4s, v9.s[2]
|
||||
ldr d3, [x1], #8
|
||||
fmla v25.4s, v6.4s, v9.s[2]
|
||||
ldr x8, [x1], #8
|
||||
fmla v26.4s, v7.4s, v9.s[2]
|
||||
ins v0.d[1], x5
|
||||
fmla v27.4s, v8.4s, v9.s[2]
|
||||
ins v1.d[1], x6
|
||||
fmla v28.4s, v5.4s, v9.s[3]
|
||||
ins v2.d[1], x7
|
||||
fmla v29.4s, v6.4s, v9.s[3]
|
||||
ins v3.d[1], x8
|
||||
fmla v30.4s, v7.4s, v9.s[3]
|
||||
ins v4.d[1], x9
|
||||
fmla v31.4s, v8.4s, v9.s[3]
|
||||
+174
@@ -0,0 +1,174 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs: v16 to v31, (scratch)
|
||||
// - x19-x29 to preserve (but x19, x28, x29 not used)
|
||||
// - d8..d15 to preserve
|
||||
// - v16 to v31, no need to preserve
|
||||
|
||||
// v8 is used, d8 (lower half) must preserved
|
||||
// v0-v7 (scratch registers)
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd
|
||||
.global {{G}}arm64simd_mmm_f32_16x4_{{core}}_{{suffix}}
|
||||
{{G}}arm64simd_mmm_f32_16x4_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
|
||||
ld1 { v4.4s }, [ x2 ], #16
|
||||
|
||||
{% set packed_packed_loop1 %}
|
||||
{% if core == "a53" %}
|
||||
{% include "arm64simd_mmm_f32_16x4/packed_packed_loop1/cortex_a53.S.raw" %}
|
||||
{% else %}
|
||||
{% include "arm64simd_mmm_f32_16x4/packed_packed_loop1/naive.S.raw" %}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
{% set packed_packed_loop2 %}
|
||||
{% if core == "a55" %}
|
||||
{% include "arm64simd_mmm_f32_16x4/packed_packed_loop2/cortex_a55.S.raw" %}
|
||||
{% else %}
|
||||
{{ packed_packed_loop1 }}
|
||||
{{ packed_packed_loop1 }}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
cmp x3, #4
|
||||
blt .packed_packed_loop_1
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_4:
|
||||
{{ packed_packed_loop2 }}
|
||||
{{ packed_packed_loop2 }}
|
||||
|
||||
sub x3, x3, #4
|
||||
cmp x3, #4
|
||||
bge .packed_packed_loop_4
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
{{ packed_packed_loop1 }}
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_scalars.j2" %}
|
||||
{% set mr = 16 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_rows.j2" %}
|
||||
{% set mr = 16 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_cols.j2" %}
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
{% for col in range(0, 4) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 4) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.s}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{ col * 4 + 16 + reg }}.4s, v{{ col * 4 + 16 + reg }}.4s, v0.4s
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x2, [x0, #8]
|
||||
ldr x3, [x0, #16]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x2 ]
|
||||
ld1 { v4.4s }, [ x3 ]
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
fmla v18.4s, v2.4s, v4.s[0]
|
||||
fmla v19.4s, v3.4s, v4.s[0]
|
||||
fmla v20.4s, v0.4s, v4.s[1]
|
||||
fmla v21.4s, v1.4s, v4.s[1]
|
||||
fmla v22.4s, v2.4s, v4.s[1]
|
||||
fmla v23.4s, v3.4s, v4.s[1]
|
||||
|
||||
fmla v24.4s, v0.4s, v4.s[2]
|
||||
fmla v25.4s, v1.4s, v4.s[2]
|
||||
fmla v26.4s, v2.4s, v4.s[2]
|
||||
fmla v27.4s, v3.4s, v4.s[2]
|
||||
fmla v28.4s, v0.4s, v4.s[3]
|
||||
fmla v29.4s, v1.4s, v4.s[3]
|
||||
fmla v30.4s, v2.4s, v4.s[3]
|
||||
fmla v31.4s, v3.4s, v4.s[3]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x6, #4
|
||||
bne .store_strides_generic
|
||||
|
||||
{% for col in range(0, 4) %}
|
||||
str q{{ col * 4 + 16 + 0 }}, [ x5 ]
|
||||
str q{{ col * 4 + 16 + 1 }}, [ x5, #16 ]
|
||||
str q{{ col * 4 + 16 + 2 }}, [ x5, #32 ]
|
||||
str q{{ col * 4 + 16 + 3 }}, [ x5, #48 ]
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_generic:
|
||||
|
||||
{% for col in range(0, 4) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 4) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{ col * 4 + 16 + reg }}.s }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+73
@@ -0,0 +1,73 @@
|
||||
// mul a: v0, v1, v2, v3, v4, v5 b: v7
|
||||
// load a: v5(d5/x5), v6(d6,x6), v7(d7,x7), v8(d8, x8)
|
||||
// load b: v9(d9/x9)
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
ldr d5, [x1], #8
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
ldr d9, [x2], #8
|
||||
fmla v18.4s, v2.4s, v4.s[0]
|
||||
ldr x5, [x1], #8
|
||||
fmla v19.4s, v3.4s, v4.s[0]
|
||||
ldr x9, [x2], #8
|
||||
fmla v20.4s, v0.4s, v4.s[1]
|
||||
ldr d6, [x1], #8
|
||||
fmla v21.4s, v1.4s, v4.s[1]
|
||||
ldr x6, [x1], #8
|
||||
fmla v22.4s, v2.4s, v4.s[1]
|
||||
ldr d7, [x1], #8
|
||||
fmla v23.4s, v3.4s, v4.s[1]
|
||||
ldr x7, [x1], #8
|
||||
|
||||
fmla v24.4s, v0.4s, v4.s[2]
|
||||
ldr d8, [x1], #8
|
||||
fmla v25.4s, v1.4s, v4.s[2]
|
||||
ldr x8, [x1], #8
|
||||
fmla v26.4s, v2.4s, v4.s[2]
|
||||
ins v5.d[1], x5
|
||||
fmla v27.4s, v3.4s, v4.s[2]
|
||||
ins v6.d[1], x6
|
||||
fmla v28.4s, v0.4s, v4.s[3]
|
||||
ins v7.d[1], x7
|
||||
fmla v29.4s, v1.4s, v4.s[3]
|
||||
ins v8.d[1], x8
|
||||
fmla v30.4s, v2.4s, v4.s[3]
|
||||
ins v9.d[1], x9
|
||||
fmla v31.4s, v3.4s, v4.s[3]
|
||||
|
||||
// mul a: v5, v6, v7, v8 b: v9
|
||||
// load a: v0(d0/x5), v1(d1,x6), v2(d2,x7), v3(d3, x8)
|
||||
// load b: v4(d4/x9)
|
||||
|
||||
fmla v16.4s, v5.4s, v9.s[0]
|
||||
ldr d0, [x1], #8
|
||||
fmla v17.4s, v6.4s, v9.s[0]
|
||||
ldr d4, [x2], #8
|
||||
fmla v18.4s, v7.4s, v9.s[0]
|
||||
ldr x5, [x1], #8
|
||||
fmla v19.4s, v8.4s, v9.s[0]
|
||||
ldr x9, [x2], #8
|
||||
fmla v20.4s, v5.4s, v9.s[1]
|
||||
ldr d1, [x1], #8
|
||||
fmla v21.4s, v6.4s, v9.s[1]
|
||||
ldr x6, [x1], #8
|
||||
fmla v22.4s, v7.4s, v9.s[1]
|
||||
ldr d2, [x1], #8
|
||||
fmla v23.4s, v8.4s, v9.s[1]
|
||||
ldr x7, [x1], #8
|
||||
|
||||
fmla v24.4s, v5.4s, v9.s[2]
|
||||
ldr d3, [x1], #8
|
||||
fmla v25.4s, v6.4s, v9.s[2]
|
||||
ldr x8, [x1], #8
|
||||
fmla v26.4s, v7.4s, v9.s[2]
|
||||
ins v0.d[1], x5
|
||||
fmla v27.4s, v8.4s, v9.s[2]
|
||||
ins v1.d[1], x6
|
||||
fmla v28.4s, v5.4s, v9.s[3]
|
||||
ins v2.d[1], x7
|
||||
fmla v29.4s, v6.4s, v9.s[3]
|
||||
ins v3.d[1], x8
|
||||
fmla v30.4s, v7.4s, v9.s[3]
|
||||
ins v4.d[1], x9
|
||||
fmla v31.4s, v8.4s, v9.s[3]
|
||||
+68
@@ -0,0 +1,68 @@
|
||||
|
||||
fmla v8.4s, v0.4s, v7.s[0]
|
||||
ldr x4, [x1]
|
||||
fmla v9.4s, v1.4s, v7.s[0]
|
||||
ldr x5, [x1, #8]
|
||||
fmla v10.4s, v2.4s, v7.s[0]
|
||||
ldr x6, [x1, #16]
|
||||
fmla v11.4s, v3.4s, v7.s[0]
|
||||
ldr x7, [x1, #24]
|
||||
fmla v12.4s, v4.4s, v7.s[0]
|
||||
ldr x8, [x1, #32]
|
||||
fmla v13.4s, v5.4s, v7.s[0]
|
||||
ldr x9, [x1, #40]
|
||||
|
||||
fmla v14.4s, v0.4s, v7.s[1]
|
||||
ldr x10, [x1, #48]
|
||||
fmla v15.4s, v1.4s, v7.s[1]
|
||||
ldr x11, [x1, #56]
|
||||
fmla v16.4s, v2.4s, v7.s[1]
|
||||
ldr x12, [x1, #64]
|
||||
fmla v17.4s, v3.4s, v7.s[1]
|
||||
ldr x13, [x1, #72]
|
||||
fmla v18.4s, v4.4s, v7.s[1]
|
||||
ldr x14, [x1, #80]
|
||||
fmla v19.4s, v5.4s, v7.s[1]
|
||||
ldr x15, [x1, #88]
|
||||
|
||||
fmla v20.4s, v0.4s, v7.s[2]
|
||||
ldr x20, [x2]
|
||||
fmla v21.4s, v1.4s, v7.s[2]
|
||||
ldr x21, [x2, #8]
|
||||
fmla v22.4s, v2.4s, v7.s[2]
|
||||
add x1, x1, #96
|
||||
fmla v23.4s, v3.4s, v7.s[2]
|
||||
add x2, x2, #16
|
||||
fmla v24.4s, v4.4s, v7.s[2]
|
||||
prfm pldl1keep, [x1, #256]
|
||||
fmla v25.4s, v5.4s, v7.s[2]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
|
||||
fmla v26.4s, v0.4s, v7.s[3]
|
||||
prfm pldl1keep, [x1, #320]
|
||||
fmla v27.4s, v1.4s, v7.s[3]
|
||||
prfm pldl1keep, [x1, #384]
|
||||
fmla v28.4s, v2.4s, v7.s[3]
|
||||
prfm pldl1keep, [x1, #448]
|
||||
fmla v29.4s, v3.4s, v7.s[3]
|
||||
prfm pldl1keep, [x2, #320]
|
||||
fmla v30.4s, v4.4s, v7.s[3]
|
||||
prfm pldl1keep, [x2, #384]
|
||||
fmla v31.4s, v5.4s, v7.s[3]
|
||||
prfm pldl1keep, [x2, #448]
|
||||
|
||||
ins v0.d[0], x4
|
||||
ins v1.d[0], x6
|
||||
ins v2.d[0], x8
|
||||
ins v3.d[0], x10
|
||||
ins v4.d[0], x12
|
||||
ins v5.d[0], x14
|
||||
ins v7.d[0], x20
|
||||
|
||||
ins v0.d[1], x5
|
||||
ins v1.d[1], x7
|
||||
ins v2.d[1], x9
|
||||
ins v3.d[1], x11
|
||||
ins v4.d[1], x13
|
||||
ins v5.d[1], x15
|
||||
ins v7.d[1], x21
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
fmla v8.4s, v0.4s, v7.s[0]
|
||||
fmla v14.4s, v0.4s, v7.s[1]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
fmla v20.4s, v0.4s, v7.s[2]
|
||||
fmla v26.4s, v0.4s, v7.s[3]
|
||||
ldr d0, [x1], #8
|
||||
fmla v9.4s, v1.4s, v7.s[0]
|
||||
ldr x5, [x1], #8
|
||||
fmla v15.4s, v1.4s, v7.s[1]
|
||||
ldr x20, [x2], #8
|
||||
fmla v21.4s, v1.4s, v7.s[2]
|
||||
ldr x21, [x2], #8
|
||||
fmla v27.4s, v1.4s, v7.s[3]
|
||||
ldr d1, [x1], #8
|
||||
fmla v10.4s, v2.4s, v7.s[0]
|
||||
ldr x7, [x1], #8
|
||||
fmla v16.4s, v2.4s, v7.s[1]
|
||||
prfm pldl1keep, [x1, #256]
|
||||
fmla v22.4s, v2.4s, v7.s[2]
|
||||
prfm pldl1keep, [x1, #320]
|
||||
fmla v28.4s, v2.4s, v7.s[3]
|
||||
ldr d2, [x1], #8
|
||||
fmla v11.4s, v3.4s, v7.s[0]
|
||||
ldr x9, [x1], #8
|
||||
fmla v17.4s, v3.4s, v7.s[1]
|
||||
ins v0.d[1], x5
|
||||
fmla v23.4s, v3.4s, v7.s[2]
|
||||
ins v1.d[1], x7
|
||||
fmla v29.4s, v3.4s, v7.s[3]
|
||||
ldr d3, [x1], #8
|
||||
fmla v12.4s, v4.4s, v7.s[0]
|
||||
ldr x11, [x1], #8
|
||||
fmla v18.4s, v4.4s, v7.s[1]
|
||||
ins v2.d[1], x9
|
||||
fmla v24.4s, v4.4s, v7.s[2]
|
||||
fmla v30.4s, v4.4s, v7.s[3]
|
||||
ldr d4, [x1], #8
|
||||
fmla v13.4s, v5.4s, v7.s[0]
|
||||
ldr x13, [x1], #8
|
||||
fmla v19.4s, v5.4s, v7.s[1]
|
||||
ldr x14, [x1], #8
|
||||
fmla v25.4s, v5.4s, v7.s[2]
|
||||
ldr x15, [x1], #8
|
||||
fmla v31.4s, v5.4s, v7.s[3]
|
||||
|
||||
ins v7.d[0], x20
|
||||
ins v7.d[1], x21
|
||||
|
||||
ins v5.d[0], x14
|
||||
ins v5.d[1], x15
|
||||
|
||||
ins v3.d[1], x11
|
||||
ins v4.d[1], x13
|
||||
+31
@@ -0,0 +1,31 @@
|
||||
fmla v8.4s, v0.4s, v7.s[0]
|
||||
fmla v9.4s, v1.4s, v7.s[0]
|
||||
fmla v10.4s, v2.4s, v7.s[0]
|
||||
fmla v11.4s, v3.4s, v7.s[0]
|
||||
fmla v12.4s, v4.4s, v7.s[0]
|
||||
fmla v13.4s, v5.4s, v7.s[0]
|
||||
|
||||
fmla v14.4s, v0.4s, v7.s[1]
|
||||
fmla v15.4s, v1.4s, v7.s[1]
|
||||
fmla v16.4s, v2.4s, v7.s[1]
|
||||
fmla v17.4s, v3.4s, v7.s[1]
|
||||
fmla v18.4s, v4.4s, v7.s[1]
|
||||
fmla v19.4s, v5.4s, v7.s[1]
|
||||
|
||||
fmla v20.4s, v0.4s, v7.s[2]
|
||||
fmla v21.4s, v1.4s, v7.s[2]
|
||||
fmla v22.4s, v2.4s, v7.s[2]
|
||||
fmla v23.4s, v3.4s, v7.s[2]
|
||||
fmla v24.4s, v4.4s, v7.s[2]
|
||||
fmla v25.4s, v5.4s, v7.s[2]
|
||||
|
||||
fmla v26.4s, v0.4s, v7.s[3]
|
||||
fmla v27.4s, v1.4s, v7.s[3]
|
||||
fmla v28.4s, v2.4s, v7.s[3]
|
||||
fmla v29.4s, v3.4s, v7.s[3]
|
||||
fmla v30.4s, v4.4s, v7.s[3]
|
||||
fmla v31.4s, v5.4s, v7.s[3]
|
||||
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [ x1 ], #64
|
||||
ld1 {{ v4.4s, v5.4s }}, [ x1 ], #32
|
||||
ld1 {{ v7.4s }}, [ x2 ], #16
|
||||
+185
@@ -0,0 +1,185 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// x20..x27 are used, callee-preserved
|
||||
|
||||
// C tile regs: v8 to v31, (scratch)
|
||||
// - x19-x29 to preserve (but x19, x28, x29 not used)
|
||||
// - d8..d15 to preserve
|
||||
// - v16 to v31, no need to preserve
|
||||
|
||||
// v8 is used, d8 (lower half) must preserved
|
||||
// v0-v7 (scratch registers)
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd
|
||||
.global {{G}}arm64simd_mmm_f32_24x4_{{core}}_{{suffix}}
|
||||
{{G}}arm64simd_mmm_f32_24x4_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
|
||||
ld1 { v4.4s, v5.4s }, [ x1 ], #32
|
||||
ld1 { v7.4s }, [ x2 ], #16
|
||||
|
||||
{% set packed_packed_loop1 %}
|
||||
{% if core == "a53" %}
|
||||
{% include "arm64simd_mmm_f32_24x4/packed_packed_loop1/cortex_a53.S.raw" %}
|
||||
{% elif core == "a55" %}
|
||||
{% include "arm64simd_mmm_f32_24x4/packed_packed_loop1/cortex_a55.S.raw" %}
|
||||
{% else %}
|
||||
{% include "arm64simd_mmm_f32_24x4/packed_packed_loop1/naive.S.raw" %}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
cmp x3, #4
|
||||
blt .packed_packed_loop_1
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_4:
|
||||
{{ packed_packed_loop1 }}
|
||||
{{ packed_packed_loop1 }}
|
||||
{{ packed_packed_loop1 }}
|
||||
{{ packed_packed_loop1 }}
|
||||
|
||||
sub x3, x3, #4
|
||||
cmp x3, #4
|
||||
bge .packed_packed_loop_4
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
{{ packed_packed_loop1 }}
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_scalars.j2" %}
|
||||
{% set mr = 24 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_rows.j2" %}
|
||||
{% set mr = 24 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_cols.j2" %}
|
||||
{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
{% for col in range(0, 4) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 6) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.s}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{ col * 6 + 8 + reg }}.4s, v{{ col * 6 + 8 + reg }}.4s, v0.4s
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x2, [x0, #8]
|
||||
ldr x3, [x0, #16]
|
||||
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x2 ], #64
|
||||
ld1 { v7.4s }, [ x3 ]
|
||||
ld1 { v4.4s, v5.4s }, [ x2 ]
|
||||
|
||||
fmla v8.4s, v0.4s, v7.s[0]
|
||||
fmla v9.4s, v1.4s, v7.s[0]
|
||||
fmla v10.4s, v2.4s, v7.s[0]
|
||||
fmla v11.4s, v3.4s, v7.s[0]
|
||||
fmla v12.4s, v4.4s, v7.s[0]
|
||||
fmla v13.4s, v5.4s, v7.s[0]
|
||||
|
||||
fmla v14.4s, v0.4s, v7.s[1]
|
||||
fmla v15.4s, v1.4s, v7.s[1]
|
||||
fmla v16.4s, v2.4s, v7.s[1]
|
||||
fmla v17.4s, v3.4s, v7.s[1]
|
||||
fmla v18.4s, v4.4s, v7.s[1]
|
||||
fmla v19.4s, v5.4s, v7.s[1]
|
||||
|
||||
fmla v20.4s, v0.4s, v7.s[2]
|
||||
fmla v21.4s, v1.4s, v7.s[2]
|
||||
fmla v22.4s, v2.4s, v7.s[2]
|
||||
fmla v23.4s, v3.4s, v7.s[2]
|
||||
fmla v24.4s, v4.4s, v7.s[2]
|
||||
fmla v25.4s, v5.4s, v7.s[2]
|
||||
|
||||
fmla v26.4s, v0.4s, v7.s[3]
|
||||
fmla v27.4s, v1.4s, v7.s[3]
|
||||
fmla v28.4s, v2.4s, v7.s[3]
|
||||
fmla v29.4s, v3.4s, v7.s[3]
|
||||
fmla v30.4s, v4.4s, v7.s[3]
|
||||
fmla v31.4s, v5.4s, v7.s[3]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x6, #4
|
||||
bne .store_strides_generic
|
||||
|
||||
{% for col in range(0, 4) %}
|
||||
str q{{ col * 6 + 8 + 0 }}, [ x5 ]
|
||||
str q{{ col * 6 + 8 + 1 }}, [ x5, #16 ]
|
||||
str q{{ col * 6 + 8 + 2 }}, [ x5, #32 ]
|
||||
str q{{ col * 6 + 8 + 3 }}, [ x5, #48 ]
|
||||
str q{{ col * 6 + 8 + 4 }}, [ x5, #64 ]
|
||||
str q{{ col * 6 + 8 + 5 }}, [ x5, #80 ]
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_generic:
|
||||
|
||||
{% for col in range(0, 4) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 6) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{ col * 6 + 8 + reg }}.s }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+403
@@ -0,0 +1,403 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs:
|
||||
// - x19-x29 to preserve (but x19, x28, x29 not used)
|
||||
// - d8..d15 to preserve
|
||||
// - v16 to v31, no need to preserve
|
||||
//
|
||||
// v16[0] v18[0] v20[0] v22[0] v24[0] v26[0] v28[0] v30[0]
|
||||
// v16[1] v18[1]
|
||||
// v16[2] v18[2]
|
||||
// v16[3] v18[3]
|
||||
//
|
||||
// v17[0] v19[0] v21[0] v23[0] v25[0] v27[0] v29[0] v31[0]
|
||||
// v17[1] v19[1]
|
||||
// v17[2] v19[2]
|
||||
// v17[3] v19[3]
|
||||
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd
|
||||
.global {{G}}arm64simd_mmm_f32_32x1_{{core}}_{{suffix}}
|
||||
{{G}}arm64simd_mmm_f32_32x1_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldp x2, x4, [x0, #24] // b, packing
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
cmp x4, #1
|
||||
beq .q4f16se
|
||||
cmp x4, #2
|
||||
beq .q4f32se
|
||||
cmp x4, #3
|
||||
beq .f16f16
|
||||
cmp x4, #4
|
||||
beq .f32f16
|
||||
cmp x4, #5
|
||||
beq .f16f32
|
||||
|
||||
sub x3, x3, #1
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
ld1 { v8.s }[0], [ x2 ], #4
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
|
||||
ld1 { v4.4s, v5.4s, v6.4s, v7.4s }, [ x1 ], #64
|
||||
|
||||
fmla v24.4s, v0.4s, v8.s[0]
|
||||
fmla v25.4s, v1.4s, v8.s[0]
|
||||
fmla v26.4s, v2.4s, v8.s[0]
|
||||
fmla v27.4s, v3.4s, v8.s[0]
|
||||
fmla v28.4s, v4.4s, v8.s[0]
|
||||
fmla v29.4s, v5.4s, v8.s[0]
|
||||
fmla v30.4s, v6.4s, v8.s[0]
|
||||
fmla v31.4s, v7.4s, v8.s[0]
|
||||
|
||||
subs x3, x3, #1
|
||||
bge .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.p2align 8
|
||||
.q40f16_const:
|
||||
.byte 0xc8, 0xc7, 0xc6, 0xc5, 0xc4, 0xc2, 0xc0, 0xbc
|
||||
.byte 0x00, 0x3c, 0x40, 0x42, 0x44, 0x45, 0x46, 0x47
|
||||
|
||||
.q4f16se:
|
||||
adr x4, .q40f16_const
|
||||
movi v15.16b, 15
|
||||
ld1 {v13.16b}, [ x4 ]
|
||||
eor v12.16b, v12.16b, v12.16b
|
||||
|
||||
.q4f16se_outerloop:
|
||||
{% for i in range(0, 8) %}
|
||||
eor v{{ i + 16 }}.16b, v{{ i + 16 }}.16b, v{{ i + 16 }}.16b
|
||||
{% endfor %}
|
||||
mov x4, #32
|
||||
|
||||
.p2align 4
|
||||
.q4f16se_innerloop:
|
||||
ld1 { v10.16b }, [ x1 ], #16
|
||||
ld1 { v11.h }[0], [ x2 ], #2
|
||||
|
||||
and v9.16b, v10.16b, v15.16b
|
||||
ushr v10.16b, v10.16b, 4
|
||||
|
||||
tbl v9.16b, { v13.16b }, v9.16b
|
||||
tbl v10.16b, { v13.16b }, v10.16b
|
||||
|
||||
zip1 v0.16b, v12.16b, v9.16b
|
||||
zip2 v2.16b, v12.16b, v9.16b
|
||||
zip1 v4.16b, v12.16b, v10.16b
|
||||
zip2 v6.16b, v12.16b, v10.16b
|
||||
|
||||
fcvtl v11.4s, v11.4h
|
||||
|
||||
fcvtl2 v1.4s, v0.8h
|
||||
fcvtl2 v3.4s, v2.8h
|
||||
fcvtl2 v5.4s, v4.8h
|
||||
fcvtl2 v7.4s, v6.8h
|
||||
fcvtl v0.4s, v0.4h
|
||||
fcvtl v2.4s, v2.4h
|
||||
fcvtl v4.4s, v4.4h
|
||||
fcvtl v6.4s, v6.4h
|
||||
|
||||
{% for i in range(0, 8) %}
|
||||
fmla v{{ i + 16 }}.4s, v{{i}}.4s, v11.s[0]
|
||||
{% endfor %}
|
||||
|
||||
subs x4, x4, #1
|
||||
bne .q4f16se_innerloop
|
||||
|
||||
// scales
|
||||
ld1 { v0.8h-v3.8h }, [ x1 ], #64
|
||||
|
||||
fcvtl v4.4s, v0.4h
|
||||
fcvtl2 v5.4s, v0.8h
|
||||
fcvtl v6.4s, v1.4h
|
||||
fcvtl2 v7.4s, v1.8h
|
||||
fcvtl v8.4s, v2.4h
|
||||
fcvtl2 v9.4s, v2.8h
|
||||
fcvtl v10.4s, v3.4h
|
||||
fcvtl2 v11.4s, v3.8h
|
||||
|
||||
{% for i in range(0, 8) %}
|
||||
fmla v{{ i + 24 }}.4s, v{{ i + 4 }}.4s, v{{ i + 16 }}.4s
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #32
|
||||
bne .q4f16se_outerloop
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.q4f32se:
|
||||
adr x4, .q40f16_const
|
||||
movi v15.16b, 15
|
||||
ld1 {v13.16b}, [ x4 ]
|
||||
eor v12.16b, v12.16b, v12.16b
|
||||
|
||||
.q4f32se_outerloop:
|
||||
{% for i in range(0, 8) %}
|
||||
eor v{{ i + 16 }}.16b, v{{ i + 16 }}.16b, v{{ i + 16 }}.16b
|
||||
{% endfor %}
|
||||
mov x4, #32
|
||||
|
||||
.p2align 4
|
||||
.q4f32se_innerloop:
|
||||
ld1 { v10.16b }, [ x1 ], #16
|
||||
ld1 { v11.s }[0], [ x2 ], #4
|
||||
|
||||
and v9.16b, v10.16b, v15.16b
|
||||
ushr v10.16b, v10.16b, 4
|
||||
|
||||
tbl v9.16b, { v13.16b }, v9.16b
|
||||
tbl v10.16b, { v13.16b }, v10.16b
|
||||
|
||||
zip1 v0.16b, v12.16b, v9.16b
|
||||
zip2 v2.16b, v12.16b, v9.16b
|
||||
zip1 v4.16b, v12.16b, v10.16b
|
||||
zip2 v6.16b, v12.16b, v10.16b
|
||||
|
||||
fcvtl2 v1.4s, v0.8h
|
||||
fcvtl2 v3.4s, v2.8h
|
||||
fcvtl2 v5.4s, v4.8h
|
||||
fcvtl2 v7.4s, v6.8h
|
||||
fcvtl v0.4s, v0.4h
|
||||
fcvtl v2.4s, v2.4h
|
||||
fcvtl v4.4s, v4.4h
|
||||
fcvtl v6.4s, v6.4h
|
||||
|
||||
{% for i in range(0, 8) %}
|
||||
fmla v{{ i + 16 }}.4s, v{{i}}.4s, v11.s[0]
|
||||
{% endfor %}
|
||||
|
||||
subs x4, x4, #1
|
||||
bne .q4f32se_innerloop
|
||||
|
||||
// scales
|
||||
ld1 { v0.8h-v3.8h }, [ x1 ], #64
|
||||
|
||||
fcvtl v4.4s, v0.4h
|
||||
fcvtl2 v5.4s, v0.8h
|
||||
fcvtl v6.4s, v1.4h
|
||||
fcvtl2 v7.4s, v1.8h
|
||||
fcvtl v8.4s, v2.4h
|
||||
fcvtl2 v9.4s, v2.8h
|
||||
fcvtl v10.4s, v3.4h
|
||||
fcvtl2 v11.4s, v3.8h
|
||||
|
||||
{% for i in range(0, 8) %}
|
||||
fmla v{{ i + 24 }}.4s, v{{ i + 4 }}.4s, v{{ i + 16 }}.4s
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #32
|
||||
bne .q4f32se_outerloop
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.f16f16:
|
||||
sub x3, x3, #1
|
||||
.f16f16_loop:
|
||||
ld1 { v9.h }[0], [ x2 ], #2
|
||||
ld1 { v10.8h-v13.8h }, [ x1 ], #64
|
||||
|
||||
fcvtl v8.4s, v9.4h
|
||||
{% for reg in range(0, 4) %}
|
||||
fcvtl v{{ reg * 2 }}.4s, v{{ reg + 10 }}.4h
|
||||
fcvtl2 v{{ reg * 2 + 1 }}.4s, v{{ reg + 10 }}.8h
|
||||
{% endfor %}
|
||||
|
||||
fmla v24.4s, v0.4s, v8.s[0]
|
||||
fmla v25.4s, v1.4s, v8.s[0]
|
||||
fmla v26.4s, v2.4s, v8.s[0]
|
||||
fmla v27.4s, v3.4s, v8.s[0]
|
||||
fmla v28.4s, v4.4s, v8.s[0]
|
||||
fmla v29.4s, v5.4s, v8.s[0]
|
||||
fmla v30.4s, v6.4s, v8.s[0]
|
||||
fmla v31.4s, v7.4s, v8.s[0]
|
||||
|
||||
subs x3, x3, #1
|
||||
bge .f16f16_loop
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.f32f16:
|
||||
sub x3, x3, #1
|
||||
.f32f16_loop:
|
||||
ld1 { v9.h }[0], [ x2 ], #2
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
|
||||
ld1 { v4.4s, v5.4s, v6.4s, v7.4s }, [ x1 ], #64
|
||||
|
||||
fcvtl v8.4s, v9.4h
|
||||
|
||||
fmla v24.4s, v0.4s, v8.s[0]
|
||||
fmla v25.4s, v1.4s, v8.s[0]
|
||||
fmla v26.4s, v2.4s, v8.s[0]
|
||||
fmla v27.4s, v3.4s, v8.s[0]
|
||||
fmla v28.4s, v4.4s, v8.s[0]
|
||||
fmla v29.4s, v5.4s, v8.s[0]
|
||||
fmla v30.4s, v6.4s, v8.s[0]
|
||||
fmla v31.4s, v7.4s, v8.s[0]
|
||||
|
||||
subs x3, x3, #1
|
||||
bge .f32f16_loop
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.f16f32:
|
||||
sub x3, x3, #1
|
||||
.f16f32_loop:
|
||||
ld1 { v8.s }[0], [ x2 ], #4
|
||||
ld1 { v10.8h-v13.8h }, [ x1 ], #64
|
||||
|
||||
{% for reg in range(0, 4) %}
|
||||
fcvtl v{{ reg * 2 }}.4s, v{{ reg + 10 }}.4h
|
||||
fcvtl2 v{{ reg * 2 + 1 }}.4s, v{{ reg + 10 }}.8h
|
||||
{% endfor %}
|
||||
|
||||
fmla v24.4s, v0.4s, v8.s[0]
|
||||
fmla v25.4s, v1.4s, v8.s[0]
|
||||
fmla v26.4s, v2.4s, v8.s[0]
|
||||
fmla v27.4s, v3.4s, v8.s[0]
|
||||
fmla v28.4s, v4.4s, v8.s[0]
|
||||
fmla v29.4s, v5.4s, v8.s[0]
|
||||
fmla v30.4s, v6.4s, v8.s[0]
|
||||
fmla v31.4s, v7.4s, v8.s[0]
|
||||
|
||||
subs x3, x3, #1
|
||||
bge .f16f32_loop
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 24 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_scalars.j2" %}
|
||||
{% set mr = 32 %}{% set from = 24 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_rows.j2" %}
|
||||
{% set mr = 32 %}{% set from = 24 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_cols.j2" %}
|
||||
{% set from = 24 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
cmp x6, #4
|
||||
beq .do_per_row_add
|
||||
|
||||
{% for reg in range(24, 32) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.s}[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{reg}}.4s, v{{reg}}.4s, v0.4s
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.do_per_row_add:
|
||||
ld1 {v0.4s-v3.4s}, [x5], #64
|
||||
ld1 {v4.4s-v7.4s}, [x5], #64
|
||||
|
||||
{% for r in range(0, 8) %}
|
||||
fadd v{{ r + 24 }}.4s, v{{ r + 24 }}.4s, v{{r}}.4s
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x3, [x0, #16]
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 {v8.s}[0], [ x3 ]
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x2 ], #64
|
||||
ld1 { v4.4s, v5.4s, v6.4s, v7.4s }, [ x2 ], #64
|
||||
|
||||
fmla v24.4s, v0.4s, v8.s[0]
|
||||
fmla v25.4s, v1.4s, v8.s[0]
|
||||
fmla v26.4s, v2.4s, v8.s[0]
|
||||
fmla v27.4s, v3.4s, v8.s[0]
|
||||
fmla v28.4s, v4.4s, v8.s[0]
|
||||
fmla v29.4s, v5.4s, v8.s[0]
|
||||
fmla v30.4s, v6.4s, v8.s[0]
|
||||
fmla v31.4s, v7.4s, v8.s[0]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x8, #2
|
||||
beq .store_f16
|
||||
|
||||
cmp x6, #4
|
||||
beq .store_strides_contig
|
||||
|
||||
{% for reg in range(24, 32) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{reg}}.s }[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_contig:
|
||||
|
||||
{% for reg in range(24, 32) %}
|
||||
st1 { v{{reg}}.4s }, [ x5 ], #16
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.store_f16:
|
||||
{% for reg in range(0, 4) %}
|
||||
fcvtn v{{reg}}.4h, v{{ reg * 2 + 24 }}.4s
|
||||
fcvtn2 v{{reg}}.8h, v{{ reg * 2 + 25 }}.4s
|
||||
{% endfor %}
|
||||
|
||||
cmp x6, #2
|
||||
beq .store_strides_contig_f16
|
||||
|
||||
{% for reg in range(0, 4) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
st1 { v{{reg}}.h }[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_contig_f16:
|
||||
|
||||
{% for reg in range(0, 4) %}
|
||||
st1 { v{{reg}}.8h }, [ x5 ], #16
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+307
@@ -0,0 +1,307 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs: v16 to v31, no need to preserve
|
||||
|
||||
// no preservation either for v0-v7...
|
||||
// v8..v15 are callee-preserved
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.global {{G}}arm64simd_mmm_f32_32x3_{{core}}_{{suffix}}
|
||||
{{G}}arm64simd_mmm_f32_32x3_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldp x2, x4, [x0, #24] // b, packing
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
cmp x4, #1
|
||||
beq .f32f16
|
||||
cmp x4, #2
|
||||
beq .f16f32
|
||||
cmp x4, #3
|
||||
beq .f16f16
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
ld1 { v7.4s }, [ x2 ]
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
|
||||
ld1 { v4.4s, v5.4s, v6.4s }, [ x1 ], #48
|
||||
add x2, x2, #12
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 }}.4s, v0.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
ld1 { v0.4s }, [ x1 ], #16
|
||||
|
||||
{% for row in range(1, 7) %}
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 + row }}.4s, v{{row}}.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 15 }}.4s, v0.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.f32f16:
|
||||
ld1 { v7.4h }, [ x2 ]
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
|
||||
ld1 { v4.4s, v5.4s, v6.4s }, [ x1 ], #48
|
||||
fcvtl v7.4s, v7.4h
|
||||
add x2, x2, #6
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 }}.4s, v0.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
ld1 { v0.4s }, [ x1 ], #16
|
||||
|
||||
{% for row in range(1, 7) %}
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 + row }}.4s, v{{row}}.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 15 }}.4s, v0.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .f32f16
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.f16f32:
|
||||
ld1 { v7.4s }, [ x2 ]
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [ x1 ], #64
|
||||
add x2, x2, #12
|
||||
|
||||
fcvtl v4.4s, v0.4h
|
||||
fcvtl2 v5.4s, v0.8h
|
||||
fcvtl v6.4s, v1.4h
|
||||
fcvtl2 v0.4s, v1.8h
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 }}.4s, v4.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 9 }}.4s, v5.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 10 }}.4s, v6.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 11 }}.4s, v0.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
fcvtl v4.4s, v2.4h
|
||||
fcvtl2 v5.4s, v2.8h
|
||||
fcvtl v6.4s, v3.4h
|
||||
fcvtl2 v1.4s, v3.8h
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 12 }}.4s, v4.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 13 }}.4s, v5.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 14 }}.4s, v6.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 15 }}.4s, v1.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .f16f32
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.f16f16:
|
||||
ld1 { v7.4h }, [ x2 ]
|
||||
ld1 { v0.8h, v1.8h, v2.8h, v3.8h }, [ x1 ], #64
|
||||
add x2, x2, #6
|
||||
|
||||
fcvtl v7.4s, v7.4h
|
||||
|
||||
fcvtl v4.4s, v0.4h
|
||||
fcvtl2 v5.4s, v0.8h
|
||||
fcvtl v6.4s, v1.4h
|
||||
fcvtl2 v0.4s, v1.8h
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 }}.4s, v4.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 9 }}.4s, v5.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 10 }}.4s, v6.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 11 }}.4s, v0.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
fcvtl v4.4s, v2.4h
|
||||
fcvtl2 v5.4s, v2.8h
|
||||
fcvtl v6.4s, v3.4h
|
||||
fcvtl2 v1.4s, v3.8h
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 12 }}.4s, v4.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 13 }}.4s, v5.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 14 }}.4s, v6.4s, v7.s[{{ col }}]
|
||||
fmla v{{ col * 8 + 15 }}.4s, v1.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .f16f16
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
|
||||
{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_scalars.j2" %}
|
||||
{% set mr = 32 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_rows.j2" %}
|
||||
{% set mr = 32 %}{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_cols.j2" %}
|
||||
{% set from = 8 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 8) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.s}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{ col * 8 + 8 + reg }}.4s, v{{ col * 8 + 8 + reg }}.4s, v0.4s
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldp x2, x3, [x0, #8]
|
||||
|
||||
ld1 { v7.d }[0], [ x3 ], #8
|
||||
ld1 { v7.s }[2], [ x3 ], #4
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x2 ], #64
|
||||
ld1 { v4.4s, v5.4s, v6.4s }, [ x2 ], #48
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 }}.4s, v0.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
ld1 { v0.4s }, [ x2 ], #16
|
||||
|
||||
{% for row in range(1, 7) %}
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 8 + row }}.4s, v{{row}}.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
fmla v{{ col * 8 + 15 }}.4s, v0.4s, v7.s[{{ col }}]
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x8, #2
|
||||
beq .store_f16
|
||||
|
||||
cmp x6, #4
|
||||
beq .store_strides_contig
|
||||
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 8) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{ col * 8 + 8 + reg }}.s }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_contig:
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
mov x4, x5
|
||||
{% for r in range(0, 8) %}
|
||||
st1 { v{{ col * 8 + 8 + r }}.4s }, [ x4 ], 16
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_f16:
|
||||
|
||||
cmp x6, #2
|
||||
beq .store_strides_contig_f16
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
{% for reg in range(0, 4) %}
|
||||
fcvtn v{{reg}}.4h, v{{ (col * 4 + reg) * 2 + 8 }}.4s
|
||||
fcvtn2 v{{reg}}.8h, v{{ (col * 4 + reg) * 2 + 9 }}.4s
|
||||
{% endfor %}
|
||||
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 4) %}
|
||||
{% for lane in range(0, 8) %}
|
||||
st1 { v{{reg}}.h }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
|
||||
{% endfor %}
|
||||
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_contig_f16:
|
||||
|
||||
{% for col in range(0, 3) %}
|
||||
{% for reg in range(0, 4) %}
|
||||
fcvtn v{{reg}}.4h, v{{ (col * 4 + reg) * 2 + 8 }}.4s
|
||||
fcvtn2 v{{reg}}.8h, v{{ (col * 4 + reg) * 2 + 9 }}.4s
|
||||
{% endfor %}
|
||||
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 4) %}
|
||||
st1 { v{{reg}}.4s }, [ x4 ], #16
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
|
||||
.return:
|
||||
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+65
@@ -0,0 +1,65 @@
|
||||
fmla v16.4s, v0.4s, v8.s[0]
|
||||
ldr x5, [x1, #128]
|
||||
fmla v17.4s, v1.4s, v8.s[0]
|
||||
ldr x6, [x1, #136]
|
||||
fmla v18.4s, v2.4s, v8.s[0]
|
||||
ldr x7, [x1, #144]
|
||||
fmla v19.4s, v3.4s, v8.s[0]
|
||||
ldr x9, [x1, #152]
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [ x1 ], #64
|
||||
|
||||
fmla v20.4s, v4.4s, v8.s[0]
|
||||
ldr x10, [x1, #96]
|
||||
fmla v21.4s, v5.4s, v8.s[0]
|
||||
ldr x11, [x1, #104]
|
||||
fmla v22.4s, v6.4s, v8.s[0]
|
||||
ldr x12, [x1, #112]
|
||||
fmla v23.4s, v7.4s, v8.s[0]
|
||||
ldr x13, [x1, #120]
|
||||
|
||||
ld1 {{ v4.4s, v5.4s, v6.4s, v7.4s }}, [ x1 ]
|
||||
|
||||
fmla v24.4s, v0.4s, v8.s[0]
|
||||
ldr x14, [x1, #128]
|
||||
fmla v25.4s, v1.4s, v8.s[0]
|
||||
ldr x15, [x1, #136]
|
||||
fmla v26.4s, v2.4s, v8.s[0]
|
||||
ldr x20, [x1, #144]
|
||||
fmla v27.4s, v3.4s, v8.s[0]
|
||||
ldr x21, [x1, #152]
|
||||
fmla v28.4s, v4.4s, v8.s[0]
|
||||
ldr x22, [x1, #160]
|
||||
fmla v29.4s, v5.4s, v8.s[0]
|
||||
ldr x23, [x1, #168]
|
||||
fmla v30.4s, v6.4s, v8.s[0]
|
||||
ldr x24, [x1, #176]
|
||||
fmla v31.4s, v7.4s, v8.s[0]
|
||||
ldr x25, [x1, #184]
|
||||
|
||||
ld1 {{ v8.s }}[0], [ x2 ], #4
|
||||
|
||||
prfm pldl1keep, [x1, #1024]
|
||||
prfm pldl1keep, [x1, #1088]
|
||||
prfm pldl1keep, [x1, #1152]
|
||||
prfm pldl1keep, [x1, #1216]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
|
||||
ins v0.d[0], x5
|
||||
ins v1.d[0], x7
|
||||
ins v2.d[0], x10
|
||||
ins v3.d[0], x12
|
||||
ins v4.d[0], x14
|
||||
ins v5.d[0], x20
|
||||
ins v6.d[0], x22
|
||||
ins v7.d[0], x24
|
||||
|
||||
ins v0.d[1], x6
|
||||
ins v1.d[1], x9
|
||||
ins v2.d[1], x11
|
||||
ins v3.d[1], x13
|
||||
ins v4.d[1], x15
|
||||
ins v5.d[1], x21
|
||||
ins v6.d[1], x23
|
||||
ins v7.d[1], x25
|
||||
|
||||
add x1, x1, #192
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
ld1 {{ v9.4s, v10.4s, v11.4s, v12.4s }}, [x1], #64
|
||||
ld1 {{ v13.4s, v14.4s, v15.4s }}, [x1], #48
|
||||
|
||||
fmla v16.4s, v0.4s, v8.s[0]
|
||||
fmla v17.4s, v1.4s, v8.s[0]
|
||||
fmla v18.4s, v2.4s, v8.s[0]
|
||||
fmla v19.4s, v3.4s, v8.s[0]
|
||||
fmla v20.4s, v4.4s, v8.s[0]
|
||||
fmla v21.4s, v5.4s, v8.s[0]
|
||||
fmla v22.4s, v6.4s, v8.s[0]
|
||||
fmla v23.4s, v7.4s, v8.s[0]
|
||||
fmla v24.4s, v9.4s, v8.s[0]
|
||||
ld1 {{ v9.4s }}, [ x1 ], #16
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [x1], #64
|
||||
ld1 {{ v4.4s, v5.4s, v6.4s, v7.4s }}, [x1], #64
|
||||
fmla v25.4s, v10.4s, v8.s[0]
|
||||
fmla v26.4s, v11.4s, v8.s[0]
|
||||
fmla v27.4s, v12.4s, v8.s[0]
|
||||
fmla v28.4s, v13.4s, v8.s[0]
|
||||
fmla v29.4s, v14.4s, v8.s[0]
|
||||
fmla v30.4s, v15.4s, v8.s[0]
|
||||
|
||||
fmla v31.4s, v9.4s, v8.s[0]
|
||||
|
||||
ld1 {{ v8.s }}[0], [ x2 ], #4
|
||||
|
||||
prfm pldl1keep, [x1, #1024]
|
||||
prfm pldl1keep, [x1, #1088]
|
||||
prfm pldl1keep, [x1, #1152]
|
||||
prfm pldl1keep, [x1, #1216]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
|
||||
+85
@@ -0,0 +1,85 @@
|
||||
ld1 {{ v9.4s, v10.4s, v11.4s }}, [x1], #48
|
||||
|
||||
fmla v16.4s, v0.4s, v8.s[0]
|
||||
ldr x8, [x2], #8
|
||||
fmla v17.4s, v1.4s, v8.s[0]
|
||||
ldr d12, [x1], #8
|
||||
fmla v18.4s, v2.4s, v8.s[0]
|
||||
ldr x12, [x1], #8
|
||||
fmla v19.4s, v3.4s, v8.s[0]
|
||||
ldr d13, [x1], #8
|
||||
fmla v20.4s, v4.4s, v8.s[0]
|
||||
ldr x13, [x1], #8
|
||||
fmla v21.4s, v5.4s, v8.s[0]
|
||||
ldr d14, [x1], #8
|
||||
fmla v22.4s, v6.4s, v8.s[0]
|
||||
ldr x14, [x1], #8
|
||||
fmla v23.4s, v7.4s, v8.s[0]
|
||||
ldr d15, [x1], #8
|
||||
fmla v24.4s, v9.4s, v8.s[0]
|
||||
ldr x15, [x1], #8
|
||||
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [x1], #64
|
||||
ins v8.d[1], x8
|
||||
ld1 {{ v4.4s, v5.4s, v6.4s, v7.4s }}, [x1], #64
|
||||
|
||||
fmla v25.4s, v10.4s, v8.s[0]
|
||||
ins v12.d[1], x12
|
||||
fmla v26.4s, v11.4s, v8.s[0]
|
||||
ins v13.d[1], x13
|
||||
fmla v27.4s, v12.4s, v8.s[0]
|
||||
ins v14.d[1], x14
|
||||
fmla v28.4s, v13.4s, v8.s[0]
|
||||
ins v15.d[1], x15
|
||||
|
||||
ld1 {{ v9.4s, v10.4s, v11.4s, v12.4s }}, [x1], #64
|
||||
|
||||
fmla v29.4s, v14.4s, v8.s[0]
|
||||
ldr d13, [x1], #8
|
||||
fmla v30.4s, v15.4s, v8.s[0]
|
||||
ldr x13, [x1], #8
|
||||
fmla v31.4s, v0.4s, v8.s[0]
|
||||
ldr d14, [x1], #8
|
||||
|
||||
fmla v16.4s, v1.4s, v8.s[2]
|
||||
ldr x14, [x1], #8
|
||||
fmla v17.4s, v2.4s, v8.s[2]
|
||||
ldr d15, [x1], #8
|
||||
fmla v18.4s, v3.4s, v8.s[2]
|
||||
ldr x15, [x1], #8
|
||||
fmla v19.4s, v4.4s, v8.s[2]
|
||||
|
||||
ld1 {{ v0.4s }}, [x1], #16
|
||||
|
||||
fmla v20.4s, v5.4s, v8.s[2]
|
||||
ldr d1, [x1], #8
|
||||
fmla v21.4s, v6.4s, v8.s[2]
|
||||
ldr x10, [x1], #8
|
||||
|
||||
fmla v22.4s, v7.4s, v8.s[2]
|
||||
|
||||
fmla v23.4s, v9.4s, v8.s[2]
|
||||
ins v13.d[1], x13
|
||||
fmla v24.4s, v10.4s, v8.s[2]
|
||||
ins v14.d[1], x14
|
||||
fmla v25.4s, v11.4s, v8.s[2]
|
||||
ins v15.d[1], x15
|
||||
|
||||
fmla v26.4s, v12.4s, v8.s[2]
|
||||
prfm pldl1keep, [x1, #1024]
|
||||
fmla v27.4s, v13.4s, v8.s[2]
|
||||
ins v1.d[1], x10
|
||||
fmla v28.4s, v14.4s, v8.s[2]
|
||||
prfm pldl1keep, [x1, #1088]
|
||||
fmla v29.4s, v15.4s, v8.s[2]
|
||||
prfm pldl1keep, [x1, #1152]
|
||||
fmla v30.4s, v0.4s, v8.s[2]
|
||||
prfm pldl1keep, [x1, #1216]
|
||||
fmla v31.4s, v1.4s, v8.s[2]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [x1], #64
|
||||
ins v8.s[0], v8.s[3]
|
||||
ld1 {{ v4.4s, v5.4s, v6.4s, v7.4s }}, [x1], #64
|
||||
|
||||
|
||||
+66
@@ -0,0 +1,66 @@
|
||||
// load a: v9, v10, v11, v12, v13, v14, v15
|
||||
// load a: v0, v1, v2, v3, v4, v4, v6, v7
|
||||
|
||||
ld1 {{ v9.4s, v10.4s, v11.4s, v12.4s }}, [x1], #64
|
||||
ld1 {{ v13.4s, v14.4s, v15.4s }}, [x1], #48
|
||||
|
||||
fmla v16.4s, v0.4s, v8.s[0]
|
||||
fmla v17.4s, v1.4s, v8.s[0]
|
||||
fmla v18.4s, v2.4s, v8.s[0]
|
||||
fmla v19.4s, v3.4s, v8.s[0]
|
||||
|
||||
ld1 {{ v0.4s, v1.4s }}, [x1], #32
|
||||
|
||||
fmla v20.4s, v4.4s, v8.s[0]
|
||||
fmla v21.4s, v5.4s, v8.s[0]
|
||||
|
||||
ld1 {{ v2.4s, v3.4s, v4.4s, v5.4s }}, [x1], #64
|
||||
fmla v22.4s, v6.4s, v8.s[0]
|
||||
fmla v23.4s, v7.4s, v8.s[0]
|
||||
|
||||
ld1 {{ v6.4s, v7.4s }}, [x1], #32
|
||||
|
||||
fmla v24.4s, v9.4s, v8.s[0]
|
||||
fmla v25.4s, v10.4s, v8.s[0]
|
||||
fmla v26.4s, v11.4s, v8.s[0]
|
||||
fmla v27.4s, v12.4s, v8.s[0]
|
||||
fmla v28.4s, v13.4s, v8.s[0]
|
||||
fmla v29.4s, v14.4s, v8.s[0]
|
||||
fmla v30.4s, v15.4s, v8.s[0]
|
||||
|
||||
ld1 {{ v9.4s, v10.4s, v11.4s, v12.4s }}, [x1], #64
|
||||
ld1 {{ v13.4s, v14.4s, v15.4s }}, [x1], #48
|
||||
|
||||
fmla v31.4s, v0.4s, v8.s[0]
|
||||
ld1 {{ v8.s }}[0], [ x2 ], #4
|
||||
|
||||
fmla v16.4s, v1.4s, v8.s[0]
|
||||
ld1 {{ v0.4s, v1.4s }}, [x1], #32
|
||||
fmla v17.4s, v2.4s, v8.s[0]
|
||||
fmla v18.4s, v3.4s, v8.s[0]
|
||||
fmla v19.4s, v4.4s, v8.s[0]
|
||||
|
||||
fmla v20.4s, v5.4s, v8.s[0]
|
||||
fmla v21.4s, v6.4s, v8.s[0]
|
||||
fmla v22.4s, v7.4s, v8.s[0]
|
||||
fmla v23.4s, v9.4s, v8.s[0]
|
||||
|
||||
fmla v24.4s, v10.4s, v8.s[0]
|
||||
fmla v25.4s, v11.4s, v8.s[0]
|
||||
fmla v26.4s, v12.4s, v8.s[0]
|
||||
fmla v27.4s, v13.4s, v8.s[0]
|
||||
fmla v28.4s, v14.4s, v8.s[0]
|
||||
fmla v29.4s, v15.4s, v8.s[0]
|
||||
fmla v30.4s, v0.4s, v8.s[0]
|
||||
fmla v31.4s, v1.4s, v8.s[0]
|
||||
ld1 {{ v8.s }}[0], [ x2 ], #4
|
||||
|
||||
ld1 {{ v0.4s, v1.4s, v2.4s, v3.4s }}, [x1], #64
|
||||
ld1 {{ v4.4s, v5.4s, v6.4s, v7.4s }}, [x1], #64
|
||||
|
||||
prfm pldl1keep, [x1, #1024]
|
||||
prfm pldl1keep, [x1, #1088]
|
||||
prfm pldl1keep, [x1, #1152]
|
||||
prfm pldl1keep, [x1, #1216]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
|
||||
+225
@@ -0,0 +1,225 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs:
|
||||
// - x19-x29 to preserve (but x19, x28, x29 not used)
|
||||
// - d8..d15 to preserve
|
||||
// - v16 to v31, no need to preserve
|
||||
//
|
||||
// v16[0] v18[0] v20[0] v22[0] v24[0] v26[0] v28[0] v30[0]
|
||||
// v16[1] v18[1]
|
||||
// v16[2] v18[2]
|
||||
// v16[3] v18[3]
|
||||
//
|
||||
// v17[0] v19[0] v21[0] v23[0] v25[0] v27[0] v29[0] v31[0]
|
||||
// v17[1] v19[1]
|
||||
// v17[2] v19[2]
|
||||
// v17[3] v19[3]
|
||||
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd
|
||||
.global {{G}}arm64simd_mmm_f32_64x1_{{core}}_{{suffix}}
|
||||
{{G}}arm64simd_mmm_f32_64x1_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
sub x3, x3, #1
|
||||
|
||||
|
||||
ld1 { v8.s }[0], [ x2 ], #4
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [ x1 ], #64
|
||||
ld1 { v4.4s, v5.4s, v6.4s, v7.4s }, [ x1 ], #64
|
||||
|
||||
cmp x3, #0
|
||||
beq .packed_packed_loop_1_last
|
||||
|
||||
cmp x3, #4
|
||||
blt .packed_packed_loop_1
|
||||
|
||||
{% set packed_packed_loop1 %}
|
||||
{% if core == "a53" %}
|
||||
{% include "arm64simd_mmm_f32_64x1/loop1/cortex_a53.S.raw" %}
|
||||
{% else %}
|
||||
{% include "arm64simd_mmm_f32_64x1/loop1/naive.S.raw" %}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
{% set packed_packed_loop2 %}
|
||||
{% if core == "a53" %}
|
||||
{{ packed_packed_loop1 }}
|
||||
{{ packed_packed_loop1 }}
|
||||
{% elif core == "a55" %}
|
||||
{% include "arm64simd_mmm_f32_64x1/loop2/cortex_a55.S.raw" %}
|
||||
{% else %}
|
||||
{% include "arm64simd_mmm_f32_64x1/loop2/naive.S.raw" %}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_4:
|
||||
{{ packed_packed_loop2 }}
|
||||
{{ packed_packed_loop2 }}
|
||||
|
||||
sub x3, x3, #4
|
||||
cmp x3, #4
|
||||
bge .packed_packed_loop_4
|
||||
|
||||
cmp x3, #0
|
||||
beq .packed_packed_loop_1_last
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
{{ packed_packed_loop1 }}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
// last loop can't read beyond actual input as it's likely not packed and padded
|
||||
.packed_packed_loop_1_last:
|
||||
ld1 { v9.4s, v10.4s, v11.4s, v12.4s }, [x1], #64
|
||||
ld1 { v13.4s, v14.4s, v15.4s }, [x1], #48
|
||||
|
||||
fmla v16.4s, v0.4s, v8.s[0]
|
||||
fmla v17.4s, v1.4s, v8.s[0]
|
||||
ld1 { v0.4s }, [ x1 ]
|
||||
fmla v18.4s, v2.4s, v8.s[0]
|
||||
fmla v19.4s, v3.4s, v8.s[0]
|
||||
fmla v20.4s, v4.4s, v8.s[0]
|
||||
fmla v21.4s, v5.4s, v8.s[0]
|
||||
fmla v22.4s, v6.4s, v8.s[0]
|
||||
fmla v23.4s, v7.4s, v8.s[0]
|
||||
|
||||
fmla v24.4s, v9.4s, v8.s[0]
|
||||
fmla v25.4s, v10.4s, v8.s[0]
|
||||
fmla v26.4s, v11.4s, v8.s[0]
|
||||
fmla v27.4s, v12.4s, v8.s[0]
|
||||
fmla v28.4s, v13.4s, v8.s[0]
|
||||
fmla v29.4s, v14.4s, v8.s[0]
|
||||
fmla v30.4s, v15.4s, v8.s[0]
|
||||
fmla v31.4s, v0.4s, v8.s[0]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_scalars.j2" %}
|
||||
{% set mr = 64 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_rows.j2" %}
|
||||
{% set mr = 64 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_cols.j2" %}
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
cmp x6, #4
|
||||
beq .do_per_row_add
|
||||
|
||||
{% for reg in range(16, 32) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.s}[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{reg}}.4s, v{{reg}}.4s, v0.4s
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.do_per_row_add:
|
||||
ld1 {v0.4s-v3.4s}, [x5], #64
|
||||
ld1 {v4.4s-v7.4s}, [x5], #64
|
||||
ld1 {v8.4s-v11.4s}, [x5], #64
|
||||
ld1 {v12.4s-v15.4s}, [x5], #64
|
||||
|
||||
{% for r in range(0, 16) %}
|
||||
fadd v{{ r + 16 }}.4s, v{{ r + 16 }}.4s, v{{r}}.4s
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x3, [x0, #16]
|
||||
ldr x2, [x0, #8]
|
||||
|
||||
ld1 {v8.s}[0], [ x3 ]
|
||||
|
||||
{% for r in range(0, 8) %}
|
||||
ldr q{{r}}, [x2], #16
|
||||
{% endfor %}
|
||||
|
||||
fmla v16.4s, v0.4s, v8.s[0]
|
||||
ldr q0, [x2], #16
|
||||
fmla v17.4s, v1.4s, v8.s[0]
|
||||
ldr q1, [x2], #16
|
||||
fmla v18.4s, v2.4s, v8.s[0]
|
||||
ldr q2, [x2], #16
|
||||
fmla v19.4s, v3.4s, v8.s[0]
|
||||
ldr q3, [x2], #16
|
||||
fmla v20.4s, v4.4s, v8.s[0]
|
||||
ldr q4, [x2], #16
|
||||
fmla v21.4s, v5.4s, v8.s[0]
|
||||
ldr q5, [x2], #16
|
||||
fmla v22.4s, v6.4s, v8.s[0]
|
||||
ldr q6, [x2], #16
|
||||
fmla v23.4s, v7.4s, v8.s[0]
|
||||
ldr q7, [x2], #16
|
||||
|
||||
fmla v24.4s, v0.4s, v8.s[0]
|
||||
fmla v25.4s, v1.4s, v8.s[0]
|
||||
fmla v26.4s, v2.4s, v8.s[0]
|
||||
fmla v27.4s, v3.4s, v8.s[0]
|
||||
fmla v28.4s, v4.4s, v8.s[0]
|
||||
fmla v29.4s, v5.4s, v8.s[0]
|
||||
fmla v30.4s, v6.4s, v8.s[0]
|
||||
fmla v31.4s, v7.4s, v8.s[0]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc$
|
||||
|
||||
cmp x6, #4
|
||||
beq .store_strides_contig
|
||||
|
||||
{% for reg in range(16, 32) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{reg}}.s }[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_contig:
|
||||
|
||||
{% for reg in range(16, 32) %}
|
||||
st1 { v{{reg}}.4s }, [ x5 ], #16
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
ld1 {{ v2.4s, v3.4s }}, [x1], #32
|
||||
ld1 {{ v6.4s, v7.4s }}, [x2], #32
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
fmla v18.4s, v0.4s, v4.s[1]
|
||||
fmla v19.4s, v1.4s, v4.s[1]
|
||||
fmla v20.4s, v0.4s, v4.s[2]
|
||||
fmla v21.4s, v1.4s, v4.s[2]
|
||||
fmla v22.4s, v0.4s, v4.s[3]
|
||||
fmla v23.4s, v1.4s, v4.s[3]
|
||||
|
||||
fmla v24.4s, v0.4s, v5.s[0]
|
||||
fmla v25.4s, v1.4s, v5.s[0]
|
||||
fmla v26.4s, v0.4s, v5.s[1]
|
||||
fmla v27.4s, v1.4s, v5.s[1]
|
||||
fmla v28.4s, v0.4s, v5.s[2]
|
||||
fmla v29.4s, v1.4s, v5.s[2]
|
||||
fmla v30.4s, v0.4s, v5.s[3]
|
||||
fmla v31.4s, v1.4s, v5.s[3]
|
||||
|
||||
and v0.16b, v2.16b, v2.16b
|
||||
and v1.16b, v3.16b, v3.16b
|
||||
and v4.16b, v6.16b, v6.16b
|
||||
and v5.16b, v7.16b, v7.16b
|
||||
+51
@@ -0,0 +1,51 @@
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
ldr w5, [x1], #4
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
ldr w20, [x2], #4
|
||||
fmla v18.4s, v0.4s, v4.s[1]
|
||||
ldr w6, [x1], #4
|
||||
fmla v20.4s, v1.4s, v4.s[1]
|
||||
ldr w21, [x2], #4
|
||||
fmla v20.4s, v0.4s, v4.s[2]
|
||||
ldr w7, [x1], #4
|
||||
fmla v21.4s, v1.4s, v4.s[2]
|
||||
ldr w22, [x2], #4
|
||||
fmla v22.4s, v0.4s, v4.s[3]
|
||||
ldr w8, [x1], #4
|
||||
fmla v23.4s, v1.4s, v4.s[3]
|
||||
ldr w23, [x2], #4
|
||||
|
||||
fmla v24.4s, v0.4s, v5.s[0]
|
||||
ldr w9, [x1], #4
|
||||
fmla v25.4s, v1.4s, v5.s[0]
|
||||
ldr w24, [x2], #4
|
||||
fmla v26.4s, v0.4s, v5.s[1]
|
||||
ldr w10, [x1], #4
|
||||
fmla v27.4s, v1.4s, v5.s[1]
|
||||
ldr w25, [x2], #4
|
||||
fmla v28.4s, v0.4s, v5.s[2]
|
||||
ldr w11, [x1], #4
|
||||
fmla v29.4s, v1.4s, v5.s[2]
|
||||
ldr w26, [x2], #4
|
||||
fmla v30.4s, v0.4s, v5.s[3]
|
||||
ldr w12, [x1], #4
|
||||
fmla v31.4s, v1.4s, v5.s[3]
|
||||
ldr w27, [x2], #4
|
||||
|
||||
ins v0.s[0], w5
|
||||
ins v4.s[0], w20
|
||||
ins v1.s[0], w9
|
||||
ins v5.s[0], w24
|
||||
ins v0.s[2], w7
|
||||
ins v4.s[2], w22
|
||||
ins v1.s[2], w11
|
||||
ins v5.s[2], w26
|
||||
ins v0.s[1], w6
|
||||
ins v4.s[1], w21
|
||||
ins v1.s[1], w10
|
||||
ins v5.s[1], w25
|
||||
ins v0.s[3], w8
|
||||
ins v4.s[3], w23
|
||||
ins v1.s[3], w12
|
||||
ins v5.s[3], w27
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
ldr w5, [x1], #4
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
ldr w20, [x2], #4
|
||||
fmla v18.4s, v0.4s, v4.s[1]
|
||||
ldr w6, [x1], #4
|
||||
fmla v19.4s, v1.4s, v4.s[1]
|
||||
ldr w21, [x2], #4
|
||||
fmla v20.4s, v0.4s, v4.s[2]
|
||||
ldr w7, [x1], #4
|
||||
fmla v21.4s, v1.4s, v4.s[2]
|
||||
ldr w22, [x2], #4
|
||||
fmla v22.4s, v0.4s, v4.s[3]
|
||||
ldr w8, [x1], #4
|
||||
fmla v23.4s, v1.4s, v4.s[3]
|
||||
ldr w23, [x2], #4
|
||||
|
||||
fmla v24.4s, v0.4s, v5.s[0]
|
||||
ldr w9, [x1], #4
|
||||
fmla v25.4s, v1.4s, v5.s[0]
|
||||
ldr w24, [x2], #4
|
||||
fmla v26.4s, v0.4s, v5.s[1]
|
||||
ldr w10, [x1], #4
|
||||
fmla v27.4s, v1.4s, v5.s[1]
|
||||
ldr w25, [x2], #4
|
||||
fmla v28.4s, v0.4s, v5.s[2]
|
||||
ldr w11, [x1], #4
|
||||
fmla v29.4s, v1.4s, v5.s[2]
|
||||
ldr w26, [x2], #4
|
||||
fmla v30.4s, v0.4s, v5.s[3]
|
||||
ldr w12, [x1], #4
|
||||
fmla v31.4s, v1.4s, v5.s[3]
|
||||
ldr w27, [x2], #4
|
||||
|
||||
prfm pldl1keep, [x1, #256]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
|
||||
ins v0.s[0], w5
|
||||
ins v4.s[0], w20
|
||||
ins v1.s[0], w9
|
||||
ins v5.s[0], w24
|
||||
ins v0.s[2], w7
|
||||
ins v4.s[2], w22
|
||||
ins v1.s[2], w11
|
||||
ins v5.s[2], w26
|
||||
ins v0.s[1], w6
|
||||
ins v4.s[1], w21
|
||||
ins v1.s[1], w10
|
||||
ins v5.s[1], w25
|
||||
ins v0.s[3], w8
|
||||
ins v4.s[3], w23
|
||||
ins v1.s[3], w12
|
||||
ins v5.s[3], w27
|
||||
|
||||
+35
@@ -0,0 +1,35 @@
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
ldr x5, [x1], #8
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
ldr x9, [x2], #8
|
||||
fmla v18.4s, v0.4s, v4.s[1]
|
||||
ldr x6, [x1], #8
|
||||
fmla v19.4s, v1.4s, v4.s[1]
|
||||
ldr x10, [x2], #8
|
||||
fmla v20.4s, v0.4s, v4.s[2]
|
||||
ldr x7, [x1], #8
|
||||
fmla v21.4s, v1.4s, v4.s[2]
|
||||
ldr x11, [x2], #8
|
||||
fmla v22.4s, v0.4s, v4.s[3]
|
||||
ldr x8, [x1], #8
|
||||
fmla v23.4s, v1.4s, v4.s[3]
|
||||
ldr x12, [x2], #8
|
||||
|
||||
fmla v24.4s, v0.4s, v5.s[0]
|
||||
fmla v25.4s, v1.4s, v5.s[0]
|
||||
fmla v26.4s, v0.4s, v5.s[1]
|
||||
fmla v27.4s, v1.4s, v5.s[1]
|
||||
fmla v28.4s, v0.4s, v5.s[2]
|
||||
fmla v29.4s, v1.4s, v5.s[2]
|
||||
fmla v30.4s, v0.4s, v5.s[3]
|
||||
fmla v31.4s, v1.4s, v5.s[3]
|
||||
|
||||
ins v2.d[0], x5
|
||||
ins v6.d[0], x9
|
||||
ins v3.d[0], x7
|
||||
ins v7.d[0], x11
|
||||
ins v2.d[1], x6
|
||||
ins v6.d[1], x10
|
||||
ins v3.d[1], x8
|
||||
ins v7.d[1], x12
|
||||
+43
@@ -0,0 +1,43 @@
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
ldr x5, [x1], #8
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
ldr x9, [x2], #8
|
||||
fmla v18.4s, v0.4s, v4.s[1]
|
||||
ldr x6, [x1], #8
|
||||
fmla v19.4s, v1.4s, v4.s[1]
|
||||
ldr x10, [x2], #8
|
||||
fmla v20.4s, v0.4s, v4.s[2]
|
||||
ldr x7, [x1], #8
|
||||
fmla v21.4s, v1.4s, v4.s[2]
|
||||
ldr x11, [x2], #8
|
||||
fmla v22.4s, v0.4s, v4.s[3]
|
||||
ldr x8, [x1], #8
|
||||
fmla v23.4s, v1.4s, v4.s[3]
|
||||
ldr x12, [x2], #8
|
||||
|
||||
fmla v24.4s, v0.4s, v5.s[0]
|
||||
prfm pldl1keep, [x1, #256]
|
||||
fmla v25.4s, v1.4s, v5.s[0]
|
||||
prfm pldl1keep, [x1, #320]
|
||||
fmla v26.4s, v0.4s, v5.s[1]
|
||||
prfm pldl1keep, [x1, #384]
|
||||
fmla v27.4s, v1.4s, v5.s[1]
|
||||
prfm pldl1keep, [x1, #448]
|
||||
fmla v28.4s, v0.4s, v5.s[2]
|
||||
prfm pldl1keep, [x2, #256]
|
||||
fmla v29.4s, v1.4s, v5.s[2]
|
||||
prfm pldl1keep, [x2, #320]
|
||||
fmla v30.4s, v0.4s, v5.s[3]
|
||||
prfm pldl1keep, [x2, #384]
|
||||
fmla v31.4s, v1.4s, v5.s[3]
|
||||
prfm pldl1keep, [x2, #448]
|
||||
|
||||
ins v0.d[0], x5
|
||||
ins v4.d[0], x9
|
||||
ins v1.d[0], x7
|
||||
ins v5.d[0], x11
|
||||
ins v0.d[1], x6
|
||||
ins v4.d[1], x10
|
||||
ins v1.d[1], x8
|
||||
ins v5.d[1], x12
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
fmla v18.4s, v0.4s, v4.s[1]
|
||||
fmla v19.4s, v1.4s, v4.s[1]
|
||||
fmla v20.4s, v0.4s, v4.s[2]
|
||||
fmla v21.4s, v1.4s, v4.s[2]
|
||||
fmla v22.4s, v0.4s, v4.s[3]
|
||||
fmla v23.4s, v1.4s, v4.s[3]
|
||||
|
||||
fmla v24.4s, v0.4s, v5.s[0]
|
||||
fmla v25.4s, v1.4s, v5.s[0]
|
||||
fmla v26.4s, v0.4s, v5.s[1]
|
||||
fmla v27.4s, v1.4s, v5.s[1]
|
||||
fmla v28.4s, v0.4s, v5.s[2]
|
||||
fmla v29.4s, v1.4s, v5.s[2]
|
||||
fmla v30.4s, v0.4s, v5.s[3]
|
||||
fmla v31.4s, v1.4s, v5.s[3]
|
||||
|
||||
ld1 {{ v0.4s, v1.4s }}, [x1], #32
|
||||
ld1 {{ v4.4s, v5.4s }}, [x2], #32
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
ld1 {{ v2.4s, v3.4s }}, [x1], #32
|
||||
ld1 {{ v6.4s, v7.4s }}, [x2], #32
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
fmla v18.4s, v0.4s, v4.s[1]
|
||||
fmla v19.4s, v1.4s, v4.s[1]
|
||||
fmla v20.4s, v0.4s, v4.s[2]
|
||||
fmla v21.4s, v1.4s, v4.s[2]
|
||||
fmla v22.4s, v0.4s, v4.s[3]
|
||||
fmla v23.4s, v1.4s, v4.s[3]
|
||||
|
||||
fmla v24.4s, v0.4s, v5.s[0]
|
||||
fmla v25.4s, v1.4s, v5.s[0]
|
||||
fmla v26.4s, v0.4s, v5.s[1]
|
||||
fmla v27.4s, v1.4s, v5.s[1]
|
||||
fmla v28.4s, v0.4s, v5.s[2]
|
||||
fmla v29.4s, v1.4s, v5.s[2]
|
||||
fmla v30.4s, v0.4s, v5.s[3]
|
||||
fmla v31.4s, v1.4s, v5.s[3]
|
||||
|
||||
ld1 {{ v0.4s, v1.4s }}, [x1], #32
|
||||
ld1 {{ v4.4s, v5.4s }}, [x2], #32
|
||||
|
||||
fmla v16.4s, v2.4s, v6.s[0]
|
||||
fmla v17.4s, v3.4s, v6.s[0]
|
||||
fmla v18.4s, v2.4s, v6.s[1]
|
||||
fmla v19.4s, v3.4s, v6.s[1]
|
||||
fmla v20.4s, v2.4s, v6.s[2]
|
||||
fmla v21.4s, v3.4s, v6.s[2]
|
||||
fmla v22.4s, v2.4s, v6.s[3]
|
||||
fmla v23.4s, v3.4s, v6.s[3]
|
||||
|
||||
fmla v24.4s, v2.4s, v7.s[0]
|
||||
fmla v25.4s, v3.4s, v7.s[0]
|
||||
fmla v26.4s, v2.4s, v7.s[1]
|
||||
fmla v27.4s, v3.4s, v7.s[1]
|
||||
fmla v28.4s, v2.4s, v7.s[2]
|
||||
fmla v29.4s, v3.4s, v7.s[2]
|
||||
fmla v30.4s, v2.4s, v7.s[3]
|
||||
fmla v31.4s, v3.4s, v7.s[3]
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
ldr d2, [x1], #8
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
ldr d6, [x2], #8
|
||||
fmla v18.4s, v0.4s, v4.s[1]
|
||||
ldr x5, [x1], #8
|
||||
fmla v19.4s, v1.4s, v4.s[1]
|
||||
ldr x7, [x2], #8
|
||||
fmla v20.4s, v0.4s, v4.s[2]
|
||||
ldr d3, [x1], #8
|
||||
fmla v21.4s, v1.4s, v4.s[2]
|
||||
ldr d7, [x2], #8
|
||||
fmla v22.4s, v0.4s, v4.s[3]
|
||||
ldr x6, [x1], #8
|
||||
fmla v23.4s, v1.4s, v4.s[3]
|
||||
ldr x8, [x2], #8
|
||||
|
||||
fmla v24.4s, v0.4s, v5.s[0]
|
||||
fmla v25.4s, v1.4s, v5.s[0]
|
||||
fmla v26.4s, v0.4s, v5.s[1]
|
||||
fmla v27.4s, v1.4s, v5.s[1]
|
||||
fmla v28.4s, v0.4s, v5.s[2]
|
||||
ins v2.d[1], x5
|
||||
fmla v29.4s, v1.4s, v5.s[2]
|
||||
ins v6.d[1], x7
|
||||
fmla v30.4s, v0.4s, v5.s[3]
|
||||
ins v3.d[1], x6
|
||||
fmla v31.4s, v1.4s, v5.s[3]
|
||||
ins v7.d[1], x8
|
||||
|
||||
fmla v16.4s, v2.4s, v6.s[0]
|
||||
ldr d0, [x1], #8
|
||||
fmla v17.4s, v3.4s, v6.s[0]
|
||||
ldr d4, [x2], #8
|
||||
fmla v18.4s, v2.4s, v6.s[1]
|
||||
ldr x5, [x1], #8
|
||||
fmla v19.4s, v3.4s, v6.s[1]
|
||||
ldr x7, [x2], #8
|
||||
fmla v20.4s, v2.4s, v6.s[2]
|
||||
ldr d1, [x1], #8
|
||||
fmla v21.4s, v3.4s, v6.s[2]
|
||||
ldr d5, [x2], #8
|
||||
fmla v22.4s, v2.4s, v6.s[3]
|
||||
ldr x6, [x1], #8
|
||||
fmla v23.4s, v3.4s, v6.s[3]
|
||||
ldr x8, [x2], #8
|
||||
|
||||
fmla v24.4s, v2.4s, v7.s[0]
|
||||
fmla v25.4s, v3.4s, v7.s[0]
|
||||
fmla v26.4s, v2.4s, v7.s[1]
|
||||
fmla v27.4s, v3.4s, v7.s[1]
|
||||
fmla v28.4s, v2.4s, v7.s[2]
|
||||
ins v0.d[1], x5
|
||||
fmla v29.4s, v3.4s, v7.s[2]
|
||||
ins v4.d[1], x7
|
||||
fmla v30.4s, v2.4s, v7.s[3]
|
||||
ins v1.d[1], x6
|
||||
fmla v31.4s, v3.4s, v7.s[3]
|
||||
ins v5.d[1], x8
|
||||
|
||||
+182
@@ -0,0 +1,182 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs: v16 to v31, (scratch)
|
||||
// - x19-x29 to preserve (but x19, x28, x29 not used)
|
||||
// - d8..d15 to preserve
|
||||
// - v16 to v31, no need to preserve
|
||||
//
|
||||
// v16[0] v18[0] v20[0] v22[0] v24[0] v26[0] v28[0] v30[0]
|
||||
// v16[1] v18[1]
|
||||
// v16[2] v18[2]
|
||||
// v16[3] v18[3]
|
||||
//
|
||||
// v17[0] v19[0] v21[0] v23[0] v25[0] v27[0] v29[0] v31[0]
|
||||
// v17[1] v19[1]
|
||||
// v17[2] v19[2]
|
||||
// v17[3] v19[3]
|
||||
|
||||
// v0-v7 (scratch registers)
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd
|
||||
.global {{G}}arm64simd_mmm_f32_8x8_{{core}}_{{suffix}}
|
||||
{{G}}arm64simd_mmm_f32_8x8_{{core}}_{{suffix}}:
|
||||
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.packed_packed:
|
||||
ld1 { v0.4s, v1.4s }, [ x1 ], #32
|
||||
ld1 { v4.4s, v5.4s }, [ x2 ], #32
|
||||
|
||||
{% set packed_packed_loop1 %}
|
||||
{% if core == "a53" %}
|
||||
{% include "arm64simd_mmm_f32_8x8/packed_packed_loop1/ldr_x_preload.S.raw" %}
|
||||
{% else %}
|
||||
{% include "arm64simd_mmm_f32_8x8/packed_packed_loop1/naive.S.raw" %}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
{% set packed_packed_loop2 %}
|
||||
{% if core == "a55" %}
|
||||
{% include "arm64simd_mmm_f32_8x8/packed_packed_loop2/cortex_a55.S.raw" %}
|
||||
{% else %}
|
||||
{{ packed_packed_loop1 }}
|
||||
{{ packed_packed_loop1 }}
|
||||
{% endif %}
|
||||
{% endset %}
|
||||
|
||||
cmp x3, #4
|
||||
blt .packed_packed_loop_1
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_4:
|
||||
{{ packed_packed_loop2 }}
|
||||
{{ packed_packed_loop2 }}
|
||||
|
||||
sub x3, x3, #4
|
||||
cmp x3, #4
|
||||
bge .packed_packed_loop_4
|
||||
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
.p2align 4
|
||||
.packed_packed_loop_1:
|
||||
{{ packed_packed_loop1 }}
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_scalars.j2" %}
|
||||
{% set mr = 8 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_rows.j2" %}
|
||||
{% set mr = 8 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_f32_per_cols.j2" %}
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.s}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
fadd v{{ col * 2 + reg }}.4s, v{{ col * 2 + reg }}.4s, v0.4s
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x2, [x0, #8]
|
||||
ldr x3, [x0, #16]
|
||||
|
||||
ld1 { v0.4s, v1.4s }, [ x2 ], #32
|
||||
ld1 { v4.4s, v5.4s }, [ x3 ], #32
|
||||
|
||||
fmla v16.4s, v0.4s, v4.s[0]
|
||||
fmla v17.4s, v1.4s, v4.s[0]
|
||||
fmla v18.4s, v0.4s, v4.s[1]
|
||||
fmla v19.4s, v1.4s, v4.s[1]
|
||||
fmla v20.4s, v0.4s, v4.s[2]
|
||||
fmla v21.4s, v1.4s, v4.s[2]
|
||||
fmla v22.4s, v0.4s, v4.s[3]
|
||||
fmla v23.4s, v1.4s, v4.s[3]
|
||||
|
||||
fmla v24.4s, v0.4s, v5.s[0]
|
||||
fmla v25.4s, v1.4s, v5.s[0]
|
||||
fmla v26.4s, v0.4s, v5.s[1]
|
||||
fmla v27.4s, v1.4s, v5.s[1]
|
||||
fmla v28.4s, v0.4s, v5.s[2]
|
||||
fmla v29.4s, v1.4s, v5.s[2]
|
||||
fmla v30.4s, v0.4s, v5.s[3]
|
||||
fmla v31.4s, v1.4s, v5.s[3]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x6, #4
|
||||
bne .store_strides_generic
|
||||
|
||||
{% for col in range(8, 16) %}
|
||||
str q{{ col * 2 }}, [ x5 ]
|
||||
str q{{ col * 2 + 1 }}, [ x5, #16 ]
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_generic:
|
||||
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{ col * 2 + reg }}.s }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
// vim: ft=arm
|
||||
|
||||
{% from "arm64simd_mmm_4s_ops.j2" import per_col %}
|
||||
|
||||
{{ per_col(label="per_col_min", op="fmin", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_max", op="fmax", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_mul", op="fmul", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_add", op="fadd", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_sub", op="fsub", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_sub_flipped", op="fsub", mr=mr, from=from, to=to, flipped=true) }}
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
// vim: ft=arm
|
||||
|
||||
{% from "arm64simd_mmm_4s_ops.j2" import per_row %}
|
||||
|
||||
{{ per_row(label="per_row_min", op="fmin", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_max", op="fmax", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_mul", op="fmul", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_add", op="fadd", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_sub", op="fsub", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_sub_flipped", op="fsub", mr=mr, from=from, to=to, flipped=true) }}
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
// vim: ft=arm
|
||||
|
||||
{% from "arm64simd_mmm_4s_ops.j2" import scalar %}
|
||||
|
||||
{{ scalar(label="scalar_min", op="fmin", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_max", op="fmax", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_mul", op="fmul", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_add", op="fadd", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_sub", op="fsub", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_sub_flipped", op="fsub", from=from, to=to, flipped=true) }}
|
||||
|
||||
|
||||
.clear:
|
||||
{% for r in range(from, to + 1) %}
|
||||
eor v{{r}}.8b, v{{r}}.8b, v{{r}}.8b
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.leaky_relu:
|
||||
add x2, x0, #8
|
||||
ld1 {v4.s}[0], [ x2 ]
|
||||
dup v4.4s, v4.s[0]
|
||||
|
||||
// bsl cond/dst, then, else
|
||||
// fcmge dst, src, #0.0
|
||||
{% for r in range(from, to + 1) %}
|
||||
fmul v0.4s, v{{r}}.4s, v4.4s
|
||||
fcmge v1.4s, v{{r}}.4s, #0.0
|
||||
bsl v1.16b, v{{r}}.16b, v0.16b
|
||||
and v{{r}}.16b, v1.16b, v1.16b
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.q_scale:
|
||||
.q_shl:
|
||||
.q_shr:
|
||||
b .unsupported
|
||||
+180
@@ -0,0 +1,180 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs:
|
||||
// - x19-x29 to preserve (but x19, x28, x29 not used)
|
||||
// - d8..d15 to preserve
|
||||
// - v16 to v31, no need to preserve
|
||||
|
||||
// no preservation either for v0-v7...
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd
|
||||
.global {{G}}arm64simd_mmm_i32_64x1_{{suffix}}
|
||||
{{G}}arm64simd_mmm_i32_64x1_{{suffix}}:
|
||||
|
||||
/*
|
||||
prfm pldl1keep, [x1]
|
||||
prfm pldl1keep, [x2]
|
||||
*/
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldp x2, x4, [x0, #24] // b, packing
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
cmp x4, #1
|
||||
beq .packed_packed_loop_1_i8i8
|
||||
|
||||
.packed_packed_loop_1:
|
||||
ld1 {v9.s}[0], [ x2 ], 4
|
||||
|
||||
ld1 { v0.4s-v3.4s }, [ x1 ], #64
|
||||
ld1 { v4.4s-v7.4s }, [ x1 ], #64
|
||||
{% for reg in range(0, 4) %}
|
||||
mla v{{ reg * 2 + 16 }}.4s, v{{ reg * 2 }}.4s, v9.s[0]
|
||||
mla v{{ reg * 2 + 17 }}.4s, v{{ reg * 2 + 1 }}.4s, v9.s[0]
|
||||
{% endfor %}
|
||||
|
||||
ld1 { v0.4s-v3.4s }, [ x1 ], #64
|
||||
ld1 { v4.4s-v7.4s }, [ x1 ], #64
|
||||
{% for reg in range(0, 4) %}
|
||||
mla v{{ reg * 2 + 24 }}.4s, v{{ reg * 2 }}.4s, v9.s[0]
|
||||
mla v{{ reg * 2 + 25 }}.4s, v{{ reg * 2 + 1 }}.4s, v9.s[0]
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.packed_packed_loop_1_i8i8:
|
||||
ld1 {v9.b}[0], [ x2 ], 1
|
||||
sshll v9.8h, v9.8b, 0
|
||||
|
||||
ld1 { v0.8b-v3.8b }, [ x1 ], #32
|
||||
ld1 { v4.8b-v7.8b }, [ x1 ], #32
|
||||
|
||||
{% for reg in range(0, 8) %}
|
||||
sshll v10.8h, v{{reg}}.8b, 0
|
||||
smlal v{{ reg * 2 + 16 }}.4s, v10.4h, v9.h[0]
|
||||
smlal2 v{{ reg * 2 + 17 }}.4s, v10.8h, v9.h[0]
|
||||
{% endfor %}
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1_i8i8
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
cmp x8, #4
|
||||
beq non_linear_addc_i32
|
||||
|
||||
{% for reg in range(16, 32) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.b}[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
sshll v0.8h, v0.8b, 0
|
||||
sshll v0.4s, v0.4h, 0
|
||||
add v{{reg}}.4s, v{{reg}}.4s, v0.4s
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
non_linear_addc_i32:
|
||||
{% for reg in range(16, 32) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.s}[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
add v{{reg}}.4s, v{{reg}}.4s, v0.4s
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x2, [x0, #8]
|
||||
ldr x3, [x0, #16]
|
||||
|
||||
ld1 { v15.s }[0], [ x3 ]
|
||||
xtn v15.4h, v15.4s
|
||||
|
||||
ld1 { v0.4s-v3.4s }, [ x2 ], #64
|
||||
ld1 { v4.4s-v7.4s }, [ x2 ], #64
|
||||
|
||||
{% for reg in range(0, 8) %}
|
||||
xtn v{{reg}}.4h, v{{reg}}.4s
|
||||
smlal v{{ reg + 16 }}.4s, v{{reg}}.4h, v15.h[0]
|
||||
{% endfor %}
|
||||
|
||||
ld1 { v0.4s-v3.4s }, [ x2 ], #64
|
||||
ld1 { v4.4s-v7.4s }, [ x2 ], #64
|
||||
|
||||
{% for reg in range(0, 8) %}
|
||||
xtn v{{reg}}.4h, v{{reg}}.4s
|
||||
smlal v{{ reg + 24 }}.4s, v{{reg}}.4h, v15.h[0]
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_i32_scalars.j2" %}
|
||||
{% set mr = 64 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_i32_per_rows.j2" %}
|
||||
{% set mr = 64 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_i32_per_cols.j2" %}
|
||||
{% include "arm64simd_mmm_i32_scale_q16_q31.j2" %}
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x8, #4
|
||||
beq .store_strides_i32
|
||||
|
||||
{% for reg in range(16, 32) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{reg}}.b }[{{ lane * 4 }}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_i32:
|
||||
{% for reg in range(16, 32) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{reg}}.s }[{{lane}}], [ x5 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+234
@@ -0,0 +1,234 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs:
|
||||
// - x19-x29 to preserve (but x19, x28, x29 not used)
|
||||
// - d8..d15 to preserve
|
||||
// - v16 to v31, no need to preserve
|
||||
//
|
||||
// v16[0] v18[0] v20[0] v22[0] v24[0] v26[0] v28[0] v30[0]
|
||||
// v16[1] v18[1]
|
||||
// v16[2] v18[2]
|
||||
// v16[3] v18[3]
|
||||
//
|
||||
// v17[0] v19[0] v21[0] v23[0] v25[0] v27[0] v29[0] v31[0]
|
||||
// v17[1] v19[1]
|
||||
// v17[2] v19[2]
|
||||
// v17[3] v19[3]
|
||||
|
||||
// no preservation either for v0-v7...
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd
|
||||
.global {{G}}arm64simd_mmm_i32_8x8_{{suffix}}
|
||||
{{G}}arm64simd_mmm_i32_8x8_{{suffix}}:
|
||||
|
||||
/*
|
||||
prfm pldl1keep, [x1]
|
||||
prfm pldl1keep, [x2]
|
||||
*/
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldp x2, x4, [x0, #24] // b, packing
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
cmp x4, #1
|
||||
beq .packed_packed_loop_1_i8i8
|
||||
|
||||
.packed_packed_loop_1:
|
||||
|
||||
ld1 { v0.4s, v1.4s }, [ x1 ], #32
|
||||
ld1 { v4.4s, v5.4s }, [ x2 ], #32
|
||||
|
||||
mla v16.4s, v0.4s, v4.s[0]
|
||||
mla v17.4s, v1.4s, v4.s[0]
|
||||
mla v18.4s, v0.4s, v4.s[1]
|
||||
mla v19.4s, v1.4s, v4.s[1]
|
||||
|
||||
mla v20.4s, v0.4s, v4.s[2]
|
||||
mla v21.4s, v1.4s, v4.s[2]
|
||||
mla v22.4s, v0.4s, v4.s[3]
|
||||
mla v23.4s, v1.4s, v4.s[3]
|
||||
|
||||
mla v24.4s, v0.4s, v5.s[0]
|
||||
mla v25.4s, v1.4s, v5.s[0]
|
||||
mla v26.4s, v0.4s, v5.s[1]
|
||||
mla v27.4s, v1.4s, v5.s[1]
|
||||
|
||||
mla v28.4s, v0.4s, v5.s[2]
|
||||
mla v29.4s, v1.4s, v5.s[2]
|
||||
mla v30.4s, v0.4s, v5.s[3]
|
||||
mla v31.4s, v1.4s, v5.s[3]
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.packed_packed_loop_1_i8i8:
|
||||
|
||||
ld1 { v0.8b }, [ x1 ], #8
|
||||
sshll v0.8h, v0.8b, 0
|
||||
ld1 { v4.8b }, [ x2 ], #8
|
||||
sshll v4.8h, v4.8b, 0
|
||||
|
||||
smlal v16.4s, v0.4h, v4.h[0]
|
||||
smlal2 v17.4s, v0.8h, v4.h[0]
|
||||
smlal v18.4s, v0.4h, v4.h[1]
|
||||
smlal2 v19.4s, v0.8h, v4.h[1]
|
||||
smlal v20.4s, v0.4h, v4.h[2]
|
||||
smlal2 v21.4s, v0.8h, v4.h[2]
|
||||
smlal v22.4s, v0.4h, v4.h[3]
|
||||
smlal2 v23.4s, v0.8h, v4.h[3]
|
||||
|
||||
smlal v24.4s, v0.4h, v4.h[4]
|
||||
smlal2 v25.4s, v0.8h, v4.h[4]
|
||||
smlal v26.4s, v0.4h, v4.h[5]
|
||||
smlal2 v27.4s, v0.8h, v4.h[5]
|
||||
smlal v28.4s, v0.4h, v4.h[6]
|
||||
smlal2 v29.4s, v0.8h, v4.h[6]
|
||||
smlal v30.4s, v0.4h, v4.h[7]
|
||||
smlal2 v31.4s, v0.8h, v4.h[7]
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1_i8i8
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_i32_scalars.j2" %}
|
||||
{% set mr = 8 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_i32_per_rows.j2" %}
|
||||
{% set mr = 8 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_i32_per_cols.j2" %}
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
cmp x8, #4
|
||||
beq non_linear_addc_i32
|
||||
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.b}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
sshll v0.8h, v0.8b, 0
|
||||
sshll v0.4s, v0.4h, 0
|
||||
add v{{ col * 2 + reg }}.4s, v{{ col * 2 + reg }}.4s, v0.4s
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
non_linear_addc_i32:
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.s}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
add v{{ col * 2 + reg }}.4s, v{{ col * 2 + reg }}.4s, v0.4s
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x2, [x0, #8]
|
||||
ldr x3, [x0, #16]
|
||||
|
||||
ld1 { v0.4s, v1.4s }, [ x2 ]
|
||||
ld1 { v4.4s, v5.4s }, [ x3 ]
|
||||
|
||||
xtn v0.4h, v0.4s
|
||||
xtn v1.4h, v1.4s
|
||||
xtn v4.4h, v4.4s
|
||||
xtn v5.4h, v5.4s
|
||||
|
||||
smlal v16.4s, v0.4h, v4.h[0]
|
||||
smlal v17.4s, v1.4h, v4.h[0]
|
||||
smlal v18.4s, v0.4h, v4.h[1]
|
||||
smlal v19.4s, v1.4h, v4.h[1]
|
||||
smlal v20.4s, v0.4h, v4.h[2]
|
||||
smlal v21.4s, v1.4h, v4.h[2]
|
||||
smlal v22.4s, v0.4h, v4.h[3]
|
||||
smlal v23.4s, v1.4h, v4.h[3]
|
||||
|
||||
smlal v24.4s, v0.4h, v5.h[0]
|
||||
smlal v25.4s, v1.4h, v5.h[0]
|
||||
smlal v26.4s, v0.4h, v5.h[1]
|
||||
smlal v27.4s, v1.4h, v5.h[1]
|
||||
smlal v28.4s, v0.4h, v5.h[2]
|
||||
smlal v29.4s, v1.4h, v5.h[2]
|
||||
smlal v30.4s, v0.4h, v5.h[3]
|
||||
smlal v31.4s, v1.4h, v5.h[3]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% include "arm64simd_mmm_i32_scale_q16_q31.j2" %}
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x8, #4
|
||||
beq .store_strides_i32
|
||||
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{ col * 2 + reg }}.b }[{{ lane * 4 }}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_i32:
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{ col * 2 + reg }}.s }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+235
@@ -0,0 +1,235 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// C tile regs:
|
||||
// - x19-x29 to preserve (but x19, x28, x29 not used)
|
||||
// - d8..d15 to preserve
|
||||
// - v16 to v31, no need to preserve
|
||||
//
|
||||
// v16[0] v18[0] v20[0] v22[0] v24[0] v26[0] v28[0] v30[0]
|
||||
// v16[1] v18[1]
|
||||
// v16[2] v18[2]
|
||||
// v16[3] v18[3]
|
||||
//
|
||||
// v17[0] v19[0] v21[0] v23[0] v25[0] v27[0] v29[0] v31[0]
|
||||
// v17[1] v19[1]
|
||||
// v17[2] v19[2]
|
||||
// v17[3] v19[3]
|
||||
|
||||
// no preservation either for v0-v7...
|
||||
// packed A buffering (2x8 values): alternating v0, v1 with v2, v3
|
||||
// packed B buffering (2x8 values): alternating v4, v5 with v6, v7
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd+dotprod
|
||||
.global {{G}}arm64simd_mmm_i32_8x8_dot_{{suffix}}
|
||||
{{G}}arm64simd_mmm_i32_8x8_dot_{{suffix}}:
|
||||
|
||||
/*
|
||||
prfm pldl1keep, [x1]
|
||||
prfm pldl1keep, [x2]
|
||||
*/
|
||||
stp x20, x21, [sp, #-16]!
|
||||
stp x22, x23, [sp, #-16]!
|
||||
stp x24, x25, [sp, #-16]!
|
||||
stp x26, x27, [sp, #-16]!
|
||||
|
||||
stp d8, d9, [sp, #-16]!
|
||||
stp d10, d11, [sp, #-16]!
|
||||
stp d12, d13, [sp, #-16]!
|
||||
stp d14, d15, [sp, #-16]!
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
.add_mat_mul:
|
||||
ldp x2, x4, [x0, #24] // b, packing
|
||||
ldp x3, x1, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
beq .non_linear_loop
|
||||
|
||||
cmp x4, #1
|
||||
beq .packed_packed_loop_1_i8i8
|
||||
|
||||
.packed_packed_loop_1:
|
||||
|
||||
ld1 { v0.4s, v1.4s }, [ x1 ], #32
|
||||
ld1 { v4.4s, v5.4s }, [ x2 ], #32
|
||||
|
||||
mla v16.4s, v0.4s, v4.s[0]
|
||||
mla v17.4s, v1.4s, v4.s[0]
|
||||
mla v18.4s, v0.4s, v4.s[1]
|
||||
mla v19.4s, v1.4s, v4.s[1]
|
||||
|
||||
mla v20.4s, v0.4s, v4.s[2]
|
||||
mla v21.4s, v1.4s, v4.s[2]
|
||||
mla v22.4s, v0.4s, v4.s[3]
|
||||
mla v23.4s, v1.4s, v4.s[3]
|
||||
|
||||
mla v24.4s, v0.4s, v5.s[0]
|
||||
mla v25.4s, v1.4s, v5.s[0]
|
||||
mla v26.4s, v0.4s, v5.s[1]
|
||||
mla v27.4s, v1.4s, v5.s[1]
|
||||
|
||||
mla v28.4s, v0.4s, v5.s[2]
|
||||
mla v29.4s, v1.4s, v5.s[2]
|
||||
mla v30.4s, v0.4s, v5.s[3]
|
||||
mla v31.4s, v1.4s, v5.s[3]
|
||||
|
||||
subs x3, x3, #1
|
||||
bne .packed_packed_loop_1
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.packed_packed_loop_1_i8i8:
|
||||
// PackedI8K4 (K=4-inner, r=8): per 4-K block, A is m0-3 (v0) / m4-7 (v1),
|
||||
// B is n0-3 (v4) / n4-7 (v5), each lane a 4xi8 group. SDOT by-element dots
|
||||
// a B column's 4 K against all 4 m rows of an A half. Same v16..v31 tile
|
||||
// layout as the SMLAL kernel: v[16 + n*2 + m_half] = C[m_half*4..][n].
|
||||
ld1 { v0.16b, v1.16b }, [ x1 ], #32
|
||||
ld1 { v4.16b, v5.16b }, [ x2 ], #32
|
||||
|
||||
sdot v16.4s, v0.16b, v4.4b[0]
|
||||
sdot v17.4s, v1.16b, v4.4b[0]
|
||||
sdot v18.4s, v0.16b, v4.4b[1]
|
||||
sdot v19.4s, v1.16b, v4.4b[1]
|
||||
sdot v20.4s, v0.16b, v4.4b[2]
|
||||
sdot v21.4s, v1.16b, v4.4b[2]
|
||||
sdot v22.4s, v0.16b, v4.4b[3]
|
||||
sdot v23.4s, v1.16b, v4.4b[3]
|
||||
|
||||
sdot v24.4s, v0.16b, v5.4b[0]
|
||||
sdot v25.4s, v1.16b, v5.4b[0]
|
||||
sdot v26.4s, v0.16b, v5.4b[1]
|
||||
sdot v27.4s, v1.16b, v5.4b[1]
|
||||
sdot v28.4s, v0.16b, v5.4b[2]
|
||||
sdot v29.4s, v1.16b, v5.4b[2]
|
||||
sdot v30.4s, v0.16b, v5.4b[3]
|
||||
sdot v31.4s, v1.16b, v5.4b[3]
|
||||
|
||||
subs x3, x3, #4
|
||||
bgt .packed_packed_loop_1_i8i8
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_i32_scalars.j2" %}
|
||||
{% set mr = 8 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_i32_per_rows.j2" %}
|
||||
{% set mr = 8 %}{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_i32_per_cols.j2" %}
|
||||
{% set from = 16 %}{% set to = 31 %}{% include "arm64simd_mmm_load_tile.j2" %}
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8]
|
||||
ldp x7, x8, [x0, #24]
|
||||
|
||||
cmp x8, #4
|
||||
beq non_linear_addc_i32
|
||||
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.b}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
sshll v0.8h, v0.8b, 0
|
||||
sshll v0.4s, v0.4h, 0
|
||||
add v{{ col * 2 + reg }}.4s, v{{ col * 2 + reg }}.4s, v0.4s
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
non_linear_addc_i32:
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
ld1 {v0.s}[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
add v{{ col * 2 + reg }}.4s, v{{ col * 2 + reg }}.4s, v0.4s
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.add_row_col_products:
|
||||
ldr x2, [x0, #8]
|
||||
ldr x3, [x0, #16]
|
||||
|
||||
ld1 { v0.4s, v1.4s }, [ x2 ]
|
||||
ld1 { v4.4s, v5.4s }, [ x3 ]
|
||||
|
||||
xtn v0.4h, v0.4s
|
||||
xtn v1.4h, v1.4s
|
||||
xtn v4.4h, v4.4s
|
||||
xtn v5.4h, v5.4s
|
||||
|
||||
smlal v16.4s, v0.4h, v4.h[0]
|
||||
smlal v17.4s, v1.4h, v4.h[0]
|
||||
smlal v18.4s, v0.4h, v4.h[1]
|
||||
smlal v19.4s, v1.4h, v4.h[1]
|
||||
smlal v20.4s, v0.4h, v4.h[2]
|
||||
smlal v21.4s, v1.4h, v4.h[2]
|
||||
smlal v22.4s, v0.4h, v4.h[3]
|
||||
smlal v23.4s, v1.4h, v4.h[3]
|
||||
|
||||
smlal v24.4s, v0.4h, v5.h[0]
|
||||
smlal v25.4s, v1.4h, v5.h[0]
|
||||
smlal v26.4s, v0.4h, v5.h[1]
|
||||
smlal v27.4s, v1.4h, v5.h[1]
|
||||
smlal v28.4s, v0.4h, v5.h[2]
|
||||
smlal v29.4s, v1.4h, v5.h[2]
|
||||
smlal v30.4s, v0.4h, v5.h[3]
|
||||
smlal v31.4s, v1.4h, v5.h[3]
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
{% include "arm64simd_mmm_i32_scale_q16_q31.j2" %}
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // c base ptr, rsc
|
||||
ldp x7, x8, [x0, #24] // csc, item_size
|
||||
|
||||
cmp x8, #4
|
||||
beq .store_strides_i32
|
||||
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{ col * 2 + reg }}.b }[{{ lane * 4 }}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.store_strides_i32:
|
||||
{% for col in range(8, 16) %}
|
||||
mov x4, x5
|
||||
{% for reg in range(0, 2) %}
|
||||
{% for lane in range(0, 4) %}
|
||||
st1 { v{{ col * 2 + reg }}.s }[{{lane}}], [ x4 ], x6
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
add x5, x5, x7
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.return:
|
||||
ldp d14, d15, [sp], #16
|
||||
ldp d12, d13, [sp], #16
|
||||
ldp d10, d11, [sp], #16
|
||||
ldp d8, d9, [sp], #16
|
||||
|
||||
ldp x26, x27, [sp], #16
|
||||
ldp x24, x25, [sp], #16
|
||||
ldp x22, x23, [sp], #16
|
||||
ldp x20, x21, [sp], #16
|
||||
|
||||
ret
|
||||
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
// vim: ft=arm
|
||||
|
||||
{% from "arm64simd_mmm_4s_ops.j2" import per_col %}
|
||||
|
||||
{{ per_col(label="per_col_min", op="smin", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_max", op="smax", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_mul", op="mul", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_add", op="add", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_sub", op="sub", mr=mr, from=from, to=to) }}
|
||||
{{ per_col(label="per_col_sub_flipped", op="sub", mr=mr, from=from, to=to, flipped=true) }}
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
// vim: ft=arm
|
||||
|
||||
{% from "arm64simd_mmm_4s_ops.j2" import per_row %}
|
||||
|
||||
{{ per_row(label="per_row_min", op="smin", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_max", op="smax", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_mul", op="mul", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_add", op="add", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_sub", op="sub", mr=mr, from=from, to=to) }}
|
||||
{{ per_row(label="per_row_sub_flipped", op="sub", mr=mr, from=from, to=to, flipped=true) }}
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
// vim: ft=arm
|
||||
|
||||
{% from "arm64simd_mmm_4s_ops.j2" import scalar %}
|
||||
|
||||
{{ scalar(label="scalar_min", op="smin", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_max", op="smax", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_mul", op="mul", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_add", op="add", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_sub", op="sub", from=from, to=to) }}
|
||||
{{ scalar(label="scalar_sub_flipped", op="sub", from=from, to=to, flipped=true) }}
|
||||
|
||||
.clear:
|
||||
{% for r in range(from, to + 1) %}
|
||||
eor v{{r}}.8b, v{{r}}.8b, v{{r}}.8b
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.leaky_relu:
|
||||
add x2, x0, #8
|
||||
ld1 {v4.s}[0], [ x2 ]
|
||||
dup v4.4s, v4.s[0]
|
||||
|
||||
// bsl cond/dst, then, else
|
||||
// fcmge dst, src, #0.0
|
||||
{% for r in range(from, to + 1) %}
|
||||
mul v0.4s, v{{r}}.4s, v4.4s
|
||||
cmge v1.4s, v{{r}}.4s, #0
|
||||
bsl v1.16b, v{{r}}.16b, v0.16b
|
||||
and v{{r}}.16b, v1.16b, v1.16b
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
+267
@@ -0,0 +1,267 @@
|
||||
|
||||
// vim: ft=arm
|
||||
|
||||
.q_scale:
|
||||
ldp x5, x6, [x0, #8] // x5: shift, x6: policy
|
||||
add x2, x0, #24
|
||||
ld1r { v2.4s }, [x2] // v2.4s <- multiplier
|
||||
|
||||
mov w3, #1
|
||||
ins v4.d[0], x3
|
||||
dup v4.2d, v4.d[0] // v4.2d <- 1
|
||||
|
||||
add x5, x5, #32 // add 32 to shift
|
||||
neg x5, x5 // broadcast shift
|
||||
ins v1.d[0], x5
|
||||
dup v1.2d, v1.d[0] // v1.2s <- -(shift + 32)
|
||||
|
||||
cmp x6, 1
|
||||
beq .q_scale_rounding_zero
|
||||
cmp x6, 2
|
||||
beq .q_scale_rounding_away
|
||||
cmp x6, 3
|
||||
beq .q_scale_rounding_minus_inf
|
||||
cmp x6, 4
|
||||
beq .q_scale_rounding_plus_inf
|
||||
cmp x6, 5
|
||||
beq .q_scale_rounding_even
|
||||
cmp x6, 6
|
||||
beq .q_scale_rounding_odd
|
||||
|
||||
b .unsupported
|
||||
|
||||
.q_scale_rounding_zero:
|
||||
// rust: signum * ((abs + nudge2) >> shift
|
||||
// asm: signum * (2*abs - 1) >>r (shift + 1)
|
||||
|
||||
{% for q in range(16, 32) %}
|
||||
cmlt v0.4s, v{{q}}.4s, #0
|
||||
abs v{{q}}.4s, v{{q}}.4s
|
||||
sqdmull v8.2d, v{{q}}.2s, v2.2s
|
||||
sqdmull2 v9.2d, v{{q}}.4s, v2.4s //mul without shift and store results in v8 and v9
|
||||
|
||||
sub v8.2d, v8.2d, v4.2d
|
||||
sqrshl v8.2d, v8.2d, v1.2d
|
||||
|
||||
sub v9.2d, v9.2d, v4.2d
|
||||
sqrshl v9.2d, v9.2d, v1.2d
|
||||
|
||||
uzp1 v{{q}}.4s, v8.4s, v9.4s //combine back
|
||||
|
||||
neg v3.4s, v{{q}}.4s
|
||||
bit v{{q}}.16b, v3.16b, v0.16b
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.q_scale_rounding_away: // signum * (abs >> (shift-1) + 1 >> 1)
|
||||
|
||||
{% for q in range(16, 32) %}
|
||||
cmlt v0.4s, v{{q}}.4s, #0
|
||||
abs v{{q}}.4s, v{{q}}.4s
|
||||
sqdmull v8.2d, v{{q}}.2s, v2.2s
|
||||
sqdmull2 v9.2d, v{{q}}.4s, v2.4s //mul without shift and store results in v8 and v9
|
||||
|
||||
sqrshl v8.2d, v8.2d, v1.2d
|
||||
sqrshl v9.2d, v9.2d, v1.2d
|
||||
|
||||
uzp1 v{{q}}.4s, v8.4s, v9.4s //combine back
|
||||
|
||||
neg v3.4s, v{{q}}.4s
|
||||
bit v{{q}}.16b, v3.16b, v0.16b
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.q_scale_rounding_minus_inf: // val >> shift
|
||||
|
||||
{% for q in range(16, 32) %}
|
||||
sqdmull v8.2d, v{{q}}.2s, v2.2s
|
||||
sqdmull2 v9.2d, v{{q}}.4s, v2.4s //mul without shift and store results in v8 and v9
|
||||
|
||||
sub v8.2d, v8.2d, v4.2d
|
||||
sqrshl v8.2d, v8.2d, v1.2d
|
||||
|
||||
sub v9.2d, v9.2d, v4.2d
|
||||
sqrshl v9.2d, v9.2d, v1.2d
|
||||
|
||||
uzp1 v{{q}}.4s, v8.4s, v9.4s //combine back
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.q_scale_rounding_plus_inf: // (val >> shift-1)+1 >>1
|
||||
|
||||
{% for q in range(16, 32) %}
|
||||
sqdmull v8.2d, v{{q}}.2s, v2.2s
|
||||
sqdmull2 v9.2d, v{{q}}.4s, v2.4s //mul without shift and store results in v8 and v9
|
||||
|
||||
sqrshl v8.2d, v8.2d, v1.2d
|
||||
sqrshl v9.2d, v9.2d, v1.2d
|
||||
|
||||
uzp1 v{{q}}.4s, v8.4s, v9.4s //combine back
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.q_scale_rounding_even: // signum * ((abs >> shift-1) + (abs & 0x1) - 1 >> 1)
|
||||
|
||||
{% for q in range(16, 32) %}
|
||||
cmlt v0.4s, v{{q}}.4s, #0
|
||||
abs v{{q}}.4s, v{{q}}.4s
|
||||
sqdmull v8.2d, v{{q}}.2s, v2.2s
|
||||
sqdmull2 v9.2d, v{{q}}.4s, v2.4s //mul without shift and store results in v8 and v9
|
||||
|
||||
sqshl v3.2d, v8.2d, v1.2d // abs >> shift - 1
|
||||
and v3.16b, v3.16b, v4.16b // abs & 0x1
|
||||
sub v3.2d, v3.2d, v4.2d //nudge : -1 if we want to round down, 0 if up
|
||||
|
||||
add v8.2d, v8.2d, v3.2d
|
||||
sqrshl v8.2d, v8.2d, v1.2d
|
||||
|
||||
sqshl v3.2d, v9.2d, v1.2d
|
||||
and v3.16b, v3.16b, v4.16b
|
||||
sub v3.2d, v3.2d, v4.2d //nudge : -1 if we want to round down, 0 if up
|
||||
|
||||
add v9.2d, v9.2d, v3.2d
|
||||
sqrshl v9.2d, v9.2d, v1.2d
|
||||
|
||||
uzp1 v{{q}}.4s, v8.4s, v9.4s //combine back
|
||||
|
||||
neg v3.4s, v{{q}}.4s
|
||||
bit v{{q}}.16b, v3.16b, v0.16b
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.q_scale_rounding_odd: // signum * ((abs >> shift-1) - (abs & 0x1) >> 1)
|
||||
|
||||
{% for q in range(16, 32) %}
|
||||
cmlt v0.4s, v{{q}}.4s, #0
|
||||
abs v{{q}}.4s, v{{q}}.4s
|
||||
sqdmull v8.2d, v{{q}}.2s, v2.2s
|
||||
sqdmull2 v9.2d, v{{q}}.4s, v2.4s //mul without shift and store results in v8 and v9
|
||||
|
||||
sqshl v3.2d, v8.2d, v1.2d
|
||||
and v3.16b, v3.16b, v4.16b //nudge : -1 if we want to round down, 0 if up
|
||||
|
||||
sub v8.2d, v8.2d, v3.2d
|
||||
sqrshl v8.2d, v8.2d, v1.2d
|
||||
|
||||
sqshl v3.2d, v9.2d, v1.2d
|
||||
and v3.16b, v3.16b, v4.16b //nudge : -1 if we want to round down, 0 if up
|
||||
|
||||
sub v9.2d, v9.2d, v3.2d
|
||||
sqrshl v9.2d, v9.2d, v1.2d
|
||||
|
||||
uzp1 v{{q}}.4s, v8.4s, v9.4s //combine back
|
||||
|
||||
neg v3.4s, v{{q}}.4s
|
||||
bit v{{q}}.16b, v3.16b, v0.16b
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
.q_shl:
|
||||
ldr x5, [x0, #8] // x5: shift
|
||||
ins v1.s[0], w5
|
||||
dup v1.4s, v1.s[0] // v1.4s <- shift
|
||||
|
||||
{% for q in range(16, 32) %}
|
||||
sqrshl v{{q}}.4s, v{{q}}.4s, v1.4s
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.q_shr:
|
||||
ldp x5, x6, [x0, #8] // x5: shift, x6: policy
|
||||
|
||||
mov w3, #1
|
||||
ins v4.s[0], w3
|
||||
dup v4.4s, v4.s[0] // v4.4d <- 1
|
||||
|
||||
neg w5, w5 // broadcast shift
|
||||
ins v1.s[0], w5
|
||||
dup v1.4s, v1.s[0] // v1.4s <- -shift
|
||||
|
||||
cmp x6, 1
|
||||
beq .q_shr_rounding_zero
|
||||
cmp x6, 2
|
||||
beq .q_shr_rounding_away
|
||||
cmp x6, 3
|
||||
beq .q_shr_rounding_minus_inf
|
||||
cmp x6, 4
|
||||
beq .q_shr_rounding_plus_inf
|
||||
cmp x6, 5
|
||||
beq .q_shr_rounding_even
|
||||
cmp x6, 6
|
||||
beq .q_shr_rounding_odd
|
||||
|
||||
b .unsupported
|
||||
|
||||
.q_shr_rounding_zero:
|
||||
// asm: signum * (abs >>r shift)
|
||||
{% for q in range(16, 32) %}
|
||||
cmlt v0.4s, v{{q}}.4s, #0
|
||||
abs v{{q}}.4s, v{{q}}.4s
|
||||
|
||||
sub v{{q}}.4s, v{{q}}.4s, v4.4s
|
||||
sqrshl v{{q}}.4s, v{{q}}.4s, v1.4s
|
||||
|
||||
neg v3.4s, v{{q}}.4s
|
||||
bit v{{q}}.16b, v3.16b, v0.16b
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.q_shr_rounding_away:
|
||||
{% for q in range(16, 32) %}
|
||||
cmlt v0.4s, v{{q}}.4s, #0
|
||||
abs v{{q}}.4s, v{{q}}.4s
|
||||
|
||||
sqrshl v{{q}}.4s, v{{q}}.4s, v1.4s
|
||||
|
||||
neg v3.4s, v{{q}}.4s
|
||||
bit v{{q}}.16b, v3.16b, v0.16b
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.q_shr_rounding_minus_inf:
|
||||
{% for q in range(16, 32) %}
|
||||
sqneg v{{q}}.4s, v{{q}}.4s
|
||||
sqrshl v{{q}}.4s, v{{q}}.4s, v1.4s
|
||||
sqneg v{{q}}.4s, v{{q}}.4s
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.q_shr_rounding_plus_inf:
|
||||
{% for q in range(16, 32) %}
|
||||
sqrshl v{{q}}.4s, v{{q}}.4s, v1.4s
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.q_shr_rounding_even:
|
||||
// sqrshl is round(+inf), sqshl trauncates
|
||||
// we look at parity of result by truncation: if it's odd, we have nothing more to do, we go towards +inf
|
||||
// if it's even, we need to nudge towards 0 by adding -1
|
||||
// => nudge = (x >>l shift) & 0x1 - 1 (>>l is sqshl)
|
||||
// => result is (x + nudge) >>r shift (with sqrshl)
|
||||
{% for q in range(16, 32) %}
|
||||
sqshl v3.4s, v{{q}}.4s, v1.4s // trunc
|
||||
and v3.16b, v3.16b, v4.16b
|
||||
sub v3.4s, v3.4s, v4.4s
|
||||
add v{{q}}.4s, v{{q}}.4s, v3.4s
|
||||
|
||||
sqrshl v{{q}}.4s, v{{q}}.4s, v1.4s
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
|
||||
.q_shr_rounding_odd:
|
||||
// here: nudge is -((x >>l shift) & 0x1)
|
||||
{% for q in range(16, 32) %}
|
||||
sqshl v3.4s, v{{q}}.4s, v1.4s // trunc
|
||||
and v3.16b, v3.16b, v4.16b
|
||||
neg v3.4s, v3.4s
|
||||
add v{{q}}.4s, v{{q}}.4s, v3.4s
|
||||
|
||||
sqrshl v{{q}}.4s, v{{q}}.4s, v1.4s
|
||||
{% endfor %}
|
||||
b .non_linear_loop
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
// vim: ft=arm
|
||||
|
||||
.load_tile:
|
||||
ldr x2, [ x0, #8 ]
|
||||
{% for reg in range(from, to + 1) %}
|
||||
ld1 { v{{reg}}.4s }, [ x2 ], #16
|
||||
{% endfor %}
|
||||
|
||||
b .non_linear_loop
|
||||
|
||||
+206
@@ -0,0 +1,206 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// no preservation either for v0-v7 and v16-v31
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd
|
||||
.global {{G}}arm64simd_sigmoid_f32_4n_{{suffix}}
|
||||
{{G}}arm64simd_sigmoid_f32_4n_{{suffix}}:
|
||||
|
||||
cmp x1, #0
|
||||
beq .return
|
||||
|
||||
adr x2, .coeffs_num
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
|
||||
dup v5.4s, v0.s[0] // v5 <- low, broadcasted
|
||||
dup v6.4s, v0.s[1] // v6 <- high, broadcasted
|
||||
dup v7.4s, v3.s[1] // v7 <- 0.5, broadcasted
|
||||
|
||||
cmp x1, #16
|
||||
blt .loop
|
||||
|
||||
.loop4:
|
||||
ld1 { v16.4s, v17.4s, v18.4s, v19.4s }, [x0]
|
||||
|
||||
fmax v16.4s, v16.4s, v5.4s
|
||||
fmax v17.4s, v17.4s, v5.4s
|
||||
fmax v18.4s, v18.4s, v5.4s
|
||||
fmax v19.4s, v19.4s, v5.4s
|
||||
|
||||
fmin v16.4s, v16.4s, v6.4s
|
||||
fmin v17.4s, v17.4s, v6.4s
|
||||
fmin v18.4s, v18.4s, v6.4s
|
||||
fmin v19.4s, v19.4s, v6.4s // v16 <- x
|
||||
|
||||
fmul v20.4s, v16.4s, v16.4s
|
||||
fmul v21.4s, v17.4s, v17.4s
|
||||
fmul v22.4s, v18.4s, v18.4s
|
||||
fmul v23.4s, v19.4s, v19.4s // v20 <- x2
|
||||
|
||||
dup v24.4s, v0.s[3]
|
||||
fmla v24.4s, v20.4s, v0.s[2]
|
||||
dup v25.4s, v0.s[3]
|
||||
fmla v25.4s, v21.4s, v0.s[2]
|
||||
dup v26.4s, v0.s[3]
|
||||
fmla v26.4s, v22.4s, v0.s[2]
|
||||
dup v27.4s, v0.s[3]
|
||||
fmla v27.4s, v23.4s, v0.s[2]
|
||||
|
||||
dup v28.4s, v1.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v1.s[0]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v1.s[0]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v1.s[0]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v1.s[1]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v1.s[1]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v1.s[1]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v1.s[1]
|
||||
fmla v27.4s, v23.4s, v31.4s
|
||||
|
||||
dup v28.4s, v1.s[2]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v1.s[2]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v1.s[2]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v1.s[2]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v1.s[3]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v1.s[3]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v1.s[3]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v1.s[3]
|
||||
fmla v27.4s, v23.4s, v31.4s
|
||||
|
||||
dup v28.4s, v2.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v2.s[0]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v2.s[0]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v2.s[0]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
fmul v16.4s, v16.4s, v28.4s
|
||||
fmul v17.4s, v17.4s, v29.4s
|
||||
fmul v18.4s, v18.4s, v30.4s
|
||||
fmul v19.4s, v19.4s, v31.4s // v16 <- numerator
|
||||
|
||||
dup v24.4s, v2.s[2]
|
||||
fmla v24.4s, v20.4s, v2.s[1]
|
||||
dup v25.4s, v2.s[2]
|
||||
fmla v25.4s, v21.4s, v2.s[1]
|
||||
dup v26.4s, v2.s[2]
|
||||
fmla v26.4s, v22.4s, v2.s[1]
|
||||
dup v27.4s, v2.s[2]
|
||||
fmla v27.4s, v23.4s, v2.s[1]
|
||||
|
||||
dup v28.4s, v2.s[3]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v2.s[3]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v2.s[3]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v2.s[3]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v3.s[0]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v3.s[0]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v3.s[0]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v3.s[0]
|
||||
fmla v27.4s, v23.4s, v31.4s // v24 denum
|
||||
|
||||
fdiv v16.4s, v16.4s, v24.4s
|
||||
fdiv v17.4s, v17.4s, v25.4s
|
||||
fdiv v18.4s, v18.4s, v26.4s
|
||||
fdiv v19.4s, v19.4s, v27.4s
|
||||
|
||||
fadd v16.4s, v16.4s, v7.4s
|
||||
fadd v17.4s, v17.4s, v7.4s
|
||||
fadd v18.4s, v18.4s, v7.4s
|
||||
fadd v19.4s, v19.4s, v7.4s
|
||||
|
||||
st1 { v16.4s, v17.4s, v18.4s, v19.4s }, [x0], #64
|
||||
|
||||
subs x1, x1, #16
|
||||
cmp x1, #16
|
||||
bge .loop4
|
||||
|
||||
cmp x1, #0
|
||||
beq .return
|
||||
|
||||
.loop:
|
||||
ld1 { v16.4s }, [x0]
|
||||
|
||||
fmax v16.4s, v16.4s, v5.4s
|
||||
fmin v16.4s, v16.4s, v6.4s // v16 <- x
|
||||
fmul v20.4s, v16.4s, v16.4s // v20 <- x2
|
||||
|
||||
dup v24.4s, v0.s[3]
|
||||
fmla v24.4s, v20.4s, v0.s[2]
|
||||
dup v28.4s, v1.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v1.s[1]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v28.4s, v1.s[2]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v1.s[3]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v28.4s, v2.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
fmul v16.4s, v16.4s, v28.4s // v16 <- numerator
|
||||
|
||||
dup v24.4s, v2.s[2]
|
||||
fmla v24.4s, v20.4s, v2.s[1]
|
||||
dup v28.4s, v2.s[3]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v3.s[0]
|
||||
fmla v24.4s, v20.4s, v28.4s // v24 <- denum
|
||||
|
||||
fdiv v16.4s, v16.4s, v24.4s
|
||||
fadd v16.4s, v16.4s, v7.4s
|
||||
|
||||
st1 { v16.4s }, [x0], #16
|
||||
|
||||
subs x1, x1, #4
|
||||
bne .loop
|
||||
|
||||
.return:
|
||||
ret
|
||||
|
||||
.coeffs_num:
|
||||
.float -18.6 // low
|
||||
.float 18.6 // high
|
||||
.float -4.433153405e-18 // alpha_13
|
||||
.float 1.169974371e-14
|
||||
|
||||
.float -1.875289645e-11
|
||||
.float 4.257889523e-8
|
||||
.float 0.00004811817576
|
||||
.float 0.008163842030
|
||||
|
||||
.float 0.2499999971
|
||||
.float 3.922935744e-6 // beta_6
|
||||
.float 0.001524872358
|
||||
.float 0.1159886749
|
||||
|
||||
.float 1.0
|
||||
.float 0.5 //
|
||||
.float 0.0 // padding
|
||||
.float 0.0
|
||||
|
||||
+198
@@ -0,0 +1,198 @@
|
||||
// vim: ft=arm
|
||||
|
||||
// no preservation either for v0-v7 and v16-v31
|
||||
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.cpu generic+fp+simd
|
||||
.global {{G}}arm64simd_tanh_f32_4n_{{suffix}}
|
||||
{{G}}arm64simd_tanh_f32_4n_{{suffix}}:
|
||||
|
||||
cmp x1, #0
|
||||
beq .return
|
||||
|
||||
adr x2, .coeffs_num
|
||||
ld1 { v0.4s, v1.4s, v2.4s, v3.4s }, [x2]
|
||||
dup v5.4s, v0.s[0] // v5 <- low, broadcasted
|
||||
dup v6.4s, v0.s[1] // v6 <- high, broadcasted
|
||||
|
||||
cmp x1, #16
|
||||
blt .loop
|
||||
|
||||
.loop4:
|
||||
ld1 { v16.4s, v17.4s, v18.4s, v19.4s }, [x0]
|
||||
|
||||
fmax v16.4s, v16.4s, v5.4s
|
||||
fmax v17.4s, v17.4s, v5.4s
|
||||
fmax v18.4s, v18.4s, v5.4s
|
||||
fmax v19.4s, v19.4s, v5.4s
|
||||
|
||||
fmin v16.4s, v16.4s, v6.4s
|
||||
fmin v17.4s, v17.4s, v6.4s
|
||||
fmin v18.4s, v18.4s, v6.4s
|
||||
fmin v19.4s, v19.4s, v6.4s // v16 <- x
|
||||
|
||||
fmul v20.4s, v16.4s, v16.4s
|
||||
fmul v21.4s, v17.4s, v17.4s
|
||||
fmul v22.4s, v18.4s, v18.4s
|
||||
fmul v23.4s, v19.4s, v19.4s // v20 <- x2
|
||||
|
||||
dup v24.4s, v0.s[3]
|
||||
fmla v24.4s, v20.4s, v0.s[2]
|
||||
dup v25.4s, v0.s[3]
|
||||
fmla v25.4s, v21.4s, v0.s[2]
|
||||
dup v26.4s, v0.s[3]
|
||||
fmla v26.4s, v22.4s, v0.s[2]
|
||||
dup v27.4s, v0.s[3]
|
||||
fmla v27.4s, v23.4s, v0.s[2]
|
||||
|
||||
dup v28.4s, v1.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v1.s[0]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v1.s[0]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v1.s[0]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v1.s[1]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v1.s[1]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v1.s[1]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v1.s[1]
|
||||
fmla v27.4s, v23.4s, v31.4s
|
||||
|
||||
dup v28.4s, v1.s[2]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v1.s[2]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v1.s[2]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v1.s[2]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v1.s[3]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v1.s[3]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v1.s[3]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v1.s[3]
|
||||
fmla v27.4s, v23.4s, v31.4s
|
||||
|
||||
dup v28.4s, v2.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v2.s[0]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v2.s[0]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v2.s[0]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
fmul v16.4s, v16.4s, v28.4s
|
||||
fmul v17.4s, v17.4s, v29.4s
|
||||
fmul v18.4s, v18.4s, v30.4s
|
||||
fmul v19.4s, v19.4s, v31.4s // v16 <- numerator
|
||||
|
||||
dup v24.4s, v2.s[2]
|
||||
fmla v24.4s, v20.4s, v2.s[1]
|
||||
dup v25.4s, v2.s[2]
|
||||
fmla v25.4s, v21.4s, v2.s[1]
|
||||
dup v26.4s, v2.s[2]
|
||||
fmla v26.4s, v22.4s, v2.s[1]
|
||||
dup v27.4s, v2.s[2]
|
||||
fmla v27.4s, v23.4s, v2.s[1]
|
||||
|
||||
dup v28.4s, v2.s[3]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v29.4s, v2.s[3]
|
||||
fmla v29.4s, v21.4s, v25.4s
|
||||
dup v30.4s, v2.s[3]
|
||||
fmla v30.4s, v22.4s, v26.4s
|
||||
dup v31.4s, v2.s[3]
|
||||
fmla v31.4s, v23.4s, v27.4s
|
||||
|
||||
dup v24.4s, v3.s[0]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v25.4s, v3.s[0]
|
||||
fmla v25.4s, v21.4s, v29.4s
|
||||
dup v26.4s, v3.s[0]
|
||||
fmla v26.4s, v22.4s, v30.4s
|
||||
dup v27.4s, v3.s[0]
|
||||
fmla v27.4s, v23.4s, v31.4s // v24 denum
|
||||
|
||||
fdiv v16.4s, v16.4s, v24.4s
|
||||
fdiv v17.4s, v17.4s, v25.4s
|
||||
fdiv v18.4s, v18.4s, v26.4s
|
||||
fdiv v19.4s, v19.4s, v27.4s
|
||||
|
||||
st1 { v16.4s, v17.4s, v18.4s, v19.4s }, [x0], #64
|
||||
|
||||
subs x1, x1, #16
|
||||
cmp x1, #16
|
||||
bge .loop4
|
||||
|
||||
cmp x1, #0
|
||||
beq .return
|
||||
|
||||
.loop:
|
||||
ld1 { v16.4s }, [x0]
|
||||
|
||||
fmax v16.4s, v16.4s, v5.4s
|
||||
fmin v16.4s, v16.4s, v6.4s // v16 <- x
|
||||
fmul v20.4s, v16.4s, v16.4s // v20 <- x2
|
||||
|
||||
dup v24.4s, v0.s[3]
|
||||
fmla v24.4s, v20.4s, v0.s[2]
|
||||
dup v28.4s, v1.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v1.s[1]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v28.4s, v1.s[2]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v1.s[3]
|
||||
fmla v24.4s, v20.4s, v28.4s
|
||||
dup v28.4s, v2.s[0]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
fmul v16.4s, v16.4s, v28.4s // v16 <- numerator
|
||||
|
||||
dup v24.4s, v2.s[2]
|
||||
fmla v24.4s, v20.4s, v2.s[1]
|
||||
dup v28.4s, v2.s[3]
|
||||
fmla v28.4s, v20.4s, v24.4s
|
||||
dup v24.4s, v3.s[0]
|
||||
fmla v24.4s, v20.4s, v28.4s // v24 <- denum
|
||||
|
||||
fdiv v16.4s, v16.4s, v24.4s
|
||||
|
||||
st1 { v16.4s }, [x0], #16
|
||||
|
||||
subs x1, x1, #4
|
||||
bne .loop
|
||||
|
||||
.return:
|
||||
ret
|
||||
|
||||
.coeffs_num:
|
||||
.float -8.9 // low
|
||||
.float 8.9 // high
|
||||
.float -8.488492677e-14 // alpha_13
|
||||
.float 5.277853000e-11
|
||||
|
||||
.float -2.022500419e-8
|
||||
.float 0.00001115424833
|
||||
.float 0.003103950131
|
||||
.float 0.1308400453
|
||||
|
||||
.float 0.9999999934
|
||||
.float 0.0002546136580 // beta_6
|
||||
.float 0.02449515379
|
||||
.float 0.4641733162
|
||||
|
||||
.float 1.0
|
||||
.float 0 // padding
|
||||
.float 0 // padding
|
||||
.float 0 // padding
|
||||
Vendored
+37
@@ -0,0 +1,37 @@
|
||||
// vim: ft=arm
|
||||
|
||||
.non_linear:
|
||||
sub x0, x0, 40
|
||||
|
||||
.non_linear_loop:
|
||||
add x0, x0, 40
|
||||
ldr x2, [x0]
|
||||
|
||||
mov x4, #{{ jump_table | length }}
|
||||
|
||||
cmp x2, #{{ jump_table | length }}
|
||||
csel x2, x2, x4, lt
|
||||
cmp x2, #0
|
||||
csel x2, x4, x2, lt
|
||||
|
||||
adr x3, .jmp_table
|
||||
add x3, x3, x2, LSL#2
|
||||
br x3
|
||||
|
||||
.jmp_table:
|
||||
{% for j in jump_table %}
|
||||
b .{{j}}
|
||||
{% endfor %}
|
||||
b .unsupported
|
||||
|
||||
add x0, x2, #4000
|
||||
b .return
|
||||
|
||||
.unsupported:
|
||||
mov x0, #1
|
||||
b .return
|
||||
|
||||
.done:
|
||||
mov x0, 0
|
||||
b .return
|
||||
|
||||
Vendored
+13
@@ -0,0 +1,13 @@
|
||||
// Build-time capability probe for the assembler, used by build.rs
|
||||
// (assembler_supports_dotprod). Older binutils — notably the Debian stretch
|
||||
// aarch64 cross-toolchain in CI — predate FEAT_DotProd and cannot assemble
|
||||
// `sdot` even with `.cpu generic+fp+simd+dotprod`. If this file fails to
|
||||
// assemble, build.rs skips the SDOT kernel and the `tract_arm64_dotprod` cfg,
|
||||
// and the runtime falls back to the SMLAL 8x8 i32 kernel. Not linked into
|
||||
// anything.
|
||||
.cpu generic+fp+simd+dotprod
|
||||
.text
|
||||
.globl tract_dotprod_probe
|
||||
tract_dotprod_probe:
|
||||
sdot v0.4s, v1.16b, v2.4b[0]
|
||||
ret
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
// vim: ft=arm
|
||||
|
||||
.non_linear:
|
||||
sub x0, x0, 40
|
||||
|
||||
.non_linear_loop:
|
||||
add x0, x0, 40
|
||||
ldr x2, [x0]
|
||||
|
||||
mov x4, #{{ jump_table | length }}
|
||||
|
||||
cmp x2, #{{ jump_table | length }}
|
||||
csel x2, x2, x4, lt
|
||||
cmp x2, #0
|
||||
csel x2, x4, x2, lt
|
||||
|
||||
adr x3, .jmp_table
|
||||
add x3, x3, x2, LSL#2
|
||||
br x3
|
||||
|
||||
.jmp_table:
|
||||
{% for j in jump_table %}
|
||||
b .{{j}}
|
||||
{% endfor %}
|
||||
b .unsupported
|
||||
|
||||
add x0, x2, #4000
|
||||
b .return
|
||||
|
||||
.unsupported:
|
||||
mov x0, #1
|
||||
b .return
|
||||
|
||||
.done:
|
||||
mov x0, 0
|
||||
b .return
|
||||
|
||||
+14
@@ -0,0 +1,14 @@
|
||||
// Build-time capability probe for the assembler, used by build.rs
|
||||
// (assembler_supports_sme). Older binutils — notably the Debian stretch
|
||||
// aarch64 cross-toolchain in CI — predate SME and cannot assemble these
|
||||
// mnemonics even with `.arch armv9-a+sme2`. If this file fails to assemble,
|
||||
// build.rs skips the SME kernels and the `tract_sme` cfg, and the runtime
|
||||
// falls back to the portable path. Not linked into anything.
|
||||
.arch armv9-a+sme2
|
||||
.text
|
||||
.globl tract_sme_probe
|
||||
tract_sme_probe:
|
||||
smstart
|
||||
zero {za}
|
||||
smstop
|
||||
ret
|
||||
Vendored
+474
@@ -0,0 +1,474 @@
|
||||
// vim: ft=arm
|
||||
//
|
||||
// SME f32 32x32 matmul kernel.
|
||||
//
|
||||
// ZA tile layout (4 .S tiles, 16x16 each, indexed left/right x top/bottom):
|
||||
//
|
||||
// ZA0.S : C[0..16, 0..16] (top-left)
|
||||
// ZA1.S : C[0..16, 16..32] (top-right)
|
||||
// ZA2.S : C[16..32, 0..16] (bottom-left)
|
||||
// ZA3.S : C[16..32, 16..32] (bottom-right)
|
||||
//
|
||||
// Inner K-step: load 32 f32 of A (split z0+z2) and 32 of B (split z1+z3),
|
||||
// issue 4 FMOPAs (one per tile). All 4 tiles are independent → SME unit
|
||||
// reaches 1 fmopa/cycle = ~2 TFLOPS on M4.
|
||||
//
|
||||
// Calling convention (extern "C", AAPCS64):
|
||||
// x0 = const *FusedKerSpec<f32>, advanced 40 B per dispatcher iteration.
|
||||
// x1 = stack-resident 4 KiB scratch buffer for tile spills (Phase 1B+).
|
||||
//
|
||||
// Streaming mode: PSTATE.SM=1 from prologue smstart to epilogue smstop.
|
||||
// V0..V31 (low 128 bits = Z0..Z31 low) are destroyed by the smstart/smstop
|
||||
// pair; v8..v15 are saved/restored to stack across the streaming region per
|
||||
// AAPCS callee-save rules.
|
||||
|
||||
.arch armv9-a+sme2
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.global {{G}}sme_mmm_f32_32x32_{{suffix}}
|
||||
{{G}}sme_mmm_f32_32x32_{{suffix}}:
|
||||
|
||||
// Save callee-saved q8..q15 (AAPCS preserves low 64 bits of v8..v15;
|
||||
// we save the full 128-bit Q to keep the stack layout simple).
|
||||
stp q8, q9, [sp, #-128]!
|
||||
stp q10, q11, [sp, #32]
|
||||
stp q12, q13, [sp, #64]
|
||||
stp q14, q15, [sp, #96]
|
||||
|
||||
// Allocate 4 KiB tile-spill scratch (kept live across the whole call).
|
||||
sub sp, sp, #4096
|
||||
mov x1, sp
|
||||
|
||||
smstart
|
||||
ptrue p0.b
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
// -------- supported fuse ops ---------------------------------------------
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x2, [x0, #24] // b
|
||||
ldp x3, x4, [x0, #8] // k, a
|
||||
|
||||
cmp x3, #0
|
||||
b.eq .non_linear_loop
|
||||
|
||||
.Lmatmul_loop:
|
||||
ld1w {z0.s}, p0/z, [x4]
|
||||
ld1w {z2.s}, p0/z, [x4, #1, mul vl]
|
||||
ld1w {z1.s}, p0/z, [x2]
|
||||
ld1w {z3.s}, p0/z, [x2, #1, mul vl]
|
||||
add x4, x4, #128
|
||||
add x2, x2, #128
|
||||
|
||||
fmopa za0.s, p0/m, p0/m, z0.s, z1.s // C[0..16, 0..16]
|
||||
fmopa za1.s, p0/m, p0/m, z0.s, z3.s // C[0..16, 16..32]
|
||||
fmopa za2.s, p0/m, p0/m, z2.s, z1.s // C[16..32, 0..16]
|
||||
fmopa za3.s, p0/m, p0/m, z2.s, z3.s // C[16..32, 16..32]
|
||||
|
||||
subs x3, x3, #1
|
||||
b.ne .Lmatmul_loop
|
||||
b .non_linear_loop
|
||||
|
||||
.clear:
|
||||
zero {za}
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
// FusedKerSpec::Store(OutputStoreKer { ptr, row_byte_stride,
|
||||
// col_byte_stride, item_size })
|
||||
// [x0, #8] = ptr [x0, #16] = row_byte_stride
|
||||
// [x0, #24] = col_byte_stride [x0, #32] = item_size
|
||||
ldp x5, x6, [x0, #8] // x5 = ptr, x6 = row_byte_stride
|
||||
ldp x7, x8, [x0, #24] // x7 = col_byte_stride, x8 = item_size
|
||||
|
||||
// Fast path: contiguous f32 columns (col_stride == 4) → direct ZA→user.
|
||||
// st1w-from-ZA does not accept "[Xn, #imm, MUL VL]" offsets, so we keep
|
||||
// two parallel base pointers for the left and right halves of each row.
|
||||
cmp x7, #4
|
||||
b.ne .Lstore_generic
|
||||
cmp x8, #4
|
||||
b.ne .Lstore_generic
|
||||
|
||||
add x4, x5, #64 // right-half base
|
||||
mov w12, #0
|
||||
.Lstore_top:
|
||||
st1w {za0h.s[w12, 0]}, p0, [x5]
|
||||
st1w {za1h.s[w12, 0]}, p0, [x4]
|
||||
add x5, x5, x6
|
||||
add x4, x4, x6
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lstore_top
|
||||
mov w12, #0
|
||||
.Lstore_bot:
|
||||
st1w {za2h.s[w12, 0]}, p0, [x5]
|
||||
st1w {za3h.s[w12, 0]}, p0, [x4]
|
||||
add x5, x5, x6
|
||||
add x4, x4, x6
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lstore_bot
|
||||
b .non_linear_loop
|
||||
|
||||
.Lstore_generic:
|
||||
// Slow path: spill ZA → x1 (scratch buffer, 32x32 row-major, 128 B/row)
|
||||
// using two parallel pointers, then per-element strided scatter.
|
||||
mov x4, x1 // left-half pointer
|
||||
add x9, x1, #64 // right-half pointer
|
||||
mov w12, #0
|
||||
.Lstore_spill_top:
|
||||
st1w {za0h.s[w12, 0]}, p0, [x4]
|
||||
st1w {za1h.s[w12, 0]}, p0, [x9]
|
||||
add x4, x4, #128
|
||||
add x9, x9, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lstore_spill_top
|
||||
mov w12, #0
|
||||
.Lstore_spill_bot:
|
||||
st1w {za2h.s[w12, 0]}, p0, [x4]
|
||||
st1w {za3h.s[w12, 0]}, p0, [x9]
|
||||
add x4, x4, #128
|
||||
add x9, x9, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lstore_spill_bot
|
||||
|
||||
// Strided f32 scatter: 32 rows × 32 cols.
|
||||
mov x3, #0
|
||||
.Lstore_row:
|
||||
mov x4, x5
|
||||
mov x10, #0
|
||||
lsl x9, x3, #7 // row*128 byte offset in scratch
|
||||
add x11, x1, x9
|
||||
.Lstore_col:
|
||||
ldr w9, [x11], #4
|
||||
str w9, [x4]
|
||||
add x4, x4, x7
|
||||
add x10, x10, #1
|
||||
cmp x10, #32
|
||||
b.lt .Lstore_col
|
||||
add x5, x5, x6
|
||||
add x3, x3, #1
|
||||
cmp x3, #32
|
||||
b.lt .Lstore_row
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- scalar ops ------------------------------------------------------
|
||||
//
|
||||
// FusedKerSpec::Scalar{Add,Mul,Sub,SubF,Min,Max}(TI) — broadcast scalar from
|
||||
// [x0, #8] to all lanes, apply elementwise across the 4-tile 32x32 grid.
|
||||
//
|
||||
// Sub vs SubF semantics (matching apple_amx + tests/fuse.rs):
|
||||
// ScalarSub → result = scalar - z (mnemonic fsubr)
|
||||
// ScalarSubF → result = z - scalar (mnemonic fsub)
|
||||
//
|
||||
// Slice-op loop: for each slice index w12, extract ZA tile slice → Z reg,
|
||||
// op with broadcast-scalar in z4, insert Z back. Two halves × 4 tiles total.
|
||||
|
||||
{% macro scalar_op(label, op) %}
|
||||
{{label}}:
|
||||
ldr w2, [x0, #8]
|
||||
dup z4.s, w2
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_top:
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z4.s
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_top
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_bot:
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z4.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_bot
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{{ scalar_op('.scalar_add', 'fadd') }}
|
||||
{{ scalar_op('.scalar_mul', 'fmul') }}
|
||||
{{ scalar_op('.scalar_sub', 'fsubr') }}
|
||||
{{ scalar_op('.scalar_sub_flipped', 'fsub') }}
|
||||
{{ scalar_op('.scalar_min', 'fmin') }}
|
||||
{{ scalar_op('.scalar_max', 'fmax') }}
|
||||
|
||||
// -------- per-col ops -----------------------------------------------------
|
||||
//
|
||||
// 32-element column vector → z4 (cols 0-15) + z5 (cols 16-31).
|
||||
// Same z4/z5 is applied to every row across the 4-tile grid.
|
||||
|
||||
{% macro per_col_op(label, op) %}
|
||||
{{label}}:
|
||||
ldr x2, [x0, #8]
|
||||
ld1w {z4.s}, p0/z, [x2]
|
||||
ld1w {z5.s}, p0/z, [x2, #1, mul vl]
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_top:
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z5.s
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_top
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_bot:
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z5.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_bot
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{{ per_col_op('.per_col_add', 'fadd') }}
|
||||
{{ per_col_op('.per_col_mul', 'fmul') }}
|
||||
{{ per_col_op('.per_col_sub', 'fsubr') }}
|
||||
{{ per_col_op('.per_col_sub_flipped', 'fsub') }}
|
||||
{{ per_col_op('.per_col_min', 'fmin') }}
|
||||
{{ per_col_op('.per_col_max', 'fmax') }}
|
||||
|
||||
// -------- per-row ops -----------------------------------------------------
|
||||
//
|
||||
// 32-element row vector at x2 (top 16 rows) and x2+64 (bottom 16 rows).
|
||||
// Load one f32 per iteration and broadcast (no SVE indexed-broadcast for
|
||||
// arbitrary i across 16 lanes, so we just walk the bias pointer).
|
||||
|
||||
{% macro per_row_op(label, op) %}
|
||||
{{label}}:
|
||||
ldr x2, [x0, #8]
|
||||
add x3, x2, #64
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_top:
|
||||
ldr w4, [x2], #4
|
||||
dup z4.s, w4
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z4.s
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_top
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_bot:
|
||||
ldr w4, [x3], #4
|
||||
dup z4.s, w4
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z4.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_bot
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{{ per_row_op('.per_row_add', 'fadd') }}
|
||||
{{ per_row_op('.per_row_mul', 'fmul') }}
|
||||
{{ per_row_op('.per_row_sub', 'fsubr') }}
|
||||
{{ per_row_op('.per_row_sub_flipped', 'fsub') }}
|
||||
{{ per_row_op('.per_row_min', 'fmin') }}
|
||||
{{ per_row_op('.per_row_max', 'fmax') }}
|
||||
|
||||
// -------- AddRowColProducts: ZA += rows ⊗ cols (rank-1 update) ------------
|
||||
//
|
||||
// Same shape as a K=1 matmul step: load 32 f32 of rows + 32 of cols, four
|
||||
// FMOPAs into the 2x2 ZA grid.
|
||||
|
||||
.add_row_col_products:
|
||||
ldp x2, x3, [x0, #8] // rows ptr, cols ptr
|
||||
ld1w {z0.s}, p0/z, [x2]
|
||||
ld1w {z2.s}, p0/z, [x2, #1, mul vl]
|
||||
ld1w {z1.s}, p0/z, [x3]
|
||||
ld1w {z3.s}, p0/z, [x3, #1, mul vl]
|
||||
fmopa za0.s, p0/m, p0/m, z0.s, z1.s
|
||||
fmopa za1.s, p0/m, p0/m, z0.s, z3.s
|
||||
fmopa za2.s, p0/m, p0/m, z2.s, z1.s
|
||||
fmopa za3.s, p0/m, p0/m, z2.s, z3.s
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- AddUnicast: ZA += C[i][j] from strided buffer -------------------
|
||||
//
|
||||
// FusedKerSpec::AddUnicast(OutputStoreKer { ptr, row_byte_stride,
|
||||
// col_byte_stride, item_size })
|
||||
// Fast path: contiguous f32 cols (col_stride == 4) — load each row via
|
||||
// ld1w then in-place fadd to ZA slice.
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // ptr, row_byte_stride
|
||||
ldp x7, x8, [x0, #24] // col_byte_stride, item_size
|
||||
|
||||
cmp x7, #4
|
||||
b.ne .Laddu_generic
|
||||
cmp x8, #4
|
||||
b.ne .Laddu_generic
|
||||
|
||||
add x4, x5, #64
|
||||
mov w12, #0
|
||||
.Laddu_top:
|
||||
ld1w {z8.s}, p0/z, [x5]
|
||||
ld1w {z9.s}, p0/z, [x4]
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
fadd z6.s, p0/m, z6.s, z8.s
|
||||
fadd z7.s, p0/m, z7.s, z9.s
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add x5, x5, x6
|
||||
add x4, x4, x6
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Laddu_top
|
||||
mov w12, #0
|
||||
.Laddu_bot:
|
||||
ld1w {z8.s}, p0/z, [x5]
|
||||
ld1w {z9.s}, p0/z, [x4]
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
fadd z6.s, p0/m, z6.s, z8.s
|
||||
fadd z7.s, p0/m, z7.s, z9.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add x5, x5, x6
|
||||
add x4, x4, x6
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Laddu_bot
|
||||
b .non_linear_loop
|
||||
|
||||
.Laddu_generic:
|
||||
// Generic strided gather: walk 32 rows × 32 cols, accumulate one
|
||||
// element at a time into the spill scratch, then re-load each slice
|
||||
// into ZA and fadd.
|
||||
//
|
||||
// Phase 1B keeps this slow but correct — it triggers only for non-
|
||||
// contiguous AddUnicast which auto-tests don't exercise.
|
||||
mov x3, #0 // row idx
|
||||
mov x9, x1 // scratch ptr
|
||||
.Laddu_gen_row:
|
||||
mov x10, #0
|
||||
mov x11, x5
|
||||
.Laddu_gen_col:
|
||||
ldr w4, [x11]
|
||||
str w4, [x9], #4
|
||||
add x11, x11, x7
|
||||
add x10, x10, #1
|
||||
cmp x10, #32
|
||||
b.lt .Laddu_gen_col
|
||||
add x5, x5, x6
|
||||
add x3, x3, #1
|
||||
cmp x3, #32
|
||||
b.lt .Laddu_gen_row
|
||||
|
||||
// Now scratch holds 32x32 f32 row-major. Same loop as fast path but
|
||||
// reading from scratch (contiguous).
|
||||
mov x9, x1
|
||||
add x4, x9, #64
|
||||
mov w12, #0
|
||||
.Laddu_gen_apply_top:
|
||||
ld1w {z8.s}, p0/z, [x9]
|
||||
ld1w {z10.s}, p0/z, [x4]
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
fadd z6.s, p0/m, z6.s, z8.s
|
||||
fadd z7.s, p0/m, z7.s, z10.s
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add x9, x9, #128
|
||||
add x4, x4, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Laddu_gen_apply_top
|
||||
mov w12, #0
|
||||
.Laddu_gen_apply_bot:
|
||||
ld1w {z8.s}, p0/z, [x9]
|
||||
ld1w {z10.s}, p0/z, [x4]
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
fadd z6.s, p0/m, z6.s, z8.s
|
||||
fadd z7.s, p0/m, z7.s, z10.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add x9, x9, #128
|
||||
add x4, x4, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Laddu_gen_apply_bot
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- LoadTile: ZA := row-major tile from memory ----------------------
|
||||
//
|
||||
// FusedKerSpec::LoadTile(col_major_ptr, row_major_ptr):
|
||||
// [x0, #8] = col-major ptr (unused; AMX prefers this for its layout)
|
||||
// [x0, #16] = row-major ptr (32x32 f32, 128 B per row)
|
||||
//
|
||||
// We use the row-major pointer because the ZA H-tile store path is itself
|
||||
// row-major and matches naturally.
|
||||
|
||||
.load_tile:
|
||||
ldr x2, [x0, #16]
|
||||
add x4, x2, #64
|
||||
mov w12, #0
|
||||
.Lloadtile_top:
|
||||
ld1w {z6.s}, p0/z, [x2]
|
||||
ld1w {z7.s}, p0/z, [x4]
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add x2, x2, #128
|
||||
add x4, x4, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lloadtile_top
|
||||
mov w12, #0
|
||||
.Lloadtile_bot:
|
||||
ld1w {z6.s}, p0/z, [x2]
|
||||
ld1w {z7.s}, p0/z, [x4]
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add x2, x2, #128
|
||||
add x4, x4, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lloadtile_bot
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- still not implemented (low priority for Phase 1) ----------------
|
||||
|
||||
.leaky_relu:
|
||||
.q_scale:
|
||||
.q_shl:
|
||||
.q_shr:
|
||||
b .unsupported
|
||||
|
||||
// -------- epilogue --------------------------------------------------------
|
||||
|
||||
.return:
|
||||
smstop
|
||||
add sp, sp, #4096
|
||||
ldp q14, q15, [sp, #96]
|
||||
ldp q12, q13, [sp, #64]
|
||||
ldp q10, q11, [sp, #32]
|
||||
ldp q8, q9, [sp], #128
|
||||
ret
|
||||
Vendored
+268
@@ -0,0 +1,268 @@
|
||||
// vim: ft=arm
|
||||
//
|
||||
// SME2 f32 64x1 GEMV kernel.
|
||||
//
|
||||
// Accumulator layout: ZA tile slot rows 0..3 (one vgx4 group starting at
|
||||
// w8=0). The 64-element output column maps to these 4 slots × 16 f32 each.
|
||||
//
|
||||
// Inner K-step: load 64 f32 of A's column into {z0.s-z3.s} via one SME2
|
||||
// multi-vec LD1W, broadcast B[k] into z4 with LD1RW, issue ONE multi-vec
|
||||
// vgx4 FMLA-into-ZA. Measured peak on M4: ~125 GFLOPS (ZA-write port).
|
||||
// Plain SME single-vec predicated FMLA into Z regs caps at ~31 GFLOPS so
|
||||
// is not used. Detection therefore gates on FEAT_SME2, not just FEAT_SME.
|
||||
//
|
||||
// Calling convention (extern "C", AAPCS64):
|
||||
// x0 = const *FusedKerSpec<f32>, advanced 40 B per dispatcher iteration.
|
||||
// x1 = stack-resident 256 B scratch buffer for the strided-store path.
|
||||
// w8 = 0 throughout (vgx-group base index, set in prologue).
|
||||
//
|
||||
// Streaming-mode rules: PSTATE.SM=1 from prologue smstart to epilogue
|
||||
// smstop. v8..v15 saved/restored across the streaming region.
|
||||
|
||||
.arch armv9-a+sme2
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.global {{G}}sme_mmv_f32_64x1_{{suffix}}
|
||||
{{G}}sme_mmv_f32_64x1_{{suffix}}:
|
||||
|
||||
stp q8, q9, [sp, #-128]!
|
||||
stp q10, q11, [sp, #32]
|
||||
stp q12, q13, [sp, #64]
|
||||
stp q14, q15, [sp, #96]
|
||||
|
||||
// 256 B = 64 f32 spill buffer for the strided-store / AddUnicast paths.
|
||||
sub sp, sp, #256
|
||||
mov x1, sp
|
||||
|
||||
smstart
|
||||
ptrue p0.b
|
||||
ptrue pn8.b
|
||||
mov w8, #0
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
// -------- supported fuse ops -----------------------------------------------
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x2, [x0, #24] // b ptr
|
||||
ldp x3, x4, [x0, #8] // k, a ptr
|
||||
cmp x3, #0
|
||||
b.eq .non_linear_loop
|
||||
|
||||
.Lmmv_loop:
|
||||
ld1w {z0.s-z3.s}, pn8/z, [x4]
|
||||
add x4, x4, #256
|
||||
ld1rw {z4.s}, p0/z, [x2]
|
||||
add x2, x2, #4
|
||||
fmla za.s[w8, 0, vgx4], {z0.s-z3.s}, z4.s[0]
|
||||
subs x3, x3, #1
|
||||
b.ne .Lmmv_loop
|
||||
b .non_linear_loop
|
||||
|
||||
.clear:
|
||||
zero {za}
|
||||
b .non_linear_loop
|
||||
|
||||
.store:
|
||||
// FusedKerSpec::Store(OutputStoreKer { ptr, row_byte_stride,
|
||||
// col_byte_stride, item_size })
|
||||
// [x0, #8] = ptr [x0, #16] = row_byte_stride
|
||||
// [x0, #24] = col_byte_stride [x0, #32] = item_size
|
||||
// x8 must NOT be touched (it's the vgx-base index, set to 0 in prologue).
|
||||
ldp x5, x6, [x0, #8] // ptr, row_byte_stride
|
||||
ldp x7, x9, [x0, #24] // col_byte_stride, item_size
|
||||
|
||||
// At NR=1 the output column is one element per row; the fast path
|
||||
// triggers when (row_byte_stride==4 AND item_size==4) i.e. the 64
|
||||
// outputs are contiguous in memory.
|
||||
cmp x6, #4
|
||||
b.ne .Lstore_generic
|
||||
cmp x9, #4
|
||||
b.ne .Lstore_generic
|
||||
|
||||
mov {z0.s-z3.s}, za.s[w8, 0, vgx4]
|
||||
st1w {z0.s}, p0, [x5]
|
||||
st1w {z1.s}, p0, [x5, #1, mul vl]
|
||||
st1w {z2.s}, p0, [x5, #2, mul vl]
|
||||
st1w {z3.s}, p0, [x5, #3, mul vl]
|
||||
b .non_linear_loop
|
||||
|
||||
.Lstore_generic:
|
||||
// Spill ZA → 256 B scratch buffer x1, then per-element strided write.
|
||||
mov {z0.s-z3.s}, za.s[w8, 0, vgx4]
|
||||
st1w {z0.s}, p0, [x1]
|
||||
st1w {z1.s}, p0, [x1, #1, mul vl]
|
||||
st1w {z2.s}, p0, [x1, #2, mul vl]
|
||||
st1w {z3.s}, p0, [x1, #3, mul vl]
|
||||
|
||||
mov x3, #0
|
||||
mov x9, x1
|
||||
.Lstore_scatter:
|
||||
ldr w10, [x9], #4
|
||||
str w10, [x5]
|
||||
add x5, x5, x6
|
||||
add x3, x3, #1
|
||||
cmp x3, #64
|
||||
b.lt .Lstore_scatter
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- LoadTile: ZA := tile from row-major source -----------------------
|
||||
//
|
||||
// FusedKerSpec::LoadTile(col_major_ptr, row_major_ptr) — same as Phase 1's
|
||||
// 32x32 LoadTile. NR=1 collapses both pointers to the same 64-element vec;
|
||||
// we use the row-major form at [x0, #16].
|
||||
|
||||
.load_tile:
|
||||
ldr x2, [x0, #16]
|
||||
ld1w {z0.s-z3.s}, pn8/z, [x2]
|
||||
mov za.s[w8, 0, vgx4], {z0.s-z3.s}
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- AddRowColProducts: ZA += rows ⊗ cols (rank-1 K=1) ---------------
|
||||
//
|
||||
// NR=1: cols is a single f32, rows is a 64-element vector. Effectively one
|
||||
// K-step of add_mat_mul with K=1.
|
||||
|
||||
.add_row_col_products:
|
||||
ldp x2, x3, [x0, #8] // rows ptr, cols ptr
|
||||
ld1w {z0.s-z3.s}, pn8/z, [x2]
|
||||
ld1rw {z4.s}, p0/z, [x3]
|
||||
fmla za.s[w8, 0, vgx4], {z0.s-z3.s}, z4.s[0]
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- AddUnicast: ZA += C from strided buffer --------------------------
|
||||
//
|
||||
// NR=1 implies a 64-element column vec layout. Fast path = contiguous f32
|
||||
// rows (row_stride == 4); generic path gathers strided.
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // ptr, row_byte_stride
|
||||
ldp x7, x9, [x0, #24] // col_byte_stride, item_size
|
||||
|
||||
cmp x6, #4
|
||||
b.ne .Laddu_generic
|
||||
cmp x9, #4
|
||||
b.ne .Laddu_generic
|
||||
|
||||
// Fast path: contiguous load via 4-vec LD1W.
|
||||
ld1w {z16.s-z19.s}, pn8/z, [x5]
|
||||
mov {z0.s-z3.s}, za.s[w8, 0, vgx4]
|
||||
fadd z0.s, p0/m, z0.s, z16.s
|
||||
fadd z1.s, p0/m, z1.s, z17.s
|
||||
fadd z2.s, p0/m, z2.s, z18.s
|
||||
fadd z3.s, p0/m, z3.s, z19.s
|
||||
mov za.s[w8, 0, vgx4], {z0.s-z3.s}
|
||||
b .non_linear_loop
|
||||
|
||||
.Laddu_generic:
|
||||
// Per-element strided gather into scratch, then contiguous accumulate.
|
||||
mov x3, #0
|
||||
mov x9, x1
|
||||
.Laddu_gather:
|
||||
ldr w10, [x5]
|
||||
str w10, [x9], #4
|
||||
add x5, x5, x6
|
||||
add x3, x3, #1
|
||||
cmp x3, #64
|
||||
b.lt .Laddu_gather
|
||||
ld1w {z16.s-z19.s}, pn8/z, [x1]
|
||||
mov {z0.s-z3.s}, za.s[w8, 0, vgx4]
|
||||
fadd z0.s, p0/m, z0.s, z16.s
|
||||
fadd z1.s, p0/m, z1.s, z17.s
|
||||
fadd z2.s, p0/m, z2.s, z18.s
|
||||
fadd z3.s, p0/m, z3.s, z19.s
|
||||
mov za.s[w8, 0, vgx4], {z0.s-z3.s}
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- scalar / per_col ops (degenerate at NR=1; per_col == scalar) -----
|
||||
//
|
||||
// Per Phase 1's mapping:
|
||||
// ScalarSub → result = scalar - z (fsubr)
|
||||
// ScalarSubF → result = z - scalar (fsub)
|
||||
// Same convention applies to PerCol*.
|
||||
|
||||
{% macro scalar_op(label, op) %}
|
||||
{{label}}:
|
||||
ldr w2, [x0, #8]
|
||||
dup z4.s, w2
|
||||
mov {z0.s-z3.s}, za.s[w8, 0, vgx4]
|
||||
{{op}} z0.s, p0/m, z0.s, z4.s
|
||||
{{op}} z1.s, p0/m, z1.s, z4.s
|
||||
{{op}} z2.s, p0/m, z2.s, z4.s
|
||||
{{op}} z3.s, p0/m, z3.s, z4.s
|
||||
mov za.s[w8, 0, vgx4], {z0.s-z3.s}
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{{ scalar_op('.scalar_add', 'fadd') }}
|
||||
{{ scalar_op('.scalar_mul', 'fmul') }}
|
||||
{{ scalar_op('.scalar_sub', 'fsubr') }}
|
||||
{{ scalar_op('.scalar_sub_flipped', 'fsub') }}
|
||||
{{ scalar_op('.scalar_min', 'fmin') }}
|
||||
{{ scalar_op('.scalar_max', 'fmax') }}
|
||||
|
||||
// per_col at NR=1 takes a *pointer* to 1 f32 at [x0, #8]; dereference
|
||||
// and broadcast. Result is functionally identical to scalar but the
|
||||
// load path differs.
|
||||
|
||||
{% macro per_col_op(label, op) %}
|
||||
{{label}}:
|
||||
ldr x2, [x0, #8]
|
||||
ld1rw {z4.s}, p0/z, [x2]
|
||||
mov {z0.s-z3.s}, za.s[w8, 0, vgx4]
|
||||
{{op}} z0.s, p0/m, z0.s, z4.s
|
||||
{{op}} z1.s, p0/m, z1.s, z4.s
|
||||
{{op}} z2.s, p0/m, z2.s, z4.s
|
||||
{{op}} z3.s, p0/m, z3.s, z4.s
|
||||
mov za.s[w8, 0, vgx4], {z0.s-z3.s}
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{{ per_col_op('.per_col_add', 'fadd') }}
|
||||
{{ per_col_op('.per_col_mul', 'fmul') }}
|
||||
{{ per_col_op('.per_col_sub', 'fsubr') }}
|
||||
{{ per_col_op('.per_col_sub_flipped', 'fsub') }}
|
||||
{{ per_col_op('.per_col_min', 'fmin') }}
|
||||
{{ per_col_op('.per_col_max', 'fmax') }}
|
||||
|
||||
// -------- per_row ops: 64-element bias, lane-wise op against accumulator --
|
||||
|
||||
{% macro per_row_op(label, op) %}
|
||||
{{label}}:
|
||||
ldr x2, [x0, #8]
|
||||
ld1w {z16.s-z19.s}, pn8/z, [x2]
|
||||
mov {z0.s-z3.s}, za.s[w8, 0, vgx4]
|
||||
{{op}} z0.s, p0/m, z0.s, z16.s
|
||||
{{op}} z1.s, p0/m, z1.s, z17.s
|
||||
{{op}} z2.s, p0/m, z2.s, z18.s
|
||||
{{op}} z3.s, p0/m, z3.s, z19.s
|
||||
mov za.s[w8, 0, vgx4], {z0.s-z3.s}
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{{ per_row_op('.per_row_add', 'fadd') }}
|
||||
{{ per_row_op('.per_row_mul', 'fmul') }}
|
||||
{{ per_row_op('.per_row_sub', 'fsubr') }}
|
||||
{{ per_row_op('.per_row_sub_flipped', 'fsub') }}
|
||||
{{ per_row_op('.per_row_min', 'fmin') }}
|
||||
{{ per_row_op('.per_row_max', 'fmax') }}
|
||||
|
||||
// -------- not yet implemented ----------------------------------------------
|
||||
|
||||
.leaky_relu:
|
||||
.q_scale:
|
||||
.q_shl:
|
||||
.q_shr:
|
||||
b .unsupported
|
||||
|
||||
// -------- epilogue ---------------------------------------------------------
|
||||
|
||||
.return:
|
||||
smstop
|
||||
add sp, sp, #256
|
||||
ldp q14, q15, [sp, #96]
|
||||
ldp q12, q13, [sp, #64]
|
||||
ldp q10, q11, [sp, #32]
|
||||
ldp q8, q9, [sp], #128
|
||||
ret
|
||||
Vendored
+681
@@ -0,0 +1,681 @@
|
||||
// vim: ft=arm
|
||||
//
|
||||
// SME2 i32 32x32 quantized matmul kernel.
|
||||
//
|
||||
// ZA tile layout (4 .S tiles, 16x16 i32 each):
|
||||
// ZA0.S : C[0..16, 0..16] (top-left)
|
||||
// ZA1.S : C[0..16, 16..32] (top-right)
|
||||
// ZA2.S : C[16..32, 0..16] (bottom-left)
|
||||
// ZA3.S : C[16..32, 16..32] (bottom-right)
|
||||
//
|
||||
// Inner K-step (K decrements by 4 per iter, since SMOPA at i8 reduces 4):
|
||||
// ld1b {z0, z1}, pn8/z, [A] ; 32 M × 4 K = 128 i8 of A
|
||||
// ld1b {z2, z3}, pn8/z, [B] ; 32 N × 4 K = 128 i8 of B
|
||||
// smopa za0.s, p0/m, p0/m, z0.b, z2.b ; ZA0 += A[0..16] × B[0..16]
|
||||
// smopa za1.s, p0/m, p0/m, z0.b, z3.b
|
||||
// smopa za2.s, p0/m, p0/m, z1.b, z2.b
|
||||
// smopa za3.s, p0/m, p0/m, z1.b, z3.b
|
||||
//
|
||||
// SMOPA at i8 throughput: 4-way K reduction per insn × 16x16 cells = 1024
|
||||
// MACs per insn. With 4-tile rotation we approach 4 SMOPAs/cycle = 4 K
|
||||
// reduction × 16x16 = 4096 MACs/cycle ≈ ~16 TOPS theoretical peak.
|
||||
//
|
||||
// Calling convention (extern "C", AAPCS64):
|
||||
// x0 = const *FusedKerSpec<i32>, advanced 40 B per dispatcher iteration
|
||||
// x1 = 4 KiB scratch buffer for tile spills (used by store-generic / q_scale)
|
||||
//
|
||||
// Tract packing requirement: i8 inputs packed with K_alignment=4 (SMOPA
|
||||
// requires K%4=0). The PackedFormat::with_k_alignment(4) handles this.
|
||||
|
||||
.arch armv9-a+sme2
|
||||
.text
|
||||
.align 4
|
||||
|
||||
.global {{G}}sme_qmmm_i32_32x32_{{suffix}}
|
||||
{{G}}sme_qmmm_i32_32x32_{{suffix}}:
|
||||
|
||||
stp q8, q9, [sp, #-128]!
|
||||
stp q10, q11, [sp, #32]
|
||||
stp q12, q13, [sp, #64]
|
||||
stp q14, q15, [sp, #96]
|
||||
|
||||
sub sp, sp, #4096
|
||||
mov x1, sp
|
||||
|
||||
smstart
|
||||
ptrue p0.b
|
||||
ptrue pn8.b
|
||||
mov w8, #0
|
||||
|
||||
{% include "dispatcher.j2" %}
|
||||
|
||||
// -------- AddMatMul: ZA += A·B at i8 with K=4 reduction per SMOPA ----------
|
||||
|
||||
.add_mat_mul:
|
||||
ldr x9, [x0, #32] // packing index
|
||||
ldr x2, [x0, #24] // b ptr
|
||||
ldp x3, x4, [x0, #8] // k, a ptr
|
||||
cmp x3, #0
|
||||
b.eq .non_linear_loop
|
||||
cmp x9, #1
|
||||
b.eq .Lmatmul_loop
|
||||
// i32i32 fallback (packing != 1, auto-test path): ZA += A[:,k] (x) B[k,:], one
|
||||
// K-step at a time via predicated MLA rank-1 updates. One instruction per line:
|
||||
// the Apple/LLVM AArch64 assembler treats `;` as a COMMENT, so semicolon-packed
|
||||
// statements silently drop everything after the first `;`.
|
||||
.Lk32:
|
||||
ld1w {z2.s}, p0/z, [x2] // B[k, 0..16]
|
||||
ld1w {z3.s}, p0/z, [x2, #1, mul vl] // B[k, 16..32]
|
||||
mov w12, #0
|
||||
.Lkt:
|
||||
ldr w10, [x4, w12, uxtw #2] // A[k, w12]
|
||||
dup z4.s, w10
|
||||
mov z16.s, p0/m, za0h.s[w12, 0]
|
||||
mov z17.s, p0/m, za1h.s[w12, 0]
|
||||
mla z16.s, p0/m, z2.s, z4.s // C[w12, 0..16] += A[w12] * B[0..16]
|
||||
mla z17.s, p0/m, z3.s, z4.s // C[w12, 16..32] += A[w12] * B[16..32]
|
||||
mov za0h.s[w12, 0], p0/m, z16.s
|
||||
mov za1h.s[w12, 0], p0/m, z17.s
|
||||
add w10, w12, #16
|
||||
ldr w10, [x4, w10, uxtw #2] // A[k, w12+16]
|
||||
dup z4.s, w10
|
||||
mov z18.s, p0/m, za2h.s[w12, 0]
|
||||
mov z19.s, p0/m, za3h.s[w12, 0]
|
||||
mla z18.s, p0/m, z2.s, z4.s // C[w12+16, 0..16] += A[w12+16] * B[0..16]
|
||||
mla z19.s, p0/m, z3.s, z4.s // C[w12+16, 16..32] += A[w12+16] * B[16..32]
|
||||
mov za2h.s[w12, 0], p0/m, z18.s
|
||||
mov za3h.s[w12, 0], p0/m, z19.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lkt
|
||||
add x4, x4, #128
|
||||
add x2, x2, #128
|
||||
subs x3, x3, #1
|
||||
b.ne .Lk32
|
||||
b .non_linear_loop
|
||||
|
||||
.Lmatmul_loop:
|
||||
ld1b {z0.b, z1.b}, pn8/z, [x4]
|
||||
ld1b {z2.b, z3.b}, pn8/z, [x2]
|
||||
add x4, x4, #128
|
||||
add x2, x2, #128
|
||||
smopa za0.s, p0/m, p0/m, z0.b, z2.b
|
||||
smopa za1.s, p0/m, p0/m, z0.b, z3.b
|
||||
smopa za2.s, p0/m, p0/m, z1.b, z2.b
|
||||
smopa za3.s, p0/m, p0/m, z1.b, z3.b
|
||||
subs x3, x3, #4
|
||||
b.gt .Lmatmul_loop
|
||||
b .non_linear_loop
|
||||
|
||||
.clear:
|
||||
zero {za}
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- Store: i32 tile -> memory (port of Phase 1 f32 store) -----------
|
||||
|
||||
.store:
|
||||
ldp x5, x6, [x0, #8] // ptr, row_byte_stride
|
||||
ldp x7, x9, [x0, #24] // col_byte_stride, item_size
|
||||
|
||||
cmp x7, #4
|
||||
b.ne .Lstore_generic
|
||||
cmp x9, #4
|
||||
b.ne .Lstore_generic
|
||||
|
||||
add x4, x5, #64
|
||||
mov w12, #0
|
||||
.Lstore_top:
|
||||
st1w {za0h.s[w12, 0]}, p0, [x5]
|
||||
st1w {za1h.s[w12, 0]}, p0, [x4]
|
||||
add x5, x5, x6
|
||||
add x4, x4, x6
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lstore_top
|
||||
mov w12, #0
|
||||
.Lstore_bot:
|
||||
st1w {za2h.s[w12, 0]}, p0, [x5]
|
||||
st1w {za3h.s[w12, 0]}, p0, [x4]
|
||||
add x5, x5, x6
|
||||
add x4, x4, x6
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lstore_bot
|
||||
b .non_linear_loop
|
||||
|
||||
.Lstore_generic:
|
||||
mov x13, x9 // preserve item_size before x9 is reused as a ptr
|
||||
mov x4, x1
|
||||
add x9, x1, #64
|
||||
mov w12, #0
|
||||
.Lstore_spill_top:
|
||||
st1w {za0h.s[w12, 0]}, p0, [x4]
|
||||
st1w {za1h.s[w12, 0]}, p0, [x9]
|
||||
add x4, x4, #128
|
||||
add x9, x9, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lstore_spill_top
|
||||
mov w12, #0
|
||||
.Lstore_spill_bot:
|
||||
st1w {za2h.s[w12, 0]}, p0, [x4]
|
||||
st1w {za3h.s[w12, 0]}, p0, [x9]
|
||||
add x4, x4, #128
|
||||
add x9, x9, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lstore_spill_bot
|
||||
|
||||
mov x3, #0
|
||||
.Lstore_row:
|
||||
mov x4, x5
|
||||
mov x10, #0
|
||||
lsl x9, x3, #7
|
||||
add x11, x1, x9
|
||||
.Lstore_col:
|
||||
ldr w9, [x11], #4
|
||||
cmp x13, #1 // item_size: 1 -> strb, 2 -> strh, else (4) -> str
|
||||
b.eq .Lstore_b1
|
||||
cmp x13, #2
|
||||
b.eq .Lstore_b2
|
||||
str w9, [x4]
|
||||
b .Lstore_cnext
|
||||
.Lstore_b1:
|
||||
strb w9, [x4]
|
||||
b .Lstore_cnext
|
||||
.Lstore_b2:
|
||||
strh w9, [x4]
|
||||
.Lstore_cnext:
|
||||
add x4, x4, x7
|
||||
add x10, x10, #1
|
||||
cmp x10, #32
|
||||
b.lt .Lstore_col
|
||||
add x5, x5, x6
|
||||
add x3, x3, #1
|
||||
cmp x3, #32
|
||||
b.lt .Lstore_row
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- LoadTile: ZA := row-major i32 tile from memory -------------------
|
||||
|
||||
.load_tile:
|
||||
ldr x2, [x0, #16]
|
||||
add x4, x2, #64
|
||||
mov w12, #0
|
||||
.Lloadtile_top:
|
||||
ld1w {z6.s}, p0/z, [x2]
|
||||
ld1w {z7.s}, p0/z, [x4]
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add x2, x2, #128
|
||||
add x4, x4, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lloadtile_top
|
||||
mov w12, #0
|
||||
.Lloadtile_bot:
|
||||
ld1w {z6.s}, p0/z, [x2]
|
||||
ld1w {z7.s}, p0/z, [x4]
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add x2, x2, #128
|
||||
add x4, x4, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lloadtile_bot
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- AddUnicast: ZA += C (strided load + add) ------------------------
|
||||
|
||||
.add_unicast:
|
||||
ldp x5, x6, [x0, #8] // ptr, row_byte_stride
|
||||
ldp x7, x9, [x0, #24] // col_byte_stride, item_size
|
||||
|
||||
cmp x7, #4
|
||||
b.ne .Laddu_generic
|
||||
cmp x9, #4
|
||||
b.ne .Laddu_generic
|
||||
|
||||
add x4, x5, #64
|
||||
mov w12, #0
|
||||
.Laddu_top:
|
||||
ld1w {z8.s}, p0/z, [x5]
|
||||
ld1w {z9.s}, p0/z, [x4]
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
add z6.s, p0/m, z6.s, z8.s
|
||||
add z7.s, p0/m, z7.s, z9.s
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add x5, x5, x6
|
||||
add x4, x4, x6
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Laddu_top
|
||||
mov w12, #0
|
||||
.Laddu_bot:
|
||||
ld1w {z8.s}, p0/z, [x5]
|
||||
ld1w {z9.s}, p0/z, [x4]
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
add z6.s, p0/m, z6.s, z8.s
|
||||
add z7.s, p0/m, z7.s, z9.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add x5, x5, x6
|
||||
add x4, x4, x6
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Laddu_bot
|
||||
b .non_linear_loop
|
||||
|
||||
.Laddu_generic:
|
||||
// Strided gather to scratch, then contig accumulate (mirrors Phase 1).
|
||||
mov x3, #0
|
||||
mov x10, x1
|
||||
.Laddu_gather_row:
|
||||
mov x11, x5
|
||||
mov x4, #0
|
||||
.Laddu_gather_col:
|
||||
ldr w9, [x11]
|
||||
str w9, [x10], #4
|
||||
add x11, x11, x7
|
||||
add x4, x4, #1
|
||||
cmp x4, #32
|
||||
b.lt .Laddu_gather_col
|
||||
add x5, x5, x6
|
||||
add x3, x3, #1
|
||||
cmp x3, #32
|
||||
b.lt .Laddu_gather_row
|
||||
|
||||
mov x4, x1
|
||||
add x9, x1, #64
|
||||
mov w12, #0
|
||||
.Laddu_apply_top:
|
||||
ld1w {z8.s}, p0/z, [x4]
|
||||
ld1w {z10.s}, p0/z, [x9]
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
add z6.s, p0/m, z6.s, z8.s
|
||||
add z7.s, p0/m, z7.s, z10.s
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add x4, x4, #128
|
||||
add x9, x9, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Laddu_apply_top
|
||||
mov w12, #0
|
||||
.Laddu_apply_bot:
|
||||
ld1w {z8.s}, p0/z, [x4]
|
||||
ld1w {z10.s}, p0/z, [x9]
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
add z6.s, p0/m, z6.s, z8.s
|
||||
add z7.s, p0/m, z7.s, z10.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add x4, x4, #128
|
||||
add x9, x9, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Laddu_apply_bot
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- AddRowColProducts: ZA += rows ⊗ cols (i32 outer product) --------
|
||||
//
|
||||
// rows: 32 i32 (broadcast per M-row), cols: 32 i32 (lane vector per N-col).
|
||||
// Per ZA row, we need: ZA[i, j] += rows[i] * cols[j]. Slice-by-slice.
|
||||
|
||||
.add_row_col_products:
|
||||
ldp x2, x3, [x0, #8] // rows ptr, cols ptr
|
||||
ld1w {z4.s}, p0/z, [x3] // cols[0..16]
|
||||
ld1w {z5.s}, p0/z, [x3, #1, mul vl] // cols[16..32]
|
||||
|
||||
// Top 16 rows
|
||||
mov w12, #0
|
||||
.Larcp_top:
|
||||
ldr w9, [x2], #4
|
||||
dup z16.s, w9 // broadcast rows[i] to z16
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
mla z6.s, p0/m, z16.s, z4.s // z6 += z16 * cols[0..16]
|
||||
mla z7.s, p0/m, z16.s, z5.s // z7 += z16 * cols[16..32]
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Larcp_top
|
||||
// Bottom 16 rows
|
||||
mov w12, #0
|
||||
.Larcp_bot:
|
||||
ldr w9, [x2], #4
|
||||
dup z16.s, w9
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
mla z6.s, p0/m, z16.s, z4.s
|
||||
mla z7.s, p0/m, z16.s, z5.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Larcp_bot
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- scalar fuse ops: broadcast scalar, apply lane-wise --------------
|
||||
//
|
||||
// Sub vs SubF (matches Phase 1's f32 convention):
|
||||
// ScalarSub → result = scalar - z (mnemonic: subr)
|
||||
// ScalarSubF → result = z - scalar (mnemonic: sub)
|
||||
|
||||
{% macro scalar_op_i32(label, op) %}
|
||||
{{label}}:
|
||||
ldr w2, [x0, #8]
|
||||
dup z4.s, w2
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_top:
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z4.s
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_top
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_bot:
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z4.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_bot
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{{ scalar_op_i32('.scalar_add', 'add') }}
|
||||
{{ scalar_op_i32('.scalar_mul', 'mul') }}
|
||||
{{ scalar_op_i32('.scalar_sub', 'subr') }}
|
||||
{{ scalar_op_i32('.scalar_sub_flipped', 'sub') }}
|
||||
{{ scalar_op_i32('.scalar_min', 'smin') }}
|
||||
{{ scalar_op_i32('.scalar_max', 'smax') }}
|
||||
|
||||
// -------- per_col fuse ops: 32-elem vector, broadcast across M rows ------
|
||||
|
||||
{% macro per_col_op_i32(label, op) %}
|
||||
{{label}}:
|
||||
ldr x2, [x0, #8]
|
||||
ld1w {z4.s}, p0/z, [x2]
|
||||
ld1w {z5.s}, p0/z, [x2, #1, mul vl]
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_top:
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z5.s
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_top
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_bot:
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z5.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_bot
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{{ per_col_op_i32('.per_col_add', 'add') }}
|
||||
{{ per_col_op_i32('.per_col_mul', 'mul') }}
|
||||
{{ per_col_op_i32('.per_col_sub', 'subr') }}
|
||||
{{ per_col_op_i32('.per_col_sub_flipped', 'sub') }}
|
||||
{{ per_col_op_i32('.per_col_min', 'smin') }}
|
||||
{{ per_col_op_i32('.per_col_max', 'smax') }}
|
||||
|
||||
// -------- per_row fuse ops: 32-elem vector, one scalar per M row ---------
|
||||
|
||||
{% macro per_row_op_i32(label, op) %}
|
||||
{{label}}:
|
||||
ldr x2, [x0, #8]
|
||||
add x3, x2, #64
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_top:
|
||||
ldr w4, [x2], #4
|
||||
dup z4.s, w4
|
||||
mov z6.s, p0/m, za0h.s[w12, 0]
|
||||
mov z7.s, p0/m, za1h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z4.s
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_top
|
||||
mov w12, #0
|
||||
.L{{label|replace('.', '')}}_bot:
|
||||
ldr w4, [x3], #4
|
||||
dup z4.s, w4
|
||||
mov z6.s, p0/m, za2h.s[w12, 0]
|
||||
mov z7.s, p0/m, za3h.s[w12, 0]
|
||||
{{op}} z6.s, p0/m, z6.s, z4.s
|
||||
{{op}} z7.s, p0/m, z7.s, z4.s
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .L{{label|replace('.', '')}}_bot
|
||||
b .non_linear_loop
|
||||
{% endmacro %}
|
||||
|
||||
{{ per_row_op_i32('.per_row_add', 'add') }}
|
||||
{{ per_row_op_i32('.per_row_mul', 'mul') }}
|
||||
{{ per_row_op_i32('.per_row_sub', 'subr') }}
|
||||
{{ per_row_op_i32('.per_row_sub_flipped', 'sub') }}
|
||||
{{ per_row_op_i32('.per_row_min', 'smin') }}
|
||||
{{ per_row_op_i32('.per_row_max', 'smax') }}
|
||||
|
||||
// -------- Quantization fuse ops (bit-exact port of generic/rounding.rs) ----
|
||||
//
|
||||
// Strategy: spill the 32x32 i32 ZA tile to the 4 KiB scratch (x1), quantize
|
||||
// element-wise in SCALAR GP registers (streaming-mode legal: smull/lsr/asr/
|
||||
// cneg/cset/... are base A64 and unaffected by PSTATE.SM), then reload to ZA.
|
||||
// Quant is not the hot path; this mirrors the scalar approach already proven
|
||||
// in arm64/sve/sve_mmm_i32.c. Everything is inlined (no `bl` — a nested call
|
||||
// would clobber x30 and corrupt the final `ret`).
|
||||
//
|
||||
// Bit-exactness: the reference forms the FULL i64 product (mult*v) and does a
|
||||
// single magnitude-rounding shift by (shift+31) with a per-policy nudge. A
|
||||
// vector sqdmulh+srshl truncates the low 31 bits before the second shift, so
|
||||
// it is NOT equivalent — hence the i64 scalar port.
|
||||
//
|
||||
// RoundingPolicy: Native=0 Zero=1 Away=2 MinusInf=3 PlusInf=4 Even=5 Odd=6.
|
||||
|
||||
// Spill ZA0..ZA3 -> scratch[x1] as a contiguous 32x32 row-major i32 matrix
|
||||
// (same layout the generic store path uses). Clobbers x4, x9, w12.
|
||||
{% macro za_spill(sfx) %}
|
||||
mov x4, x1
|
||||
add x9, x1, #64
|
||||
mov w12, #0
|
||||
.Lspt_{{sfx}}:
|
||||
st1w {za0h.s[w12, 0]}, p0, [x4]
|
||||
st1w {za1h.s[w12, 0]}, p0, [x9]
|
||||
add x4, x4, #128
|
||||
add x9, x9, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lspt_{{sfx}}
|
||||
mov w12, #0
|
||||
.Lspb_{{sfx}}:
|
||||
st1w {za2h.s[w12, 0]}, p0, [x4]
|
||||
st1w {za3h.s[w12, 0]}, p0, [x9]
|
||||
add x4, x4, #128
|
||||
add x9, x9, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lspb_{{sfx}}
|
||||
{% endmacro %}
|
||||
|
||||
// Reload scratch[x1] (32x32 row-major i32) -> ZA0..ZA3. Clobbers x4,x9,w12,z6,z7.
|
||||
{% macro za_reload(sfx) %}
|
||||
mov x4, x1
|
||||
add x9, x1, #64
|
||||
mov w12, #0
|
||||
.Lrlt_{{sfx}}:
|
||||
ld1w {z6.s}, p0/z, [x4]
|
||||
ld1w {z7.s}, p0/z, [x9]
|
||||
mov za0h.s[w12, 0], p0/m, z6.s
|
||||
mov za1h.s[w12, 0], p0/m, z7.s
|
||||
add x4, x4, #128
|
||||
add x9, x9, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lrlt_{{sfx}}
|
||||
mov w12, #0
|
||||
.Lrlb_{{sfx}}:
|
||||
ld1w {z6.s}, p0/z, [x4]
|
||||
ld1w {z7.s}, p0/z, [x9]
|
||||
mov za2h.s[w12, 0], p0/m, z6.s
|
||||
mov za3h.s[w12, 0], p0/m, z7.s
|
||||
add x4, x4, #128
|
||||
add x9, x9, #128
|
||||
add w12, w12, #1
|
||||
cmp w12, #16
|
||||
b.lt .Lrlb_{{sfx}}
|
||||
{% endmacro %}
|
||||
|
||||
// Magnitude-rounding shared by q_scale and q_shr (mirrors `Mul<i32> for Scaler`
|
||||
// / `i32::q_shr`). In: x12 = val (i64), x5 = shift, x6 = policy. Out: w14 (i32).
|
||||
// Clobbers x13,x15,x16,x17. Preserves x5,x6,x7,x10,x11,x12.
|
||||
{% macro round_mag(sfx) %}
|
||||
cmp x5, #0
|
||||
b.gt .Lrpos_{{sfx}}
|
||||
neg x13, x5
|
||||
lsl x13, x12, x13 // val << (-shift)
|
||||
mov w14, w13
|
||||
b .Lrend_{{sfx}}
|
||||
.Lrpos_{{sfx}}:
|
||||
cmp x12, #0
|
||||
cneg x15, x12, mi // x15 = |val|
|
||||
sub x13, x5, #1
|
||||
mov x16, #1
|
||||
lsl x16, x16, x13 // x16 = half = 1 << (shift-1)
|
||||
cmp x6, #2 // Away -> nudge 0
|
||||
b.eq .Lrn0_{{sfx}}
|
||||
cmp x6, #1 // Zero -> nudge -1
|
||||
b.ne .Lrna_{{sfx}}
|
||||
mov x17, #-1
|
||||
b .Lrnd_{{sfx}}
|
||||
.Lrna_{{sfx}}:
|
||||
cmp x6, #3 // MinusInf -> -(val >= 0)
|
||||
b.ne .Lrnb_{{sfx}}
|
||||
cmp x12, #0
|
||||
cset x17, ge
|
||||
neg x17, x17
|
||||
b .Lrnd_{{sfx}}
|
||||
.Lrnb_{{sfx}}:
|
||||
cmp x6, #4 // PlusInf -> -(val <= 0)
|
||||
b.ne .Lrnc_{{sfx}}
|
||||
cmp x12, #0
|
||||
cset x17, le
|
||||
neg x17, x17
|
||||
b .Lrnd_{{sfx}}
|
||||
.Lrnc_{{sfx}}:
|
||||
cmp x6, #5 // Even -> ((|val|>>shift)&1) - 1
|
||||
b.ne .Lrno_{{sfx}}
|
||||
lsr x17, x15, x5
|
||||
and x17, x17, #1
|
||||
sub x17, x17, #1
|
||||
b .Lrnd_{{sfx}}
|
||||
.Lrno_{{sfx}}: // Odd -> -((|val|>>shift)&1)
|
||||
lsr x17, x15, x5
|
||||
and x17, x17, #1
|
||||
neg x17, x17
|
||||
b .Lrnd_{{sfx}}
|
||||
.Lrn0_{{sfx}}:
|
||||
mov x17, #0
|
||||
.Lrnd_{{sfx}}:
|
||||
add x15, x15, x16
|
||||
add x15, x15, x17
|
||||
lsr x15, x15, x5 // (|val| + half + nudge) >> shift
|
||||
cmp x12, #0
|
||||
cneg x14, x15, mi // signum(val) * mag
|
||||
.Lrend_{{sfx}}:
|
||||
{% endmacro %}
|
||||
|
||||
// QScale(shift, policy, mult): val = mult*v (i64); shift += 31; magnitude round.
|
||||
.q_scale:
|
||||
ldr x5, [x0, #8] // shift (isize)
|
||||
ldr x6, [x0, #16] // policy
|
||||
ldr w7, [x0, #24] // mult (i32)
|
||||
add x5, x5, #31
|
||||
{{ za_spill('qsc') }}
|
||||
mov x10, x1
|
||||
mov x11, #1024
|
||||
.Lqsc_loop:
|
||||
ldr w9, [x10]
|
||||
smull x12, w7, w9 // val = (i64)mult * (i64)v
|
||||
{{ round_mag('qsc') }}
|
||||
str w14, [x10], #4
|
||||
subs x11, x11, #1
|
||||
b.ne .Lqsc_loop
|
||||
{{ za_reload('qsc') }}
|
||||
b .non_linear_loop
|
||||
|
||||
// RoundingShiftRight(shift, policy): val = v (i64); magnitude round (shift>0).
|
||||
.q_shr:
|
||||
ldr x5, [x0, #8] // shift (usize, >= 1)
|
||||
ldr x6, [x0, #16] // policy
|
||||
{{ za_spill('qsr') }}
|
||||
mov x10, x1
|
||||
mov x11, #1024
|
||||
.Lqsr_loop:
|
||||
ldr w9, [x10]
|
||||
sxtw x12, w9 // val = (i64)v
|
||||
{{ round_mag('qsr') }}
|
||||
str w14, [x10], #4
|
||||
subs x11, x11, #1
|
||||
b.ne .Lqsr_loop
|
||||
{{ za_reload('qsr') }}
|
||||
b .non_linear_loop
|
||||
|
||||
// ShiftLeft(shift): result = v << shift (32-bit wrapping, matches i32::q_shl).
|
||||
.q_shl:
|
||||
ldr x5, [x0, #8] // shift (usize)
|
||||
{{ za_spill('qsl') }}
|
||||
mov x10, x1
|
||||
mov x11, #1024
|
||||
.Lqsl_loop:
|
||||
ldr w9, [x10]
|
||||
lsl w9, w9, w5
|
||||
str w9, [x10], #4
|
||||
subs x11, x11, #1
|
||||
b.ne .Lqsl_loop
|
||||
{{ za_reload('qsl') }}
|
||||
b .non_linear_loop
|
||||
|
||||
// -------- LeakyRelu (excluded via CAN_FUSE_I32) ---------------------------
|
||||
|
||||
.leaky_relu:
|
||||
b .unsupported
|
||||
|
||||
// -------- epilogue --------------------------------------------------------
|
||||
|
||||
.return:
|
||||
smstop
|
||||
add sp, sp, #4096
|
||||
ldp q14, q15, [sp, #96]
|
||||
ldp q12, q13, [sp, #64]
|
||||
ldp q10, q11, [sp, #32]
|
||||
ldp q8, q9, [sp], #128
|
||||
ret
|
||||
+153
@@ -0,0 +1,153 @@
|
||||
// SVE f16 GEMM kernel for tract's MMM framework (the mmm_f16 slot).
|
||||
//
|
||||
// Tile MR=8 x NR=8 with native f16 accumulation, the f16 sibling of
|
||||
// sve_mmm_f32.c. The hot AddMatMul is the same vector-length-agnostic
|
||||
// broadcast-A rank-1 update, but over f16 lanes: the NR columns are walked in
|
||||
// svcnth() chunks with whilelt predication and folded with svmla_n_f16 (native
|
||||
// f16 fused multiply-add, as the NEON arm64fp16 kernels do), so one binary is
|
||||
// correct and full-width at any SVE VL 128..2048-bit.
|
||||
//
|
||||
// Gated on FEAT_SVE2 AND FEAT_FP16 (Rust side). Built with +fp16. Consumes
|
||||
// tract's native f16 K-major packing. Fuse ops act on the MRxNR tile in memory
|
||||
// (scalar; not the hot path); as with the f32 kernel, LeakyRelu and the i32
|
||||
// quantization ops are excluded by CAN_FUSE. Returns 0 on success, 1 otherwise.
|
||||
|
||||
#include <arm_sve.h>
|
||||
#include <stdint.h>
|
||||
#include <string.h>
|
||||
|
||||
#define MR 8
|
||||
#define NR 8
|
||||
|
||||
enum {
|
||||
DONE = 0, CLEAR, LOAD_TILE,
|
||||
SCALAR_MIN, SCALAR_MAX, SCALAR_ADD, SCALAR_MUL, SCALAR_SUB, SCALAR_SUBF,
|
||||
LEAKY_RELU,
|
||||
PER_ROW_MIN, PER_ROW_MAX, PER_ROW_ADD, PER_ROW_MUL, PER_ROW_SUB, PER_ROW_SUBF,
|
||||
PER_COL_MIN, PER_COL_MAX, PER_COL_ADD, PER_COL_MUL, PER_COL_SUB, PER_COL_SUBF,
|
||||
Q_SCALE, Q_SHR, Q_SHL,
|
||||
ADD_UNICAST, ADD_ROW_COL_PRODUCTS, STORE, ADD_MAT_MUL
|
||||
};
|
||||
|
||||
typedef struct {
|
||||
uint64_t disc;
|
||||
uint64_t f0, f1, f2, f3;
|
||||
} spec_t;
|
||||
|
||||
static inline __fp16 f16_of(uint64_t bits) {
|
||||
__fp16 f;
|
||||
uint16_t lo = (uint16_t)bits;
|
||||
memcpy(&f, &lo, 2);
|
||||
return f;
|
||||
}
|
||||
|
||||
// AddMatMul: ab[m][n] += sum_k pa[k*MR+m] * pb[k*NR+n]. VLA over NR (f16 lanes).
|
||||
static void add_mat_mul(__fp16 ab[MR][NR], const __fp16 *pa, const __fp16 *pb, long k) {
|
||||
for (long n0 = 0; n0 < NR; n0 += svcnth()) {
|
||||
svbool_t pg = svwhilelt_b16((uint64_t)n0, (uint64_t)NR);
|
||||
svfloat16_t a0 = svld1_f16(pg, &ab[0][n0]), a1 = svld1_f16(pg, &ab[1][n0]);
|
||||
svfloat16_t a2 = svld1_f16(pg, &ab[2][n0]), a3 = svld1_f16(pg, &ab[3][n0]);
|
||||
svfloat16_t a4 = svld1_f16(pg, &ab[4][n0]), a5 = svld1_f16(pg, &ab[5][n0]);
|
||||
svfloat16_t a6 = svld1_f16(pg, &ab[6][n0]), a7 = svld1_f16(pg, &ab[7][n0]);
|
||||
for (long kk = 0; kk < k; kk++) {
|
||||
svfloat16_t b = svld1_f16(pg, &pb[kk * NR + n0]);
|
||||
const __fp16 *arow = &pa[kk * MR];
|
||||
a0 = svmla_n_f16_x(pg, a0, b, arow[0]);
|
||||
a1 = svmla_n_f16_x(pg, a1, b, arow[1]);
|
||||
a2 = svmla_n_f16_x(pg, a2, b, arow[2]);
|
||||
a3 = svmla_n_f16_x(pg, a3, b, arow[3]);
|
||||
a4 = svmla_n_f16_x(pg, a4, b, arow[4]);
|
||||
a5 = svmla_n_f16_x(pg, a5, b, arow[5]);
|
||||
a6 = svmla_n_f16_x(pg, a6, b, arow[6]);
|
||||
a7 = svmla_n_f16_x(pg, a7, b, arow[7]);
|
||||
}
|
||||
svst1_f16(pg, &ab[0][n0], a0); svst1_f16(pg, &ab[1][n0], a1);
|
||||
svst1_f16(pg, &ab[2][n0], a2); svst1_f16(pg, &ab[3][n0], a3);
|
||||
svst1_f16(pg, &ab[4][n0], a4); svst1_f16(pg, &ab[5][n0], a5);
|
||||
svst1_f16(pg, &ab[6][n0], a6); svst1_f16(pg, &ab[7][n0], a7);
|
||||
}
|
||||
}
|
||||
|
||||
// Store the MRxNR f16 tile with arbitrary row/col byte strides.
|
||||
static void store_tile(__fp16 ab[MR][NR], const spec_t *s) {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, cstride = (long)s->f2, isz = (long)s->f3;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) {
|
||||
uint8_t *p = ptr + i * rstride + j * cstride;
|
||||
if (isz == 2)
|
||||
*(__fp16 *)p = ab[i][j];
|
||||
else if (isz == 4)
|
||||
*(float *)p = (float)ab[i][j];
|
||||
else
|
||||
memcpy(p, &ab[i][j], isz);
|
||||
}
|
||||
}
|
||||
|
||||
intptr_t sve_mmm_f16_kernel(const spec_t *ops) {
|
||||
__fp16 ab[MR][NR];
|
||||
memset(ab, 0, sizeof(ab));
|
||||
for (const spec_t *s = ops;; s++) {
|
||||
switch (s->disc) {
|
||||
case DONE:
|
||||
return 0;
|
||||
case CLEAR:
|
||||
memset(ab, 0, sizeof(ab));
|
||||
break;
|
||||
case ADD_MAT_MUL: {
|
||||
long k = (long)s->f0;
|
||||
add_mat_mul(ab, (const __fp16 *)s->f1, (const __fp16 *)s->f2, k);
|
||||
break;
|
||||
}
|
||||
case STORE:
|
||||
store_tile(ab, s);
|
||||
break;
|
||||
case LOAD_TILE: {
|
||||
const __fp16 *src = (const __fp16 *)s->f1;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) ab[i][j] = src[i * NR + j];
|
||||
break;
|
||||
}
|
||||
case ADD_UNICAST: {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, cstride = (long)s->f2, isz = (long)s->f3;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) {
|
||||
const uint8_t *p = ptr + i * rstride + j * cstride;
|
||||
if (isz == 2)
|
||||
ab[i][j] += *(const __fp16 *)p;
|
||||
else
|
||||
ab[i][j] += (__fp16) * (const float *)p;
|
||||
}
|
||||
break;
|
||||
}
|
||||
case ADD_ROW_COL_PRODUCTS: {
|
||||
const __fp16 *rows = (const __fp16 *)s->f0;
|
||||
const __fp16 *cols = (const __fp16 *)s->f1;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) ab[i][j] += rows[i] * cols[j];
|
||||
break;
|
||||
}
|
||||
case SCALAR_MIN: { __fp16 v=f16_of(s->f0); for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]<v?ab[i][j]:v; break; }
|
||||
case SCALAR_MAX: { __fp16 v=f16_of(s->f0); for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]>v?ab[i][j]:v; break; }
|
||||
case SCALAR_ADD: { __fp16 v=f16_of(s->f0); for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]+=v; break; }
|
||||
case SCALAR_MUL: { __fp16 v=f16_of(s->f0); for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]*=v; break; }
|
||||
case SCALAR_SUB: { __fp16 v=f16_of(s->f0); for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=v-ab[i][j]; break; }
|
||||
case SCALAR_SUBF:{ __fp16 v=f16_of(s->f0); for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]-v; break; }
|
||||
case PER_ROW_MIN: { const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]<m[i]?ab[i][j]:m[i]; break; }
|
||||
case PER_ROW_MAX: { const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]>m[i]?ab[i][j]:m[i]; break; }
|
||||
case PER_ROW_ADD: { const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]+=m[i]; break; }
|
||||
case PER_ROW_MUL: { const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]*=m[i]; break; }
|
||||
case PER_ROW_SUB: { const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=m[i]-ab[i][j]; break; }
|
||||
case PER_ROW_SUBF:{ const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]-m[i]; break; }
|
||||
case PER_COL_MIN: { const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]<m[j]?ab[i][j]:m[j]; break; }
|
||||
case PER_COL_MAX: { const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]>m[j]?ab[i][j]:m[j]; break; }
|
||||
case PER_COL_ADD: { const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]+=m[j]; break; }
|
||||
case PER_COL_MUL: { const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]*=m[j]; break; }
|
||||
case PER_COL_SUB: { const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=m[j]-ab[i][j]; break; }
|
||||
case PER_COL_SUBF:{ const __fp16*m=(const __fp16*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]-m[j]; break; }
|
||||
default:
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
+159
@@ -0,0 +1,159 @@
|
||||
// SVE f32 GEMM kernel for tract's MMM framework.
|
||||
//
|
||||
// Tile MR=8 x NR=8. The hot AddMatMul uses the vector-length-agnostic
|
||||
// broadcast-A rank-1 update: the NR columns are walked in svcntw() chunks with
|
||||
// whilelt predication, so the SAME code runs correctly (and uses the full
|
||||
// vector) at any SVE vector length — 128 to 2048-bit. The MR=8 accumulators are
|
||||
// held in SVE registers across the K loop. Fuse ops operate on an MRxNR tile in
|
||||
// memory (scalar C — they are not the hot path), mirroring the generic kernel.
|
||||
//
|
||||
// ABI: the kernel walks a *const FusedKerSpec<f32> array (40 bytes / entry,
|
||||
// discriminant u64 at offset 0, fields at 8/16/24/32) until Done, exactly like
|
||||
// the asm dispatcher. f32 GEMM consumes tract's native K-major packing
|
||||
// (pa[k*MR+m], pb[k*NR+n]) so no custom packing format is required.
|
||||
//
|
||||
// Returns 0 on success, 1 if asked to do an unsupported fused op.
|
||||
|
||||
#include <arm_sve.h>
|
||||
#include <stdint.h>
|
||||
#include <string.h>
|
||||
|
||||
#define MR 8
|
||||
#define NR 8
|
||||
|
||||
// FusedKerSpec discriminants (must match frame/mmm/fuse.rs enum order).
|
||||
enum {
|
||||
DONE = 0, CLEAR, LOAD_TILE,
|
||||
SCALAR_MIN, SCALAR_MAX, SCALAR_ADD, SCALAR_MUL, SCALAR_SUB, SCALAR_SUBF,
|
||||
LEAKY_RELU,
|
||||
PER_ROW_MIN, PER_ROW_MAX, PER_ROW_ADD, PER_ROW_MUL, PER_ROW_SUB, PER_ROW_SUBF,
|
||||
PER_COL_MIN, PER_COL_MAX, PER_COL_ADD, PER_COL_MUL, PER_COL_SUB, PER_COL_SUBF,
|
||||
Q_SCALE, Q_SHR, Q_SHL,
|
||||
ADD_UNICAST, ADD_ROW_COL_PRODUCTS, STORE, ADD_MAT_MUL
|
||||
};
|
||||
|
||||
typedef struct {
|
||||
uint64_t disc;
|
||||
uint64_t f0, f1, f2, f3; // fields at byte offsets 8, 16, 24, 32
|
||||
} spec_t;
|
||||
|
||||
// AddMatMul: ab[m][n] += sum_k pa[k*MR+m] * pb[k*NR+n]. VLA over NR.
|
||||
static void add_mat_mul(float ab[MR][NR], const float *pa, const float *pb, long k) {
|
||||
for (long n0 = 0; n0 < NR; n0 += svcntw()) {
|
||||
svbool_t pg = svwhilelt_b32((uint64_t)n0, (uint64_t)NR);
|
||||
svfloat32_t a0 = svld1_f32(pg, &ab[0][n0]), a1 = svld1_f32(pg, &ab[1][n0]);
|
||||
svfloat32_t a2 = svld1_f32(pg, &ab[2][n0]), a3 = svld1_f32(pg, &ab[3][n0]);
|
||||
svfloat32_t a4 = svld1_f32(pg, &ab[4][n0]), a5 = svld1_f32(pg, &ab[5][n0]);
|
||||
svfloat32_t a6 = svld1_f32(pg, &ab[6][n0]), a7 = svld1_f32(pg, &ab[7][n0]);
|
||||
for (long kk = 0; kk < k; kk++) {
|
||||
svfloat32_t b = svld1_f32(pg, &pb[kk * NR + n0]);
|
||||
const float *arow = &pa[kk * MR];
|
||||
a0 = svmla_n_f32_x(pg, a0, b, arow[0]);
|
||||
a1 = svmla_n_f32_x(pg, a1, b, arow[1]);
|
||||
a2 = svmla_n_f32_x(pg, a2, b, arow[2]);
|
||||
a3 = svmla_n_f32_x(pg, a3, b, arow[3]);
|
||||
a4 = svmla_n_f32_x(pg, a4, b, arow[4]);
|
||||
a5 = svmla_n_f32_x(pg, a5, b, arow[5]);
|
||||
a6 = svmla_n_f32_x(pg, a6, b, arow[6]);
|
||||
a7 = svmla_n_f32_x(pg, a7, b, arow[7]);
|
||||
}
|
||||
svst1_f32(pg, &ab[0][n0], a0); svst1_f32(pg, &ab[1][n0], a1);
|
||||
svst1_f32(pg, &ab[2][n0], a2); svst1_f32(pg, &ab[3][n0], a3);
|
||||
svst1_f32(pg, &ab[4][n0], a4); svst1_f32(pg, &ab[5][n0], a5);
|
||||
svst1_f32(pg, &ab[6][n0], a6); svst1_f32(pg, &ab[7][n0], a7);
|
||||
}
|
||||
}
|
||||
|
||||
static inline float f32_of(uint64_t bits) {
|
||||
float f;
|
||||
uint32_t lo = (uint32_t)bits;
|
||||
memcpy(&f, &lo, 4);
|
||||
return f;
|
||||
}
|
||||
|
||||
// Store the MRxNR tile to memory with arbitrary row/col byte strides.
|
||||
static void store_tile(float ab[MR][NR], const spec_t *s) {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, cstride = (long)s->f2, isz = (long)s->f3;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) {
|
||||
uint8_t *p = ptr + i * rstride + j * cstride;
|
||||
if (isz == 4)
|
||||
*(float *)p = ab[i][j];
|
||||
else
|
||||
memcpy(p, &ab[i][j], isz);
|
||||
}
|
||||
}
|
||||
|
||||
// Returns isize (64-bit) to match tract's kernel ABI — NOT int (would leave the
|
||||
// upper 32 bits of x0 undefined).
|
||||
intptr_t sve_mmm_f32_kernel(const spec_t *ops) {
|
||||
float ab[MR][NR];
|
||||
memset(ab, 0, sizeof(ab));
|
||||
for (const spec_t *s = ops;; s++) {
|
||||
switch (s->disc) {
|
||||
case DONE:
|
||||
return 0;
|
||||
case CLEAR:
|
||||
memset(ab, 0, sizeof(ab));
|
||||
break;
|
||||
case ADD_MAT_MUL: {
|
||||
long k = (long)s->f0;
|
||||
const float *pa = (const float *)s->f1;
|
||||
const float *pb = (const float *)s->f2;
|
||||
add_mat_mul(ab, pa, pb, k);
|
||||
break;
|
||||
}
|
||||
case STORE:
|
||||
store_tile(ab, s);
|
||||
break;
|
||||
case LOAD_TILE: {
|
||||
// LoadTile(col_major_ptr, row_major_ptr); use the row-major one.
|
||||
const float *src = (const float *)s->f1;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) ab[i][j] = src[i * NR + j];
|
||||
break;
|
||||
}
|
||||
case ADD_UNICAST: {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, cstride = (long)s->f2;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++)
|
||||
ab[i][j] += *(const float *)(ptr + i * rstride + j * cstride);
|
||||
break;
|
||||
}
|
||||
case ADD_ROW_COL_PRODUCTS: {
|
||||
const float *rows = (const float *)s->f0;
|
||||
const float *cols = (const float *)s->f1;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) ab[i][j] += rows[i] * cols[j];
|
||||
break;
|
||||
}
|
||||
// ---- scalar fuse ops ----
|
||||
case SCALAR_MIN: { float v = f32_of(s->f0); for (long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]<v?ab[i][j]:v; break; }
|
||||
case SCALAR_MAX: { float v = f32_of(s->f0); for (long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]>v?ab[i][j]:v; break; }
|
||||
case SCALAR_ADD: { float v = f32_of(s->f0); for (long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]+=v; break; }
|
||||
case SCALAR_MUL: { float v = f32_of(s->f0); for (long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]*=v; break; }
|
||||
case SCALAR_SUB: { float v = f32_of(s->f0); for (long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=v-ab[i][j]; break; }
|
||||
case SCALAR_SUBF:{ float v = f32_of(s->f0); for (long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]-v; break; }
|
||||
// ---- per-row fuse ops (one value per m-row) ----
|
||||
case PER_ROW_MIN: { const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]<m[i]?ab[i][j]:m[i]; break; }
|
||||
case PER_ROW_MAX: { const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]>m[i]?ab[i][j]:m[i]; break; }
|
||||
case PER_ROW_ADD: { const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]+=m[i]; break; }
|
||||
case PER_ROW_MUL: { const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]*=m[i]; break; }
|
||||
case PER_ROW_SUB: { const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=m[i]-ab[i][j]; break; }
|
||||
case PER_ROW_SUBF:{ const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]-m[i]; break; }
|
||||
// ---- per-col fuse ops (one value per n-col) ----
|
||||
case PER_COL_MIN: { const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]<m[j]?ab[i][j]:m[j]; break; }
|
||||
case PER_COL_MAX: { const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]>m[j]?ab[i][j]:m[j]; break; }
|
||||
case PER_COL_ADD: { const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]+=m[j]; break; }
|
||||
case PER_COL_MUL: { const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]*=m[j]; break; }
|
||||
case PER_COL_SUB: { const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=m[j]-ab[i][j]; break; }
|
||||
case PER_COL_SUBF:{ const float*m=(const float*)s->f0; for(long i=0;i<MR;i++) for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]-m[j]; break; }
|
||||
default:
|
||||
// LeakyRelu / QScale / RoundingShiftRight / ShiftLeft: excluded
|
||||
// by CAN_FUSE for f32 — should never arrive. Anything else: error.
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
+261
@@ -0,0 +1,261 @@
|
||||
// SVE int8 -> int32 GEMM kernel for tract's MMM framework (the qmmm_i32 slot).
|
||||
//
|
||||
// Tile MR=8 x NR=8, i32 accumulator. The hot AddMatMul is the vector-length-
|
||||
// agnostic *widening* broadcast-A rank-1 update: per K-step it loads NR signed
|
||||
// bytes of a B row and sign-extends them to i32 with `svld1sb_s32`, then folds
|
||||
// MR `svmla_n_s32` updates with the (sign-extended) scalar from the A column.
|
||||
// The NR columns are walked in svcntw() chunks with whilelt predication, so the
|
||||
// SAME code is correct and full-width at any SVE vector length (128..2048-bit).
|
||||
//
|
||||
// Why widening MLA and not SDOT: SDOT reduces 4 K-contiguous i8 per i32 lane, so
|
||||
// it needs A and B packed K-contiguous within a lane. tract's PackedFormat is
|
||||
// K-major (mn-inner: for each k, r contiguous mn values), which is exactly what
|
||||
// the per-k widening update consumes — and what arm64simd's i32 kernel uses via
|
||||
// NEON SMLAL. A SDOT/SMMLA path would need a custom interleaved packer; that is
|
||||
// a separate (max-throughput) kernel, not this one.
|
||||
//
|
||||
// int8 inputs arrive via tract's native i8i8 packing (AddMatMul packing == 1).
|
||||
// The default i32i32 packing (packing == 0) is also handled (scalar) so the
|
||||
// generic auto-test surface (mmm_packed_packed_tests i32i32:0) passes.
|
||||
//
|
||||
// ABI: identical 40-byte FusedKerSpec<i32> walk as the f32 kernel (discriminant
|
||||
// u64 at offset 0, fields at 8/16/24/32). Fuse ops act on the MRxNR i32 tile in
|
||||
// memory (scalar C — not the hot path), including the quantization ops
|
||||
// q_scale / q_shr (rounding) / q_shl, ported bit-exact from
|
||||
// linalg/src/generic/rounding.rs.
|
||||
//
|
||||
// Returns 0 on success, 1 if asked to do an unsupported fused op / packing.
|
||||
|
||||
#include <arm_sve.h>
|
||||
#include <stdint.h>
|
||||
#include <string.h>
|
||||
|
||||
#define MR 8
|
||||
#define NR 8
|
||||
|
||||
// FusedKerSpec discriminants (must match frame/mmm/fuse.rs enum order).
|
||||
enum {
|
||||
DONE = 0, CLEAR, LOAD_TILE,
|
||||
SCALAR_MIN, SCALAR_MAX, SCALAR_ADD, SCALAR_MUL, SCALAR_SUB, SCALAR_SUBF,
|
||||
LEAKY_RELU,
|
||||
PER_ROW_MIN, PER_ROW_MAX, PER_ROW_ADD, PER_ROW_MUL, PER_ROW_SUB, PER_ROW_SUBF,
|
||||
PER_COL_MIN, PER_COL_MAX, PER_COL_ADD, PER_COL_MUL, PER_COL_SUB, PER_COL_SUBF,
|
||||
Q_SCALE, Q_SHR, Q_SHL,
|
||||
ADD_UNICAST, ADD_ROW_COL_PRODUCTS, STORE, ADD_MAT_MUL
|
||||
};
|
||||
|
||||
// RoundingPolicy is #[repr(usize)] in fuse.rs: Native=0, Zero=1, Away=2,
|
||||
// MinusInf=3, PlusInf=4, Even=5, Odd=6.
|
||||
enum { RP_NATIVE = 0, RP_ZERO, RP_AWAY, RP_MINUSINF, RP_PLUSINF, RP_EVEN, RP_ODD };
|
||||
|
||||
typedef struct {
|
||||
uint64_t disc;
|
||||
uint64_t f0, f1, f2, f3; // fields at byte offsets 8, 16, 24, 32
|
||||
} spec_t;
|
||||
|
||||
// AddMatMul, i8 x i8 -> i32 (packing 1): ab[m][n] += sum_k pa[k*MR+m]*pb[k*NR+n].
|
||||
// VLA widening rank-1 update over NR.
|
||||
static void add_mat_mul_i8(int32_t ab[MR][NR], const int8_t *pa, const int8_t *pb, long k) {
|
||||
for (long n0 = 0; n0 < NR; n0 += svcntw()) {
|
||||
svbool_t pg = svwhilelt_b32((uint64_t)n0, (uint64_t)NR);
|
||||
svint32_t a0 = svld1_s32(pg, &ab[0][n0]), a1 = svld1_s32(pg, &ab[1][n0]);
|
||||
svint32_t a2 = svld1_s32(pg, &ab[2][n0]), a3 = svld1_s32(pg, &ab[3][n0]);
|
||||
svint32_t a4 = svld1_s32(pg, &ab[4][n0]), a5 = svld1_s32(pg, &ab[5][n0]);
|
||||
svint32_t a6 = svld1_s32(pg, &ab[6][n0]), a7 = svld1_s32(pg, &ab[7][n0]);
|
||||
for (long kk = 0; kk < k; kk++) {
|
||||
// Load NR int8 of B row kk, sign-extending each lane to i32.
|
||||
svint32_t b = svld1sb_s32(pg, &pb[kk * NR + n0]);
|
||||
const int8_t *arow = &pa[kk * MR];
|
||||
a0 = svmla_n_s32_x(pg, a0, b, (int32_t)arow[0]);
|
||||
a1 = svmla_n_s32_x(pg, a1, b, (int32_t)arow[1]);
|
||||
a2 = svmla_n_s32_x(pg, a2, b, (int32_t)arow[2]);
|
||||
a3 = svmla_n_s32_x(pg, a3, b, (int32_t)arow[3]);
|
||||
a4 = svmla_n_s32_x(pg, a4, b, (int32_t)arow[4]);
|
||||
a5 = svmla_n_s32_x(pg, a5, b, (int32_t)arow[5]);
|
||||
a6 = svmla_n_s32_x(pg, a6, b, (int32_t)arow[6]);
|
||||
a7 = svmla_n_s32_x(pg, a7, b, (int32_t)arow[7]);
|
||||
}
|
||||
svst1_s32(pg, &ab[0][n0], a0); svst1_s32(pg, &ab[1][n0], a1);
|
||||
svst1_s32(pg, &ab[2][n0], a2); svst1_s32(pg, &ab[3][n0], a3);
|
||||
svst1_s32(pg, &ab[4][n0], a4); svst1_s32(pg, &ab[5][n0], a5);
|
||||
svst1_s32(pg, &ab[6][n0], a6); svst1_s32(pg, &ab[7][n0], a7);
|
||||
}
|
||||
}
|
||||
|
||||
// AddMatMul, i32 x i32 -> i32 (packing 0, default): only used by the auto-test
|
||||
// surface, never in production (quantized matmul uses the i8i8 packing). Scalar.
|
||||
static void add_mat_mul_i32(int32_t ab[MR][NR], const int32_t *pa, const int32_t *pb, long k) {
|
||||
for (long kk = 0; kk < k; kk++) {
|
||||
const int32_t *arow = &pa[kk * MR], *brow = &pb[kk * NR];
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) ab[i][j] += arow[i] * brow[j];
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Quantization helpers, ported bit-exact from generic/rounding.rs ----
|
||||
|
||||
// i32::q_shr(shift, rp): rounding arithmetic shift right.
|
||||
static int32_t q_shr_i32(int32_t v, long shift, int rp) {
|
||||
int32_t half = (int32_t)1 << (shift - 1);
|
||||
int32_t a = v < 0 ? -v : v; // abs (test inputs are small; matches Rust .abs())
|
||||
int32_t nudge;
|
||||
switch (rp) {
|
||||
case RP_ZERO: nudge = -1; break;
|
||||
case RP_MINUSINF: nudge = -(int32_t)(v >= 0); break;
|
||||
case RP_PLUSINF: nudge = -(int32_t)(v <= 0); break;
|
||||
case RP_AWAY: nudge = 0; break;
|
||||
case RP_EVEN: nudge = ((a >> shift) & 0x1) - 1; break;
|
||||
case RP_ODD: nudge = -((a >> shift) & 0x1); break;
|
||||
default: nudge = 0; break; // Native: unreachable for q ops
|
||||
}
|
||||
int32_t sign = (v > 0) - (v < 0); // signum: -1 / 0 / 1
|
||||
return sign * ((a + half + nudge) >> shift);
|
||||
}
|
||||
|
||||
// i32::q_scale(Scaler{mult, shift, policy}) with mult always present (the QScale
|
||||
// fused op carries an explicit multiplier). Mirrors `Mul<i32> for Scaler`.
|
||||
static int32_t q_scale_i32(int32_t v, long shift_in, int policy, int32_t mult) {
|
||||
int64_t val = (int64_t)mult * (int64_t)v;
|
||||
long shift = shift_in + 31;
|
||||
if (shift > 0) {
|
||||
int64_t half = (int64_t)1 << (shift - 1);
|
||||
int64_t a = val < 0 ? -val : val;
|
||||
int64_t nudge;
|
||||
switch (policy) {
|
||||
case RP_ZERO: nudge = -1; break;
|
||||
case RP_MINUSINF: nudge = -(int64_t)(val >= 0); break;
|
||||
case RP_PLUSINF: nudge = -(int64_t)(val <= 0); break;
|
||||
case RP_AWAY: nudge = 0; break;
|
||||
case RP_EVEN: nudge = ((a >> shift) & 0x1) - 1; break;
|
||||
case RP_ODD: nudge = -((a >> shift) & 0x1); break;
|
||||
default: nudge = 0; break;
|
||||
}
|
||||
int64_t sign = (val > 0) - (val < 0);
|
||||
return (int32_t)(sign * ((a + half + nudge) >> shift));
|
||||
} else {
|
||||
return (int32_t)(val << (-shift));
|
||||
}
|
||||
}
|
||||
|
||||
// Store the MRxNR i32 tile to memory with arbitrary row/col byte strides,
|
||||
// truncating to the destination item size (matches generic store_t semantics
|
||||
// for the tested widths 1 and 4).
|
||||
static void store_tile(int32_t ab[MR][NR], const spec_t *s) {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, cstride = (long)s->f2, isz = (long)s->f3;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) {
|
||||
uint8_t *p = ptr + i * rstride + j * cstride;
|
||||
int32_t v = ab[i][j];
|
||||
switch (isz) {
|
||||
case 1: *(uint8_t *)p = (uint8_t)v; break;
|
||||
case 2: *(uint16_t *)p = (uint16_t)v; break;
|
||||
case 4: *(int32_t *)p = v; break;
|
||||
case 8: { int64_t w = v; memcpy(p, &w, 8); break; }
|
||||
default: memcpy(p, &v, isz < 4 ? (size_t)isz : 4); break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Returns isize (64-bit) to match tract's kernel ABI.
|
||||
intptr_t sve_mmm_i32_kernel(const spec_t *ops) {
|
||||
int32_t ab[MR][NR];
|
||||
memset(ab, 0, sizeof(ab));
|
||||
for (const spec_t *s = ops;; s++) {
|
||||
switch (s->disc) {
|
||||
case DONE:
|
||||
return 0;
|
||||
case CLEAR:
|
||||
memset(ab, 0, sizeof(ab));
|
||||
break;
|
||||
case ADD_MAT_MUL: {
|
||||
long k = (long)s->f0;
|
||||
long packing = (long)s->f3;
|
||||
if (packing == 1) {
|
||||
add_mat_mul_i8(ab, (const int8_t *)s->f1, (const int8_t *)s->f2, k);
|
||||
} else if (packing == 0) {
|
||||
add_mat_mul_i32(ab, (const int32_t *)s->f1, (const int32_t *)s->f2, k);
|
||||
} else {
|
||||
return 1;
|
||||
}
|
||||
break;
|
||||
}
|
||||
case STORE:
|
||||
store_tile(ab, s);
|
||||
break;
|
||||
case LOAD_TILE: {
|
||||
// LoadTile(col_major_ptr, row_major_ptr); use the row-major one.
|
||||
const int32_t *src = (const int32_t *)s->f1;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) ab[i][j] = src[i * NR + j];
|
||||
break;
|
||||
}
|
||||
case ADD_UNICAST: {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, cstride = (long)s->f2, isz = (long)s->f3;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) {
|
||||
const uint8_t *p = ptr + i * rstride + j * cstride;
|
||||
if (isz == 1)
|
||||
ab[i][j] += *(const int8_t *)p; // sign-extend
|
||||
else
|
||||
ab[i][j] += *(const int32_t *)p;
|
||||
}
|
||||
break;
|
||||
}
|
||||
case ADD_ROW_COL_PRODUCTS: {
|
||||
const int32_t *rows = (const int32_t *)s->f0;
|
||||
const int32_t *cols = (const int32_t *)s->f1;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) ab[i][j] += rows[i] * cols[j];
|
||||
break;
|
||||
}
|
||||
// ---- quantization fuse ops ----
|
||||
case Q_SCALE: {
|
||||
long shift = (long)s->f0;
|
||||
int policy = (int)s->f1;
|
||||
int32_t mult = (int32_t)s->f2;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) ab[i][j] = q_scale_i32(ab[i][j], shift, policy, mult);
|
||||
break;
|
||||
}
|
||||
case Q_SHR: {
|
||||
long shift = (long)s->f0;
|
||||
int policy = (int)s->f1;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) ab[i][j] = q_shr_i32(ab[i][j], shift, policy);
|
||||
break;
|
||||
}
|
||||
case Q_SHL: {
|
||||
long shift = (long)s->f0;
|
||||
for (long i = 0; i < MR; i++)
|
||||
for (long j = 0; j < NR; j++) ab[i][j] = ab[i][j] << shift;
|
||||
break;
|
||||
}
|
||||
// ---- scalar fuse ops (value is an i32 in the low 32 bits of f0) ----
|
||||
case SCALAR_MIN: { int32_t v=(int32_t)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]<v?ab[i][j]:v; break; }
|
||||
case SCALAR_MAX: { int32_t v=(int32_t)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]>v?ab[i][j]:v; break; }
|
||||
case SCALAR_ADD: { int32_t v=(int32_t)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]+=v; break; }
|
||||
case SCALAR_MUL: { int32_t v=(int32_t)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]*=v; break; }
|
||||
case SCALAR_SUB: { int32_t v=(int32_t)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=v-ab[i][j]; break; }
|
||||
case SCALAR_SUBF:{ int32_t v=(int32_t)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]-v; break; }
|
||||
// ---- per-row fuse ops (one i32 per m-row) ----
|
||||
case PER_ROW_MIN: { const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]<m[i]?ab[i][j]:m[i]; break; }
|
||||
case PER_ROW_MAX: { const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]>m[i]?ab[i][j]:m[i]; break; }
|
||||
case PER_ROW_ADD: { const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]+=m[i]; break; }
|
||||
case PER_ROW_MUL: { const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]*=m[i]; break; }
|
||||
case PER_ROW_SUB: { const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=m[i]-ab[i][j]; break; }
|
||||
case PER_ROW_SUBF:{ const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]-m[i]; break; }
|
||||
// ---- per-col fuse ops (one i32 per n-col) ----
|
||||
case PER_COL_MIN: { const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]<m[j]?ab[i][j]:m[j]; break; }
|
||||
case PER_COL_MAX: { const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]>m[j]?ab[i][j]:m[j]; break; }
|
||||
case PER_COL_ADD: { const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]+=m[j]; break; }
|
||||
case PER_COL_MUL: { const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]*=m[j]; break; }
|
||||
case PER_COL_SUB: { const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=m[j]-ab[i][j]; break; }
|
||||
case PER_COL_SUBF:{ const int32_t*m=(const int32_t*)s->f0; for(long i=0;i<MR;i++)for(long j=0;j<NR;j++) ab[i][j]=ab[i][j]-m[j]; break; }
|
||||
default:
|
||||
// LeakyRelu excluded by CAN_FUSE_I32; anything else is an error.
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+210
@@ -0,0 +1,210 @@
|
||||
// SVE int8 -> int32 GEMV kernel for tract's MMM framework (the qmmv_i32 slot,
|
||||
// dispatched when N == 1: matrix x int8 column vector).
|
||||
//
|
||||
// Tile MR=64 x NR=1, i32 accumulator. The hot AddMatMul is the vector-length-
|
||||
// agnostic widening update vectorized over M: per K-step it loads MR signed
|
||||
// bytes of the A-panel column, sign-extends to i32 (svld1sb_s32), and folds a
|
||||
// single svmla_n_s32 with the (sign-extended) scalar B[k]. The MR rows are
|
||||
// walked in svcntw() chunks with whilelt predication, so the SAME code is
|
||||
// correct and full-width at any SVE vector length (128..2048-bit).
|
||||
//
|
||||
// Same rationale as the 8x8 kernel: widening MLA (not SDOT) consumes tract's
|
||||
// native K-major i8i8 packing directly. int8 inputs arrive via that packing
|
||||
// (AddMatMul packing == 1); the default i32i32 packing (packing == 0) is handled
|
||||
// scalar for the auto-test surface.
|
||||
//
|
||||
// ABI: identical 40-byte FusedKerSpec<i32> walk. At NR=1, per_col / scalar fuse
|
||||
// ops degenerate to a single broadcast value and per_row is element-wise over
|
||||
// the MR outputs. Quantization ops q_scale / q_shr / q_shl are ported bit-exact
|
||||
// from linalg/src/generic/rounding.rs. Returns 0 on success, 1 on an
|
||||
// unsupported fused op / packing.
|
||||
|
||||
#include <arm_sve.h>
|
||||
#include <stdint.h>
|
||||
#include <string.h>
|
||||
|
||||
#define MR 64
|
||||
#define NR 1
|
||||
|
||||
enum {
|
||||
DONE = 0, CLEAR, LOAD_TILE,
|
||||
SCALAR_MIN, SCALAR_MAX, SCALAR_ADD, SCALAR_MUL, SCALAR_SUB, SCALAR_SUBF,
|
||||
LEAKY_RELU,
|
||||
PER_ROW_MIN, PER_ROW_MAX, PER_ROW_ADD, PER_ROW_MUL, PER_ROW_SUB, PER_ROW_SUBF,
|
||||
PER_COL_MIN, PER_COL_MAX, PER_COL_ADD, PER_COL_MUL, PER_COL_SUB, PER_COL_SUBF,
|
||||
Q_SCALE, Q_SHR, Q_SHL,
|
||||
ADD_UNICAST, ADD_ROW_COL_PRODUCTS, STORE, ADD_MAT_MUL
|
||||
};
|
||||
|
||||
enum { RP_NATIVE = 0, RP_ZERO, RP_AWAY, RP_MINUSINF, RP_PLUSINF, RP_EVEN, RP_ODD };
|
||||
|
||||
typedef struct {
|
||||
uint64_t disc;
|
||||
uint64_t f0, f1, f2, f3;
|
||||
} spec_t;
|
||||
|
||||
// AddMatMul, i8 x i8 -> i32 (packing 1): ab[m] += sum_k pa[k*MR+m]*pb[k].
|
||||
// VLA widening update over MR.
|
||||
static void add_mat_mul_i8(int32_t ab[MR], const int8_t *pa, const int8_t *pb, long k) {
|
||||
for (long m0 = 0; m0 < MR; m0 += svcntw()) {
|
||||
svbool_t pg = svwhilelt_b32((uint64_t)m0, (uint64_t)MR);
|
||||
svint32_t acc = svld1_s32(pg, &ab[m0]);
|
||||
for (long kk = 0; kk < k; kk++) {
|
||||
svint32_t a = svld1sb_s32(pg, &pa[kk * MR + m0]); // load i8 col, sign-extend
|
||||
acc = svmla_n_s32_x(pg, acc, a, (int32_t)pb[kk]);
|
||||
}
|
||||
svst1_s32(pg, &ab[m0], acc);
|
||||
}
|
||||
}
|
||||
|
||||
// AddMatMul, i32 x i32 -> i32 (packing 0, default): auto-test surface only.
|
||||
static void add_mat_mul_i32(int32_t ab[MR], const int32_t *pa, const int32_t *pb, long k) {
|
||||
for (long kk = 0; kk < k; kk++) {
|
||||
int32_t b = pb[kk];
|
||||
const int32_t *acol = &pa[kk * MR];
|
||||
for (long m = 0; m < MR; m++) ab[m] += acol[m] * b;
|
||||
}
|
||||
}
|
||||
|
||||
// ---- quantization helpers, ported bit-exact from generic/rounding.rs ----
|
||||
|
||||
static int32_t q_shr_i32(int32_t v, long shift, int rp) {
|
||||
int32_t half = (int32_t)1 << (shift - 1);
|
||||
int32_t a = v < 0 ? -v : v;
|
||||
int32_t nudge;
|
||||
switch (rp) {
|
||||
case RP_ZERO: nudge = -1; break;
|
||||
case RP_MINUSINF: nudge = -(int32_t)(v >= 0); break;
|
||||
case RP_PLUSINF: nudge = -(int32_t)(v <= 0); break;
|
||||
case RP_AWAY: nudge = 0; break;
|
||||
case RP_EVEN: nudge = ((a >> shift) & 0x1) - 1; break;
|
||||
case RP_ODD: nudge = -((a >> shift) & 0x1); break;
|
||||
default: nudge = 0; break;
|
||||
}
|
||||
int32_t sign = (v > 0) - (v < 0);
|
||||
return sign * ((a + half + nudge) >> shift);
|
||||
}
|
||||
|
||||
static int32_t q_scale_i32(int32_t v, long shift_in, int policy, int32_t mult) {
|
||||
int64_t val = (int64_t)mult * (int64_t)v;
|
||||
long shift = shift_in + 31;
|
||||
if (shift > 0) {
|
||||
int64_t half = (int64_t)1 << (shift - 1);
|
||||
int64_t a = val < 0 ? -val : val;
|
||||
int64_t nudge;
|
||||
switch (policy) {
|
||||
case RP_ZERO: nudge = -1; break;
|
||||
case RP_MINUSINF: nudge = -(int64_t)(val >= 0); break;
|
||||
case RP_PLUSINF: nudge = -(int64_t)(val <= 0); break;
|
||||
case RP_AWAY: nudge = 0; break;
|
||||
case RP_EVEN: nudge = ((a >> shift) & 0x1) - 1; break;
|
||||
case RP_ODD: nudge = -((a >> shift) & 0x1); break;
|
||||
default: nudge = 0; break;
|
||||
}
|
||||
int64_t sign = (val > 0) - (val < 0);
|
||||
return (int32_t)(sign * ((a + half + nudge) >> shift));
|
||||
} else {
|
||||
return (int32_t)(val << (-shift));
|
||||
}
|
||||
}
|
||||
|
||||
intptr_t sve_mmm_i32_64x1_kernel(const spec_t *ops) {
|
||||
int32_t ab[MR];
|
||||
memset(ab, 0, sizeof(ab));
|
||||
for (const spec_t *s = ops;; s++) {
|
||||
switch (s->disc) {
|
||||
case DONE:
|
||||
return 0;
|
||||
case CLEAR:
|
||||
memset(ab, 0, sizeof(ab));
|
||||
break;
|
||||
case ADD_MAT_MUL: {
|
||||
long k = (long)s->f0, packing = (long)s->f3;
|
||||
if (packing == 1)
|
||||
add_mat_mul_i8(ab, (const int8_t *)s->f1, (const int8_t *)s->f2, k);
|
||||
else if (packing == 0)
|
||||
add_mat_mul_i32(ab, (const int32_t *)s->f1, (const int32_t *)s->f2, k);
|
||||
else
|
||||
return 1;
|
||||
break;
|
||||
}
|
||||
case STORE: {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, isz = (long)s->f3;
|
||||
for (long m = 0; m < MR; m++) {
|
||||
uint8_t *p = ptr + m * rstride;
|
||||
int32_t v = ab[m];
|
||||
switch (isz) {
|
||||
case 1: *(uint8_t *)p = (uint8_t)v; break;
|
||||
case 2: *(uint16_t *)p = (uint16_t)v; break;
|
||||
case 4: *(int32_t *)p = v; break;
|
||||
case 8: { int64_t w = v; memcpy(p, &w, 8); break; }
|
||||
default: memcpy(p, &v, isz < 4 ? (size_t)isz : 4); break;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case LOAD_TILE: {
|
||||
const int32_t *src = (const int32_t *)s->f1; // row-major MR values
|
||||
for (long m = 0; m < MR; m++) ab[m] = src[m];
|
||||
break;
|
||||
}
|
||||
case ADD_UNICAST: {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, isz = (long)s->f3;
|
||||
for (long m = 0; m < MR; m++) {
|
||||
const uint8_t *p = ptr + m * rstride;
|
||||
if (isz == 1)
|
||||
ab[m] += *(const int8_t *)p;
|
||||
else
|
||||
ab[m] += *(const int32_t *)p;
|
||||
}
|
||||
break;
|
||||
}
|
||||
case ADD_ROW_COL_PRODUCTS: {
|
||||
const int32_t *rows = (const int32_t *)s->f0;
|
||||
const int32_t *cols = (const int32_t *)s->f1;
|
||||
for (long m = 0; m < MR; m++) ab[m] += rows[m] * cols[0];
|
||||
break;
|
||||
}
|
||||
case Q_SCALE: {
|
||||
long shift = (long)s->f0; int policy = (int)s->f1; int32_t mult = (int32_t)s->f2;
|
||||
for (long m = 0; m < MR; m++) ab[m] = q_scale_i32(ab[m], shift, policy, mult);
|
||||
break;
|
||||
}
|
||||
case Q_SHR: {
|
||||
long shift = (long)s->f0; int policy = (int)s->f1;
|
||||
for (long m = 0; m < MR; m++) ab[m] = q_shr_i32(ab[m], shift, policy);
|
||||
break;
|
||||
}
|
||||
case Q_SHL: {
|
||||
long shift = (long)s->f0;
|
||||
for (long m = 0; m < MR; m++) ab[m] = ab[m] << shift;
|
||||
break;
|
||||
}
|
||||
// scalar fuse ops (single i32 in low 32 bits of f0)
|
||||
case SCALAR_MIN: { int32_t v=(int32_t)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]<v?ab[m]:v; break; }
|
||||
case SCALAR_MAX: { int32_t v=(int32_t)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]>v?ab[m]:v; break; }
|
||||
case SCALAR_ADD: { int32_t v=(int32_t)s->f0; for(long m=0;m<MR;m++) ab[m]+=v; break; }
|
||||
case SCALAR_MUL: { int32_t v=(int32_t)s->f0; for(long m=0;m<MR;m++) ab[m]*=v; break; }
|
||||
case SCALAR_SUB: { int32_t v=(int32_t)s->f0; for(long m=0;m<MR;m++) ab[m]=v-ab[m]; break; }
|
||||
case SCALAR_SUBF:{ int32_t v=(int32_t)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]-v; break; }
|
||||
// per-row fuse ops (one i32 per m-row)
|
||||
case PER_ROW_MIN: { const int32_t*m_=(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]<m_[m]?ab[m]:m_[m]; break; }
|
||||
case PER_ROW_MAX: { const int32_t*m_=(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]>m_[m]?ab[m]:m_[m]; break; }
|
||||
case PER_ROW_ADD: { const int32_t*m_=(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]+=m_[m]; break; }
|
||||
case PER_ROW_MUL: { const int32_t*m_=(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]*=m_[m]; break; }
|
||||
case PER_ROW_SUB: { const int32_t*m_=(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]=m_[m]-ab[m]; break; }
|
||||
case PER_ROW_SUBF:{ const int32_t*m_=(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]-m_[m]; break; }
|
||||
// per-col fuse ops degenerate to a single broadcast value at NR=1
|
||||
case PER_COL_MIN: { int32_t v=*(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]<v?ab[m]:v; break; }
|
||||
case PER_COL_MAX: { int32_t v=*(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]>v?ab[m]:v; break; }
|
||||
case PER_COL_ADD: { int32_t v=*(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]+=v; break; }
|
||||
case PER_COL_MUL: { int32_t v=*(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]*=v; break; }
|
||||
case PER_COL_SUB: { int32_t v=*(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]=v-ab[m]; break; }
|
||||
case PER_COL_SUBF:{ int32_t v=*(const int32_t*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]-v; break; }
|
||||
default:
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+132
@@ -0,0 +1,132 @@
|
||||
// SVE f16 GEMV kernel for tract's MMM framework (the mmv_f16 slot, dispatched
|
||||
// when N == 1: matrix x f16 column vector).
|
||||
//
|
||||
// Tile MR=64 x NR=1 with native f16 accumulation, the f16 sibling of
|
||||
// sve_mmv_f32_64x1.c. The hot AddMatMul is vectorized over M: per K-step it
|
||||
// loads MR f16 of the A-panel column and folds a single svmla_n_f16 (native f16
|
||||
// fused multiply-add) with the scalar B[k]. The MR rows are walked in svcnth()
|
||||
// chunks with whilelt predication, so one binary is correct and full-width at
|
||||
// any SVE VL 128..2048-bit.
|
||||
//
|
||||
// Gated on FEAT_SVE2 AND FEAT_FP16 (Rust side). Built with +fp16. At NR=1 the
|
||||
// per_col/scalar fuse ops degenerate to a broadcast and per_row is element-wise
|
||||
// over the MR outputs. LeakyRelu and the i32 quantization ops are excluded by
|
||||
// CAN_FUSE. Returns 0 on success, 1 otherwise.
|
||||
|
||||
#include <arm_sve.h>
|
||||
#include <stdint.h>
|
||||
#include <string.h>
|
||||
|
||||
#define MR 64
|
||||
#define NR 1
|
||||
|
||||
enum {
|
||||
DONE = 0, CLEAR, LOAD_TILE,
|
||||
SCALAR_MIN, SCALAR_MAX, SCALAR_ADD, SCALAR_MUL, SCALAR_SUB, SCALAR_SUBF,
|
||||
LEAKY_RELU,
|
||||
PER_ROW_MIN, PER_ROW_MAX, PER_ROW_ADD, PER_ROW_MUL, PER_ROW_SUB, PER_ROW_SUBF,
|
||||
PER_COL_MIN, PER_COL_MAX, PER_COL_ADD, PER_COL_MUL, PER_COL_SUB, PER_COL_SUBF,
|
||||
Q_SCALE, Q_SHR, Q_SHL,
|
||||
ADD_UNICAST, ADD_ROW_COL_PRODUCTS, STORE, ADD_MAT_MUL
|
||||
};
|
||||
|
||||
typedef struct {
|
||||
uint64_t disc;
|
||||
uint64_t f0, f1, f2, f3;
|
||||
} spec_t;
|
||||
|
||||
static inline __fp16 f16_of(uint64_t bits) {
|
||||
__fp16 f;
|
||||
uint16_t lo = (uint16_t)bits;
|
||||
memcpy(&f, &lo, 2);
|
||||
return f;
|
||||
}
|
||||
|
||||
// AddMatMul: ab[m] += sum_k pa[k*MR+m] * pb[k]. VLA over MR (f16 lanes).
|
||||
static void add_mat_mul(__fp16 ab[MR], const __fp16 *pa, const __fp16 *pb, long k) {
|
||||
for (long m0 = 0; m0 < MR; m0 += svcnth()) {
|
||||
svbool_t pg = svwhilelt_b16((uint64_t)m0, (uint64_t)MR);
|
||||
svfloat16_t acc = svld1_f16(pg, &ab[m0]);
|
||||
for (long kk = 0; kk < k; kk++) {
|
||||
svfloat16_t a = svld1_f16(pg, &pa[kk * MR + m0]);
|
||||
acc = svmla_n_f16_x(pg, acc, a, pb[kk]);
|
||||
}
|
||||
svst1_f16(pg, &ab[m0], acc);
|
||||
}
|
||||
}
|
||||
|
||||
intptr_t sve_mmv_f16_64x1_kernel(const spec_t *ops) {
|
||||
__fp16 ab[MR];
|
||||
memset(ab, 0, sizeof(ab));
|
||||
for (const spec_t *s = ops;; s++) {
|
||||
switch (s->disc) {
|
||||
case DONE:
|
||||
return 0;
|
||||
case CLEAR:
|
||||
memset(ab, 0, sizeof(ab));
|
||||
break;
|
||||
case ADD_MAT_MUL: {
|
||||
long k = (long)s->f0;
|
||||
add_mat_mul(ab, (const __fp16 *)s->f1, (const __fp16 *)s->f2, k);
|
||||
break;
|
||||
}
|
||||
case STORE: {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, isz = (long)s->f3;
|
||||
for (long m = 0; m < MR; m++) {
|
||||
uint8_t *p = ptr + m * rstride;
|
||||
if (isz == 2)
|
||||
*(__fp16 *)p = ab[m];
|
||||
else if (isz == 4)
|
||||
*(float *)p = (float)ab[m];
|
||||
else
|
||||
memcpy(p, &ab[m], isz);
|
||||
}
|
||||
break;
|
||||
}
|
||||
case LOAD_TILE: {
|
||||
const __fp16 *src = (const __fp16 *)s->f1;
|
||||
for (long m = 0; m < MR; m++) ab[m] = src[m];
|
||||
break;
|
||||
}
|
||||
case ADD_UNICAST: {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, isz = (long)s->f3;
|
||||
for (long m = 0; m < MR; m++) {
|
||||
const uint8_t *p = ptr + m * rstride;
|
||||
if (isz == 2)
|
||||
ab[m] += *(const __fp16 *)p;
|
||||
else
|
||||
ab[m] += (__fp16) * (const float *)p;
|
||||
}
|
||||
break;
|
||||
}
|
||||
case ADD_ROW_COL_PRODUCTS: {
|
||||
const __fp16 *rows = (const __fp16 *)s->f0;
|
||||
const __fp16 *cols = (const __fp16 *)s->f1;
|
||||
for (long m = 0; m < MR; m++) ab[m] += rows[m] * cols[0];
|
||||
break;
|
||||
}
|
||||
case SCALAR_MIN: { __fp16 v=f16_of(s->f0); for(long m=0;m<MR;m++) ab[m]=ab[m]<v?ab[m]:v; break; }
|
||||
case SCALAR_MAX: { __fp16 v=f16_of(s->f0); for(long m=0;m<MR;m++) ab[m]=ab[m]>v?ab[m]:v; break; }
|
||||
case SCALAR_ADD: { __fp16 v=f16_of(s->f0); for(long m=0;m<MR;m++) ab[m]+=v; break; }
|
||||
case SCALAR_MUL: { __fp16 v=f16_of(s->f0); for(long m=0;m<MR;m++) ab[m]*=v; break; }
|
||||
case SCALAR_SUB: { __fp16 v=f16_of(s->f0); for(long m=0;m<MR;m++) ab[m]=v-ab[m]; break; }
|
||||
case SCALAR_SUBF:{ __fp16 v=f16_of(s->f0); for(long m=0;m<MR;m++) ab[m]=ab[m]-v; break; }
|
||||
case PER_ROW_MIN: { const __fp16*m_=(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]<m_[m]?ab[m]:m_[m]; break; }
|
||||
case PER_ROW_MAX: { const __fp16*m_=(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]>m_[m]?ab[m]:m_[m]; break; }
|
||||
case PER_ROW_ADD: { const __fp16*m_=(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]+=m_[m]; break; }
|
||||
case PER_ROW_MUL: { const __fp16*m_=(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]*=m_[m]; break; }
|
||||
case PER_ROW_SUB: { const __fp16*m_=(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]=m_[m]-ab[m]; break; }
|
||||
case PER_ROW_SUBF:{ const __fp16*m_=(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]-m_[m]; break; }
|
||||
case PER_COL_MIN: { __fp16 v=*(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]<v?ab[m]:v; break; }
|
||||
case PER_COL_MAX: { __fp16 v=*(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]>v?ab[m]:v; break; }
|
||||
case PER_COL_ADD: { __fp16 v=*(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]+=v; break; }
|
||||
case PER_COL_MUL: { __fp16 v=*(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]*=v; break; }
|
||||
case PER_COL_SUB: { __fp16 v=*(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]=v-ab[m]; break; }
|
||||
case PER_COL_SUBF:{ __fp16 v=*(const __fp16*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]-v; break; }
|
||||
default:
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+128
@@ -0,0 +1,128 @@
|
||||
// SVE f32 GEMV kernel for tract's MMM framework (the mmv_f32 slot, dispatched
|
||||
// when N == 1: matrix x f32 column vector).
|
||||
//
|
||||
// Tile MR=64 x NR=1. The hot AddMatMul is the vector-length-agnostic update
|
||||
// vectorized over M: per K-step it loads MR f32 of the A-panel column and folds
|
||||
// a single svmla_n_f32 with the scalar B[k]. The MR rows are walked in svcntw()
|
||||
// chunks with whilelt predication, so the SAME code is correct and full-width at
|
||||
// any SVE vector length (128..2048-bit).
|
||||
//
|
||||
// Sibling of sve_mmm_f32.c (the 8x8 GEMM); shares the same FusedKerSpec<f32> ABI
|
||||
// and fuse-op surface. At NR=1, per_col / scalar fuse ops degenerate to a single
|
||||
// broadcast value and per_row is element-wise over the MR outputs. As with the
|
||||
// f32 GEMM kernel, LeakyRelu and the i32 quantization ops are excluded by
|
||||
// CAN_FUSE. Returns 0 on success, 1 on an unsupported fused op.
|
||||
|
||||
#include <arm_sve.h>
|
||||
#include <stdint.h>
|
||||
#include <string.h>
|
||||
|
||||
#define MR 64
|
||||
#define NR 1
|
||||
|
||||
enum {
|
||||
DONE = 0, CLEAR, LOAD_TILE,
|
||||
SCALAR_MIN, SCALAR_MAX, SCALAR_ADD, SCALAR_MUL, SCALAR_SUB, SCALAR_SUBF,
|
||||
LEAKY_RELU,
|
||||
PER_ROW_MIN, PER_ROW_MAX, PER_ROW_ADD, PER_ROW_MUL, PER_ROW_SUB, PER_ROW_SUBF,
|
||||
PER_COL_MIN, PER_COL_MAX, PER_COL_ADD, PER_COL_MUL, PER_COL_SUB, PER_COL_SUBF,
|
||||
Q_SCALE, Q_SHR, Q_SHL,
|
||||
ADD_UNICAST, ADD_ROW_COL_PRODUCTS, STORE, ADD_MAT_MUL
|
||||
};
|
||||
|
||||
typedef struct {
|
||||
uint64_t disc;
|
||||
uint64_t f0, f1, f2, f3;
|
||||
} spec_t;
|
||||
|
||||
static inline float f32_of(uint64_t bits) {
|
||||
float f;
|
||||
uint32_t lo = (uint32_t)bits;
|
||||
memcpy(&f, &lo, 4);
|
||||
return f;
|
||||
}
|
||||
|
||||
// AddMatMul: ab[m] += sum_k pa[k*MR+m] * pb[k]. VLA over MR.
|
||||
static void add_mat_mul(float ab[MR], const float *pa, const float *pb, long k) {
|
||||
for (long m0 = 0; m0 < MR; m0 += svcntw()) {
|
||||
svbool_t pg = svwhilelt_b32((uint64_t)m0, (uint64_t)MR);
|
||||
svfloat32_t acc = svld1_f32(pg, &ab[m0]);
|
||||
for (long kk = 0; kk < k; kk++) {
|
||||
svfloat32_t a = svld1_f32(pg, &pa[kk * MR + m0]);
|
||||
acc = svmla_n_f32_x(pg, acc, a, pb[kk]);
|
||||
}
|
||||
svst1_f32(pg, &ab[m0], acc);
|
||||
}
|
||||
}
|
||||
|
||||
intptr_t sve_mmv_f32_64x1_kernel(const spec_t *ops) {
|
||||
float ab[MR];
|
||||
memset(ab, 0, sizeof(ab));
|
||||
for (const spec_t *s = ops;; s++) {
|
||||
switch (s->disc) {
|
||||
case DONE:
|
||||
return 0;
|
||||
case CLEAR:
|
||||
memset(ab, 0, sizeof(ab));
|
||||
break;
|
||||
case ADD_MAT_MUL: {
|
||||
long k = (long)s->f0;
|
||||
add_mat_mul(ab, (const float *)s->f1, (const float *)s->f2, k);
|
||||
break;
|
||||
}
|
||||
case STORE: {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1, isz = (long)s->f3;
|
||||
for (long m = 0; m < MR; m++) {
|
||||
uint8_t *p = ptr + m * rstride;
|
||||
if (isz == 4)
|
||||
*(float *)p = ab[m];
|
||||
else
|
||||
memcpy(p, &ab[m], isz);
|
||||
}
|
||||
break;
|
||||
}
|
||||
case LOAD_TILE: {
|
||||
const float *src = (const float *)s->f1; // row-major MR values
|
||||
for (long m = 0; m < MR; m++) ab[m] = src[m];
|
||||
break;
|
||||
}
|
||||
case ADD_UNICAST: {
|
||||
uint8_t *ptr = (uint8_t *)s->f0;
|
||||
long rstride = (long)s->f1;
|
||||
for (long m = 0; m < MR; m++) ab[m] += *(const float *)(ptr + m * rstride);
|
||||
break;
|
||||
}
|
||||
case ADD_ROW_COL_PRODUCTS: {
|
||||
const float *rows = (const float *)s->f0;
|
||||
const float *cols = (const float *)s->f1;
|
||||
for (long m = 0; m < MR; m++) ab[m] += rows[m] * cols[0];
|
||||
break;
|
||||
}
|
||||
// scalar fuse ops (f32 bits in low 32 bits of f0)
|
||||
case SCALAR_MIN: { float v=f32_of(s->f0); for(long m=0;m<MR;m++) ab[m]=ab[m]<v?ab[m]:v; break; }
|
||||
case SCALAR_MAX: { float v=f32_of(s->f0); for(long m=0;m<MR;m++) ab[m]=ab[m]>v?ab[m]:v; break; }
|
||||
case SCALAR_ADD: { float v=f32_of(s->f0); for(long m=0;m<MR;m++) ab[m]+=v; break; }
|
||||
case SCALAR_MUL: { float v=f32_of(s->f0); for(long m=0;m<MR;m++) ab[m]*=v; break; }
|
||||
case SCALAR_SUB: { float v=f32_of(s->f0); for(long m=0;m<MR;m++) ab[m]=v-ab[m]; break; }
|
||||
case SCALAR_SUBF:{ float v=f32_of(s->f0); for(long m=0;m<MR;m++) ab[m]=ab[m]-v; break; }
|
||||
// per-row fuse ops (one f32 per m-row)
|
||||
case PER_ROW_MIN: { const float*m_=(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]<m_[m]?ab[m]:m_[m]; break; }
|
||||
case PER_ROW_MAX: { const float*m_=(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]>m_[m]?ab[m]:m_[m]; break; }
|
||||
case PER_ROW_ADD: { const float*m_=(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]+=m_[m]; break; }
|
||||
case PER_ROW_MUL: { const float*m_=(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]*=m_[m]; break; }
|
||||
case PER_ROW_SUB: { const float*m_=(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]=m_[m]-ab[m]; break; }
|
||||
case PER_ROW_SUBF:{ const float*m_=(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]-m_[m]; break; }
|
||||
// per-col fuse ops degenerate to a single broadcast value at NR=1
|
||||
case PER_COL_MIN: { float v=*(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]<v?ab[m]:v; break; }
|
||||
case PER_COL_MAX: { float v=*(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]>v?ab[m]:v; break; }
|
||||
case PER_COL_ADD: { float v=*(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]+=v; break; }
|
||||
case PER_COL_MUL: { float v=*(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]*=v; break; }
|
||||
case PER_COL_SUB: { float v=*(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]=v-ab[m]; break; }
|
||||
case PER_COL_SUBF:{ float v=*(const float*)s->f0; for(long m=0;m<MR;m++) ab[m]=ab[m]-v; break; }
|
||||
default:
|
||||
// LeakyRelu / QScale / RoundingShiftRight / ShiftLeft excluded by CAN_FUSE.
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user