Skip to content

Add previx_load and prefix_store! as optimized load with mask etc #161

Open
oscardssmith wants to merge 3 commits into
eschnett:masterfrom
oscardssmith:patch-1
Open

Add previx_load and prefix_store! as optimized load with mask etc #161
oscardssmith wants to merge 3 commits into
eschnett:masterfrom
oscardssmith:patch-1

Conversation

@oscardssmith

@oscardssmith oscardssmith commented May 6, 2026

Copy link
Copy Markdown
Contributor

This is helpful for generating slightly better code. PR generated with claude help, but I've reviewed it to ensure reasonability:

julia> @inline lanewise_mask(::Val{N}, n::Int64) where {N} =
           Vec{N,Bool}(ntuple(i -> i <= n, Val(N)))
lanewise_mask (generic function with 1 method)

julia> @inline vload_via_lanewise(::Type{Vec{N,T}}, p::Ptr{T}, n::Int64) where {N,T} =
           vload(Vec{N,T}, p, lanewise_mask(Val(N), n))
vload_via_lanewise (generic function with 1 method)

julia> @inline vstore_via_lanewise!(v::Vec{N,T}, p::Ptr{T}, n::Int64) where {N,T} =
           vstore(v, p, lanewise_mask(Val(N), n))
vstore_via_lanewise! (generic function with 1 method)

julia> @code_native debuginfo=:none syntax=:intel prefix_mask(Val(W), Int64(5))
	.text
	.file	"prefix_mask"
	.section	.rodata.cst16,"aM",@progbits,16
	.p2align	4, 0x0                          # -- Begin function julia_prefix_mask_10453
.LCPI0_0:
	.zero	16,1
	.section	.ltext,"axl",@progbits
	.globl	julia_prefix_mask_10453
	.p2align	4, 0x90
	.type	julia_prefix_mask_10453,@function
julia_prefix_mask_10453:                # @julia_prefix_mask_10453
; Function Signature: prefix_mask(Base.Val{16}, Int64)
# %bb.0:                                # %top
	#DEBUG_VALUE: prefix_mask:n <- $rsi
	push	rbp
	mov	rbp, rsp
	cmp	rsi, 16
	mov	ecx, 16
	cmovb	rcx, rsi
	neg	ecx
	xor	eax, eax
	mov	rdx, -1
                                        # kill: def $cl killed $cl killed $rcx
	shrd	rdx, rax, cl
	mov	rax, rdi
	kmovd	k1, edx
	movabs	rcx, offset .LCPI0_0
	vmovdqu8	xmm0 {k1} {z}, xmmword ptr [rcx]
	vmovdqa	xmmword ptr [rdi], xmm0
	pop	rbp
	ret
.Lfunc_end0:
	.size	julia_prefix_mask_10453, .Lfunc_end0-julia_prefix_mask_10453
                                        # -- End function
	.type	".L+SIMD.Vec#10455",@object     # @"+SIMD.Vec#10455"
	.section	.lrodata,"al",@progbits
	.p2align	3, 0x0
".L+SIMD.Vec#10455":
	.quad	".L+SIMD.Vec#10455.jit"
	.size	".L+SIMD.Vec#10455", 8

.set ".L+SIMD.Vec#10455.jit", 123450746584720
	.size	".L+SIMD.Vec#10455.jit", 8
	.section	".note.GNU-stack","",@progbits

julia> @code_native debuginfo=:none syntax=:intel lanewise_mask(Val(W), Int64(5))

	.text
	.file	"lanewise_mask"
	.section	.rodata,"a",@progbits
	.p2align	6, 0x0                          # -- Begin function julia_lanewise_mask_10459
.LCPI0_0:
	.quad	0                               # 0x0
	.quad	1                               # 0x1
	.quad	2                               # 0x2
	.quad	3                               # 0x3
	.quad	4                               # 0x4
	.quad	5                               # 0x5
	.quad	6                               # 0x6
	.quad	7                               # 0x7
.LCPI0_1:
	.quad	8                               # 0x8
	.quad	9                               # 0x9
	.quad	10                              # 0xa
	.quad	11                              # 0xb
	.quad	12                              # 0xc
	.quad	13                              # 0xd
	.quad	14                              # 0xe
	.quad	15                              # 0xf
	.section	.rodata.cst16,"aM",@progbits,16
	.p2align	4, 0x0
.LCPI0_2:
	.zero	16,1
	.section	.ltext,"axl",@progbits
	.globl	julia_lanewise_mask_10459
	.p2align	4, 0x90
	.type	julia_lanewise_mask_10459,@function
julia_lanewise_mask_10459:              # @julia_lanewise_mask_10459
; Function Signature: lanewise_mask(Base.Val{16}, Int64)
# %bb.0:                                # %top
	#DEBUG_VALUE: lanewise_mask:n <- $rsi
	push	rbp
	mov	rbp, rsp
	mov	rax, rdi
	vpbroadcastq	zmm0, rsi
	movabs	rcx, offset .LCPI0_0
	vpcmpgtq	k0, zmm0, zmmword ptr [rcx]
	movabs	rcx, offset .LCPI0_1
	vpcmpgtq	k1, zmm0, zmmword ptr [rcx]
	kunpckbw	k1, k1, k0
	movabs	rcx, offset .LCPI0_2
	vmovdqu8	xmm0 {k1} {z}, xmmword ptr [rcx]
	vmovdqa	xmmword ptr [rdi], xmm0
	pop	rbp
	vzeroupper
	ret
.Lfunc_end0:
	.size	julia_lanewise_mask_10459, .Lfunc_end0-julia_lanewise_mask_10459
                                        # -- End function
	.type	".L+SIMD.Vec#10461",@object     # @"+SIMD.Vec#10461"
	.section	.lrodata,"al",@progbits
	.p2align	3, 0x0
".L+SIMD.Vec#10461":
	.quad	".L+SIMD.Vec#10461.jit"
	.size	".L+SIMD.Vec#10461", 8

.set ".L+SIMD.Vec#10461.jit", 123450746584720
	.size	".L+SIMD.Vec#10461.jit", 8
	.section	".note.GNU-stack","",@progbits

julia>  @code_native debuginfo=:none syntax=:intel vload_prefix(Vec{W,Float32}, Ptr{Float32}(0), Int64(5))
	.text
	.file	"vload_prefix"
	.section	.ltext,"axl",@progbits
	.globl	julia_vload_prefix_10462        # -- Begin function julia_vload_prefix_10462
	.p2align	4, 0x90
	.type	julia_vload_prefix_10462,@function
julia_vload_prefix_10462:               # @julia_vload_prefix_10462
; Function Signature: vload_prefix(Type{SIMD.Vec{16, Float32}}, Ptr{Float32}, Int64)
# %bb.0:                                # %top
	#DEBUG_VALUE: vload_prefix:ptr <- $rsi
	#DEBUG_VALUE: vload_prefix:n <- $rdx
	push	rbp
	mov	rbp, rsp
	cmp	rdx, 16
	mov	ecx, 16
	cmovb	rcx, rdx
	neg	ecx
	xor	eax, eax
	mov	rdx, -1
                                        # kill: def $cl killed $cl killed $rcx
	shrd	rdx, rax, cl
	mov	rax, rdi
	kmovd	k1, edx
	vmovups	zmm0 {k1} {z}, zmmword ptr [rsi]
	vmovups	zmmword ptr [rdi], zmm0
	pop	rbp
	vzeroupper
	ret
.Lfunc_end0:
	.size	julia_vload_prefix_10462, .Lfunc_end0-julia_vload_prefix_10462
                                        # -- End function
	.type	".L+SIMD.Vec#10464",@object     # @"+SIMD.Vec#10464"
	.section	.lrodata,"al",@progbits
	.p2align	3, 0x0
".L+SIMD.Vec#10464":
	.quad	".L+SIMD.Vec#10464.jit"
	.size	".L+SIMD.Vec#10464", 8

.set ".L+SIMD.Vec#10464.jit", 123450787514512
	.size	".L+SIMD.Vec#10464.jit", 8
	.section	".note.GNU-stack","",@progbits

julia> @code_native debuginfo=:none syntax=:intel vload_via_lanewise(Vec{W,Float32}, Ptr{Float32}(0), Int64(5))

	.text
	.file	"vload_via_lanewise"
	.section	.rodata,"a",@progbits
	.p2align	6, 0x0                          # -- Begin function julia_vload_via_lanewise_10470
.LCPI0_0:
	.quad	0                               # 0x0
	.quad	1                               # 0x1
	.quad	2                               # 0x2
	.quad	3                               # 0x3
	.quad	4                               # 0x4
	.quad	5                               # 0x5
	.quad	6                               # 0x6
	.quad	7                               # 0x7
.LCPI0_1:
	.quad	8                               # 0x8
	.quad	9                               # 0x9
	.quad	10                              # 0xa
	.quad	11                              # 0xb
	.quad	12                              # 0xc
	.quad	13                              # 0xd
	.quad	14                              # 0xe
	.quad	15                              # 0xf
	.section	.ltext,"axl",@progbits
	.globl	julia_vload_via_lanewise_10470
	.p2align	4, 0x90
	.type	julia_vload_via_lanewise_10470,@function
julia_vload_via_lanewise_10470:         # @julia_vload_via_lanewise_10470
; Function Signature: vload_via_lanewise(Type{SIMD.Vec{16, Float32}}, Ptr{Float32}, Int64)
# %bb.0:                                # %top
	#DEBUG_VALUE: vload_via_lanewise:p <- $rsi
	#DEBUG_VALUE: vload_via_lanewise:n <- $rdx
	push	rbp
	mov	rbp, rsp
	vpbroadcastq	zmm0, rdx
	movabs	rax, offset .LCPI0_0
	vpcmpgtq	k0, zmm0, zmmword ptr [rax]
	movabs	rax, offset .LCPI0_1
	vpcmpgtq	k1, zmm0, zmmword ptr [rax]
	mov	rax, rdi
	kunpckbw	k1, k1, k0
	vmovups	zmm0 {k1} {z}, zmmword ptr [rsi]
	vmovups	zmmword ptr [rdi], zmm0
	pop	rbp
	vzeroupper
	ret
.Lfunc_end0:
	.size	julia_vload_via_lanewise_10470, .Lfunc_end0-julia_vload_via_lanewise_10470
                                        # -- End function
	.type	".L+SIMD.Vec#10472",@object     # @"+SIMD.Vec#10472"
	.section	.lrodata,"al",@progbits
	.p2align	3, 0x0
".L+SIMD.Vec#10472":
	.quad	".L+SIMD.Vec#10472.jit"
	.size	".L+SIMD.Vec#10472", 8

.set ".L+SIMD.Vec#10472.jit", 123450787514512
	.size	".L+SIMD.Vec#10472.jit", 8
	.section	".note.GNU-stack","",@progbits

julia> @code_native debuginfo=:none syntax=:intel vstore_prefix!(Vec{W,Float32}(0f0), Ptr{Float32}(0), Int64(5))

	.text
	.file	"vstore_prefix!"
	.section	.ltext,"axl",@progbits
	.globl	"julia_vstore_prefix!_10476"    # -- Begin function julia_vstore_prefix!_10476
	.p2align	4, 0x90
	.type	"julia_vstore_prefix!_10476",@function
"julia_vstore_prefix!_10476":           # @"julia_vstore_prefix!_10476"
; Function Signature: vstore_prefix!(SIMD.Vec{16, Float32}, Ptr{Float32}, Int64)
# %bb.0:                                # %top
	#DEBUG_VALUE: vstore_prefix!:v <- [$rdi+0]
	#DEBUG_VALUE: vstore_prefix!:ptr <- $rsi
	#DEBUG_VALUE: vstore_prefix!:n <- $rdx
	push	rbp
	mov	rbp, rsp
	cmp	rdx, 16
	mov	ecx, 16
	cmovb	rcx, rdx
	neg	ecx
	xor	eax, eax
	mov	rdx, -1
                                        # kill: def $cl killed $cl killed $rcx
	shrd	rdx, rax, cl
	kmovd	k1, edx
	vmovups	zmm0, zmmword ptr [rdi]
	vmovups	zmmword ptr [rsi] {k1}, zmm0
	pop	rbp
	vzeroupper
	ret
.Lfunc_end0:
	.size	"julia_vstore_prefix!_10476", .Lfunc_end0-"julia_vstore_prefix!_10476"
                                        # -- End function
	.section	".note.GNU-stack","",@progbits

julia> @code_native debuginfo=:none syntax=:intel vstore_via_lanewise!(Vec{W,Float32}(0f0), Ptr{Float32}(0), Int64(5))

	.text
	.file	"vstore_via_lanewise!"
	.section	.rodata,"a",@progbits
	.p2align	6, 0x0                          # -- Begin function julia_vstore_via_lanewise!_10483
.LCPI0_0:
	.quad	0                               # 0x0
	.quad	1                               # 0x1
	.quad	2                               # 0x2
	.quad	3                               # 0x3
	.quad	4                               # 0x4
	.quad	5                               # 0x5
	.quad	6                               # 0x6
	.quad	7                               # 0x7
.LCPI0_1:
	.quad	8                               # 0x8
	.quad	9                               # 0x9
	.quad	10                              # 0xa
	.quad	11                              # 0xb
	.quad	12                              # 0xc
	.quad	13                              # 0xd
	.quad	14                              # 0xe
	.quad	15                              # 0xf
	.section	.ltext,"axl",@progbits
	.globl	"julia_vstore_via_lanewise!_10483"
	.p2align	4, 0x90
	.type	"julia_vstore_via_lanewise!_10483",@function
"julia_vstore_via_lanewise!_10483":     # @"julia_vstore_via_lanewise!_10483"
; Function Signature: vstore_via_lanewise!(SIMD.Vec{16, Float32}, Ptr{Float32}, Int64)
# %bb.0:                                # %top
	#DEBUG_VALUE: vstore_via_lanewise!:v <- [$rdi+0]
	#DEBUG_VALUE: vstore_via_lanewise!:p <- $rsi
	#DEBUG_VALUE: vstore_via_lanewise!:n <- $rdx
	push	rbp
	mov	rbp, rsp
	vpbroadcastq	zmm0, rdx
	movabs	rax, offset .LCPI0_0
	vpcmpgtq	k0, zmm0, zmmword ptr [rax]
	movabs	rax, offset .LCPI0_1
	vpcmpgtq	k1, zmm0, zmmword ptr [rax]
	kunpckbw	k1, k1, k0
	vmovups	zmm0, zmmword ptr [rdi]
	vmovups	zmmword ptr [rsi] {k1}, zmm0
	pop	rbp
	vzeroupper
	ret
.Lfunc_end0:
	.size	"julia_vstore_via_lanewise!_10483", .Lfunc_end0-"julia_vstore_via_lanewise!_10483"
                                        # -- End function
	.section	".note.GNU-stack","",@progbits
	```

I found this while trying to write a fast SIMD.jl only matmul

@codecov-commenter

codecov-commenter commented May 6, 2026 •
edited
Loading

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 0% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.77%. Comparing base (ec78d89) to head (a8ba018).

Files with missing lines Patch % Lines
src/LLVM_intrinsics.jl 0.00% 7 Missing ⚠️
src/arrayops.jl 0.00% 4 Missing ⚠️
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.
Additional details and impacted files 
@@            Coverage Diff             @@
##           master     #161      +/-   ##
==========================================
- Coverage   83.04%   81.77%   -1.28%     
==========================================
  Files           5        5              
  Lines         643      653      +10     
==========================================
  Hits          534      534              
- Misses        109      119      +10

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

vchuravy
vchuravy reviewed May 6, 2026
View reviewed changes
Comment thread src/LLVM_intrinsics.jl
and the remaining lanes are `false`. The canonical "vectorized-loop tail"
mask: lane `i` (0-indexed) is `true` iff `i < n`.

Lowers to LLVM's `llvm.get.active.lane.mask` intrinsic, which targets

@vchuravy vchuravy May 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this part of the comment still true?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@vchuravy vchuravy vchuravy left review comments

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@oscardssmith @codecov-commenter @vchuravy