From 6b14828474b4b0d4cb317207e8aeaa47339b3969 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Tue, 30 Jun 2026 13:42:30 +0530 Subject: [PATCH 1/5] feat(intrinsics): add KI.vload / KI.vstore! wide vector memory operations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements vectorised load/store for AbstractArray across all KA backends: - CPU (Ptr path): reinterpret to Ptr{NTuple{N,VecElement{T}}} + unsafe_load/store! - GPU (LLVMPtr path): reinterpret to LLVMPtr{NTuple{N,VecElement{T}},AS} + unsafe_load/store! with N*sizeof(T) alignment; lowers to ld.global.v4 (CUDA), global_load_dwordx4 (AMDGPU), or OpLoad (SPIR-V/POCL) — avoids pointerref/pointerset which are unsupported in SPIR-V - Scalar fallback for N=1 or non-primitive element types Adds @generated helpers (_vload_ptr, _vload_lptr, _vstore_ptr!, _vstore_lptr!, _vload_arr, _vstore_arr!) and comprehensive CPU test suite (147 tests, all pass). Co-Authored-By: Claude Sonnet 4.6 --- src/intrinsics.jl | 182 ++++++++++++++++++++++++++++++++++++++++ test/intrinsics.jl | 201 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 383 insertions(+) diff --git a/src/intrinsics.jl b/src/intrinsics.jl index 79efa4cbe..5ac329986 100644 --- a/src/intrinsics.jl +++ b/src/intrinsics.jl @@ -156,6 +156,188 @@ end function _print end +# ── Internal helpers for wide vector loads / stores ─────────────────────────── +# +# CPU (Ptr{T}): cast pointer to Ptr{NTuple{N, VecElement{T}}} and call +# unsafe_load / unsafe_store!. NTuple{N, VecElement{T}} is Julia's +# representation of LLVM's ; the Julia/LLVM compiler lowers this +# to a single wide memory op. +# +# GPU / non-Ptr (Core.LLVMPtr): reinterpret the scalar LLVMPtr as a +# LLVMPtr{NTuple{N, VecElement{T}}, AS} and issue a single unsafe_load / +# unsafe_store! with vector-width alignment. The LLVM backend lowers this to: +# CUDA/NVPTX → ld.global.v4.f32 / st.global.v4.f32 +# AMDGPU → global_load_dwordx4 / global_store_dwordx4 +# POCL/SPIR-V → OpLoad / OpStore on — avoids Julia runtime calls +# (pointerref / pointerset are not supported in SPIR-V) +# @generated is required to avoid closures — closures in GPU kernels +# require heap allocation which the GPU compiler cannot handle. +# +# Fallback (N=1, non-primitive types): scalar arr[idx+k] accesses via +# @generated to avoid closures. + +# CPU path ──────────────────────────────────────────────────────────────────── + +@inline _vptr(p::Ptr{T}, ::Val{N}) where {T, N} = + Ptr{NTuple{N, Core.VecElement{T}}}(p) + +@inline function _vload_ptr(::Val{N}, p::Ptr{T}) where {N, T} + raw = unsafe_load(_vptr(p, Val(N))) + ntuple(i -> raw[i].value, Val(N)) +end + +@inline function _vstore_ptr!(p::Ptr{T}, ::Val{N}, vals::NTuple{N, T}) where {N, T} + unsafe_store!(_vptr(p, Val(N)), ntuple(i -> Core.VecElement{T}(vals[i]), Val(N))) +end + +# GPU / non-Ptr path — single vector load / store via reinterpret ───────────── +# Use reinterpret(LLVMPtr{, AS}, p) + unsafe_load/unsafe_store!. +# This is the same pattern POCL's own CLDeviceArray uses internally. + +@generated function _vload_lptr(::Val{N}, p) where {N} + T = p.parameters[1] + AS = p.parameters[2] + VT = NTuple{N, Core.VecElement{T}} + quote + p_vec = reinterpret(Core.LLVMPtr{$VT, $AS}, p) + raw = unsafe_load(p_vec, 1, Val($(N * sizeof(T)))) + $(Expr(:tuple, [:(raw[$i].value) for i in 1:N]...)) + end +end + +@generated function _vstore_lptr!(p, ::Val{N}, vals::NTuple{N}) where {N} + T = p.parameters[1] + AS = p.parameters[2] + VT = NTuple{N, Core.VecElement{T}} + quote + p_vec = reinterpret(Core.LLVMPtr{$VT, $AS}, p) + vec_vals = $(Expr(:tuple, [:(Core.VecElement{$T}(vals[$i])) for i in 1:N]...)) + unsafe_store!(p_vec, vec_vals, 1, Val($(N * sizeof(T)))) + nothing + end +end + +# Scalar fallback — used for N=1, non-primitive types ───────────────────────── + +@generated function _vload_arr(::Val{N}, arr::AbstractArray{T}, idx::Integer) where {N, T} + Expr(:tuple, [:(Base.@inbounds arr[idx + $(i - 1)]) for i in 1:N]...) +end + +@generated function _vstore_arr!(arr::AbstractArray{T}, idx::Integer, + ::Val{N}, vals::NTuple{N, T}) where {N, T} + Expr(:block, + [:(Base.@inbounds arr[idx + $(i - 1)] = vals[$i]) for i in 1:N]..., + :nothing) +end + +# ─── vload ─────────────────────────────────────────────────────────────────── + +""" + vload(::Val{N}, arr::AbstractArray{T}, idx::Integer) → NTuple{N, T} + +Load `N` consecutive elements starting at 1-based index `idx` from `arr`. + +### How it works + +On **CPU** (`Ptr{T}` path), for primitive element types this casts +`pointer(arr, idx)` to a `Ptr{NTuple{N, Core.VecElement{T}}}` and issues +a single `unsafe_load`. Julia lowers `NTuple{N, VecElement{T}}` to +LLVM's `` vector type — one wide memory op, no LoadStoreVectorizer. + +On **GPU** (and for all arrays whose `pointer` returns a non-CPU pointer), +the pointer is reinterpreted as `LLVMPtr{NTuple{N, VecElement{T}}, AS}` and a +single `unsafe_load` with `N*sizeof(T)` alignment is issued. The backend LLVM +pipeline lowers this to one wide instruction — +`ld.global.v4.f32` (CUDA/NVPTX), `global_load_dwordx4` (AMDGPU), or +`OpLoad ` (SPIR-V/POCL). + +For non-primitive element types or `N == 1` the fallback emits N plain scalar +array accesses (`arr[idx]`, `arr[idx+1]`, …). + +### Requirements + +- `N` must be a **compile-time constant** — pass `Val(4)`, not a variable. +- The array must **not** be annotated with `@Const`; `@Const` inserts `ldg` + (read-only cache) intrinsics whose pointer type differs from a plain global + load and cannot be widened into a vector load. +- `idx` must be aligned to `N * sizeof(T)` bytes for the hardware vector + instruction to fire. Julia's allocator guarantees ≥ 64-byte base alignment, + so `idx = 1 + k*N` for `k ≥ 0` is always valid. + +### Example + +```julia +function reduce_kernel(dst, src) + g = KI.get_global_id().x + idx = (g - 1) * 4 + 1 + v0, v1, v2, v3 = KI.vload(Val(4), src, idx) # single ld.global.v4.f32 + KI.vstore!(dst, idx, (v0+1f0, v1+1f0, v2+1f0, v3+1f0)) + return +end +``` + +!!! note + Backends **may** provide an override with additional alignment hints or + cache-modifier control: + ```julia + @device_override @inline KI.vload(::Val{N}, arr::AbstractArray{T}, idx::Integer) where {N, T} + ``` +""" +@inline function vload(::Val{N}, arr::AbstractArray{T}, idx::Integer) where {N, T} + if isprimitivetype(T) && N > 1 + p = pointer(arr, idx) + if p isa Ptr + _vload_ptr(Val(N), p) + else + _vload_lptr(Val(N), p) + end + else + # N=1 or non-primitive: scalar loads + # (VecElement<1x> is invalid in SPIRV/OpenCL; non-primitive structs + # don't have a simple LLVM vector representation) + _vload_arr(Val(N), arr, idx) + end +end + +# ─── vstore! ───────────────────────────────────────────────────────────────── + +""" + vstore!(arr::AbstractArray{T}, idx::Integer, vals::NTuple{N, T}) + +Store `N` consecutive elements from `vals` into `arr` starting at 1-based +index `idx`. + +The write complement of [`vload`](@ref). On CPU this wraps `vals` in a +`` LLVM vector and calls a single `unsafe_store!` (one wide store). +On GPU the pointer is reinterpreted as `LLVMPtr{NTuple{N, VecElement{T}}, AS}` +and a single `unsafe_store!` with `N*sizeof(T)` alignment issues one wide +instruction — `st.global.v4.f32` (CUDA), `global_store_dwordx4` (AMDGPU), +or `OpStore ` (SPIR-V/POCL). + +The same `N`-must-be-`Val`, no-`@Const`, and alignment requirements as +[`vload`](@ref) apply. + +!!! note + Backends **may** provide an override: + ```julia + @device_override @inline KI.vstore!(arr::AbstractArray{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} + ``` +""" +@inline function vstore!(arr::AbstractArray{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} + if isprimitivetype(T) && N > 1 + p = pointer(arr, idx) + if p isa Ptr + _vstore_ptr!(p, Val(N), vals) + else + _vstore_lptr!(p, Val(N), vals) + end + else + _vstore_arr!(arr, idx, Val(N), vals) + end + return nothing +end + + """ Kernel{Backend, Kern} diff --git a/test/intrinsics.jl b/test/intrinsics.jl index 63216d32b..bb0cafee5 100644 --- a/test/intrinsics.jl +++ b/test/intrinsics.jl @@ -24,6 +24,65 @@ function test_intrinsics_kernel(results) return end +# ── vload / vstore! kernel helpers ────────────────────────────────────────── + +function vload_copy_kernel(dst, src, ::Val{N}) where {N} + g = KI.get_global_id().x + idx = (g - 1) * N + 1 + KI.vstore!(dst, idx, KI.vload(Val(N), src, idx)) + return +end + +function vload_scale_kernel(dst, src, scale::T, ::Val{N}) where {T, N} + g = KI.get_global_id().x + idx = (g - 1) * N + 1 + vals = KI.vload(Val(N), src, idx) + KI.vstore!(dst, idx, ntuple(i -> vals[i] * scale, Val(N))) + return +end + +function vload_reduce_kernel(dst, src, ::Val{N}) where {N} + g = KI.get_global_id().x + idx = (g - 1) * N + 1 + vals = KI.vload(Val(N), src, idx) + s = vals[1] + for i in 2:N + s += vals[i] + end + @inbounds dst[g] = s + return +end + +function vload_tail_kernel(dst, src, n, ::Val{N}) where {N} + g = KI.get_global_id().x + idx = (g - 1) * N + 1 + if idx + N - 1 <= n + KI.vstore!(dst, idx, KI.vload(Val(N), src, idx)) + else + while idx <= n + @inbounds dst[idx] = src[idx] + idx += 1 + end + end + return +end + +function vload_store_load_kernel(arr, v0::T, v1::T, v2::T, v3::T) where {T} + KI.vstore!(arr, 1, (v0, v1, v2, v3)) + loaded = KI.vload(Val(4), arr, 1) + KI.vstore!(arr, 5, loaded) + return +end + +function vload_mt_copy_kernel(dst, src, ::Val{N}) where {N} + local_id = KI.get_local_id().x + group_id = KI.get_group_id().x + groupsize = KI.get_local_size().x + idx = ((group_id - 1) * groupsize + (local_id - 1)) * N + 1 + KI.vstore!(dst, idx, KI.vload(Val(N), src, idx)) + return +end + function intrinsics_testsuite(backend, AT) @testset "KernelIntrinsics Tests" begin @testset "Launch parameters" begin @@ -122,6 +181,148 @@ function intrinsics_testsuite(backend, AT) @test k_data.local_id == expected_local end end + + @testset "vload / vstore!" begin + + # ── 1. Roundtrip for every supported (N, T) combination ───────── + @testset "roundtrip N=$N T=$T" for N in (1, 2, 4, 8), + T in (Float32, Float64, Int32, Int64) + n = 64 + src = AT(T.(1:n)) + dst = AT(zeros(T, n)) + KI.@kernel backend() numworkgroups = n ÷ N workgroupsize = 1 vload_copy_kernel(dst, src, Val(N)) + KernelAbstractions.synchronize(backend()) + @test Array(dst) == Array(src) + end + + # ── 2. N=1 degenerate (scalar vload/vstore!) ──────────────────── + @testset "N=1 scalar degenerate" begin + src = AT(Float32[42.0f0]) + dst = AT(Float32[0.0f0]) + KI.@kernel backend() numworkgroups = 1 workgroupsize = 1 vload_copy_kernel(dst, src, Val(1)) + KernelAbstractions.synchronize(backend()) + @test Array(dst) == Array(src) + end + + # ── 3. Minimal array: size == N (one vload covers whole array) ── + @testset "minimal n=N=$N" for N in (1, 2, 4) + src = AT(Float32.(1:N)) + dst = AT(zeros(Float32, N)) + KI.@kernel backend() numworkgroups = 1 workgroupsize = 1 vload_copy_kernel(dst, src, Val(N)) + KernelAbstractions.synchronize(backend()) + @test Array(dst) == Array(src) + end + + # ── 4. Arithmetic on loaded values ────────────────────────────── + @testset "scale by 2 N=4 Float32" begin + n = 64 + src = AT(Float32.(1:n)) + dst = AT(zeros(Float32, n)) + KI.@kernel backend() numworkgroups = n ÷ 4 workgroupsize = 1 vload_scale_kernel(dst, src, 2.0f0, Val(4)) + KernelAbstractions.synchronize(backend()) + @test Array(dst) ≈ 2.0f0 .* Float32.(1:n) + end + + @testset "scale by 3 N=2 Float64" begin + n = 64 + src = AT(Float64.(1:n)) + dst = AT(zeros(Float64, n)) + KI.@kernel backend() numworkgroups = n ÷ 2 workgroupsize = 1 vload_scale_kernel(dst, src, 3.0, Val(2)) + KernelAbstractions.synchronize(backend()) + @test Array(dst) ≈ 3.0 .* Float64.(1:n) + end + + # ── 5. Per-thread reduction: each thread sums its N elements ──── + @testset "per-thread sum all-ones N=4" begin + n = 64 + src = AT(ones(Float32, n)) + result = AT(zeros(Float32, n ÷ 4)) + KI.@kernel backend() numworkgroups = n ÷ 4 workgroupsize = 1 vload_reduce_kernel(result, src, Val(4)) + KernelAbstractions.synchronize(backend()) + @test all(Array(result) .== 4.0f0) + end + + @testset "per-thread sum sequential N=4" begin + n = 64 + src = AT(Float32.(1:n)) + result = AT(zeros(Float32, n ÷ 4)) + KI.@kernel backend() numworkgroups = n ÷ 4 workgroupsize = 1 vload_reduce_kernel(result, src, Val(4)) + KernelAbstractions.synchronize(backend()) + expected = [Float32(sum(4k+1:4k+4)) for k in 0:n÷4-1] + @test Array(result) == expected + end + + # ── 6. Scalar tail: n not divisible by N ──────────────────────── + @testset "scalar tail n=$n" for n in (65, 66, 67) + src = AT(Float32.(1:n)) + dst = AT(zeros(Float32, n)) + KI.@kernel backend() numworkgroups = cld(n, 4) workgroupsize = 1 vload_tail_kernel(dst, src, n, Val(4)) + KernelAbstractions.synchronize(backend()) + @test Array(dst) == Array(src) + end + + # ── 7. vstore! → vload identity (store then read back) ────────── + @testset "vstore then vload identity" begin + arr = AT(zeros(Float32, 8)) + KI.@kernel backend() numworkgroups = 1 workgroupsize = 1 vload_store_load_kernel(arr, 10.0f0, 20.0f0, 30.0f0, 40.0f0) + KernelAbstractions.synchronize(backend()) + h = Array(arr) + @test h[1:4] == [10.0f0, 20.0f0, 30.0f0, 40.0f0] + @test h[5:8] == h[1:4] + end + + # ── 8. All-zeros and all-ones arrays ──────────────────────────── + @testset "all-zeros roundtrip N=4" begin + n = 64 + src = AT(zeros(Float32, n)) + dst = AT(ones(Float32, n)) + KI.@kernel backend() numworkgroups = n ÷ 4 workgroupsize = 1 vload_copy_kernel(dst, src, Val(4)) + KernelAbstractions.synchronize(backend()) + @test all(Array(dst) .== 0.0f0) + end + + # ── 9. Large array ─────────────────────────────────────────────── + @testset "large array n=1024 N=4" begin + n = 1024 + src = AT(Float32.(1:n)) + dst = AT(zeros(Float32, n)) + KI.@kernel backend() numworkgroups = n ÷ 4 workgroupsize = 1 vload_copy_kernel(dst, src, Val(4)) + KernelAbstractions.synchronize(backend()) + @test Array(dst) == Array(src) + end + + # ── 10. Multi-thread (workgroupsize > 1) ───────────────────────── + @testset "multi-thread workgroupsize=4 N=4" begin + wgs = 4 + n = 128 + src = AT(Float32.(1:n)) + dst = AT(zeros(Float32, n)) + KI.@kernel backend() numworkgroups = n ÷ (wgs * 4) workgroupsize = wgs vload_mt_copy_kernel(dst, src, Val(4)) + KernelAbstractions.synchronize(backend()) + @test Array(dst) == Array(src) + end + + # ── 11. Integer types ──────────────────────────────────────────── + @testset "Int32 N=4 arithmetic" begin + n = 64 + src = AT(Int32.(1:n)) + dst = AT(zeros(Int32, n)) + KI.@kernel backend() numworkgroups = n ÷ 4 workgroupsize = 1 vload_scale_kernel(dst, src, Int32(3), Val(4)) + KernelAbstractions.synchronize(backend()) + @test Array(dst) == 3 .* Int32.(1:n) + end + + # ── 12. Type stability ─────────────────────────────────────────── + @testset "type stability T=$T N=$N" for T in (Float32, Float64, Int32, Int64), N in (1, 4) + n = N + src = AT(T.(1:n)) + dst = AT(zeros(T, n)) + KI.@kernel backend() numworkgroups = 1 workgroupsize = 1 vload_copy_kernel(dst, src, Val(N)) + KernelAbstractions.synchronize(backend()) + @test eltype(Array(dst)) === T + end + + end end return nothing end From 3116d0add456bde4a8e15ad1b6524ac71739a889 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Tue, 30 Jun 2026 14:16:50 +0530 Subject: [PATCH 2/5] refactor(intrinsics): use @device_override architecture for vload/vstore! MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Base vload/vstore! now fall back to scalar indexing for all non-Ptr arrays. GPU backends inject the vectorized path by registering @device_override for their concrete device array type, following the KernelIntrinsics.jl pattern: - src/intrinsics.jl: remove _vload_lptr/_vstore_lptr! and the LLVMPtr branch; base dispatch is now: Ptr{T} → wide load, everything else → scalar fallback - src/pocl/backend.jl: add @device_override for KI.vload/KI.vstore! on CLDeviceArray, using reinterpret(LLVMPtr{NTuple{N,VecElement{T}},AS}, ptr) + unsafe_load/unsafe_store! to emit OpLoad (SPIR-V/POCL) CUDA.jl, AMDGPU.jl etc. should add analogous overrides for CuDeviceArray / ROCDeviceArray to get ld.global.v4.f32 / global_load_dwordx4. 147/147 CPU tests pass. Co-Authored-By: Claude Sonnet 4.6 --- src/intrinsics.jl | 101 ++++++++++++++------------------------------ src/pocl/backend.jl | 43 +++++++++++++++++++ 2 files changed, 74 insertions(+), 70 deletions(-) diff --git a/src/intrinsics.jl b/src/intrinsics.jl index 5ac329986..11bfc1964 100644 --- a/src/intrinsics.jl +++ b/src/intrinsics.jl @@ -163,18 +163,16 @@ function _print end # representation of LLVM's ; the Julia/LLVM compiler lowers this # to a single wide memory op. # -# GPU / non-Ptr (Core.LLVMPtr): reinterpret the scalar LLVMPtr as a -# LLVMPtr{NTuple{N, VecElement{T}}, AS} and issue a single unsafe_load / -# unsafe_store! with vector-width alignment. The LLVM backend lowers this to: -# CUDA/NVPTX → ld.global.v4.f32 / st.global.v4.f32 -# AMDGPU → global_load_dwordx4 / global_store_dwordx4 -# POCL/SPIR-V → OpLoad / OpStore on — avoids Julia runtime calls -# (pointerref / pointerset are not supported in SPIR-V) -# @generated is required to avoid closures — closures in GPU kernels -# require heap allocation which the GPU compiler cannot handle. +# GPU backends: each backend registers a @device_override for its concrete +# device array type (e.g. CLDeviceArray in src/pocl/backend.jl, +# CuDeviceArray in CUDA.jl's KA extension). The override uses +# reinterpret(LLVMPtr{NTuple{N,VecElement{T}},AS}, ptr) + unsafe_load to +# emit a single wide instruction — ld.global.v4.f32 (CUDA/NVPTX), +# global_load_dwordx4 (AMDGPU), OpLoad (SPIR-V/POCL). +# Without an override, vload falls back to scalar indexing. # -# Fallback (N=1, non-primitive types): scalar arr[idx+k] accesses via -# @generated to avoid closures. +# Fallback (N=1, non-primitive types, or any array with no override): +# scalar arr[idx+k] accesses via @generated to avoid closures. # CPU path ──────────────────────────────────────────────────────────────────── @@ -190,34 +188,8 @@ end unsafe_store!(_vptr(p, Val(N)), ntuple(i -> Core.VecElement{T}(vals[i]), Val(N))) end -# GPU / non-Ptr path — single vector load / store via reinterpret ───────────── -# Use reinterpret(LLVMPtr{, AS}, p) + unsafe_load/unsafe_store!. -# This is the same pattern POCL's own CLDeviceArray uses internally. - -@generated function _vload_lptr(::Val{N}, p) where {N} - T = p.parameters[1] - AS = p.parameters[2] - VT = NTuple{N, Core.VecElement{T}} - quote - p_vec = reinterpret(Core.LLVMPtr{$VT, $AS}, p) - raw = unsafe_load(p_vec, 1, Val($(N * sizeof(T)))) - $(Expr(:tuple, [:(raw[$i].value) for i in 1:N]...)) - end -end - -@generated function _vstore_lptr!(p, ::Val{N}, vals::NTuple{N}) where {N} - T = p.parameters[1] - AS = p.parameters[2] - VT = NTuple{N, Core.VecElement{T}} - quote - p_vec = reinterpret(Core.LLVMPtr{$VT, $AS}, p) - vec_vals = $(Expr(:tuple, [:(Core.VecElement{$T}(vals[$i])) for i in 1:N]...)) - unsafe_store!(p_vec, vec_vals, 1, Val($(N * sizeof(T)))) - nothing - end -end - -# Scalar fallback — used for N=1, non-primitive types ───────────────────────── +# Scalar fallback — used for N=1, non-primitive types, and GPU arrays +# without a @device_override ────────────────────────────────────────────────── @generated function _vload_arr(::Val{N}, arr::AbstractArray{T}, idx::Integer) where {N, T} Expr(:tuple, [:(Base.@inbounds arr[idx + $(i - 1)]) for i in 1:N]...) @@ -242,17 +214,18 @@ Load `N` consecutive elements starting at 1-based index `idx` from `arr`. On **CPU** (`Ptr{T}` path), for primitive element types this casts `pointer(arr, idx)` to a `Ptr{NTuple{N, Core.VecElement{T}}}` and issues a single `unsafe_load`. Julia lowers `NTuple{N, VecElement{T}}` to -LLVM's `` vector type — one wide memory op, no LoadStoreVectorizer. +LLVM's `` vector type — one wide memory op. -On **GPU** (and for all arrays whose `pointer` returns a non-CPU pointer), -the pointer is reinterpreted as `LLVMPtr{NTuple{N, VecElement{T}}, AS}` and a -single `unsafe_load` with `N*sizeof(T)` alignment is issued. The backend LLVM -pipeline lowers this to one wide instruction — -`ld.global.v4.f32` (CUDA/NVPTX), `global_load_dwordx4` (AMDGPU), or -`OpLoad ` (SPIR-V/POCL). +On **GPU**, each backend registers a `@device_override` for its concrete +device array type (e.g. `CLDeviceArray` for POCL, `CuDeviceArray` for +CUDA.jl). The override reinterprets the pointer as +`LLVMPtr{NTuple{N, VecElement{T}}, AS}` and issues a single `unsafe_load` +with `N*sizeof(T)` alignment, which the backend LLVM pipeline lowers to one +wide instruction — `ld.global.v4.f32` (CUDA/NVPTX), +`global_load_dwordx4` (AMDGPU), `OpLoad ` (SPIR-V/POCL). -For non-primitive element types or `N == 1` the fallback emits N plain scalar -array accesses (`arr[idx]`, `arr[idx+1]`, …). +For non-primitive element types, `N == 1`, or any array type without a +registered override, the fallback emits N plain scalar array accesses. ### Requirements @@ -277,26 +250,19 @@ end ``` !!! note - Backends **may** provide an override with additional alignment hints or - cache-modifier control: + GPU backends provide the vectorized path via: ```julia - @device_override @inline KI.vload(::Val{N}, arr::AbstractArray{T}, idx::Integer) where {N, T} + @device_override @inline function KI.vload(::Val{N}, arr::MyDeviceArray{T}, idx::Integer) where {N, T} ``` """ @inline function vload(::Val{N}, arr::AbstractArray{T}, idx::Integer) where {N, T} if isprimitivetype(T) && N > 1 p = pointer(arr, idx) if p isa Ptr - _vload_ptr(Val(N), p) - else - _vload_lptr(Val(N), p) + return _vload_ptr(Val(N), p) end - else - # N=1 or non-primitive: scalar loads - # (VecElement<1x> is invalid in SPIRV/OpenCL; non-primitive structs - # don't have a simple LLVM vector representation) - _vload_arr(Val(N), arr, idx) end + _vload_arr(Val(N), arr, idx) end # ─── vstore! ───────────────────────────────────────────────────────────────── @@ -308,19 +274,16 @@ Store `N` consecutive elements from `vals` into `arr` starting at 1-based index `idx`. The write complement of [`vload`](@ref). On CPU this wraps `vals` in a -`` LLVM vector and calls a single `unsafe_store!` (one wide store). -On GPU the pointer is reinterpreted as `LLVMPtr{NTuple{N, VecElement{T}}, AS}` -and a single `unsafe_store!` with `N*sizeof(T)` alignment issues one wide -instruction — `st.global.v4.f32` (CUDA), `global_store_dwordx4` (AMDGPU), -or `OpStore ` (SPIR-V/POCL). +`` LLVM vector and issues a single `unsafe_store!`. GPU backends +provide wide store via `@device_override` on their device array type. The same `N`-must-be-`Val`, no-`@Const`, and alignment requirements as [`vload`](@ref) apply. !!! note - Backends **may** provide an override: + GPU backends provide the vectorized path via: ```julia - @device_override @inline KI.vstore!(arr::AbstractArray{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} + @device_override @inline function KI.vstore!(arr::MyDeviceArray{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} ``` """ @inline function vstore!(arr::AbstractArray{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} @@ -328,12 +291,10 @@ The same `N`-must-be-`Val`, no-`@Const`, and alignment requirements as p = pointer(arr, idx) if p isa Ptr _vstore_ptr!(p, Val(N), vals) - else - _vstore_lptr!(p, Val(N), vals) + return nothing end - else - _vstore_arr!(arr, idx, Val(N), vals) end + _vstore_arr!(arr, idx, Val(N), vals) return nothing end diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index f23e20b0f..e23fb3e43 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -241,4 +241,47 @@ end KA.argconvert(::KA.Kernel{POCLBackend}, arg) = clconvert(arg) + +## Vectorized memory operations — @device_override for CLDeviceArray +# +# These helpers live here (not in KI) because they are POCL-specific: they +# operate on LLVMPtr with SPIR-V address spaces and use reinterpret + +# unsafe_load/unsafe_store! to emit OpLoad/OpStore on vectors. +# CUDA.jl, AMDGPU.jl etc. provide analogous overrides for their array types. + +@generated function _vload_lptr(::Val{N}, p::Core.LLVMPtr{T, A}) where {N, T, A} + VT = NTuple{N, Core.VecElement{T}} + quote + p_vec = reinterpret(Core.LLVMPtr{$VT, $A}, p) + raw = unsafe_load(p_vec, 1, Val($(N * sizeof(T)))) + $(Expr(:tuple, [:(raw[$i].value) for i in 1:N]...)) + end +end + +@generated function _vstore_lptr!(p::Core.LLVMPtr{T, A}, ::Val{N}, vals::NTuple{N}) where {N, T, A} + VT = NTuple{N, Core.VecElement{T}} + quote + p_vec = reinterpret(Core.LLVMPtr{$VT, $A}, p) + vec_vals = $(Expr(:tuple, [:(Core.VecElement{$T}(vals[$i])) for i in 1:N]...)) + unsafe_store!(p_vec, vec_vals, 1, Val($(N * sizeof(T)))) + nothing + end +end + +@device_override @inline function KI.vload(::Val{N}, arr::CLDeviceArray{T}, idx::Integer) where {N, T} + if isprimitivetype(T) && N > 1 + return _vload_lptr(Val(N), pointer(arr, idx)) + end + KI._vload_arr(Val(N), arr, idx) +end + +@device_override @inline function KI.vstore!(arr::CLDeviceArray{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} + if isprimitivetype(T) && N > 1 + _vstore_lptr!(pointer(arr, idx), Val(N), vals) + return nothing + end + KI._vstore_arr!(arr, idx, Val(N), vals) + return nothing +end + end From fb7b4a4db64bc2ce4be7e185ec47cb63a046a487 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Tue, 30 Jun 2026 15:12:43 +0530 Subject: [PATCH 3/5] fix(intrinsics): specialize vload/vstore! on Array to avoid GPU compile error MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The AbstractArray fallback called pointer(arr, idx) which routes through unsafe_convert(Ptr{T}, arr) — a path with a runtime string-formatting error that GPU compilers (CUDA, SPIR-V) reject as an unsupported dynamic function invocation. Fix: split into two methods — Array (CPU Ptr path) and AbstractArray (scalar fallback). GPU @device_overrides catch their device array types before the AbstractArray fallback, so the Ptr path is never reachable in GPU IR. Co-Authored-By: Claude Sonnet 4.6 --- src/intrinsics.jl | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/src/intrinsics.jl b/src/intrinsics.jl index 11bfc1964..a898147b0 100644 --- a/src/intrinsics.jl +++ b/src/intrinsics.jl @@ -255,13 +255,12 @@ end @device_override @inline function KI.vload(::Val{N}, arr::MyDeviceArray{T}, idx::Integer) where {N, T} ``` """ +@inline function vload(::Val{N}, arr::Array{T}, idx::Integer) where {N, T} + isprimitivetype(T) && N > 1 && return _vload_ptr(Val(N), pointer(arr, idx)) + _vload_arr(Val(N), arr, idx) +end + @inline function vload(::Val{N}, arr::AbstractArray{T}, idx::Integer) where {N, T} - if isprimitivetype(T) && N > 1 - p = pointer(arr, idx) - if p isa Ptr - return _vload_ptr(Val(N), p) - end - end _vload_arr(Val(N), arr, idx) end @@ -286,18 +285,20 @@ The same `N`-must-be-`Val`, no-`@Const`, and alignment requirements as @device_override @inline function KI.vstore!(arr::MyDeviceArray{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} ``` """ -@inline function vstore!(arr::AbstractArray{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} +@inline function vstore!(arr::Array{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} if isprimitivetype(T) && N > 1 - p = pointer(arr, idx) - if p isa Ptr - _vstore_ptr!(p, Val(N), vals) - return nothing - end + _vstore_ptr!(pointer(arr, idx), Val(N), vals) + return nothing end _vstore_arr!(arr, idx, Val(N), vals) return nothing end +@inline function vstore!(arr::AbstractArray{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} + _vstore_arr!(arr, idx, Val(N), vals) + return nothing +end + """ Kernel{Backend, Kern} From 95f990a8433ba8bdce5690f59497f7bc58e3aa88 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Tue, 30 Jun 2026 17:59:30 +0530 Subject: [PATCH 4/5] docs(intrinsics): tighten vload/vstore! docstrings and inline comments Co-Authored-By: Claude Sonnet 4.6 --- src/intrinsics.jl | 94 +++++++++++---------------------------------- src/pocl/backend.jl | 9 ++--- 2 files changed, 25 insertions(+), 78 deletions(-) diff --git a/src/intrinsics.jl b/src/intrinsics.jl index a898147b0..4bf7eb191 100644 --- a/src/intrinsics.jl +++ b/src/intrinsics.jl @@ -156,25 +156,11 @@ end function _print end -# ── Internal helpers for wide vector loads / stores ─────────────────────────── -# -# CPU (Ptr{T}): cast pointer to Ptr{NTuple{N, VecElement{T}}} and call -# unsafe_load / unsafe_store!. NTuple{N, VecElement{T}} is Julia's -# representation of LLVM's ; the Julia/LLVM compiler lowers this -# to a single wide memory op. -# -# GPU backends: each backend registers a @device_override for its concrete -# device array type (e.g. CLDeviceArray in src/pocl/backend.jl, -# CuDeviceArray in CUDA.jl's KA extension). The override uses -# reinterpret(LLVMPtr{NTuple{N,VecElement{T}},AS}, ptr) + unsafe_load to -# emit a single wide instruction — ld.global.v4.f32 (CUDA/NVPTX), -# global_load_dwordx4 (AMDGPU), OpLoad (SPIR-V/POCL). -# Without an override, vload falls back to scalar indexing. -# -# Fallback (N=1, non-primitive types, or any array with no override): -# scalar arr[idx+k] accesses via @generated to avoid closures. - -# CPU path ──────────────────────────────────────────────────────────────────── +# ── vload / vstore! helpers ─────────────────────────────────────────────────── +# Three dispatch paths (in priority order): +# Array{T} → Ptr cast to Ptr{NTuple{N,VecElement{T}}} + unsafe_load +# GPU device array → @device_override per backend (see e.g. pocl/backend.jl) +# AbstractArray{T} → N scalar reads/writes (@generated, no closures) @inline _vptr(p::Ptr{T}, ::Val{N}) where {T, N} = Ptr{NTuple{N, Core.VecElement{T}}}(p) @@ -188,8 +174,6 @@ end unsafe_store!(_vptr(p, Val(N)), ntuple(i -> Core.VecElement{T}(vals[i]), Val(N))) end -# Scalar fallback — used for N=1, non-primitive types, and GPU arrays -# without a @device_override ────────────────────────────────────────────────── @generated function _vload_arr(::Val{N}, arr::AbstractArray{T}, idx::Integer) where {N, T} Expr(:tuple, [:(Base.@inbounds arr[idx + $(i - 1)]) for i in 1:N]...) @@ -202,55 +186,27 @@ end :nothing) end -# ─── vload ─────────────────────────────────────────────────────────────────── - """ vload(::Val{N}, arr::AbstractArray{T}, idx::Integer) → NTuple{N, T} -Load `N` consecutive elements starting at 1-based index `idx` from `arr`. - -### How it works - -On **CPU** (`Ptr{T}` path), for primitive element types this casts -`pointer(arr, idx)` to a `Ptr{NTuple{N, Core.VecElement{T}}}` and issues -a single `unsafe_load`. Julia lowers `NTuple{N, VecElement{T}}` to -LLVM's `` vector type — one wide memory op. - -On **GPU**, each backend registers a `@device_override` for its concrete -device array type (e.g. `CLDeviceArray` for POCL, `CuDeviceArray` for -CUDA.jl). The override reinterprets the pointer as -`LLVMPtr{NTuple{N, VecElement{T}}, AS}` and issues a single `unsafe_load` -with `N*sizeof(T)` alignment, which the backend LLVM pipeline lowers to one -wide instruction — `ld.global.v4.f32` (CUDA/NVPTX), -`global_load_dwordx4` (AMDGPU), `OpLoad ` (SPIR-V/POCL). - -For non-primitive element types, `N == 1`, or any array type without a -registered override, the fallback emits N plain scalar array accesses. - -### Requirements +Load `N` consecutive elements starting at 1-based index `idx`. -- `N` must be a **compile-time constant** — pass `Val(4)`, not a variable. -- The array must **not** be annotated with `@Const`; `@Const` inserts `ldg` - (read-only cache) intrinsics whose pointer type differs from a plain global - load and cannot be widened into a vector load. -- `idx` must be aligned to `N * sizeof(T)` bytes for the hardware vector - instruction to fire. Julia's allocator guarantees ≥ 64-byte base alignment, - so `idx = 1 + k*N` for `k ≥ 0` is always valid. +Emits a single wide memory instruction for primitive `T` and `N > 1`: on CPU +by casting the pointer to `Ptr{NTuple{N,VecElement{T}}}`, on GPU via a backend +`@device_override` that reinterprets the `LLVMPtr` with `N*sizeof(T)` alignment +(lowered to `ld.global.v4.f32` on NVPTX, `OpLoad ` on SPIR-V, etc.). +Falls back to `N` scalar reads when no override is registered, `N == 1`, or `T` +is not a primitive type. -### Example - -```julia -function reduce_kernel(dst, src) - g = KI.get_global_id().x - idx = (g - 1) * 4 + 1 - v0, v1, v2, v3 = KI.vload(Val(4), src, idx) # single ld.global.v4.f32 - KI.vstore!(dst, idx, (v0+1f0, v1+1f0, v2+1f0, v3+1f0)) - return -end -``` +**Constraints:** +- `N` must be a compile-time constant (`Val(4)`, not a variable). +- `idx` must be `N*sizeof(T)`-aligned; `idx = 1 + k*N` for `k ≥ 0` satisfies + this for Julia-allocated arrays (≥ 64-byte base alignment). +- Do not pass `@Const`-annotated arrays; the read-only-cache pointer type is + incompatible with wide loads. !!! note - GPU backends provide the vectorized path via: + GPU backends register the vectorized path as: ```julia @device_override @inline function KI.vload(::Val{N}, arr::MyDeviceArray{T}, idx::Integer) where {N, T} ``` @@ -264,23 +220,17 @@ end _vload_arr(Val(N), arr, idx) end -# ─── vstore! ───────────────────────────────────────────────────────────────── - """ vstore!(arr::AbstractArray{T}, idx::Integer, vals::NTuple{N, T}) Store `N` consecutive elements from `vals` into `arr` starting at 1-based index `idx`. -The write complement of [`vload`](@ref). On CPU this wraps `vals` in a -`` LLVM vector and issues a single `unsafe_store!`. GPU backends -provide wide store via `@device_override` on their device array type. - -The same `N`-must-be-`Val`, no-`@Const`, and alignment requirements as -[`vload`](@ref) apply. +Write counterpart of [`vload`](@ref); same dispatch logic and alignment +constraints apply. !!! note - GPU backends provide the vectorized path via: + GPU backends register the vectorized path as: ```julia @device_override @inline function KI.vstore!(arr::MyDeviceArray{T}, idx::Integer, vals::NTuple{N, T}) where {N, T} ``` diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index e23fb3e43..bbf7ce735 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -242,12 +242,9 @@ end KA.argconvert(::KA.Kernel{POCLBackend}, arg) = clconvert(arg) -## Vectorized memory operations — @device_override for CLDeviceArray -# -# These helpers live here (not in KI) because they are POCL-specific: they -# operate on LLVMPtr with SPIR-V address spaces and use reinterpret + -# unsafe_load/unsafe_store! to emit OpLoad/OpStore on vectors. -# CUDA.jl, AMDGPU.jl etc. provide analogous overrides for their array types. +## KI.vload / KI.vstore! — SPIR-V vector memory ops for CLDeviceArray +# Reinterprets the LLVMPtr as LLVMPtr{NTuple{N,VecElement{T}},AS} with +# N*sizeof(T) alignment to emit a single OpLoad/OpStore instruction. @generated function _vload_lptr(::Val{N}, p::Core.LLVMPtr{T, A}) where {N, T, A} VT = NTuple{N, Core.VecElement{T}} From 2543a3e51e858786e2442057ff33eeaef2e2c4c5 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Tue, 30 Jun 2026 18:54:01 +0530 Subject: [PATCH 5/5] fix(test): skip Float64 vload tests on backends without fp64 support oneAPI (Intel GPU) throws "Float64 is not supported on this device" when allocating a oneArray{Float64}, so guard Float64 test cases behind KernelAbstractions.supports_float64(backend()). Co-Authored-By: Claude Sonnet 4.6 --- test/intrinsics.jl | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/test/intrinsics.jl b/test/intrinsics.jl index bb0cafee5..d8b130d84 100644 --- a/test/intrinsics.jl +++ b/test/intrinsics.jl @@ -184,9 +184,12 @@ function intrinsics_testsuite(backend, AT) @testset "vload / vstore!" begin + fp64 = KernelAbstractions.supports_float64(backend()) + # ── 1. Roundtrip for every supported (N, T) combination ───────── + vload_types = fp64 ? (Float32, Float64, Int32, Int64) : (Float32, Int32, Int64) @testset "roundtrip N=$N T=$T" for N in (1, 2, 4, 8), - T in (Float32, Float64, Int32, Int64) + T in vload_types n = 64 src = AT(T.(1:n)) dst = AT(zeros(T, n)) @@ -223,13 +226,15 @@ function intrinsics_testsuite(backend, AT) @test Array(dst) ≈ 2.0f0 .* Float32.(1:n) end - @testset "scale by 3 N=2 Float64" begin - n = 64 - src = AT(Float64.(1:n)) - dst = AT(zeros(Float64, n)) - KI.@kernel backend() numworkgroups = n ÷ 2 workgroupsize = 1 vload_scale_kernel(dst, src, 3.0, Val(2)) - KernelAbstractions.synchronize(backend()) - @test Array(dst) ≈ 3.0 .* Float64.(1:n) + if fp64 + @testset "scale by 3 N=2 Float64" begin + n = 64 + src = AT(Float64.(1:n)) + dst = AT(zeros(Float64, n)) + KI.@kernel backend() numworkgroups = n ÷ 2 workgroupsize = 1 vload_scale_kernel(dst, src, 3.0, Val(2)) + KernelAbstractions.synchronize(backend()) + @test Array(dst) ≈ 3.0 .* Float64.(1:n) + end end # ── 5. Per-thread reduction: each thread sums its N elements ──── @@ -313,7 +318,7 @@ function intrinsics_testsuite(backend, AT) end # ── 12. Type stability ─────────────────────────────────────────── - @testset "type stability T=$T N=$N" for T in (Float32, Float64, Int32, Int64), N in (1, 4) + @testset "type stability T=$T N=$N" for T in vload_types, N in (1, 4) n = N src = AT(T.(1:n)) dst = AT(zeros(T, n))