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:
+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
|
||||
Reference in New Issue
Block a user