forked from tinygrad/tinygrad
Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 George HotzandGitHub
|
26f2049d43 | ||
 geohot
|
a8d13380f9 | ||
|
geohot
|
afb0463f10 | ||
|
geohot
|
947d6a3c69 |
@@ -41,6 +41,10 @@ inputs:
|
||||
description: "Install LLVM?"
|
||||
required: false
|
||||
default: 'false'
|
||||
tinydreno:
|
||||
description: "Install tinydreno"
|
||||
required: false
|
||||
default: 'false'
|
||||
qemu:
|
||||
description: "Install qemu"
|
||||
required: false
|
||||
@@ -273,6 +277,12 @@ runs:
|
||||
shell: bash
|
||||
run: brew install llvm@20
|
||||
|
||||
# *** tinydreno ***
|
||||
- name: Install tinydreno (linux)
|
||||
if: inputs.tinydreno == 'true' && runner.os == 'Linux'
|
||||
shell: bash
|
||||
run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
|
||||
|
||||
# *** OpenCL ***
|
||||
- name: Install rusticl
|
||||
if: inputs.opencl == 'true'
|
||||
|
||||
@@ -94,6 +94,7 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -148,6 +149,7 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -200,6 +202,7 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -249,6 +252,7 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
|
||||
@@ -179,3 +179,35 @@ jobs:
|
||||
- name: Run test_tiny
|
||||
shell: bash
|
||||
run: python -m pytest -n=auto test/test_tiny.py --durations=20
|
||||
|
||||
|
||||
qcomclcompiletests:
|
||||
name: Compile-only (QCOM CL)
|
||||
runs-on: ubuntu-24.04-arm
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: compile-qcomcl
|
||||
deps: testing_unit
|
||||
tinydreno: 'true'
|
||||
- name: Set env
|
||||
shell: bash
|
||||
run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV
|
||||
- name: Run test_ops
|
||||
shell: bash
|
||||
run: |
|
||||
python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
|
||||
DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
|
||||
python -m pytest -n=auto test/backend/test_ops.py --durations=20
|
||||
- name: Run test_ops (IMAGE)
|
||||
shell: bash
|
||||
env:
|
||||
IMAGE: 1
|
||||
DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
|
||||
run: |
|
||||
DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
|
||||
python -m pytest -n=auto test/backend/test_ops.py --durations=20
|
||||
|
||||
+35
-35
@@ -21,7 +21,7 @@ concurrency:
|
||||
jobs:
|
||||
docs:
|
||||
name: Docs
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 10
|
||||
env:
|
||||
CHECK_OOB: 0
|
||||
@@ -61,7 +61,7 @@ jobs:
|
||||
|
||||
torchbackend:
|
||||
name: Torch Backend Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -88,7 +88,7 @@ jobs:
|
||||
|
||||
bepython:
|
||||
name: Python Backend
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -126,7 +126,7 @@ jobs:
|
||||
|
||||
linter:
|
||||
name: Linters
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 10
|
||||
|
||||
steps:
|
||||
@@ -157,7 +157,7 @@ jobs:
|
||||
|
||||
nulltest:
|
||||
name: Null Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
@@ -191,7 +191,7 @@ jobs:
|
||||
|
||||
unittest:
|
||||
name: Unit Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
@@ -228,7 +228,7 @@ jobs:
|
||||
matrix:
|
||||
group: [1, 2]
|
||||
name: SPEC=2 (${{ matrix.group }})
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -244,7 +244,7 @@ jobs:
|
||||
|
||||
fuzzing:
|
||||
name: Fuzzing
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -260,7 +260,7 @@ jobs:
|
||||
|
||||
testopenclimage:
|
||||
name: CL IMAGE Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -280,7 +280,7 @@ jobs:
|
||||
|
||||
testopenpilot:
|
||||
name: openpilot Compile Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -309,7 +309,7 @@ jobs:
|
||||
|
||||
testonnxcpu:
|
||||
name: ONNX (CPU) Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 20
|
||||
|
||||
steps:
|
||||
@@ -328,7 +328,7 @@ jobs:
|
||||
|
||||
testoptim:
|
||||
name: Optimization Tests
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -360,7 +360,7 @@ jobs:
|
||||
|
||||
testllm:
|
||||
name: Test LLM
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
CHECK_OOB: 0
|
||||
@@ -387,7 +387,7 @@ jobs:
|
||||
|
||||
testmodels:
|
||||
name: Models
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -407,7 +407,7 @@ jobs:
|
||||
|
||||
testdsp:
|
||||
name: Linux (DSP)
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -435,7 +435,7 @@ jobs:
|
||||
- 'WEBGPU'
|
||||
|
||||
name: Linux (DEV=${{ matrix.dev }})
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -461,7 +461,7 @@ jobs:
|
||||
|
||||
testamdasm:
|
||||
name: AMD ASM IDE
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
DEV: MOCKKFD+AMD
|
||||
@@ -507,7 +507,7 @@ jobs:
|
||||
|
||||
hcq2:
|
||||
name: hcq2
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -529,7 +529,7 @@ jobs:
|
||||
|
||||
testmockam:
|
||||
name: Linux (am)
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
DEV: MOCKPCI+AMD
|
||||
@@ -565,7 +565,7 @@ jobs:
|
||||
arch: [gfx1100, gfx1201, gfx950]
|
||||
|
||||
name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
|
||||
@@ -603,7 +603,7 @@ jobs:
|
||||
backend: [ptx, nv]
|
||||
|
||||
name: Linux (${{ matrix.backend }})
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: *linux
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
FORWARD_ONLY: 1
|
||||
@@ -637,17 +637,10 @@ jobs:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
dev:
|
||||
- 'NULL:IR3:a630'
|
||||
- 'NULL:QCOMCL:a630'
|
||||
- 'NULL:NAK:sm_120'
|
||||
name: Compile-only (DEV=${{ matrix.dev }})
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
backend: [ir3, nak]
|
||||
name: Compile-only (${{ matrix.backend }})
|
||||
runs-on: *linux
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
NULL_ALLOW_COPYOUT: 1
|
||||
DEV: ${{ matrix.dev }}${{ contains(matrix.dev, 'a630') && ',IMAGE_PITCH_ALIGNMENT=64' || '' }}
|
||||
IMAGE: ${{ contains(matrix.dev, 'a630') && '1' || '0' }}
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v6
|
||||
@@ -656,14 +649,21 @@ jobs:
|
||||
with:
|
||||
key: compile-${{ matrix.backend }}
|
||||
deps: "testing_unit mesa"
|
||||
qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
|
||||
- name: Test IMAGE
|
||||
- name: Set env
|
||||
shell: bash
|
||||
if: contains(matrix.dev, 'a630')
|
||||
run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
|
||||
run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV
|
||||
- name: Run test_ops
|
||||
shell: bash
|
||||
run: |
|
||||
python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
|
||||
DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
|
||||
python -m pytest -n=auto test/backend/test_ops.py --durations=20
|
||||
- name: Run test_ops (IMAGE)
|
||||
if: matrix.backend == 'ir3'
|
||||
shell: bash
|
||||
env:
|
||||
IMAGE: 1
|
||||
DEV: "NULL:IR3:a630,IMAGE_PITCH_ALIGNMENT=64"
|
||||
run: |
|
||||
DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
|
||||
python -m pytest -n=auto test/backend/test_ops.py --durations=20
|
||||
|
||||
@@ -182,12 +182,12 @@ class GPTOSS:
|
||||
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
|
||||
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) # (B,N,H,D)/(B,N,KV,D)
|
||||
|
||||
if getenv("HK_FLASH_ATTENTION"):
|
||||
from extra.thunder.amd.fa import flash_attention
|
||||
attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
|
||||
attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
|
||||
elif sliding:
|
||||
if sliding:
|
||||
attn = self._sliding_attention(xq, xk, xv, sinks)
|
||||
elif getenv("HK_FLASH_ATTENTION"):
|
||||
from extra.thunder.amd.fa import flash_attention
|
||||
attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
|
||||
attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
|
||||
else:
|
||||
xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
|
||||
xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
|
||||
|
||||
@@ -122,8 +122,7 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
|
||||
groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
|
||||
lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
|
||||
sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
|
||||
arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
|
||||
estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
|
||||
arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
|
||||
insts = build_kernel(M, N, K, tile_m, tile_n)
|
||||
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
|
||||
|
||||
|
||||
+9
-55
@@ -110,49 +110,7 @@ def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
|
||||
return _sharded_empty(ref.shape, ref, axis)
|
||||
|
||||
@functools.cache
|
||||
def _windowed_lse(xq:Tensor, xk:Tensor, sinks, W:int) -> Tensor:
|
||||
B, N, H, hd = xq.shape
|
||||
H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
|
||||
q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
|
||||
k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
|
||||
k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
|
||||
sc_d = (q @ k.transpose(-1, -2)) * sm
|
||||
sc_p = (q @ k_prev.transpose(-1, -2)) * sm
|
||||
li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
|
||||
pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
|
||||
sc_d = (lj <= li).where(sc_d, -float("inf"))
|
||||
sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
|
||||
m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
|
||||
if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
|
||||
denom = (sc_d - m).exp().sum(-1, keepdim=True) + (sc_p - m).exp().sum(-1, keepdim=True)
|
||||
if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
|
||||
return (m + denom.log()).reshape(B, H, N).unsqueeze(2) # (B, H, 1, N), matches saved l_vec
|
||||
|
||||
def _windowed_delta(xq:Tensor, xk:Tensor, xv:Tensor, do:Tensor, sinks, W:int) -> Tensor:
|
||||
B, N, H, hd = xq.shape
|
||||
H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
|
||||
q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
|
||||
k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
|
||||
v = xv.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
|
||||
dob = do.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
|
||||
k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
|
||||
v_prev = v.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
|
||||
sc_d = (q @ k.transpose(-1, -2)) * sm
|
||||
sc_p = (q @ k_prev.transpose(-1, -2)) * sm
|
||||
li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
|
||||
pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
|
||||
sc_d = (lj <= li).where(sc_d, -float("inf"))
|
||||
sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
|
||||
m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
|
||||
if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
|
||||
e_d, e_p = (sc_d - m).exp(), (sc_p - m).exp()
|
||||
denom = e_d.sum(-1, keepdim=True) + e_p.sum(-1, keepdim=True)
|
||||
if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
|
||||
o = ((e_d / denom) @ v) + ((e_p / denom) @ v_prev)
|
||||
delta = (dob * o).sum(-1)
|
||||
return delta.reshape(B, H, N).unsqueeze(2)
|
||||
|
||||
def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
|
||||
def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink):
|
||||
def grad(dou:UOp, ker:UOp) -> tuple:
|
||||
do = Tensor(dou, device=dou.device)
|
||||
attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
|
||||
@@ -160,8 +118,6 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
|
||||
xq = Tensor(ker.src[3], device=ker.src[3].device)
|
||||
xk = Tensor(ker.src[4], device=ker.src[4].device)
|
||||
xv = Tensor(ker.src[5], device=ker.src[5].device)
|
||||
if window:
|
||||
l_vec = _windowed_lse(xq, xk, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
|
||||
|
||||
dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
|
||||
GROUP_SIZE = H_local // H_KV_local
|
||||
@@ -172,10 +128,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
|
||||
# delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
|
||||
delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
|
||||
delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
|
||||
if window:
|
||||
delta_vec = _windowed_delta(xq, xk, xv, do, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
|
||||
|
||||
dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, window=window))[:3]
|
||||
dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
|
||||
|
||||
if D == 64:
|
||||
dq = dq.reshape(B, H, N//16, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2).permute(0, 1, 2, 8, 9, 10, 11, 3, 4, 6, 7, 5, 12).reshape(B, H, N, D).transpose(1, 2)
|
||||
@@ -195,7 +149,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
|
||||
return grad
|
||||
|
||||
# TODO: remove write_flat once scheduler can remove reshapes between custom_kernel. TestCustomKernel.test_simple_reshape
|
||||
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None, window:int=0):
|
||||
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None):
|
||||
assert attn_mask is None, "attn_mask not supported"
|
||||
assert is_causal, "only causal attention supported"
|
||||
|
||||
@@ -222,18 +176,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
|
||||
attn = _sharded_empty((B, N, H * D), xq, axis=shard_axis) if write_flat else _sharded_empty_like(xq, axis=shard_axis)
|
||||
l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
|
||||
|
||||
grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=window)
|
||||
grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink)
|
||||
|
||||
fwd_inputs = (attn, l_vec, xq, xk, xv) + ((sinks,) if has_sink else ())
|
||||
attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink, window=window), grad_fxn=grad)[:2]
|
||||
attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink), grad_fxn=grad)[:2]
|
||||
|
||||
return attn, attn, l_vec
|
||||
|
||||
@functools.cache
|
||||
def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True, window:int=0):
|
||||
def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True):
|
||||
code = (pathlib.Path(__file__).parent / "fa_fwd_causal.cpp").read_text()
|
||||
compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
|
||||
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}", f"-DWINDOW={window}"]
|
||||
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}"]
|
||||
|
||||
Q_BLOCK_SIZE = 32
|
||||
NUM_WARPS = 8
|
||||
@@ -293,10 +247,10 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
|
||||
src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
|
||||
|
||||
@functools.cache
|
||||
def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, window:int=0):
|
||||
def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
|
||||
code = (pathlib.Path(__file__).parent / "fa_bwd_causal.cpp").read_text()
|
||||
compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
|
||||
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DWINDOW={window}"]
|
||||
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]
|
||||
|
||||
BLOCK_SIZE_KV = 256
|
||||
GROUP_SIZE = H // H_KV
|
||||
|
||||
-17
@@ -1,17 +0,0 @@
|
||||
from tinygrad import Device, Tensor, TinyJit, dtypes
|
||||
from tinygrad.helpers import Timing, Context
|
||||
|
||||
GPUS, DEPTH, SZ = 8, 4, 128 * 2**20
|
||||
WARMUP, ITERS = 3, 5
|
||||
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
|
||||
bufs = tuple(Tensor.empty(SZ, dtype=dtypes.uint8, device=dev).contiguous().realize() for _ in range(DEPTH) for dev in devs)
|
||||
|
||||
@TinyJit
|
||||
def all_to_all(*srcs:Tensor): return Tensor.realize(*(src.to(dst) for i,src in enumerate(srcs) for j,dst in enumerate(devs) if i % GPUS != j))
|
||||
|
||||
if __name__ == "__main__":
|
||||
with Context(ALL2ALL=1, JIT_BATCH_SIZE=0):
|
||||
for i in range(-WARMUP, ITERS):
|
||||
with Timing("ALL2ALL ", lambda ns: f" {SZ*GPUS*(GPUS-1)*DEPTH/ns:.2f} GB/s", enabled=i>=0):
|
||||
all_to_all(*bufs)
|
||||
for dev in devs: Device[dev].synchronize()
|
||||
@@ -0,0 +1,45 @@
|
||||
import unittest
|
||||
from tinygrad import UOp, dtypes
|
||||
from tinygrad.uop.ops import shape_to_shape_arg, ParamArg, Ops, AddrSpace
|
||||
|
||||
def placeholder(shape, dtype, slot):
|
||||
return UOp(Ops.PARAM, dtype, (shape_to_shape_arg(shape),), arg=ParamArg(slot, AddrSpace.GLOBAL))
|
||||
|
||||
class TestBitcastSpec(unittest.TestCase):
|
||||
def test_bitcast_no_shape_change(self):
|
||||
pl = placeholder((10,10), dtypes.int, 0)
|
||||
out = pl.bitcast(dtypes.float)
|
||||
self.assertEqual(out.shape, (10,10))
|
||||
|
||||
def test_bitcast_increase_shape(self):
|
||||
pl = placeholder((10,10), dtypes.int, 0)
|
||||
out = pl.bitcast(dtypes.short)
|
||||
self.assertEqual(out.shape, (10,20))
|
||||
|
||||
def test_bitcast_decrease_shape(self):
|
||||
pl = placeholder((10,10), dtypes.int, 0)
|
||||
out = pl.bitcast(dtypes.long)
|
||||
self.assertEqual(out.shape, (10,5))
|
||||
|
||||
def test_bitcast_remove_ones(self):
|
||||
pl = placeholder((10,2), dtypes.int, 0)
|
||||
out = pl.bitcast(dtypes.long)
|
||||
self.assertEqual(out.shape, (10,1))
|
||||
|
||||
def test_bitcast_remove_ones_full(self):
|
||||
pl = placeholder((2,), dtypes.int, 0)
|
||||
out = pl.bitcast(dtypes.long)
|
||||
self.assertEqual(out.shape, (1,))
|
||||
|
||||
def test_bitcast_add_ones_full(self):
|
||||
pl = placeholder((), dtypes.long, 0)
|
||||
out = pl.bitcast(dtypes.int)
|
||||
self.assertEqual(out.shape, (2,))
|
||||
|
||||
def test_bitcast_add_ones_full_uchar(self):
|
||||
pl = placeholder((), dtypes.long, 0)
|
||||
out = pl.bitcast(dtypes.uchar)
|
||||
self.assertEqual(out.shape, (8,))
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -13,18 +13,12 @@ V_TOKS = UOp.variable("toks", 1, 32) # 32 is the default chunk_size in generate
|
||||
class TestTransformerGenerate(unittest.TestCase):
|
||||
def test_warmup(self):
|
||||
model, calls = Transformer(TEST_CONFIG), []
|
||||
def generate(tokens, **kwargs):
|
||||
def generate(tokens):
|
||||
calls.append(tokens)
|
||||
yield from (1, 2)
|
||||
with patch.object(model, "generate", generate): model.warmup()
|
||||
self.assertEqual(calls, [[0], [0]])
|
||||
|
||||
def test_warmup_then_generate_with_default_chunk(self):
|
||||
# warmup must not capture JIT graphs that generate()'s default chunk_size then rejects
|
||||
model = Transformer(TEST_CONFIG)
|
||||
model.warmup()
|
||||
self.assertIsInstance(next(model.generate([5, 6, 7, 8])), int)
|
||||
|
||||
def test_first_recurrent_generate_before_state_init(self):
|
||||
model = Transformer(TEST_CONFIG)
|
||||
model.has_recurrent_block = True
|
||||
@@ -44,15 +38,6 @@ class TestTransformerGenerate(unittest.TestCase):
|
||||
next(model.generate([1, 2, 3, 4, 5, 42, 10]))
|
||||
self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])
|
||||
|
||||
def test_recurrent_divergent_prompt_restarts(self):
|
||||
model, calls = Transformer(TEST_CONFIG), []
|
||||
model.has_recurrent_block, model._cached_tokens = True, [1, 2, 9]
|
||||
def mock_call(self, tokens, start_pos, temperature):
|
||||
calls.append(start_pos)
|
||||
return Tensor([[42]])
|
||||
with patch.object(Transformer, '__call__', mock_call): next(model.generate([1, 2, 10, 11]))
|
||||
self.assertEqual(calls[0], V_START_POS.bind(0))
|
||||
|
||||
def test_template_starts_reasoning(self):
|
||||
router = StreamRouter(reasoning=True)
|
||||
self.assertEqual(list(router.route("reasoning</think>answer")),
|
||||
@@ -193,12 +178,6 @@ class TestTransformerGenerate(unittest.TestCase):
|
||||
# with temperature=2.0, we should see at least 2 distinct outputs across 5 runs
|
||||
self.assertGreater(len(runs), 1, "high temperature should produce varied outputs")
|
||||
|
||||
def test_recurrent_temperature_high_produces_variety(self):
|
||||
model = Transformer(TEST_CONFIG)
|
||||
model.has_recurrent_block = True
|
||||
outputs = {model.forward(Tensor([[1]]), 0, Tensor([2.0])).item() for _ in range(5)}
|
||||
self.assertGreater(len(outputs), 1)
|
||||
|
||||
def test_temperature_passed_to_forward(self):
|
||||
"""Temperature from generate should be passed through to __call__."""
|
||||
model = Transformer(TEST_CONFIG)
|
||||
|
||||
+8
-13
@@ -1,7 +1,7 @@
|
||||
from __future__ import annotations
|
||||
import functools, itertools, pathlib
|
||||
from dataclasses import dataclass, replace
|
||||
from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
|
||||
from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
|
||||
from tinygrad.nn import Linear
|
||||
from tinygrad.llm.gguf import gguf_load
|
||||
from tinygrad.uop.ops import resolve
|
||||
@@ -187,8 +187,8 @@ class TransformerBlock(FFNBlock):
|
||||
|
||||
def _init_state(self, x:Tensor):
|
||||
if not hasattr(self, "cache_kv"):
|
||||
self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
|
||||
dtype=dtypes.default_float, device=x.device)
|
||||
# TODO: how is the dtype of this determined?
|
||||
self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim, device=x.device)
|
||||
self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
|
||||
|
||||
class MLATransformerBlock(FFNBlock):
|
||||
@@ -261,14 +261,13 @@ class GatedDeltaNetBlock(FFNBlock):
|
||||
def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
|
||||
B, T, _ = x.shape
|
||||
assert T == 1, "GatedDeltaNetBlock currently only supports T=1"
|
||||
is_kda = hasattr(self, "ssm_g_a")
|
||||
|
||||
# input processing
|
||||
x = x.half()
|
||||
out_gate = self.ssm_g_b(self.ssm_g_a(x)) if is_kda else self.attn_gate(x)
|
||||
out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
|
||||
out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
|
||||
beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
|
||||
alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x)
|
||||
alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
|
||||
alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
|
||||
self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
|
||||
|
||||
@@ -292,13 +291,14 @@ class GatedDeltaNetBlock(FFNBlock):
|
||||
|
||||
# output
|
||||
core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
|
||||
out_gate = out_gate.sigmoid() if is_kda else out_gate.silu()
|
||||
out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
|
||||
return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
|
||||
|
||||
# recurrent state can't be partially reused after divergence, force a full rebuild
|
||||
def _state_reset_ops(self):
|
||||
return [self.conv_state.assign(self.conv_state.const_like(0)),
|
||||
self.recurrent_state.assign(self.recurrent_state.const_like(0))] if hasattr(self, "conv_state") else []
|
||||
def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return 0 if prefix_len != cached_len else prefix_len
|
||||
|
||||
def _init_state(self, x):
|
||||
if not hasattr(self, "conv_state"):
|
||||
@@ -326,8 +326,7 @@ class Transformer:
|
||||
def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
|
||||
x = self.token_embd(tokens).float() # (B, T, D)
|
||||
for block in self.blk: x = block(x, start_pos)
|
||||
# only run the output projection on the last token
|
||||
logits = self.output(self.output_norm(x[:, -1:]))[:, -1, :]
|
||||
logits = self.output(self.output_norm(x))[:, -1, :]
|
||||
# Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp)
|
||||
return (logits / temperature.maximum(1e-12) - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True)
|
||||
|
||||
@@ -421,10 +420,6 @@ class Transformer:
|
||||
for _ in range(2): list(zip(range(2), self.generate([0])))
|
||||
|
||||
def get_start_pos(self, tokens:list[int]) -> int:
|
||||
# recurrent state can't be partially reused after divergence: reuse it only when tokens extend the cached prefix
|
||||
if self.has_recurrent_block:
|
||||
return len(self._cached_tokens) if self._cached_tokens and len(self._cached_tokens) < len(tokens) \
|
||||
and tokens[:len(self._cached_tokens)] == self._cached_tokens else 0
|
||||
prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
|
||||
return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
|
||||
|
||||
|
||||
@@ -790,14 +790,10 @@ class KFDIface:
|
||||
|
||||
def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
|
||||
xcc_id=0, idx=0):
|
||||
if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA and idx and self.kfd_ver >= (1,17):
|
||||
queue_type = kfd.KFD_IOC_QUEUE_TYPE_SDMA_BY_ENG_ID
|
||||
sdma_engine_id = idx % (self.props['num_sdma_engines'] + self.props.get('num_sdma_xgmi_engines', 0))
|
||||
else: sdma_engine_id = 0
|
||||
queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring.va_addr, ring_size=ring.size, gpu_id=self.gpu_id,
|
||||
queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
|
||||
eop_buffer_address=eop_buffer.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer.size if eop_buffer else 0, ctl_stack_size=ctl_stack_size,
|
||||
ctx_save_restore_address=cwsr_buffer.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size, sdma_engine_id=sdma_engine_id,
|
||||
ctx_save_restore_address=cwsr_buffer.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size,
|
||||
write_pointer_address=gart.va_addr+wptr, read_pointer_address=gart.va_addr+rptr+8*xcc_id)
|
||||
|
||||
if not hasattr(self, 'doorbells'):
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import ctypes, struct, platform, pathlib, shutil, subprocess, sys, tarfile, tempfile
|
||||
import ctypes, struct
|
||||
from tinygrad.device import Compiler
|
||||
from tinygrad.helpers import DEBUG, system, fetch, unwrap
|
||||
from tinygrad.helpers import DEBUG, system
|
||||
from tinygrad.runtime.support.compiler_mesa import disas_adreno
|
||||
# see https://github.com/sirhcm/tinydreno
|
||||
from tinygrad.runtime.autogen import llvm_qcom
|
||||
@@ -10,17 +10,10 @@ def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
|
||||
class QCOMCompiler(Compiler):
|
||||
def __init__(self, arch:str):
|
||||
assert arch.split(',')[0] == "a630", "only a630 supported"
|
||||
if platform.machine() == "aarch64": self.arch, self.chip_id, self.llvm_inst = arch, 0x6030001, llvm_qcom.cl_compiler_create_llvm_instance()
|
||||
else:
|
||||
self.arch, self.chip_id, self.fs = arch, 0x6030001, tempfile.TemporaryDirectory()
|
||||
with tarfile.open(fetch('https://git.tinygrad.win/sirhcm/images/releases/download/v2/qcomcl.tar.gz')) as t: t.extractall(fs:=self.fs.name)
|
||||
if (qemu:=shutil.which("qemu-aarch64-static")): argv = f"{qemu} -cpu max,pauth=off -L {fs} {fs}/usr/bin/python3 {__file__} {arch}"
|
||||
else: argv = (f"docker run --rm -i --platform linux/aarch64 -v {fs}/usr:/usr -v {pathlib.Path(__file__).parents[2]}:/tinygrad "
|
||||
f"-e PYTHONPATH=/ -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3 /tinygrad/runtime/support/compiler_qcom.py {arch}")
|
||||
self.compiler_process = subprocess.Popen(argv.split(), stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
|
||||
self.arch, self.chip_id, self.llvm_inst = arch, 0x6030001, llvm_qcom.cl_compiler_create_llvm_instance()
|
||||
super().__init__(f"compile_qcomcl_{arch}")
|
||||
|
||||
def __del__(self): llvm_qcom.cl_compiler_destroy_llvm_instance(self.llvm_inst) if platform.machine() == "aarch64" else self.compiler_process.kill()
|
||||
def __del__(self): llvm_qcom.cl_compiler_destroy_llvm_instance(self.llvm_inst)
|
||||
|
||||
def __reduce__(self): return QCOMCompiler, (self.arch,)
|
||||
|
||||
@@ -32,10 +25,6 @@ class QCOMCompiler(Compiler):
|
||||
return handle
|
||||
|
||||
def compile(self, src) -> bytes:
|
||||
if platform.machine() != "aarch64":
|
||||
unwrap(self.compiler_process.stdin).write(struct.pack("I", len(src.encode())) + src.encode())
|
||||
if (lib:=unwrap(self.compiler_process.stdout).read(struct.unpack("I", unwrap(self.compiler_process.stdout).read(4))[0])): return lib
|
||||
raise RuntimeError("QCOM Compilation Error")
|
||||
ch = self.checked(llvm_qcom.cl_compiler_compile_source(self.llvm_inst, self.chip_id, llvm_qcom.CL_MODE_64BIT, b"", 0, 0, 0, src.encode(), 0,
|
||||
llvm_qcom.CL_SRC_STR, None))
|
||||
if DEBUG >= 8: print(system("llvm-dis", input=ctypes.string_at((comp:=ch.contents.compiled.contents).llvm_bitcode, comp.llvm_bitcode_size)))
|
||||
@@ -47,13 +36,3 @@ class QCOMCompiler(Compiler):
|
||||
return ret
|
||||
|
||||
def disassemble(self, lib: bytes): disas_adreno(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)], self.chip_id)
|
||||
|
||||
if __name__ == "__main__":
|
||||
compiler = QCOMCompiler(sys.argv[1])
|
||||
while (amt:=sys.stdin.buffer.read(4)):
|
||||
try: lib = compiler.compile(sys.stdin.buffer.read(struct.unpack("I", amt)[0]).decode())
|
||||
except Exception as e:
|
||||
lib = b""
|
||||
print(e, file=sys.stderr, flush=True)
|
||||
sys.stdout.buffer.write(struct.pack("I", len(lib)) + lib)
|
||||
sys.stdout.buffer.flush()
|
||||
|
||||
Reference in New Issue
Block a user