Skip to content

Use GPUCompiler's workaround for nested insertvalues #646

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
maleadt merged 1 commit into from
Sep 26, 2026
+4 −115
Merged

Use GPUCompiler's workaround for nested insertvalues #646

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ Adapt = "4"
CEnum = "0.4, 0.5"
ExprTools = "0.1"
GPUArrays = "11.5.14"
GPUCompiler = "2.8"
GPUCompiler = "2.9"
GPUToolbox = "3.1"
KernelAbstractions = "0.9.39"
LLVM = "6, 7, 8, 9"
Expand Down
115 changes: 2 additions & 113 deletions src/compiler/compilation.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,6 @@ function GPUCompiler.finish_module!(job::oneAPICompilerJob, mod::LLVM.Module,
return entry
end

# finish_ir! runs later in the pipeline, after optimizations that create nested insertvalue
function GPUCompiler.finish_ir!(job::oneAPICompilerJob, mod::LLVM.Module,
entry::LLVM.Function)
# Initialize the heap before SPIR-V lowering converts the state to a reference.
Expand All @@ -81,12 +80,6 @@ function GPUCompiler.finish_ir!(job::oneAPICompilerJob, mod::LLVM.Module,
Tuple{CompilerJob{SPIRVCompilerTarget}, typeof(mod), typeof(entry)},
job, mod, entry)

# FIX: Flatten nested insertvalue instructions to work around SPIR-V bug
# See: https://github.com/JuliaGPU/oneAPI.jl/issues/259
# Intel's SPIR-V runtime has a bug where OpCompositeInsert with nested
# indices (e.g., "1 0") corrupts adjacent struct fields.
flatten_nested_insertvalue!(mod)

# When the device supports BFloat16 but the SPIR-V runtime doesn't accept
# SPV_KHR_bfloat16, lower all bfloat types to i16 so the back-end/translator can
# handle the module without the extension. Both conditions are read from the
Expand Down Expand Up @@ -154,111 +147,6 @@ function uses_heap(mod::LLVM.Module, T_state::LLVMType, heap_field::Integer)
return false
end

# Flatten nested insertvalue instructions
# This works around a bug in Intel's SPIR-V runtime where OpCompositeInsert
# with nested array indices corrupts adjacent struct fields.
function flatten_nested_insertvalue!(mod::LLVM.Module)
changed = false
count = 0

for f in functions(mod)
isempty(blocks(f)) && continue

for bb in blocks(f)
# Collect instructions to process (can't modify while iterating)
to_process = LLVM.Instruction[]

for inst in instructions(bb)
# Check if this is an insertvalue with nested indices
if LLVM.API.LLVMGetInstructionOpcode(inst) == LLVM.API.LLVMInsertValue
num_indices = LLVM.API.LLVMGetNumIndices(inst)
if num_indices > 1
push!(to_process, inst)
end
end
end

# Flatten each nested insertvalue
for inst in to_process
try
flatten_insert!(inst)
changed = true
count += 1
catch e
@warn "Failed to flatten nested insertvalue" exception=(e, catch_backtrace())
end
end
end
end

return changed
end

function flatten_insert!(inst::LLVM.Instruction)
# Transform: insertvalue %base, %val, i, j, k...
# Into: extractvalue %base, i
# insertvalue %extracted, %val, j, k...
# insertvalue %base, %modified, i

composite = LLVM.operands(inst)[1]
value = LLVM.operands(inst)[2]

num_indices = LLVM.API.LLVMGetNumIndices(inst)
idx_ptr = LLVM.API.LLVMGetIndices(inst)
indices = unsafe_wrap(Array, idx_ptr, num_indices)

builder = LLVM.IRBuilder()
LLVM.position!(builder, inst)

# Strategy: Recursively extract and insert for each nesting level
# For insertvalue %base, %val, i, j, k
# Do: %tmp1 = extractvalue %base, i
# %tmp2 = extractvalue %tmp1, j
# %tmp3 = insertvalue %tmp2, %val, k
# %tmp4 = insertvalue %tmp1, %tmp3, j
# %result = insertvalue %base, %tmp4, i

# But that's complex. Simpler approach for 2-3 levels:
# Just do one level of flattening at a time
first_idx = indices[1]
rest_indices = indices[2:end]

# Extract the first level
extracted = LLVM.extract_value!(builder, composite, first_idx)

# Now insert into the extracted value using remaining indices
# The LLVM IR builder will handle this correctly
inserted = extracted
if length(rest_indices) == 1
# Simple case: just one more level
inserted = LLVM.insert_value!(builder, extracted, value, rest_indices[1])
else
# Multiple levels: need to extract down, insert, then insert back up
# For now, recursively extract to the deepest level
temps = [extracted]
for i in 1:(length(rest_indices)-1)
temp = LLVM.extract_value!(builder, temps[end], rest_indices[i])
push!(temps, temp)
end

# Insert the value at the deepest level
inserted = LLVM.insert_value!(builder, temps[end], value, rest_indices[end])

# Insert back up the chain
for i in (length(rest_indices)-1):-1:1
inserted = LLVM.insert_value!(builder, temps[i], inserted, rest_indices[i])
end
end

# Insert the modified structure back into the original
result = LLVM.insert_value!(builder, composite, inserted, first_idx)

LLVM.replace_uses!(inst, result)
LLVM.API.LLVMInstructionEraseFromParent(inst)
LLVM.dispose(builder)
end


# Lower bfloat types to i16 in the LLVM IR.
# This is needed when the device supports BFloat16 but the SPIR-V runtime/translator
# doesn't support SPV_KHR_bfloat16. Since sizeof(bfloat)==sizeof(i16)==2, the memory
Expand Down Expand Up @@ -421,7 +309,8 @@ end
extensions_str = join(map(ext -> "+$ext", extensions), ",")

# create GPUCompiler objects
target = SPIRVCompilerTarget(; backend, extensions = extensions_str, supports_fp16, supports_fp64, supports_bfloat16, kwargs...)
target = SPIRVCompilerTarget(; backend, extensions = extensions_str, supports_fp16, supports_fp64, supports_bfloat16,
driver = :intel, kwargs...)
params = oneAPICompilerParams()
CompilerConfig(target, params; kernel, name, always_inline)
end
Expand Down
2 changes: 1 addition & 1 deletion src/compiler/precompile.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ if _precompile_lts ? SPIRV_LLVM_Translator_jll.is_available() :
target = SPIRVCompilerTarget(;
backend = (_precompile_lts ? :khronos : :llvm),
extensions = "", supports_fp16 = true,
supports_fp64 = true, supports_bfloat16 = false,
supports_fp64 = true, supports_bfloat16 = false, driver = :intel,
)
params = oneAPICompilerParams()
config = CompilerConfig(target, params; kernel=true, name=nothing,
Expand Down
Loading