From d963762c7d337b3c54efd211468eb061fe04e4b8 Mon Sep 17 00:00:00 2001 From: ejmeitz Date: Wed, 8 Apr 2026 11:48:05 -0500 Subject: [PATCH 01/12] copy over prior fusion work --- src/fusion.jl | 385 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 385 insertions(+) create mode 100644 src/fusion.jl diff --git a/src/fusion.jl b/src/fusion.jl new file mode 100644 index 00000000..b59ac69f --- /dev/null +++ b/src/fusion.jl @@ -0,0 +1,385 @@ +export @fuse + +const ALLOWED_FUSION_TYPES = Union{<:NDArray,<:Number} +const FUSED_KERNEL_CACHE = Dict{UInt64,CUDATask}() +const FUSE_KWARGS = [:blocks, :threads, :outputs] + +struct FusedKernelData{I,O,S} + ct::CUDATask + inputs::I + outputs::O + scalars::S +end + +#! ASSUME OUTPUT_INDICIES ARE THE ARRAYS +#! RETURNED FROM THE FUNCTION AS WELL?? +#! HOW TO HANDLE CASE WHEN output_indicies = () +#! HOW TO HANDLE KWARGS?? +function fuse_function( + fn::Function, + output_indices, + args...; + kwargs..., +) + println("Compiling $(Symbol(fn)) to PTX") + + N_args = length(args) + if length(output_indices) > N_args || max(output_indices) > N_args + error( + ArgumentError( + "Marked arguments $(output_indices) as outputs, but there are only $(N_args) arguments" + ), + ) + end + + input_indices = setdiff(1:length(args), Set(output_indices)) + inputs = [args[i] for i in input_indices] + outputs = [args[i] for i in output_indices] + + #! ADD KWARGS TO THESE ?? + scalars = filter(x -> isa(x, Number), inputs) + ndarray_inputs = filter(x -> !isa(x, NDArray), inputs) + other_inputs = filter(x -> !isa(x, Number) && !isa(x, NDArray), inputs) + + #! Parse out the isbtis structs and unwrap them?? + + if any(isa.(outputs, Number)) + @error "Scalar outputs are not allowed, use 0D store." + end + + # Dummy types so we can use CUDA.jl to compile + converted_types = ndarray_cuda_type.(args) # enforces everything is an `isbitstype` + ct = CUDATask(fn, converted_types) + + return FusedKernelData(ct, other_inputs, outputs, scalars) +end + +# For @fuse on user defined function +function maybe_fuse_kernel(fn::Function, output_indices, args::T; kwargs...) where {T} + cache_key = hash((fn, T)) #! ADD KWRAGS TO THIS? + + if !haskey(FUSED_KERNEL_CACHE, cache_key) + FUSED_KERNEL_CACHE[cache_key] = fuse_function( + fn, + output_indices, + args...; + kwargs..., + ) + end + + return cache_key +end + +function run_fused_kernel(cache_key; blocks, threads) + fkd = FUSED_KERNEL_CACHE[cache_key] + cuNumeric.launch(fkd.ct, fkd.inputs, fkd.outputs, fkd.scalars; blocks=blocks, threads=threads) +end + +macro fuse(ex...) + call = ex[end] + kwargs = map(ex[1:(end - 1)]) do kwarg + if kwarg in FUSE_KWARGS + :($kwarg = $kwarg) + else + throw(ArgumentError("Invalid keyword argument '$kwarg', expected one of $(FUSE_KWARGS)")) + end + end + + #! TODO SET DEFAULT BLOCKS/THREADS + #! HOW TO KNOW WHATS THE RIGHT DIMENSION?? + + blocks = get(kwargs[:blocks], DEFAULT_BLOCKS) + threads = get(kwargs[:threads], DEFAULT_THREADS) + output_indices = get(kwargs[:output_indices], ()) + + code = quote end + + if Meta.isexpr(call, :function) + # Parses function epxression into args, name etc. + data = splitdef(longdef(call)) + # splitarg parses args into (name, type, is_slurp, default) + arg_data = map(splitarg, data[:args]) + kwarg_data = map(splitarg, data[:kwargs]) + + arg_names = getindex.(arg_data, 1) + + @gensym wrapper_name, cache_key + + push!(code.args, + quote + + # The fused function + function $(wrapper_name)() where {$(dict[:whereparams]...)} + #! NOT SURE THIS IS RIGHT WAY TO INTERPOLATE ARGS + #! WANT IT INTERPOLATE AS TUPLE OF ARGS + #! TODO PASS KWARGS + $cache_key = maybe_fuse_kernel($(data[:name]), $output_indices, $(arg_names...)) + run_fused_kernel($cache_key; blocks=($blocks), threads=($threads)) + end + + # Replace original function with call to fused + function $(data[:name])( + $(data[:args]...); $(data[:kwargs]...) + ) where {$(dict[:whereparams]...)} + $(wrapper_name)() + end + end, + ) + else + throw(ArgumentError("fuse expected Broadcasted object or function, got $(call.head)")) + end + + return esc(quote + let + $code + end + end) +end + +# Collects operators and function calls from an expression. +# # E.g. x .+ y .* z will return (.+, .*) +# function collect_ops(expr::Expr) +# ops = Symbol[] + +# postwalk(expr) do node +# if node isa Expr && (node.head === :call || node.head === :.) +# op = node.args[1] +# if op isa Symbol +# push!(ops, op) +# # strip one leading '.' so :.+ → :+, but keep names like :sqrt +# # op_sym = startswith(op, ".") ? Symbol(string(op)[2:end]) : op +# # push!(ops, op_sym) +# end +# end +# return node +# end + +# return Tuple(ops) +# end + +# This version requires all args are the same type +# @generated function maybe_fuse_kernel( +# :: NTuple{NARGS, NDArray{T,DIM}}, +# rhs_ops :: NTuple{NOPS, Symbol}, +# wrapper :: Function +# ) where {NARGS,T, DIM, NOPS} + +# # converted_types is constructed purely from the types of the arguments +# # so a generated function will allow us to pre-compute most of this function +# # at compile time! +# converted_types = ntuple(_ -> CuDeviceArray{T, DIM, 1}, Val(NARGS)) +# _hash = xor(hash(rhs_ops), hash(converted_types)) + +# quote +# if !haskey($FUSED_KERNEL_CACHE, $_hash) +# println("Compiling kernel to PTX") +# local ptx = ptx_as_string(wrapper, $converted_types) #! DO NOT interpolate wrapper +# f = open("/pool/emeitz/ptx2.txt", "w") +# println(f, ptx) +# close(f) +# local ptx_f_name = cuNumeric.extract_kernel_name(ptx) +# cuNumeric.ptx_task(ptx, ptx_f_name) +# $FUSED_KERNEL_CACHE[$_hash] = FusedKernelData(ptx_f_name, $converted_types) +# end +# $_hash +# end +# end + +# # This version does not assume args have same types +# # Manually enforces that all args are NDArray +# @generated function maybe_fuse_kernel(rhs_args :: R, +# rhs_ops :: NTuple{NOPS,Symbol}, +# wrapper :: Function) where +# {R<:Tuple, NOPS} + +# arg_types = R.parameters # e.g. (NDArray{Float32,2}, NDArray{Int,1}, …) +# println("IN HERE") + +# for T in arg_types +# T <: NDArray || throw(ArgumentError("all arguments must be NDArray, got $T")) +# end + +# # 2. Build the matching CuDeviceArray type for each element ------------- +# conv_type_exprs = Expr[] +# for T in arg_types +# elty, dim = T.parameters[1:2] # the {T, N} in NDArray{T,N} +# push!(conv_type_exprs, :(CuDeviceArray{$elty, $dim, 1})) +# end +# # an Expr(:tuple, …) is a compile‑time literal of the tuple of types +# converted_types_expr = Expr(:tuple, conv_type_exprs...) + +# const_hash = xor(hash(rhs_ops), hash(arg_types)) + +# quote +# _hash = $const_hash # value is a literal, not recomputed +# if !haskey($FUSED_KERNEL_CACHE, _hash) +# println("Compiling kernel to PTX") +# local ptx = ptx_as_string(wrapper, $converted_types_expr) +# local ptx_f_name = cuNumeric.extract_kernel_name(ptx) +# cuNumeric.ptx_task(ptx, ptx_f_name) #! UNCOMMENT LATER +# $FUSED_KERNEL_CACHE[_hash] = +# FusedKernelData(ptx_f_name, $converted_types_expr) +# end +# _hash +# end +# end + +# macro fuse(ex...) + +# call = ex[end] +# kwargs = map(ex[1:end-1]) do kwarg +# if kwarg in FUSE_KWARGS +# :($kwarg = $kwarg) +# else +# throw(ArgumentError("Invalid keyword argument '$kwarg', expected one of $(FUSE_KWARGS)")) +# end +# end + +# code = quote end + +# if Meta.isexpr(call, :function) +# #TODO use symbol_state.func_defs to fuse each definition +# throw(ArgumentError("@fuse before function definitions is not supported yet.")) +# elseif Meta.isexpr(call, :(=)) || Meta.isexpr(call, :(.=)) +# lhs, rhs = call.args + +# lhs_symbol_state = compute_symbols_state(lhs) +# rhs_symbol_state = compute_symbols_state(rhs) + +# if !isempty(intersect(lhs_symbol_state.references, rhs_symbol_state.references)) +# throw(ArgumentError("LHS and RHS of @fuse must not share variables.")) +# end + +# rhs_ops = collect_ops(rhs) + +# @gensym wrapper_name ptx ptx_f_name task _hash converted_types args_tuple + +# push!(code.args, +# quote + +# # Create wrapper function that we can fuse +# function $(wrapper_name)($(rhs_symbol_state.references...)) +# # println("Wrapper called!") +# return $(rhs) +# end + +# $_hash = $maybe_fuse_kernel(($(rhs_symbol_state.references...),), $rhs_ops, $wrapper_name) + +# println($FUSED_KERNEL_CACHE[$_hash]) + +# $task = $(cuNumeric.CUDATask)($FUSED_KERNEL_CACHE[$_hash]) + +# $(cuNumeric.launch)( +# $task, ($(rhs_symbol_state.references...),), $lhs, (); $(kwargs...) +# ) +# end) +# else +# throw(ArgumentError("fuse expected assignment operator `=` or `.=`, got $(call.head)")) +# end + +# return esc(quote +# let +# $code +# end +# end) +# end + +#* BE SURE TO TEST: +# - x .+ y .*z and x .* y .+ z +# x,y = ... +# y = to_fuse(x) +# z = unary.(x) .+ y +# z = binary.(x,y) .+ z +# macro fuse(ex...) + +# call = ex[end] +# kwargs = map(ex[1:end-1]) do kwarg +# if kwarg in FUSE_KWARGS +# :($kwarg = $kwarg) +# else +# throw(ArgumentError("Invalid keyword argument '$kwarg', expected one of $(FUSE_KWARGS)")) +# end +# end + +# code = quote end + +# if Meta.isexpr(call, :function) +# #TODO use symbol_state.func_defs to fuse each definition +# throw(ArgumentError("@fuse before function definitions is not supported yet.")) +# elseif Meta.isexpr(call, :(=)) || Meta.isexpr(call, :(.=)) +# lhs, rhs = call.args + +# lhs_symbol_state = compute_symbols_state(lhs) +# rhs_symbol_state = compute_symbols_state(rhs) + +# if !isempty(intersect(lhs_symbol_state.references, rhs_symbol_state.references)) +# throw(ArgumentError("LHS and RHS of @fuse must not share variables.")) +# end + +# rhs_ops = collect_ops(rhs) + +# # This will error if all variables are not the same type +# # E.g. NDArrays with different dimensions or el-types +# rhs_args = ntuple(i -> rhs_symbol_state.references[i], length(rhs_symbol_state.references)) + +# @gensym wrapper_name ptx ptx_f_name task _hash converted_types args_tuple + +# push!(code.args, +# quote + +# # Create wrapper function that we can fuse +# # $(wrapper_name) = ($(rhs_symbol_state.references...)) -> $(rhs) + +# #* IF @fuse is in a loop this will be called multiple times for no reason... +# function $(wrapper_name)($(rhs_symbol_state.references...)) +# return $(rhs) +# end + +# #* IF @fuse is in a loop this will be called multiple times for no reason... +# $converted_types = $to_cuda_type.(($(rhs_symbol_state.references...),)) +# $_hash = xor($hash($rhs_ops), $hash($converted_types)) + +# if haskey($FUSED_KERNEL_CACHE, $_hash) +# println("Re-using fused kernel from cache") +# else +# println("Compiling kernel to PTX") +# $ptx = $ptx_as_string($wrapper_name, $converted_types) +# $ptx_f_name = $(cuNumeric.extract_kernel_name)($ptx) +# $(cuNumeric.ptx_task)($ptx, $ptx_f_name) +# $FUSED_KERNEL_CACHE[$_hash] = $FusedKernelData($ptx_f_name, string($wrapper_name), $converted_types) +# end + +# println("Args: $(($(rhs_symbol_state.references...),)), Types: $($converted_types)") +# println($FUSED_KERNEL_CACHE[$_hash]) + +# $task = $(cuNumeric.CUDATask)($FUSED_KERNEL_CACHE[$_hash]) + +# $(cuNumeric.launch)( +# $task, ($(rhs_symbol_state.references...),), $lhs, (); $(kwargs...) +# ) +# end) +# else +# throw(ArgumentError("fuse expected assignment operator `=` or `.=`, got $(call.head)")) +# end + +# return esc(quote +# let +# $code +# end +# end) +# end + +# function to_fuse(x) +# T = eltype(x) +# return exp.(x.*x) .+ T(1.0) +# end + +# x = NDArray{Float32, 1}() +# @fuse y = to_fuse(x) + +# function kernel(x) +# tid = threadIdx().x +# if tid <= length(x) +# x[tid] = x[tid] + 1.0f0 +# end +# end From 0da8ff2a79f332db58cf3c16a17e7707b34541b7 Mon Sep 17 00:00:00 2001 From: ejmeitz Date: Fri, 24 Apr 2026 13:01:22 -0500 Subject: [PATCH 02/12] front-end for expression parsing --- src/cuNumeric.jl | 3 +- src/fusion.jl | 840 ++++++++++++++++++++++++++++++----------------- 2 files changed, 544 insertions(+), 299 deletions(-) diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index e9407207..fa5ed37f 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -152,8 +152,9 @@ include("ndarray/ndarray.jl") include("ndarray/unary.jl") include("ndarray/binary.jl") -# scoping macro +# special features include("scoping.jl") +include("fusion.jl") # Utilities include("utilities/version.jl") diff --git a/src/fusion.jl b/src/fusion.jl index b59ac69f..98022459 100644 --- a/src/fusion.jl +++ b/src/fusion.jl @@ -1,299 +1,574 @@ -export @fuse +export @fuse, @fuse_plan + +""" + ALLOWED_FUSION_OPS + +Set of scalar operations that the fusion frontend is allowed to place in the +fusion IR. Extend this set as the backend learns how to lower more operations. + +The parser normalizes dotted operators/functions, so `.+` and `+` both appear +as `:+` in the IR. Whether an operation was dotted is still preserved on each +`FuseCall` node. +""" +const ALLOWED_FUSION_OPS = Set{Symbol}([ + :+, :-, :*, :/, :^, + :abs, :abs2, + :sqrt, :cbrt, + :exp, :exp2, :expm1, + :log, :log2, :log10, :log1p, + :sin, :cos, :tan, + :asin, :acos, :atan, + :sinh, :cosh, :tanh, + :min, :max, +]) + +const FUSE_OPTION_NAMES = Set{Symbol}([:blocks, :threads]) + +# This cache is intentionally separate from any future function-fusion cache. +# The key is based on the parsed expression IR plus the runtime argument types. +const FUSED_BROADCAST_CACHE = Dict{UInt64,Any}() +const FUSED_BROADCAST_CACHE_LOCK = ReentrantLock() + +# ----------------------------------------------------------------------------- +# Fusion IR +# ----------------------------------------------------------------------------- + +abstract type FuseNode end + +""" + FuseInput(name, index) + +A symbolic input variable appearing in the RHS expression. `index` is the +runtime position of that input in the input tuple passed to `_fuse_broadcast!`. +Repeated uses of the same symbol point to the same index. +""" +struct FuseInput <: FuseNode + name::Symbol + index::Int +end -const ALLOWED_FUSION_TYPES = Union{<:NDArray,<:Number} -const FUSED_KERNEL_CACHE = Dict{UInt64,CUDATask}() -const FUSE_KWARGS = [:blocks, :threads, :outputs] +""" + FuseLiteral(value) -struct FusedKernelData{I,O,S} - ct::CUDATask - inputs::I - outputs::O - scalars::S +A literal scalar that appeared directly in the fused expression, for example +`2` in `z .= 2 .* x`. +""" +struct FuseLiteral <: FuseNode + value::Any end -#! ASSUME OUTPUT_INDICIES ARE THE ARRAYS -#! RETURNED FROM THE FUNCTION AS WELL?? -#! HOW TO HANDLE CASE WHEN output_indicies = () -#! HOW TO HANDLE KWARGS?? -function fuse_function( - fn::Function, - output_indices, - args...; - kwargs..., -) - println("Compiling $(Symbol(fn)) to PTX") +""" + FuseCall(op, args, dotted) + +A scalar operation in the fused expression tree. + +`op` is normalized, so `:.+` becomes `:+`. `dotted` records whether the user +wrote dotted broadcast syntax for this specific call/operator. +""" +struct FuseCall <: FuseNode + op::Symbol + args::Vector{FuseNode} + dotted::Bool +end + +""" + FusePlan(outputs, inputs, ops, rhs, expr_hash) + +A parsed broadcast assignment. + +For example, `@fuse z .= x .+ y .* y` becomes a plan with +`outputs == [:z]`, `inputs == [:x, :y]`, and an RHS tree rooted at `:+`. +""" +struct FusePlan + outputs::Vector{Symbol} + inputs::Vector{Symbol} + ops::Vector{Symbol} + rhs::FuseNode + expr_hash::UInt64 +end + +""" + CompiledFusionKernel + +Placeholder object returned by the default `compile_fused_broadcast` method. +Replace or overload `compile_fused_broadcast` and `launch_fused_broadcast!` when +hooking this parser up to the real CUDA/cuNumeric backend. +""" +struct CompiledFusionKernel + plan::FusePlan + output_types::Tuple + input_types::Tuple +end + +# ----------------------------------------------------------------------------- +# Small display helpers +# ----------------------------------------------------------------------------- + +function Base.show(io::IO, node::FuseInput) + print(io, "FuseInput(:", node.name, ", ", node.index, ")") +end + +function Base.show(io::IO, node::FuseLiteral) + print(io, "FuseLiteral(", repr(node.value), ")") +end + +function Base.show(io::IO, node::FuseCall) + dot = node.dotted ? "." : "" + print(io, "FuseCall(:", dot, node.op, ", ", node.args, ")") +end + +function Base.show(io::IO, plan::FusePlan) + print(io, "FusePlan(outputs = ", plan.outputs, + ", inputs = ", plan.inputs, + ", ops = ", plan.ops, + ", expr_hash = ", plan.expr_hash, + ", rhs = ", plan.rhs, + ")") +end + +# ----------------------------------------------------------------------------- +# Frontend parser +# ----------------------------------------------------------------------------- + +mutable struct ParseContext + input_indices::Dict{Symbol,Int} + inputs::Vector{Symbol} + ops::Vector{Symbol} +end + +ParseContext() = ParseContext(Dict{Symbol,Int}(), Symbol[], Symbol[]) + +function intern_input!(ctx::ParseContext, name::Symbol) + idx = get(ctx.input_indices, name, 0) + if idx == 0 + push!(ctx.inputs, name) + idx = length(ctx.inputs) + ctx.input_indices[name] = idx + end + return idx +end + +""" + build_fuse_plan(ex::Expr) -> FusePlan + +Parse a broadcast assignment expression into a reusable fusion plan. - N_args = length(args) - if length(output_indices) > N_args || max(output_indices) > N_args - error( +Supported MVP syntax: + +```julia +z .= x .+ y .* y +z .= sin.(x) .+ 2 .* y +``` + +This function intentionally does not compile or launch anything. +""" +function build_fuse_plan(ex::Expr) + if Meta.isexpr(ex, :(.=)) + lhs, rhs = ex.args + elseif Meta.isexpr(ex, :(=)) + throw( ArgumentError( - "Marked arguments $(output_indices) as outputs, but there are only $(N_args) arguments" + "@fuse currently supports broadcast assignment `.=` only; got ordinary `=`." + ), + ) + else + throw( + ArgumentError( + "@fuse expected a broadcast assignment like `z .= x .+ y .* y`; got expression head `$(ex.head)`." ), ) end - input_indices = setdiff(1:length(args), Set(output_indices)) - inputs = [args[i] for i in input_indices] - outputs = [args[i] for i in output_indices] + outputs = parse_lhs(lhs) + ctx = ParseContext() + rhs_node = parse_rhs(rhs, ctx) + + # Keep the full operator occurrence order in the plan. Use `unique_ops(plan)` + # if the backend only wants each operator once. + plan_without_hash = FusePlan(outputs, copy(ctx.inputs), copy(ctx.ops), rhs_node, UInt64(0)) + return FusePlan( + outputs, copy(ctx.inputs), copy(ctx.ops), rhs_node, hash(plan_signature(plan_without_hash)) + ) +end - #! ADD KWARGS TO THESE ?? - scalars = filter(x -> isa(x, Number), inputs) - ndarray_inputs = filter(x -> !isa(x, NDArray), inputs) - other_inputs = filter(x -> !isa(x, Number) && !isa(x, NDArray), inputs) +function build_fuse_plan(ex) + throw(ArgumentError("@fuse expected an expression; got `$(repr(ex))`.")) +end - #! Parse out the isbtis structs and unwrap them?? +function parse_lhs(lhs) + if lhs isa Symbol + return Symbol[lhs] + elseif Meta.isexpr(lhs, :tuple) + throw( + ArgumentError( + "Multiple-output broadcast fusion is not implemented yet. Use a single output, e.g. `z .= rhs`." + ), + ) + elseif Meta.isexpr(lhs, :ref) + throw( + ArgumentError( + "Indexed/view outputs are not implemented yet. Use a whole-array output, e.g. `z .= rhs`." + ), + ) + else + throw(ArgumentError("Unsupported @fuse output expression: `$(lhs)`.")) + end +end - if any(isa.(outputs, Number)) - @error "Scalar outputs are not allowed, use 0D store." +function parse_rhs(ex, ctx::ParseContext)::FuseNode + if ex isa Symbol + idx = intern_input!(ctx, ex) + return FuseInput(ex, idx) + elseif ex isa Number + return FuseLiteral(ex) + elseif ex isa Bool + return FuseLiteral(ex) + elseif ex isa Char + return FuseLiteral(ex) + elseif ex isa String + throw(ArgumentError("String literals are not valid scalar fusion literals: `$(ex)`.")) + elseif Meta.isexpr(ex, :call) + return parse_call(ex, ctx) + elseif Meta.isexpr(ex, :.) + return parse_dotted_call(ex, ctx) + elseif Meta.isexpr(ex, :ref) + throw(ArgumentError("Indexing inside fused RHS is not implemented yet: `$(ex)`.")) + elseif Meta.isexpr(ex, :tuple) + throw(ArgumentError("Tuple construction inside fused RHS is not implemented yet: `$(ex)`.")) + elseif Meta.isexpr(ex, :block) + throw( + ArgumentError( + "Block expressions are not supported in broadcast-fusion RHS expressions." + ), + ) + else + throw( + ArgumentError("Unsupported expression in fused RHS: `$(ex)` with type `$(typeof(ex))`.") + ) end +end - # Dummy types so we can use CUDA.jl to compile - converted_types = ndarray_cuda_type.(args) # enforces everything is an `isbitstype` - ct = CUDATask(fn, converted_types) +function parse_call(ex::Expr, ctx::ParseContext)::FuseNode + raw_op = ex.args[1] + op, dotted = normalize_call_operator(raw_op) + check_allowed_op!(op, raw_op) - return FusedKernelData(ct, other_inputs, outputs, scalars) + push!(ctx.ops, op) + args = FuseNode[parse_rhs(arg, ctx) for arg in ex.args[2:end]] + return FuseCall(op, args, dotted) end -# For @fuse on user defined function -function maybe_fuse_kernel(fn::Function, output_indices, args::T; kwargs...) where {T} - cache_key = hash((fn, T)) #! ADD KWRAGS TO THIS? - - if !haskey(FUSED_KERNEL_CACHE, cache_key) - FUSED_KERNEL_CACHE[cache_key] = fuse_function( - fn, - output_indices, - args...; - kwargs..., - ) +function parse_dotted_call(ex::Expr, ctx::ParseContext)::FuseNode + # Dotted function calls like `sin.(x)` usually parse as + # Expr(:., :sin, Expr(:tuple, :x)). Property access also uses head `:.`, so + # require the second argument to be an argument tuple. + if length(ex.args) != 2 || !(ex.args[2] isa Expr) || ex.args[2].head != :tuple + throw(ArgumentError("Property access is not supported in fused RHS expressions: `$(ex)`.")) end - return cache_key -end + op = function_name_symbol(ex.args[1]) + check_allowed_op!(op, ex.args[1]) -function run_fused_kernel(cache_key; blocks, threads) - fkd = FUSED_KERNEL_CACHE[cache_key] - cuNumeric.launch(fkd.ct, fkd.inputs, fkd.outputs, fkd.scalars; blocks=blocks, threads=threads) + push!(ctx.ops, op) + args = FuseNode[parse_rhs(arg, ctx) for arg in ex.args[2].args] + return FuseCall(op, args, true) end -macro fuse(ex...) - call = ex[end] - kwargs = map(ex[1:(end - 1)]) do kwarg - if kwarg in FUSE_KWARGS - :($kwarg = $kwarg) +function normalize_call_operator(raw_op) + if raw_op isa Symbol + op_string = String(raw_op) + if startswith(op_string, ".") + return Symbol(op_string[2:end]), true else - throw(ArgumentError("Invalid keyword argument '$kwarg', expected one of $(FUSE_KWARGS)")) + return raw_op, false end end - #! TODO SET DEFAULT BLOCKS/THREADS - #! HOW TO KNOW WHATS THE RIGHT DIMENSION?? - - blocks = get(kwargs[:blocks], DEFAULT_BLOCKS) - threads = get(kwargs[:threads], DEFAULT_THREADS) - output_indices = get(kwargs[:output_indices], ()) - - code = quote end - - if Meta.isexpr(call, :function) - # Parses function epxression into args, name etc. - data = splitdef(longdef(call)) - # splitarg parses args into (name, type, is_slurp, default) - arg_data = map(splitarg, data[:args]) - kwarg_data = map(splitarg, data[:kwargs]) - - arg_names = getindex.(arg_data, 1) - - @gensym wrapper_name, cache_key - - push!(code.args, - quote - - # The fused function - function $(wrapper_name)() where {$(dict[:whereparams]...)} - #! NOT SURE THIS IS RIGHT WAY TO INTERPOLATE ARGS - #! WANT IT INTERPOLATE AS TUPLE OF ARGS - #! TODO PASS KWARGS - $cache_key = maybe_fuse_kernel($(data[:name]), $output_indices, $(arg_names...)) - run_fused_kernel($cache_key; blocks=($blocks), threads=($threads)) - end - - # Replace original function with call to fused - function $(data[:name])( - $(data[:args]...); $(data[:kwargs]...) - ) where {$(dict[:whereparams]...)} - $(wrapper_name)() - end - end, - ) + # Qualified calls such as Base.sin(x). Dotted qualified calls such as + # Base.sin.(x) are handled by `parse_dotted_call`. + return function_name_symbol(raw_op), false +end + +function function_name_symbol(ex) + if ex isa Symbol + return ex + elseif ex isa QuoteNode && ex.value isa Symbol + return ex.value + elseif Meta.isexpr(ex, :.) + # Handles qualified function names. We intentionally drop the module + # qualification here because the backend's allow-list is scalar-op based. + return function_name_symbol(ex.args[end]) else - throw(ArgumentError("fuse expected Broadcasted object or function, got $(call.head)")) + throw(ArgumentError("Could not determine function/operator name from `$(ex)`.")) + end +end + +function check_allowed_op!(op::Symbol, raw_op) + if !(op in ALLOWED_FUSION_OPS) + throw( + ArgumentError( + "Operation `$(raw_op)` normalized to `$(op)`, which is not in ALLOWED_FUSION_OPS." + ), + ) end + return nothing +end + +# ----------------------------------------------------------------------------- +# Plan signatures and cache keys +# ----------------------------------------------------------------------------- - return esc(quote - let - $code +node_signature(node::FuseInput) = (:input, node.name, node.index) +node_signature(node::FuseLiteral) = (:literal, typeof(node.value), node.value) +node_signature(node::FuseCall) = (:call, node.op, node.dotted, Tuple(node_signature.(node.args))) + +function plan_signature(plan::FusePlan) + return (:broadcast_assignment, + Tuple(plan.outputs), + Tuple(plan.inputs), + Tuple(plan.ops), + node_signature(plan.rhs)) +end + +""" + unique_ops(plan::FusePlan) + +Return the operations appearing in the RHS in first-occurrence order. +""" +function unique_ops(plan::FusePlan) + seen = Set{Symbol}() + out = Symbol[] + for op in plan.ops + if !(op in seen) + push!(seen, op) + push!(out, op) end - end) + end + return out end -# Collects operators and function calls from an expression. -# # E.g. x .+ y .* z will return (.+, .*) -# function collect_ops(expr::Expr) -# ops = Symbol[] - -# postwalk(expr) do node -# if node isa Expr && (node.head === :call || node.head === :.) -# op = node.args[1] -# if op isa Symbol -# push!(ops, op) -# # strip one leading '.' so :.+ → :+, but keep names like :sqrt -# # op_sym = startswith(op, ".") ? Symbol(string(op)[2:end]) : op -# # push!(ops, op_sym) -# end -# end -# return node -# end +""" + fusion_cache_key(plan, outputs, inputs) -> UInt64 -# return Tuple(ops) -# end +Cache key for the compiled kernel. Launch options such as `blocks` and +`threads` are intentionally excluded because they should not affect codegen. +""" +function fusion_cache_key(plan::FusePlan, outputs::Tuple, inputs::Tuple) + return hash((plan.expr_hash, map(typeof, outputs), map(typeof, inputs))) +end -# This version requires all args are the same type -# @generated function maybe_fuse_kernel( -# :: NTuple{NARGS, NDArray{T,DIM}}, -# rhs_ops :: NTuple{NOPS, Symbol}, -# wrapper :: Function -# ) where {NARGS,T, DIM, NOPS} - -# # converted_types is constructed purely from the types of the arguments -# # so a generated function will allow us to pre-compute most of this function -# # at compile time! -# converted_types = ntuple(_ -> CuDeviceArray{T, DIM, 1}, Val(NARGS)) -# _hash = xor(hash(rhs_ops), hash(converted_types)) - -# quote -# if !haskey($FUSED_KERNEL_CACHE, $_hash) -# println("Compiling kernel to PTX") -# local ptx = ptx_as_string(wrapper, $converted_types) #! DO NOT interpolate wrapper -# f = open("/pool/emeitz/ptx2.txt", "w") -# println(f, ptx) -# close(f) -# local ptx_f_name = cuNumeric.extract_kernel_name(ptx) -# cuNumeric.ptx_task(ptx, ptx_f_name) -# $FUSED_KERNEL_CACHE[$_hash] = FusedKernelData(ptx_f_name, $converted_types) -# end -# $_hash -# end -# end +# ----------------------------------------------------------------------------- +# Runtime argument classification hooks +# ----------------------------------------------------------------------------- -# # This version does not assume args have same types -# # Manually enforces that all args are NDArray -# @generated function maybe_fuse_kernel(rhs_args :: R, -# rhs_ops :: NTuple{NOPS,Symbol}, -# wrapper :: Function) where -# {R<:Tuple, NOPS} +is_fusion_scalar(::Number) = true +is_fusion_scalar(_) = false +is_fusion_array(::AbstractArray) = true +is_fusion_array(_) = false +is_fusion_arg(x) = is_fusion_array(x) || is_fusion_scalar(x) -# arg_types = R.parameters # e.g. (NDArray{Float32,2}, NDArray{Int,1}, …) -# println("IN HERE") +function validate_fuse_args(plan::FusePlan, outputs::Tuple, inputs::Tuple) + if length(outputs) != length(plan.outputs) + throw( + ArgumentError("Plan expects $(length(plan.outputs)) output(s), got $(length(outputs)).") + ) + end + if length(inputs) != length(plan.inputs) + throw(ArgumentError("Plan expects $(length(plan.inputs)) input(s), got $(length(inputs)).")) + end -# for T in arg_types -# T <: NDArray || throw(ArgumentError("all arguments must be NDArray, got $T")) -# end + for (name, output) in zip(plan.outputs, outputs) + if is_fusion_scalar(output) + throw( + ArgumentError( + "Output `$(name)` is scalar-like. Scalar outputs are not supported; use a 0D store/array." + ), + ) + elseif !is_fusion_array(output) + throw(ArgumentError("Output `$(name)` has unsupported type `$(typeof(output))`.")) + end + end -# # 2. Build the matching CuDeviceArray type for each element ------------- -# conv_type_exprs = Expr[] -# for T in arg_types -# elty, dim = T.parameters[1:2] # the {T, N} in NDArray{T,N} -# push!(conv_type_exprs, :(CuDeviceArray{$elty, $dim, 1})) -# end -# # an Expr(:tuple, …) is a compile‑time literal of the tuple of types -# converted_types_expr = Expr(:tuple, conv_type_exprs...) - -# const_hash = xor(hash(rhs_ops), hash(arg_types)) - -# quote -# _hash = $const_hash # value is a literal, not recomputed -# if !haskey($FUSED_KERNEL_CACHE, _hash) -# println("Compiling kernel to PTX") -# local ptx = ptx_as_string(wrapper, $converted_types_expr) -# local ptx_f_name = cuNumeric.extract_kernel_name(ptx) -# cuNumeric.ptx_task(ptx, ptx_f_name) #! UNCOMMENT LATER -# $FUSED_KERNEL_CACHE[_hash] = -# FusedKernelData(ptx_f_name, $converted_types_expr) -# end -# _hash -# end -# end + for (name, input) in zip(plan.inputs, inputs) + if !is_fusion_arg(input) + throw( + ArgumentError( + "Input `$(name)` has unsupported type `$(typeof(input))`; expected array-like or scalar-like value." + ), + ) + end + end -# macro fuse(ex...) + return nothing +end -# call = ex[end] -# kwargs = map(ex[1:end-1]) do kwarg -# if kwarg in FUSE_KWARGS -# :($kwarg = $kwarg) -# else -# throw(ArgumentError("Invalid keyword argument '$kwarg', expected one of $(FUSE_KWARGS)")) -# end -# end +# ----------------------------------------------------------------------------- +# Runtime compilation / launch path +# ----------------------------------------------------------------------------- + +""" + _fuse_broadcast!(plan, outputs, inputs; blocks=nothing, threads=nothing) + +Runtime entrypoint emitted by `@fuse`. This validates runtime arguments, +computes a cache key, compiles if necessary, and launches the backend. +""" +function _fuse_broadcast!( + plan::FusePlan, + outputs::Tuple, + inputs::Tuple; + blocks=nothing, + threads=nothing, +) + validate_fuse_args(plan, outputs, inputs) + key = fusion_cache_key(plan, outputs, inputs) + + local kernel + lock(FUSED_BROADCAST_CACHE_LOCK) + try + kernel = get!(FUSED_BROADCAST_CACHE, key) do + compile_fused_broadcast(plan, outputs, inputs) + end + finally + unlock(FUSED_BROADCAST_CACHE_LOCK) + end -# code = quote end + return launch_fused_broadcast!(kernel, outputs, inputs; blocks=blocks, threads=threads) +end -# if Meta.isexpr(call, :function) -# #TODO use symbol_state.func_defs to fuse each definition -# throw(ArgumentError("@fuse before function definitions is not supported yet.")) -# elseif Meta.isexpr(call, :(=)) || Meta.isexpr(call, :(.=)) -# lhs, rhs = call.args +""" + compile_fused_broadcast(plan, outputs, inputs) -# lhs_symbol_state = compute_symbols_state(lhs) -# rhs_symbol_state = compute_symbols_state(rhs) +Backend hook. The default stores type information only. Replace this with the +real CUDATask/PTX generation path. +""" +function compile_fused_broadcast(plan::FusePlan, outputs::Tuple, inputs::Tuple) + return CompiledFusionKernel(plan, map(typeof, outputs), map(typeof, inputs)) +end -# if !isempty(intersect(lhs_symbol_state.references, rhs_symbol_state.references)) -# throw(ArgumentError("LHS and RHS of @fuse must not share variables.")) -# end +""" + launch_fused_broadcast!(kernel, outputs, inputs; blocks, threads) + +Backend hook. The default throws because this file only implements parsing, +validation, and caching. +""" +function launch_fused_broadcast!( + kernel::CompiledFusionKernel, + outputs::Tuple, + inputs::Tuple; + blocks=nothing, + threads=nothing, +) + throw( + ErrorException( + "No fusion backend has been installed. The expression parsed and cached successfully, " * + "but `launch_fused_broadcast!` must be overloaded to call the CUDA/cuNumeric backend.", + ), + ) +end -# rhs_ops = collect_ops(rhs) +# ----------------------------------------------------------------------------- +# Macros +# ----------------------------------------------------------------------------- + +function parse_fuse_macro_args(args) + isempty(args) && throw(ArgumentError("@fuse expected `z .= rhs`.")) + + call = args[end] + option_exprs = Expr[] + + for opt in args[1:(end - 1)] + if opt isa Symbol && opt in FUSE_OPTION_NAMES + # Supports the old style `@fuse blocks threads z .= rhs`, meaning + # `blocks = blocks, threads = threads`. + push!(option_exprs, Expr(:kw, opt, opt)) + elseif Meta.isexpr(opt, :(=)) || Meta.isexpr(opt, :kw) + name = opt.args[1] + value = opt.args[2] + if !(name isa Symbol) || !(name in FUSE_OPTION_NAMES) + throw( + ArgumentError( + "Invalid @fuse option `$(name)`. Expected one of $(collect(FUSE_OPTION_NAMES))." + ), + ) + end + push!(option_exprs, Expr(:kw, name, value)) + else + throw( + ArgumentError( + "Invalid @fuse option/expression `$(opt)`. Expected options followed by `z .= rhs`." + ), + ) + end + end -# @gensym wrapper_name ptx ptx_f_name task _hash converted_types args_tuple + return option_exprs, call +end -# push!(code.args, -# quote +""" + @fuse z .= rhs + @fuse blocks=blocks threads=threads z .= rhs -# # Create wrapper function that we can fuse -# function $(wrapper_name)($(rhs_symbol_state.references...)) -# # println("Wrapper called!") -# return $(rhs) -# end +Parse a broadcast assignment into a `FusePlan`, validate the runtime arguments, +compile/cache the corresponding kernel, and dispatch to `launch_fused_broadcast!`. +""" +macro fuse(args...) + option_exprs, call = parse_fuse_macro_args(args) + plan = build_fuse_plan(call) -# $_hash = $maybe_fuse_kernel(($(rhs_symbol_state.references...),), $rhs_ops, $wrapper_name) + output_tuple = Expr(:tuple, plan.outputs...) + input_tuple = Expr(:tuple, plan.inputs...) + runtime = GlobalRef(@__MODULE__, :_fuse_broadcast!) -# println($FUSED_KERNEL_CACHE[$_hash]) + if isempty(option_exprs) + return esc(Expr(:call, runtime, QuoteNode(plan), output_tuple, input_tuple)) + else + return esc( + Expr( + :call, + runtime, + Expr(:parameters, option_exprs...), + QuoteNode(plan), + output_tuple, + input_tuple, + ), + ) + end +end -# $task = $(cuNumeric.CUDATask)($FUSED_KERNEL_CACHE[$_hash]) +""" + @fuse_plan z .= rhs -# $(cuNumeric.launch)( -# $task, ($(rhs_symbol_state.references...),), $lhs, (); $(kwargs...) -# ) -# end) -# else -# throw(ArgumentError("fuse expected assignment operator `=` or `.=`, got $(call.head)")) +Parse only. Useful while developing the frontend. +""" +macro fuse_plan(call) + plan = build_fuse_plan(call) + return QuoteNode(plan) +end + +# For @fuse on user defined function +# function maybe_fuse_kernel(fn::Function, output_indices, args::T; kwargs...) where {T} +# cache_key = hash((fn, T)) #! ADD KWRAGS TO THIS? + +# if !haskey(FUSED_KERNEL_CACHE, cache_key) +# FUSED_KERNEL_CACHE[cache_key] = fuse_function( +# fn, +# output_indices, +# args...; +# kwargs..., +# ) # end -# return esc(quote -# let -# $code -# end -# end) +# return cache_key # end -#* BE SURE TO TEST: -# - x .+ y .*z and x .* y .+ z -# x,y = ... -# y = to_fuse(x) -# z = unary.(x) .+ y -# z = binary.(x,y) .+ z -# macro fuse(ex...) +# function run_fused_kernel(cache_key; blocks, threads) +# fkd = FUSED_KERNEL_CACHE[cache_key] +# cuNumeric.launch(fkd.ct, fkd.inputs, fkd.outputs, fkd.scalars; blocks=blocks, threads=threads) +# end +# macro fuse(ex...) # call = ex[end] -# kwargs = map(ex[1:end-1]) do kwarg +# kwargs = map(ex[1:(end - 1)]) do kwarg # if kwarg in FUSE_KWARGS # :($kwarg = $kwarg) # else @@ -301,65 +576,49 @@ end # end # end -# code = quote end +# #! TODO SET DEFAULT BLOCKS/THREADS +# #! HOW TO KNOW WHATS THE RIGHT DIMENSION?? -# if Meta.isexpr(call, :function) -# #TODO use symbol_state.func_defs to fuse each definition -# throw(ArgumentError("@fuse before function definitions is not supported yet.")) -# elseif Meta.isexpr(call, :(=)) || Meta.isexpr(call, :(.=)) -# lhs, rhs = call.args +# blocks = get(kwargs[:blocks], DEFAULT_BLOCKS) +# threads = get(kwargs[:threads], DEFAULT_THREADS) +# output_indices = get(kwargs[:output_indices], ()) -# lhs_symbol_state = compute_symbols_state(lhs) -# rhs_symbol_state = compute_symbols_state(rhs) - -# if !isempty(intersect(lhs_symbol_state.references, rhs_symbol_state.references)) -# throw(ArgumentError("LHS and RHS of @fuse must not share variables.")) -# end +# code = quote end -# rhs_ops = collect_ops(rhs) +# if Meta.isexpr(call, :function) +# # Parses function epxression into args, name etc. +# data = splitdef(longdef(call)) +# # splitarg parses args into (name, type, is_slurp, default) +# arg_data = map(splitarg, data[:args]) +# kwarg_data = map(splitarg, data[:kwargs]) -# # This will error if all variables are not the same type -# # E.g. NDArrays with different dimensions or el-types -# rhs_args = ntuple(i -> rhs_symbol_state.references[i], length(rhs_symbol_state.references)) +# arg_names = getindex.(arg_data, 1) -# @gensym wrapper_name ptx ptx_f_name task _hash converted_types args_tuple +# @gensym wrapper_name, cache_key # push!(code.args, # quote -# # Create wrapper function that we can fuse -# # $(wrapper_name) = ($(rhs_symbol_state.references...)) -> $(rhs) - -# #* IF @fuse is in a loop this will be called multiple times for no reason... -# function $(wrapper_name)($(rhs_symbol_state.references...)) -# return $(rhs) +# # The fused function +# function $(wrapper_name)() where {$(dict[:whereparams]...)} +# #! NOT SURE THIS IS RIGHT WAY TO INTERPOLATE ARGS +# #! WANT IT INTERPOLATE AS TUPLE OF ARGS +# #! TODO PASS KWARGS +# $cache_key = maybe_fuse_kernel($(data[:name]), $output_indices, $(arg_names...)) +# run_fused_kernel($cache_key; blocks=($blocks), threads=($threads)) # end -# #* IF @fuse is in a loop this will be called multiple times for no reason... -# $converted_types = $to_cuda_type.(($(rhs_symbol_state.references...),)) -# $_hash = xor($hash($rhs_ops), $hash($converted_types)) - -# if haskey($FUSED_KERNEL_CACHE, $_hash) -# println("Re-using fused kernel from cache") -# else -# println("Compiling kernel to PTX") -# $ptx = $ptx_as_string($wrapper_name, $converted_types) -# $ptx_f_name = $(cuNumeric.extract_kernel_name)($ptx) -# $(cuNumeric.ptx_task)($ptx, $ptx_f_name) -# $FUSED_KERNEL_CACHE[$_hash] = $FusedKernelData($ptx_f_name, string($wrapper_name), $converted_types) +# # Replace original function with call to fused +# #! NOT SURE OVERWRITING LIKE THIS IS GREAT +# function $(data[:name])( +# $(data[:args]...); $(data[:kwargs]...) +# ) where {$(dict[:whereparams]...)} +# $(wrapper_name)() # end - -# println("Args: $(($(rhs_symbol_state.references...),)), Types: $($converted_types)") -# println($FUSED_KERNEL_CACHE[$_hash]) - -# $task = $(cuNumeric.CUDATask)($FUSED_KERNEL_CACHE[$_hash]) - -# $(cuNumeric.launch)( -# $task, ($(rhs_symbol_state.references...),), $lhs, (); $(kwargs...) -# ) -# end) +# end +# ) # else -# throw(ArgumentError("fuse expected assignment operator `=` or `.=`, got $(call.head)")) +# throw(ArgumentError("fuse expected assignment expression or function, got $(call.head)")) # end # return esc(quote @@ -368,18 +627,3 @@ end # end # end) # end - -# function to_fuse(x) -# T = eltype(x) -# return exp.(x.*x) .+ T(1.0) -# end - -# x = NDArray{Float32, 1}() -# @fuse y = to_fuse(x) - -# function kernel(x) -# tid = threadIdx().x -# if tid <= length(x) -# x[tid] = x[tid] + 1.0f0 -# end -# end From eeeb752264684a5cbb91064dfdbb853b611096a5 Mon Sep 17 00:00:00 2001 From: ejmeitz Date: Sat, 25 Apr 2026 09:50:01 -0500 Subject: [PATCH 03/12] remove CUDA ext --- Project.toml | 11 +- ext/CUDAExt/CUDAExt.jl | 26 ---- src/cuNumeric.jl | 13 +- .../cuda.jl => src/cuda/cuda_ptx_task.jl | 112 ++++++++++++++---- src/cuda/cuda_util.jl | 19 +++ src/fusion/backends.jl | 29 +++++ src/{fusion.jl => fusion/parsing.jl} | 2 +- src/fusion/ptx_backend.jl | 7 ++ src/utilities/cuda_stubs.jl | 92 -------------- 9 files changed, 160 insertions(+), 151 deletions(-) delete mode 100644 ext/CUDAExt/CUDAExt.jl rename ext/CUDAExt/cuda.jl => src/cuda/cuda_ptx_task.jl (58%) create mode 100644 src/cuda/cuda_util.jl create mode 100644 src/fusion/backends.jl rename src/{fusion.jl => fusion/parsing.jl} (99%) create mode 100644 src/fusion/ptx_backend.jl delete mode 100644 src/utilities/cuda_stubs.jl diff --git a/Project.toml b/Project.toml index 59bc50ba..016ce917 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,8 @@ version = "0.1.0" [deps] CNPreferences = "3e078157-ea10-49d5-bf32-908f777cd46f" +CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" +CUDATools = "9ec180c6-1c07-47c7-9e6e-ebefa4d1f6d0" CxxWrap = "1f15a43c-97ca-5a2a-ae31-89f07a497df4" JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899" Legate = "1238f2cf-6593-4d60-9aca-2f5364e49909" @@ -20,15 +22,10 @@ cunumeric_jl_wrapper_jll = "49048992-29d2-5fd1-994f-9cecf112d624" cupynumeric_jll = "2862d674-414d-5b0b-a494-b21f8deca547" libcxxwrap_julia_jll = "3eaa8342-bff7-56a5-9981-c04077f7cee7" -[weakdeps] -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" - -[extensions] -CUDAExt = "CUDA" - [compat] CNPreferences = "0.1.2" -CUDA = "5.9" +CUDACore = "6.0.0" +CUDATools = "6.0.0" CxxWrap = "0.17" JuliaFormatter = "2.3.0" Legate = "0.1.0" diff --git a/ext/CUDAExt/CUDAExt.jl b/ext/CUDAExt/CUDAExt.jl deleted file mode 100644 index e991aac6..00000000 --- a/ext/CUDAExt/CUDAExt.jl +++ /dev/null @@ -1,26 +0,0 @@ -module CUDAExt - -using Random -using CUDA -using Legate: Legate -using CxxWrap: CxxWrap -using cuNumeric: cuNumeric -import cuNumeric: - @cuda_task, @launch, NDArray - -const KERNEL_OFFSET = sizeof(CUDA.KernelState) - -include("cuda.jl") - -function __init__() - if CUDA.functional() - # in cuda.jl to notify /wrapper/src/cuda.cpp about CUDA.jl kernel state size - cuNumeric.register_kernel_state_size(UInt64(KERNEL_OFFSET)) - # in /wrapper/src/cuda.cpp - cuNumeric.register_tasks(); - else - @warn "CUDA.jl is not functional; skipping CUDA kernel registration." - end -end - -end # module CUDAExt diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index fa5ed37f..36341b89 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -27,6 +27,10 @@ using Legate using Libdl using CxxWrap +using CUDATools: CUDATools +using CUDACore: CUDACore +import CUDACore: CuArray + using cupynumeric_jll using cunumeric_jl_wrapper_jll @@ -154,11 +158,12 @@ include("ndarray/binary.jl") # special features include("scoping.jl") -include("fusion.jl") +include("cuda/cuda_util.jl") +include("cuda/cuda_ptx_task.jl") +include("fusion/parsing.jl") # Utilities include("utilities/version.jl") -include("utilities/cuda_stubs.jl") include("util.jl") # From https://github.com/JuliaGraphics/QML.jl/blob/dca239404135d85fe5d4afe34ed3dc5f61736c63/src/QML.jl#L147 @@ -237,7 +242,11 @@ function __init__() _is_precompiling() && return nothing + # Start runtime, but only if not pre-compiling ensure_runtime!() + + # Requries runtime to be started + _setup_cuda_tasking() end end #module cuNumeric diff --git a/ext/CUDAExt/cuda.jl b/src/cuda/cuda_ptx_task.jl similarity index 58% rename from ext/CUDAExt/cuda.jl rename to src/cuda/cuda_ptx_task.jl index a70a5b90..327c5b3d 100644 --- a/ext/CUDAExt/cuda.jl +++ b/src/cuda/cuda_ptx_task.jl @@ -1,19 +1,12 @@ -function ndarray_cuda_type(A::NDArray{T,N}) where {T,N} - if N == 1 - CuDeviceVector{T,1} - elseif N == 2 - CuDeviceMatrix{T,1} - else - CuDeviceArray{T,N,1} - end -end +export @cuda_task, @launch, CUDATask -function ndarray_cuda_type(arg::T) where {T} - Base.isbits(arg) || throw(ArgumentError("Unsupported argument type: $(typeof(arg))")) - typeof(arg) +struct CUDATask + func::String + argtypes::NTuple{N,Type} where {N} end -cuNumeric.map_ndarray_cuda_types(args...) = tuple(map(ndarray_cuda_type, args)...) +#! JUST PASS TYPES HERE INSTEAD OF CALLING typeof() +map_ndarray_cuda_types(args...) = tuple(map(ndarray_cuda_type, typeof.(args))...) function to_stdvec(::Type{T}, vec) where {T} stdvec = CxxWrap.StdVector{T}() @@ -53,12 +46,12 @@ function add_padding(arr::NDArray, i::Int64; copy=false) end function check_sz!(arr, maxshape; copy=false) - sz = cuNumeric.size(arr) + sz = size(arr) if maxshape != nothing # currently require all ndarray inputs to be equal alligned_equal_size = sz == maxshape if !alligned_equal_size - cuNumeric.add_padding(arr, maxshape; copy=copy) + add_padding(arr, maxshape; copy=copy) new_size = padded_shape(arr) @warn "[Padding Added] $sz output is now $new_size" end @@ -74,7 +67,7 @@ function check_sz(arr, maxshape) end end -function Launch(kernel::cuNumeric.CUDATask, inputs::Tuple{Vararg{NDArray}}, +function Launch(kernel::CUDATask, inputs::Tuple{Vararg{NDArray}}, outputs::Tuple{Vararg{NDArray}}, scalars::Tuple{Vararg{Any}}; blocks, threads) # we find the largest input/output. @@ -120,7 +113,7 @@ function Launch(kernel::cuNumeric.CUDATask, inputs::Tuple{Vararg{NDArray}}, Legate.submit_auto_task(rt, task) end -function cuNumeric.launch(kernel::cuNumeric.CUDATask, inputs, outputs, scalars; blocks, threads) +function launch(kernel::CUDATask, inputs, outputs, scalars; blocks, threads) Launch(kernel, isa(inputs, Tuple) ? inputs : (inputs,), isa(outputs, Tuple) ? outputs : (outputs,), @@ -130,7 +123,7 @@ function cuNumeric.launch(kernel::cuNumeric.CUDATask, inputs, outputs, scalars; ) end -function cuNumeric.ptx_task(ptx::String, kernel_name) +function ptx_task(ptx::String, kernel_name) rt = Legate.get_runtime() lib = cuNumeric.get_lib() # grab lib of legate app # this taskid is directly tied to cpp code in our setup @@ -142,27 +135,100 @@ function cuNumeric.ptx_task(ptx::String, kernel_name) Legate.submit_auto_task(rt, task) end +""" + @cuda_task(f(args...)) + +Compile a Julia GPU kernel to PTX, register it with the Legate runtime, +and return a `CUDATask` object for later launch. + +# Arguments +- `f` — The name of the Julia CUDA.jl GPU kernel function to compile. +- `args...` — Example arguments to the kernel, used to determine the + argument type signature when generating PTX. + +# Description +This macro automates the process of: +1. Inferring the CUDA argument types for the given `args` using + `map_ndarray_cuda_types`. +2. Using `CUDA.code_ptx` to compile the specified GPU kernel + (`f`) into raw PTX text for the inferred types. +3. Extracting the kernel's function symbol name from the PTX using + `extract_kernel_name`. +4. Registering the compiled PTX and kernel name with the Legate runtime + via `ptx_task`, making it available for GPU execution. +5. Returning a `CUDATask` struct that stores the kernel name and type signature, + which can be used to configure and launch the kernel later. + +# Notes +- The `args...` are not executed; they are used solely for type inference. +- This macro is intended for use with the Legate runtime and + assumes a CUDA context is available. +- Make sure your kernel code is GPU-compatible and does not rely on + unsupported Julia features. + +# Example +```julia +mytask = @cuda_task my_kernel(A, B, C) +``` +""" macro cuda_task(call_expr) fname = call_expr.args[1] fargs = call_expr.args[2:end] esc(quote local _buf = IOBuffer() - local _types = cuNumeric.map_ndarray_cuda_types($(fargs...)) + local _types = map_ndarray_cuda_types($(fargs...)) # generate ptx using CUDA.jl - CUDA.code_ptx(_buf, $fname, _types; raw=false, kernel=true) + CUDATools.code_ptx(_buf, $fname, _types; raw=false, kernel=true) local _ptx = String(take!(_buf)) - local _func_name = cuNumeric.extract_kernel_name(_ptx) + local _func_name = extract_kernel_name(_ptx) # issue ptx_task within legate runtime to register cufunction ptr with cucontext - cuNumeric.ptx_task(_ptx, _func_name) + ptx_task(_ptx, _func_name) # create a cuNumeric.CUDAtask that stores some info for a launch config - cuNumeric.CUDATask(_func_name, _types) + CUDATask(_func_name, _types) end) end +""" + @launch(; task, blocks=(1,), threads=(256,), inputs=(), outputs=(), scalars=()) + +Launch a GPU kernel (previously registered via [`@cuda_task`](@ref)) through the Legate runtime. + +# Keywords +- `task` — A `CUDATask` object, typically returned by [`@cuda_task`](@ref). +- `blocks` — Tuple or single element specifying the CUDA grid dimensions. Defaults to `(1,)`. +- `threads` — Tuple or single element specifying the CUDA block dimensions. Defaults to `(256,)`. +- `inputs` — Tuple or single element of input NDArray objects. +- `outputs` — Tuple or single element of output NDArray objects. +- `scalars` — Tuple or single element of scalar values. + +# Description +The `@launch` macro validates the provided keywords, ensuring only +the allowed set (`:task`, `:blocks`, `:threads`, `:inputs`, `:outputs`, `:scalars`) +are present. It then expands to a call to `cuNumeric.launch`, +passing the given arguments to the Legate runtime for execution. + +This macro is meant to provide a concise, declarative syntax for +launching GPU kernels, separating kernel compilation (via `@cuda_task`) +from execution configuration. + +# Notes +- `task` **must** be a kernel registered with the runtime, usually from `@cuda_task`. +- All keyword arguments must be specified as assignments, e.g. `blocks=(2,2)` not positional arguments. +- Defaults are chosen for single-block, 256-thread 1D launches. +- The macro escapes its body so that the values of inputs/outputs/scalars are captured + from the surrounding scope at macro expansion time. + +# Example +```julia +mytask = @cuda_task my_kernel(A, B, C) + +@launch task=mytask blocks=(8,8) threads=(32,32) inputs=(A, B) outputs=(C) +``` +""" macro launch(args...) allowed_keys = Set([:task, :blocks, :threads, :inputs, :outputs, :scalars]) kwargs = Dict{Symbol,Any}() diff --git a/src/cuda/cuda_util.jl b/src/cuda/cuda_util.jl new file mode 100644 index 00000000..12384baa --- /dev/null +++ b/src/cuda/cuda_util.jl @@ -0,0 +1,19 @@ +const KERNEL_OFFSET = sizeof(CUDACore.KernelState) + +function _setup_cuda_tasking() + if CUDACore.functional() + # in cuda.jl to notify /wrapper/src/cuda.cpp about CUDA.jl kernel state size + register_kernel_state_size(UInt64(KERNEL_OFFSET)) + # in /wrapper/src/cuda.cpp + register_tasks() + else + @warn "CUDA.jl is not functional; skipping CUDA kernel registration." + end +end + +ndarray_cuda_type(::Type{<:NDArray{T,N}}) where {T,N} = CuDeviceArray{T,N,1} + +function ndarray_cuda_type(::Type{T}) where {T} + Base.isbitstype(T) || throw(ArgumentError("Unsupported argument type: $(T)")) + T +end diff --git a/src/fusion/backends.jl b/src/fusion/backends.jl new file mode 100644 index 00000000..c261bb1d --- /dev/null +++ b/src/fusion/backends.jl @@ -0,0 +1,29 @@ +abstract type FusionBackend end + +""" + PTXBackend + +The default backend for broadcasted expressions. This backend + leverages CUDA.jl to generate PTX code which is executed by the Legate runtime. +""" +struct PTXBackend <: FusionBackend end + +# Backend for expressions that return `Broadcasted` objects. +# For example, `@fuse z .= x .+ y` uses this backend. +const DEFAULT_BROADCASTED_FUSION_BACKEND = PTXBackend +const SUPPORTED_BROADCASTED_FUSION_BACKENDS = (PTXBackend,) + +#! Call this on init or in global scope +function set_broadcasted_fusion_backend() + backend = load_preference( + CNPreferences, "BCAST_FUSION_BACKEND", DEFAULT_BROADCASTED_FUSION_BACKEND + ) + if backend ∉ SUPPORTED_BROADCASTED_FUSION_BACKENDS + throw( + ArgumentError( + "Unsupported broadcasted fusion backend: $backend. Supported backends are: $(SUPPORTED_BROADCASTED_FUSION_BACKENDS)" + ), + ) + end + return backend +end diff --git a/src/fusion.jl b/src/fusion/parsing.jl similarity index 99% rename from src/fusion.jl rename to src/fusion/parsing.jl index 98022459..8f59a449 100644 --- a/src/fusion.jl +++ b/src/fusion/parsing.jl @@ -363,7 +363,7 @@ end is_fusion_scalar(::Number) = true is_fusion_scalar(_) = false -is_fusion_array(::AbstractArray) = true +is_fusion_array(::NDArray) = true is_fusion_array(_) = false is_fusion_arg(x) = is_fusion_array(x) || is_fusion_scalar(x) diff --git a/src/fusion/ptx_backend.jl b/src/fusion/ptx_backend.jl new file mode 100644 index 00000000..4de48e8b --- /dev/null +++ b/src/fusion/ptx_backend.jl @@ -0,0 +1,7 @@ + + +function + +macro ptx_debug(ex) + # PRINT GENERATED PTX TO STDOUT +end diff --git a/src/utilities/cuda_stubs.jl b/src/utilities/cuda_stubs.jl deleted file mode 100644 index caead31a..00000000 --- a/src/utilities/cuda_stubs.jl +++ /dev/null @@ -1,92 +0,0 @@ -## This file contains stubs for methods implemented in -## the CUDA package extensions not implemented -## elsewhere in the package. - -export @cuda_task, @launch -export map_ndarray_cuda_types, launch, CUDATask - -function ptx_task end -function map_ndarray_cuda_types end -function launch end - -struct CUDATask - func::String - argtypes::NTuple{N,Type} where {N} -end - -""" - @cuda_task(f(args...)) - -Compile a Julia GPU kernel to PTX, register it with the Legate runtime, -and return a `CUDATask` object for later launch. - -# Arguments -- `f` — The name of the Julia CUDA.jl GPU kernel function to compile. -- `args...` — Example arguments to the kernel, used to determine the - argument type signature when generating PTX. - -# Description -This macro automates the process of: -1. Inferring the CUDA argument types for the given `args` using - `map_ndarray_cuda_types`. -2. Using `CUDA.code_ptx` to compile the specified GPU kernel - (`f`) into raw PTX text for the inferred types. -3. Extracting the kernel's function symbol name from the PTX using - `extract_kernel_name`. -4. Registering the compiled PTX and kernel name with the Legate runtime - via `ptx_task`, making it available for GPU execution. -5. Returning a `CUDATask` struct that stores the kernel name and type signature, - which can be used to configure and launch the kernel later. - -# Notes -- The `args...` are not executed; they are used solely for type inference. -- This macro is intended for use with the Legate runtime and - assumes a CUDA context is available. -- Make sure your kernel code is GPU-compatible and does not rely on - unsupported Julia features. - -# Example -```julia -mytask = @cuda_task my_kernel(A, B, C) -``` -""" -macro cuda_task end - -""" - @launch(; task, blocks=(1,), threads=(256,), inputs=(), outputs=(), scalars=()) - -Launch a GPU kernel (previously registered via [`@cuda_task`](@ref)) through the Legate runtime. - -# Keywords -- `task` — A `CUDATask` object, typically returned by [`@cuda_task`](@ref). -- `blocks` — Tuple or single element specifying the CUDA grid dimensions. Defaults to `(1,)`. -- `threads` — Tuple or single element specifying the CUDA block dimensions. Defaults to `(256,)`. -- `inputs` — Tuple or single element of input NDArray objects. -- `outputs` — Tuple or single element of output NDArray objects. -- `scalars` — Tuple or single element of scalar values. - -# Description -The `@launch` macro validates the provided keywords, ensuring only -the allowed set (`:task`, `:blocks`, `:threads`, `:inputs`, `:outputs`, `:scalars`) -are present. It then expands to a call to `cuNumeric.launch`, -passing the given arguments to the Legate runtime for execution. - -This macro is meant to provide a concise, declarative syntax for -launching GPU kernels, separating kernel compilation (via `@cuda_task`) -from execution configuration. - -# Notes -- `task` **must** be a kernel registered with the runtime, usually from `@cuda_task`. -- All keyword arguments must be specified as assignments, e.g. `blocks=(2,2)` not positional arguments. -- Defaults are chosen for single-block, 256-thread 1D launches. -- The macro escapes its body so that the values of inputs/outputs/scalars are captured - from the surrounding scope at macro expansion time. - -# Example -```julia -mytask = @cuda_task my_kernel(A, B, C) - -@launch task=mytask blocks=(8,8) threads=(32,32) inputs=(A, B) outputs=(C) -``` -""" -macro launch end From 1b16ab10004ce8113a6b31f7bffd31d8dfa058fc Mon Sep 17 00:00:00 2001 From: ejmeitz Date: Mon, 27 Apr 2026 10:48:23 -0500 Subject: [PATCH 04/12] integrate things into broadcast machinery more --- Project.toml | 4 + src/cuNumeric.jl | 2 + src/cuda/cuda_util.jl | 2 +- src/fusion/backends.jl | 50 +-- src/fusion/parsing.jl | 796 +++++++++----------------------------- src/fusion/ptx_backend.jl | 66 +++- src/ndarray/broadcast.jl | 97 ++++- src/ndarray/ndarray.jl | 14 +- 8 files changed, 377 insertions(+), 654 deletions(-) diff --git a/Project.toml b/Project.toml index 76925176..458ff8da 100644 --- a/Project.toml +++ b/Project.toml @@ -7,7 +7,9 @@ CNPreferences = "3e078157-ea10-49d5-bf32-908f777cd46f" CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" CUDATools = "9ec180c6-1c07-47c7-9e6e-ebefa4d1f6d0" CxxWrap = "1f15a43c-97ca-5a2a-ae31-89f07a497df4" +ExpressionExplorer = "21656369-7473-754a-2065-74616d696c43" JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" Legate = "1238f2cf-6593-4d60-9aca-2f5364e49909" LegatePreferences = "8028f36a-2b64-49e9-aa04-2d0933fd2ed9" Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb" @@ -27,7 +29,9 @@ CNPreferences = "0.1.2" CUDACore = "6.0.0" CUDATools = "6.0.0" CxxWrap = "0.17" +ExpressionExplorer = "1.1.4" JuliaFormatter = "2.3.0" +KernelAbstractions = "0.9.41" Legate = "0.1.2" LegatePreferences = "0.1.6" MacroTools = "0.5.16" diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index dbfc8989..61a3d086 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -30,6 +30,7 @@ using CxxWrap using CUDATools: CUDATools using CUDACore: CUDACore import CUDACore: CuArray +import KernelAbstractions: @kernel, @index using cupynumeric_jll using cunumeric_jl_wrapper_jll @@ -150,6 +151,7 @@ include("ndarray/unary.jl") include("ndarray/binary.jl") # special features +const FUSE_BROADCAST_EXPRS = load_preference(CNPreferences, "FUSE_BROADCAST_EXPRS", true) include("scoping.jl") include("cuda/cuda_util.jl") include("cuda/cuda_ptx_task.jl") diff --git a/src/cuda/cuda_util.jl b/src/cuda/cuda_util.jl index 12384baa..5cdc4595 100644 --- a/src/cuda/cuda_util.jl +++ b/src/cuda/cuda_util.jl @@ -15,5 +15,5 @@ ndarray_cuda_type(::Type{<:NDArray{T,N}}) where {T,N} = CuDeviceArray{T,N,1} function ndarray_cuda_type(::Type{T}) where {T} Base.isbitstype(T) || throw(ArgumentError("Unsupported argument type: $(T)")) - T + return T end diff --git a/src/fusion/backends.jl b/src/fusion/backends.jl index c261bb1d..a1791f05 100644 --- a/src/fusion/backends.jl +++ b/src/fusion/backends.jl @@ -1,29 +1,29 @@ -abstract type FusionBackend end +# abstract type FusionBackend end -""" - PTXBackend +# """ +# PTXBackend -The default backend for broadcasted expressions. This backend - leverages CUDA.jl to generate PTX code which is executed by the Legate runtime. -""" -struct PTXBackend <: FusionBackend end +# The default backend for broadcasted expressions. This backend +# leverages CUDA.jl to generate PTX code which is executed by the Legate runtime. +# """ +# struct PTXBackend <: FusionBackend end -# Backend for expressions that return `Broadcasted` objects. -# For example, `@fuse z .= x .+ y` uses this backend. -const DEFAULT_BROADCASTED_FUSION_BACKEND = PTXBackend -const SUPPORTED_BROADCASTED_FUSION_BACKENDS = (PTXBackend,) +# # Backend for expressions that return `Broadcasted` objects. +# # For example, `@fuse z .= x .+ y` uses this backend. +# const DEFAULT_BROADCASTED_FUSION_BACKEND = PTXBackend +# const SUPPORTED_BROADCASTED_FUSION_BACKENDS = (PTXBackend,) -#! Call this on init or in global scope -function set_broadcasted_fusion_backend() - backend = load_preference( - CNPreferences, "BCAST_FUSION_BACKEND", DEFAULT_BROADCASTED_FUSION_BACKEND - ) - if backend ∉ SUPPORTED_BROADCASTED_FUSION_BACKENDS - throw( - ArgumentError( - "Unsupported broadcasted fusion backend: $backend. Supported backends are: $(SUPPORTED_BROADCASTED_FUSION_BACKENDS)" - ), - ) - end - return backend -end +# #! Call this on init or in global scope +# function set_broadcasted_fusion_backend() +# backend = load_preference( +# CNPreferences, "BCAST_FUSION_BACKEND", DEFAULT_BROADCASTED_FUSION_BACKEND +# ) +# if backend ∉ SUPPORTED_BROADCASTED_FUSION_BACKENDS +# throw( +# ArgumentError( +# "Unsupported broadcasted fusion backend: $backend. Supported backends are: $(SUPPORTED_BROADCASTED_FUSION_BACKENDS)" +# ), +# ) +# end +# return backend +# end diff --git a/src/fusion/parsing.jl b/src/fusion/parsing.jl index 8f59a449..2347a824 100644 --- a/src/fusion/parsing.jl +++ b/src/fusion/parsing.jl @@ -1,629 +1,219 @@ -export @fuse, @fuse_plan - -""" - ALLOWED_FUSION_OPS - -Set of scalar operations that the fusion frontend is allowed to place in the -fusion IR. Extend this set as the backend learns how to lower more operations. - -The parser normalizes dotted operators/functions, so `.+` and `+` both appear -as `:+` in the IR. Whether an operation was dotted is still preserved on each -`FuseCall` node. -""" -const ALLOWED_FUSION_OPS = Set{Symbol}([ - :+, :-, :*, :/, :^, - :abs, :abs2, - :sqrt, :cbrt, - :exp, :exp2, :expm1, - :log, :log2, :log10, :log1p, - :sin, :cos, :tan, - :asin, :acos, :atan, - :sinh, :cosh, :tanh, - :min, :max, -]) - -const FUSE_OPTION_NAMES = Set{Symbol}([:blocks, :threads]) - -# This cache is intentionally separate from any future function-fusion cache. -# The key is based on the parsed expression IR plus the runtime argument types. -const FUSED_BROADCAST_CACHE = Dict{UInt64,Any}() -const FUSED_BROADCAST_CACHE_LOCK = ReentrantLock() - -# ----------------------------------------------------------------------------- -# Fusion IR -# ----------------------------------------------------------------------------- - -abstract type FuseNode end - -""" - FuseInput(name, index) - -A symbolic input variable appearing in the RHS expression. `index` is the -runtime position of that input in the input tuple passed to `_fuse_broadcast!`. -Repeated uses of the same symbol point to the same index. -""" -struct FuseInput <: FuseNode - name::Symbol - index::Int -end - -""" - FuseLiteral(value) - -A literal scalar that appeared directly in the fused expression, for example -`2` in `z .= 2 .* x`. -""" -struct FuseLiteral <: FuseNode - value::Any -end - -""" - FuseCall(op, args, dotted) - -A scalar operation in the fused expression tree. - -`op` is normalized, so `:.+` becomes `:+`. `dotted` records whether the user -wrote dotted broadcast syntax for this specific call/operator. -""" -struct FuseCall <: FuseNode - op::Symbol - args::Vector{FuseNode} - dotted::Bool -end - -""" - FusePlan(outputs, inputs, ops, rhs, expr_hash) - -A parsed broadcast assignment. - -For example, `@fuse z .= x .+ y .* y` becomes a plan with -`outputs == [:z]`, `inputs == [:x, :y]`, and an RHS tree rooted at `:+`. -""" -struct FusePlan - outputs::Vector{Symbol} - inputs::Vector{Symbol} - ops::Vector{Symbol} - rhs::FuseNode - expr_hash::UInt64 -end - -""" - CompiledFusionKernel - -Placeholder object returned by the default `compile_fused_broadcast` method. -Replace or overload `compile_fused_broadcast` and `launch_fused_broadcast!` when -hooking this parser up to the real CUDA/cuNumeric backend. -""" -struct CompiledFusionKernel - plan::FusePlan - output_types::Tuple - input_types::Tuple -end - -# ----------------------------------------------------------------------------- -# Small display helpers -# ----------------------------------------------------------------------------- - -function Base.show(io::IO, node::FuseInput) - print(io, "FuseInput(:", node.name, ", ", node.index, ")") -end - -function Base.show(io::IO, node::FuseLiteral) - print(io, "FuseLiteral(", repr(node.value), ")") -end - -function Base.show(io::IO, node::FuseCall) - dot = node.dotted ? "." : "" - print(io, "FuseCall(:", dot, node.op, ", ", node.args, ")") -end - -function Base.show(io::IO, plan::FusePlan) - print(io, "FusePlan(outputs = ", plan.outputs, - ", inputs = ", plan.inputs, - ", ops = ", plan.ops, - ", expr_hash = ", plan.expr_hash, - ", rhs = ", plan.rhs, - ")") -end - -# ----------------------------------------------------------------------------- -# Frontend parser -# ----------------------------------------------------------------------------- - -mutable struct ParseContext - input_indices::Dict{Symbol,Int} - inputs::Vector{Symbol} - ops::Vector{Symbol} -end - -ParseContext() = ParseContext(Dict{Symbol,Int}(), Symbol[], Symbol[]) - -function intern_input!(ctx::ParseContext, name::Symbol) - idx = get(ctx.input_indices, name, 0) - if idx == 0 - push!(ctx.inputs, name) - idx = length(ctx.inputs) - ctx.input_indices[name] = idx - end - return idx -end - -""" - build_fuse_plan(ex::Expr) -> FusePlan - -Parse a broadcast assignment expression into a reusable fusion plan. - -Supported MVP syntax: - -```julia -z .= x .+ y .* y -z .= sin.(x) .+ 2 .* y -``` - -This function intentionally does not compile or launch anything. -""" -function build_fuse_plan(ex::Expr) - if Meta.isexpr(ex, :(.=)) - lhs, rhs = ex.args - elseif Meta.isexpr(ex, :(=)) - throw( - ArgumentError( - "@fuse currently supports broadcast assignment `.=` only; got ordinary `=`." - ), - ) - else - throw( - ArgumentError( - "@fuse expected a broadcast assignment like `z .= x .+ y .* y`; got expression head `$(ex.head)`." - ), - ) - end - - outputs = parse_lhs(lhs) - ctx = ParseContext() - rhs_node = parse_rhs(rhs, ctx) - - # Keep the full operator occurrence order in the plan. Use `unique_ops(plan)` - # if the backend only wants each operator once. - plan_without_hash = FusePlan(outputs, copy(ctx.inputs), copy(ctx.ops), rhs_node, UInt64(0)) - return FusePlan( - outputs, copy(ctx.inputs), copy(ctx.ops), rhs_node, hash(plan_signature(plan_without_hash)) - ) -end - -function build_fuse_plan(ex) - throw(ArgumentError("@fuse expected an expression; got `$(repr(ex))`.")) -end - -function parse_lhs(lhs) - if lhs isa Symbol - return Symbol[lhs] - elseif Meta.isexpr(lhs, :tuple) - throw( - ArgumentError( - "Multiple-output broadcast fusion is not implemented yet. Use a single output, e.g. `z .= rhs`." - ), - ) - elseif Meta.isexpr(lhs, :ref) - throw( - ArgumentError( - "Indexed/view outputs are not implemented yet. Use a whole-array output, e.g. `z .= rhs`." - ), - ) - else - throw(ArgumentError("Unsupported @fuse output expression: `$(lhs)`.")) - end -end - -function parse_rhs(ex, ctx::ParseContext)::FuseNode - if ex isa Symbol - idx = intern_input!(ctx, ex) - return FuseInput(ex, idx) - elseif ex isa Number - return FuseLiteral(ex) - elseif ex isa Bool - return FuseLiteral(ex) - elseif ex isa Char - return FuseLiteral(ex) - elseif ex isa String - throw(ArgumentError("String literals are not valid scalar fusion literals: `$(ex)`.")) - elseif Meta.isexpr(ex, :call) - return parse_call(ex, ctx) - elseif Meta.isexpr(ex, :.) - return parse_dotted_call(ex, ctx) - elseif Meta.isexpr(ex, :ref) - throw(ArgumentError("Indexing inside fused RHS is not implemented yet: `$(ex)`.")) - elseif Meta.isexpr(ex, :tuple) - throw(ArgumentError("Tuple construction inside fused RHS is not implemented yet: `$(ex)`.")) - elseif Meta.isexpr(ex, :block) - throw( - ArgumentError( - "Block expressions are not supported in broadcast-fusion RHS expressions." - ), - ) - else - throw( - ArgumentError("Unsupported expression in fused RHS: `$(ex)` with type `$(typeof(ex))`.") - ) - end -end - -function parse_call(ex::Expr, ctx::ParseContext)::FuseNode - raw_op = ex.args[1] - op, dotted = normalize_call_operator(raw_op) - check_allowed_op!(op, raw_op) - - push!(ctx.ops, op) - args = FuseNode[parse_rhs(arg, ctx) for arg in ex.args[2:end]] - return FuseCall(op, args, dotted) -end - -function parse_dotted_call(ex::Expr, ctx::ParseContext)::FuseNode - # Dotted function calls like `sin.(x)` usually parse as - # Expr(:., :sin, Expr(:tuple, :x)). Property access also uses head `:.`, so - # require the second argument to be an argument tuple. - if length(ex.args) != 2 || !(ex.args[2] isa Expr) || ex.args[2].head != :tuple - throw(ArgumentError("Property access is not supported in fused RHS expressions: `$(ex)`.")) - end - - op = function_name_symbol(ex.args[1]) - check_allowed_op!(op, ex.args[1]) - - push!(ctx.ops, op) - args = FuseNode[parse_rhs(arg, ctx) for arg in ex.args[2].args] - return FuseCall(op, args, true) -end - -function normalize_call_operator(raw_op) - if raw_op isa Symbol - op_string = String(raw_op) - if startswith(op_string, ".") - return Symbol(op_string[2:end]), true - else - return raw_op, false - end - end - - # Qualified calls such as Base.sin(x). Dotted qualified calls such as - # Base.sin.(x) are handled by `parse_dotted_call`. - return function_name_symbol(raw_op), false -end - -function function_name_symbol(ex) - if ex isa Symbol - return ex - elseif ex isa QuoteNode && ex.value isa Symbol - return ex.value - elseif Meta.isexpr(ex, :.) - # Handles qualified function names. We intentionally drop the module - # qualification here because the backend's allow-list is scalar-op based. - return function_name_symbol(ex.args[end]) - else - throw(ArgumentError("Could not determine function/operator name from `$(ex)`.")) - end -end - -function check_allowed_op!(op::Symbol, raw_op) - if !(op in ALLOWED_FUSION_OPS) - throw( - ArgumentError( - "Operation `$(raw_op)` normalized to `$(op)`, which is not in ALLOWED_FUSION_OPS." - ), - ) - end - return nothing -end - -# ----------------------------------------------------------------------------- -# Plan signatures and cache keys -# ----------------------------------------------------------------------------- - -node_signature(node::FuseInput) = (:input, node.name, node.index) -node_signature(node::FuseLiteral) = (:literal, typeof(node.value), node.value) -node_signature(node::FuseCall) = (:call, node.op, node.dotted, Tuple(node_signature.(node.args))) - -function plan_signature(plan::FusePlan) - return (:broadcast_assignment, - Tuple(plan.outputs), - Tuple(plan.inputs), - Tuple(plan.ops), - node_signature(plan.rhs)) -end - -""" - unique_ops(plan::FusePlan) - -Return the operations appearing in the RHS in first-occurrence order. -""" -function unique_ops(plan::FusePlan) - seen = Set{Symbol}() - out = Symbol[] - for op in plan.ops - if !(op in seen) - push!(seen, op) - push!(out, op) - end - end - return out -end - -""" - fusion_cache_key(plan, outputs, inputs) -> UInt64 - -Cache key for the compiled kernel. Launch options such as `blocks` and -`threads` are intentionally excluded because they should not affect codegen. -""" -function fusion_cache_key(plan::FusePlan, outputs::Tuple, inputs::Tuple) - return hash((plan.expr_hash, map(typeof, outputs), map(typeof, inputs))) -end - -# ----------------------------------------------------------------------------- -# Runtime argument classification hooks -# ----------------------------------------------------------------------------- - -is_fusion_scalar(::Number) = true -is_fusion_scalar(_) = false -is_fusion_array(::NDArray) = true -is_fusion_array(_) = false -is_fusion_arg(x) = is_fusion_array(x) || is_fusion_scalar(x) - -function validate_fuse_args(plan::FusePlan, outputs::Tuple, inputs::Tuple) - if length(outputs) != length(plan.outputs) - throw( - ArgumentError("Plan expects $(length(plan.outputs)) output(s), got $(length(outputs)).") - ) - end - if length(inputs) != length(plan.inputs) - throw(ArgumentError("Plan expects $(length(plan.inputs)) input(s), got $(length(inputs)).")) - end - - for (name, output) in zip(plan.outputs, outputs) - if is_fusion_scalar(output) - throw( - ArgumentError( - "Output `$(name)` is scalar-like. Scalar outputs are not supported; use a 0D store/array." - ), - ) - elseif !is_fusion_array(output) - throw(ArgumentError("Output `$(name)` has unsupported type `$(typeof(output))`.")) - end - end - - for (name, input) in zip(plan.inputs, inputs) - if !is_fusion_arg(input) - throw( - ArgumentError( - "Input `$(name)` has unsupported type `$(typeof(input))`; expected array-like or scalar-like value." - ), - ) - end - end - - return nothing -end - -# ----------------------------------------------------------------------------- -# Runtime compilation / launch path -# ----------------------------------------------------------------------------- - -""" - _fuse_broadcast!(plan, outputs, inputs; blocks=nothing, threads=nothing) - -Runtime entrypoint emitted by `@fuse`. This validates runtime arguments, -computes a cache key, compiles if necessary, and launches the backend. -""" -function _fuse_broadcast!( - plan::FusePlan, - outputs::Tuple, - inputs::Tuple; - blocks=nothing, - threads=nothing, -) - validate_fuse_args(plan, outputs, inputs) - key = fusion_cache_key(plan, outputs, inputs) - - local kernel - lock(FUSED_BROADCAST_CACHE_LOCK) - try - kernel = get!(FUSED_BROADCAST_CACHE, key) do - compile_fused_broadcast(plan, outputs, inputs) - end - finally - unlock(FUSED_BROADCAST_CACHE_LOCK) - end - - return launch_fused_broadcast!(kernel, outputs, inputs; blocks=blocks, threads=threads) -end - -""" - compile_fused_broadcast(plan, outputs, inputs) - -Backend hook. The default stores type information only. Replace this with the -real CUDATask/PTX generation path. -""" -function compile_fused_broadcast(plan::FusePlan, outputs::Tuple, inputs::Tuple) - return CompiledFusionKernel(plan, map(typeof, outputs), map(typeof, inputs)) -end - -""" - launch_fused_broadcast!(kernel, outputs, inputs; blocks, threads) - -Backend hook. The default throws because this file only implements parsing, -validation, and caching. -""" -function launch_fused_broadcast!( - kernel::CompiledFusionKernel, - outputs::Tuple, - inputs::Tuple; - blocks=nothing, - threads=nothing, -) - throw( - ErrorException( - "No fusion backend has been installed. The expression parsed and cached successfully, " * - "but `launch_fused_broadcast!` must be overloaded to call the CUDA/cuNumeric backend.", - ), - ) -end - -# ----------------------------------------------------------------------------- -# Macros -# ----------------------------------------------------------------------------- - -function parse_fuse_macro_args(args) - isempty(args) && throw(ArgumentError("@fuse expected `z .= rhs`.")) - - call = args[end] - option_exprs = Expr[] - - for opt in args[1:(end - 1)] - if opt isa Symbol && opt in FUSE_OPTION_NAMES - # Supports the old style `@fuse blocks threads z .= rhs`, meaning - # `blocks = blocks, threads = threads`. - push!(option_exprs, Expr(:kw, opt, opt)) - elseif Meta.isexpr(opt, :(=)) || Meta.isexpr(opt, :kw) - name = opt.args[1] - value = opt.args[2] - if !(name isa Symbol) || !(name in FUSE_OPTION_NAMES) - throw( - ArgumentError( - "Invalid @fuse option `$(name)`. Expected one of $(collect(FUSE_OPTION_NAMES))." - ), - ) - end - push!(option_exprs, Expr(:kw, name, value)) - else - throw( - ArgumentError( - "Invalid @fuse option/expression `$(opt)`. Expected options followed by `z .= rhs`." - ), - ) - end - end - - return option_exprs, call -end - -""" - @fuse z .= rhs - @fuse blocks=blocks threads=threads z .= rhs - -Parse a broadcast assignment into a `FusePlan`, validate the runtime arguments, -compile/cache the corresponding kernel, and dispatch to `launch_fused_broadcast!`. -""" -macro fuse(args...) - option_exprs, call = parse_fuse_macro_args(args) - plan = build_fuse_plan(call) +# export @fuse + +# import ExpressionExplorer: compute_symbols_state +# import MacroTools: postwalk + +# struct FusePlan +# outputs::Tuple{Vararg{Symbol}} +# inputs::Tuple{Vararg{Symbol}} +# ops::Tuple{Vararg{Symbol}} +# scalar_literals::Tuple{Vararg{Number}} +# lhs::Any +# rhs::Any +# end - output_tuple = Expr(:tuple, plan.outputs...) - input_tuple = Expr(:tuple, plan.inputs...) - runtime = GlobalRef(@__MODULE__, :_fuse_broadcast!) +# # ----------------------------------------------------------------------------- +# # Runtime argument classification hooks +# # ----------------------------------------------------------------------------- + +# is_fusion_scalar(::Number) = true +# is_fusion_scalar(_) = false +# is_fusion_array(::NDArray) = true +# is_fusion_array(_) = false +# is_fusion_arg(x) = is_fusion_array(x) || is_fusion_scalar(x) + +# function parse_fuse_call(args) +# isempty(args) && throw(ArgumentError("@fuse expected at least one argument.")) + +# call = args[end] +# kwargs = Dict{Symbol,Any}() +# for arg in args[1:(end - 1)] +# if !(arg isa Expr) || !Meta.isexpr(arg, :(=)) +# throw( +# ArgumentError( +# "Invalid @fuse argument `$(arg)`. Expected keyword form `name=value`.", +# ), +# ) +# end +# key, value = arg.args +# if !(key isa Symbol) +# throw(ArgumentError("Invalid @fuse keyword `$(arg)`; keyword name must be a symbol.")) +# end +# if key != :blocks && key != :threads +# throw( +# ArgumentError( +# "Invalid @fuse keyword `$(key)`. Supported keywords are `blocks` and `threads`.", +# ), +# ) +# end +# kwargs[key] = value +# end - if isempty(option_exprs) - return esc(Expr(:call, runtime, QuoteNode(plan), output_tuple, input_tuple)) - else - return esc( - Expr( - :call, - runtime, - Expr(:parameters, option_exprs...), - QuoteNode(plan), - output_tuple, - input_tuple, - ), - ) - end -end +# if !(call isa Expr) +# throw( +# ArgumentError( +# "@fuse expected a broadcast assignment expression as its final argument; got `$(typeof(call))`.", +# ), +# ) +# end -""" - @fuse_plan z .= rhs +# return kwargs, call +# end -Parse only. Useful while developing the frontend. -""" -macro fuse_plan(call) - plan = build_fuse_plan(call) - return QuoteNode(plan) -end +# function parse_broadcast_assignment(ex::Expr) +# lhs, rhs = if length(ex.args) == 2 +# ex.args[1], ex.args[2] +# else +# throw( +# ArgumentError( +# "@fuse expected a broadcast assignment like `z .= x .+ y .* y`; got malformed expression `$(ex)`.", +# ), +# ) +# end -# For @fuse on user defined function -# function maybe_fuse_kernel(fn::Function, output_indices, args::T; kwargs...) where {T} -# cache_key = hash((fn, T)) #! ADD KWRAGS TO THIS? +# lhs_symbol = lhs isa Symbol ? lhs : nothing +# state = compute_symbols_state(ex) +# is_plain_assignment = lhs_symbol !== nothing && in(lhs_symbol, state.assignments) -# if !haskey(FUSED_KERNEL_CACHE, cache_key) -# FUSED_KERNEL_CACHE[cache_key] = fuse_function( -# fn, -# output_indices, -# args...; -# kwargs..., +# if is_plain_assignment || Meta.isexpr(ex, :(=)) +# throw( +# ArgumentError("@fuse currently supports broadcast assignment `.=` only; got ordinary `=`."), # ) +# elseif Meta.isexpr(ex, :(.=)) +# return lhs, rhs # end -# return cache_key +# throw( +# ArgumentError( +# "@fuse expected a broadcast assignment like `z .= x .+ y .* y`; got expression head `$(ex.head)`.", +# ), +# ) # end -# function run_fused_kernel(cache_key; blocks, threads) -# fkd = FUSED_KERNEL_CACHE[cache_key] -# cuNumeric.launch(fkd.ct, fkd.inputs, fkd.outputs, fkd.scalars; blocks=blocks, threads=threads) +# function parse_lhs_outputs(lhs) +# if lhs isa Symbol +# return (lhs,), lhs +# end +# throw( +# ArgumentError( +# "Unsupported @fuse LHS form `$(lhs)`. Supported form for now is a simple symbol target, e.g. `z .= ...`.", +# ), +# ) # end -# macro fuse(ex...) -# call = ex[end] -# kwargs = map(ex[1:(end - 1)]) do kwarg -# if kwarg in FUSE_KWARGS -# :($kwarg = $kwarg) -# else -# throw(ArgumentError("Invalid keyword argument '$kwarg', expected one of $(FUSE_KWARGS)")) +# function collect_rhs_ops(rhs::Expr) +# ops = Symbol[] +# # Lifted from fusion_old's collect_ops idea: post-order walk, collect call heads. +# postwalk(rhs) do node +# if node isa Expr && (node.head === :call || node.head === :.) +# op = node.args[1] +# if op isa Symbol +# push!(ops, op) +# end # end +# return node # end +# return Tuple(ops) +# end -# #! TODO SET DEFAULT BLOCKS/THREADS -# #! HOW TO KNOW WHATS THE RIGHT DIMENSION?? +# function collect_rhs_symbol_refs(rhs::Expr) +# symbol_state = compute_symbols_state(rhs) +# referenced_symbols = symbol_state.references +# refs = Symbol[] +# seen = Set{Symbol}() +# postwalk(rhs) do node +# if node isa Symbol && +# node !== :true && +# node !== :false && +# node !== :nothing && +# in(node, referenced_symbols) && +# !in(node, seen) +# push!(refs, node) +# push!(seen, node) +# end +# return node +# end +# return Tuple(refs) +# end -# blocks = get(kwargs[:blocks], DEFAULT_BLOCKS) -# threads = get(kwargs[:threads], DEFAULT_THREADS) -# output_indices = get(kwargs[:output_indices], ()) +# function collect_rhs_scalar_literals(rhs::Expr) +# scalars = Number[] +# postwalk(rhs) do node +# if node isa Number +# push!(scalars, node) +# elseif node isa QuoteNode && node.value isa Number +# push!(scalars, node.value) +# end +# return node +# end +# return Tuple(scalars) +# end -# code = quote end +# function build_fuse_plan(ex::Expr) +# lhs, rhs = parse_broadcast_assignment(ex) +# outputs, lhs_expr = parse_lhs_outputs(lhs) +# rhs_expr = rhs isa Expr ? rhs : Expr(:block, rhs) +# lhs_symbol_state = compute_symbols_state(lhs_expr) +# rhs_symbol_state = compute_symbols_state(rhs_expr) +# inputs = collect_rhs_symbol_refs(rhs_expr) +# ops = collect_rhs_ops(rhs_expr) +# scalar_literals = collect_rhs_scalar_literals(rhs_expr) + +# shared_names = intersect(lhs_symbol_state.references, rhs_symbol_state.references) +# if !isempty(shared_names) +# throw( +# ArgumentError( +# "LHS and RHS of @fuse must not share variables. Shared symbol(s): $(collect(shared_names)).", +# ), +# ) +# end -# if Meta.isexpr(call, :function) -# # Parses function epxression into args, name etc. -# data = splitdef(longdef(call)) -# # splitarg parses args into (name, type, is_slurp, default) -# arg_data = map(splitarg, data[:args]) -# kwarg_data = map(splitarg, data[:kwargs]) +# return FusePlan(outputs, inputs, ops, scalar_literals, lhs_expr, rhs_expr) +# end -# arg_names = getindex.(arg_data, 1) +# function validate_fuse_args(plan::FusePlan, outputs::Tuple, inputs::Tuple) +# if length(outputs) != length(plan.outputs) +# throw( +# ArgumentError("Plan expects $(length(plan.outputs)) output(s), got $(length(outputs)).") +# ) +# end +# if length(inputs) != length(plan.inputs) +# throw(ArgumentError("Plan expects $(length(plan.inputs)) input(s), got $(length(inputs)).")) +# end -# @gensym wrapper_name, cache_key +# for (name, output) in zip(plan.outputs, outputs) +# if is_fusion_scalar(output) +# throw( +# ArgumentError( +# "Output `$(name)` is scalar-like. Scalar outputs are not supported; use a 0D store/array." +# ), +# ) +# elseif !is_fusion_array(output) +# throw(ArgumentError("Output `$(name)` has unsupported type `$(typeof(output))`.")) +# end +# end -# push!(code.args, -# quote +# for (name, input) in zip(plan.inputs, inputs) +# if !is_fusion_arg(input) +# throw( +# ArgumentError( +# "Input `$(name)` has unsupported type `$(typeof(input))`; expected array-like or scalar-like value." +# ), +# ) +# end +# end -# # The fused function -# function $(wrapper_name)() where {$(dict[:whereparams]...)} -# #! NOT SURE THIS IS RIGHT WAY TO INTERPOLATE ARGS -# #! WANT IT INTERPOLATE AS TUPLE OF ARGS -# #! TODO PASS KWARGS -# $cache_key = maybe_fuse_kernel($(data[:name]), $output_indices, $(arg_names...)) -# run_fused_kernel($cache_key; blocks=($blocks), threads=($threads)) -# end +# return nothing +# end -# # Replace original function with call to fused -# #! NOT SURE OVERWRITING LIKE THIS IS GREAT -# function $(data[:name])( -# $(data[:args]...); $(data[:kwargs]...) -# ) where {$(dict[:whereparams]...)} -# $(wrapper_name)() -# end -# end -# ) -# else -# throw(ArgumentError("fuse expected assignment expression or function, got $(call.head)")) -# end +# """ +# @fuse z .= rhs +# @fuse blocks=blocks threads=threads z .= rhs -# return esc(quote -# let -# $code -# end -# end) +# Parse a broadcast assignment into a `FusePlan` that captures output symbols, +# RHS symbol references, operation ordering, and scalar literals. +# """ +# macro fuse(args...) +# _, call = parse_fuse_call(args) +# return esc(:(cuNumeric.build_fuse_plan($(QuoteNode(call))))) # end diff --git a/src/fusion/ptx_backend.jl b/src/fusion/ptx_backend.jl index 4de48e8b..dc6186c5 100644 --- a/src/fusion/ptx_backend.jl +++ b/src/fusion/ptx_backend.jl @@ -1,7 +1,65 @@ +# const PTX_FUSION_CACHE = Dict{UInt64, CUDATask}() +# const PTX_FUSION_KWARGS = (:blocks, :threads) +# macro ptx_debug(ex) +# # PRINT GENERATED PTX TO STDOUT +# end -function +# function _fuse_broadcast(ex::Expr, ::Type{PTXBackend}) -macro ptx_debug(ex) - # PRINT GENERATED PTX TO STDOUT -end +# call = ex[end] +# kwargs = map(ex[1:end-1]) do kwarg +# if kwarg in FUSE_KWARGS +# :($kwarg = $kwarg) +# else +# throw(ArgumentError("Invalid keyword argument '$kwarg', expected one of $(FUSE_KWARGS)")) +# end +# end + +# #! TODO SET DEFAULT BLOCKS/THREADS +# #! HOW TO KNOW WHATS THE RIGHT DIMENSION?? + +# blocks = get(kwargs[:blocks], DEFAULT_BLOCKS) +# threads = get(kwargs[:threads], DEFAULT_THREADS) + +# # I think its safe to not put kwargs in the hash +# # as the PTX does not specialize on blocks/threads +# #! NEED TO PUT TYPE INFO OF INPUTS/OUTPUTS HERE +# cache_key = hash(call) + +# return esc(quote + +# local task::CUDATask +# if haskey(PTX_FUSION_CACHE, cache_key) +# task = PTX_FUSION_CACHE[cache_key] +# else +# local _buf = IOBuffer() +# # generate ptx using CUDA.jl +# CUDATools.@device_code_ptx io=_buf $(call) + +# local _ptx = String(take!(_buf)) +# local _func_name = extract_kernel_name(_ptx) + +# # issue ptx_task within legate runtime to register cufunction ptr with cucontext +# ptx_task(_ptx, _func_name) + +# PTX_FUSION_CACHE[cache_key] = CUDATask(_func_name, _types) +# end +# end) + +# end + +# function CUDA.code_ptx( +# io::IO, +# @nospecialize(func), +# @nospecialize(types); +# kernel=false, +# kwargs... +# ) +# compiler_kwargs, kwargs = split_kwargs_runtime(kwargs, CUDACore.COMPILER_KWARGS) +# source = methodinstance(typeof(func), Base.to_tuple_type(types)) +# config = CUDACore.compiler_config(device(); kernel, compiler_kwargs...) +# job = CompilerJob(source, config) +# # use frozen world to avoid recompiling the compiler infrastructure +# CUDACore.invoke_frozen(code_ptx, $(args...); kwargs...) +# end diff --git a/src/ndarray/broadcast.jl b/src/ndarray/broadcast.jl index fcf0a3e5..8066b3ca 100644 --- a/src/ndarray/broadcast.jl +++ b/src/ndarray/broadcast.jl @@ -38,6 +38,7 @@ function Base.similar(bc::Broadcasted{NDArrayStyle{N}}, ::Type{ElType}) where {N end function __broadcast(f::Function, _, args...) + #! WITH FUSION I THINK WE CAN SUPPORT THIS BY JUST CALLING MAP or MAP! error( """ Tried to broadcast $(f). cuNumeric.jl does not support broadcasting user-defined functions yet. Please re-define \ @@ -52,18 +53,45 @@ end bcast_depth(bc::Base.Broadcast.Broadcasted) = maximum(bcast_depth, bc.args, init=0) + 1; bcast_depth(::Any) = 0 -function Base.Broadcast.materialize(bc::Broadcasted{<:NDArrayStyle}) +# Copied from GPUArrays: https://github.com/JuliaGPU/GPUArrays.jl/blob/a9df2ba41ca2358c1de2f3cc6b020578bf6e39b1/src/host/broadcast.jl#L60-L63 +# Defined with KernelAbstractions.jl. Makes it easier to generate indexing for +# various dimensions of inputs/outputs. Assumes broadcast is `Base.Broadcast.process`ed so that +# dest/bc have singleton dimensions inserted and we can index 1-1 like this. +@kernel function broadcast_kernel_cartesian(dest, bc) + I = @index(Global, Cartesian) + @inbounds dest[I] = bc[I] +end + +@kernel function broadcast_kernel_linear(dest, bc) + I = @index(Global, Linear) + @inbounds dest[I] = bc[I] +end + +# No compilation here, just generating CUDA specific kernel. +const GPU_CARTESIAN_KERNEL = broadcast_kernel_cartesian(CUDACore.CUDAKernels.CUDABackend()) +const GPU_LINEAR_KERNEL = broadcast_kernel_linear(CUDACore.CUDAKernels.CUDABackend()) + +struct BrokenBroadcast{T} end +Base.convert(::Type{BrokenBroadcast{T}}, x) where {T} = BrokenBroadcast{T}() +Base.convert(::Type{BrokenBroadcast{T}}, x::BrokenBroadcast{T}) where {T} = x +Base.eltype(::Type{BrokenBroadcast{T}}) where {T} = T + +function Broadcast.copy(bc::Broadcasted{<:NDArrayStyle{0}}) ElType = Broadcast.combine_eltypes(bc.f, bc.args) - if ElType == Union{} || !Base.allocatedinline(ElType) - error("Cannot broadcast $(bc.f) over NDArrays with eltypes: $(eltype.(bc.args))") + if ElType == Union{} + ElType = Nothing end + dest = copyto!(similar(bc, ElType), bc) + #! CHECK THIS DOESNT CAUSE ISSUES DUE TO BLOCKING NATURE + return @allowscalar dest[CartesianIndex()] +end - #* This be the place to inject kernel fusion via CUDA.jl - #* Use the function in Base.Broadcast.flatten(bc). - #* How can we check all the funcs in this expr - #* are supported by CUDA? - - return unravel_broadcast_tree(bc) +@inline function Broadcast.copy(bc::Broadcasted{<:NDArrayStyle}) + ElType = Broadcast.combine_eltypes(bc.f, bc.args) + if ElType == Union{} || !Base.allocatedinline(ElType) + ElType = BrokenBroadcast{ElType} + end + copyto!(similar(bc, ElType), bc) end # Recursion base cases @@ -109,10 +137,51 @@ function unravel_broadcast_tree(bc::Broadcasted) return __broadcast(bc.f, out, in_args...) end -# Support .= -function Base.copyto!(dest::NDArray{T,N}, bc::Broadcasted{<:NDArrayStyle{N}}) where {T,N} - # Moves result from broadcast (src) to dest. src array is no longer valid - #! THIS ENABLES FOOT GUN IF USER SPECIFIES INTEGER ARRAY AT OUTPUT - nda_move(dest, checked_promote_arr(Base.Broadcast.materialize(bc), T)) +function fuse_broadcast_tree!(dest::NDArray, bc::Broadcasted) + bc = Base.Broadcast.preprocess(dest, bc) + + # Get proper kernel + broadcast_kernel = + if ndims(dest) == 1 || + (isa(IndexStyle(dest), IndexLinear) && + isa(IndexStyle(bc), IndexLinear)) + GPU_LINEAR_KERNEL + else + GPU_CARTESIAN_KERNEL + end + + #! DO I NEED TO DO TYPE PROMOTION CHECKS?? + # ndims check for 0D support + broadcast_kernel(dest, bc; ndrange=ndims(dest) > 0 ? size(dest) : (1,)) return dest end + +@inline function _copyto!(dest::NDArray, bc::Broadcasted) + axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc)) + isempty(dest) && return dest + if eltype(dest) <: BrokenBroadcast + throw( + ArgumentError( + "Broadcast operation resulting in $(eltype(eltype(dest))) is not NDArray compatible" + ), + ) + end + + #! IF THIS IS KNOWN AT COMPILE TIME WE CAN GENERATE + #! THIS CODE WITHOUT THE IF STATEMENT + if FUSE_BROADCAST_EXPRS + #! DO I NEED TO DO TYPE PROMOTION CHECKS BEFORE RETURNING? + #! WE MIGHT NEED TO CHECK IF ON GPU OR CPU AND FALLBACK + return fuse_broadcast_tree!(dest, bc) + else + temp_result = unravel_broadcast_tree(bc) + nda_move(dest, checked_promote_arr(temp_result, eltype(dest))) + return dest + end +end + +# Support .= +@inline Base.copyto!(dest::NDArray, bc::Broadcasted{Nothing}) = _copyto!(dest, bc) +@inline Base.copyto!(dest::NDArray, bc::Broadcasted{<:NDArrayStyle}) = _copyto!(dest, bc) + +#! TODO ADD MAP FUSED IMPLEMENTATIONS diff --git a/src/ndarray/ndarray.jl b/src/ndarray/ndarray.jl index d58f6572..edfed03f 100644 --- a/src/ndarray/ndarray.jl +++ b/src/ndarray/ndarray.jl @@ -214,6 +214,7 @@ size(arr, 2) """ Base.size(arr::NDArray{<:Any,N}) where {N} = cuNumeric.shape(arr) Base.size(arr::NDArray, dim::Int) = Base.size(arr)[dim] +Base.isempty(arr::NDArray) = any(==(0), size(arr)) @doc""" Base.firstindex(arr::NDArray, dim::Int) @@ -240,9 +241,8 @@ Base.view(arr::NDArray, inds...) = arr[inds...] # NDArray slices are views by de Base.IndexStyle(::NDArray) = IndexCartesian() function Base.show(io::IO, arr::NDArray{T,0}) where {T} - println(io, "0-dimensional NDArray{$(T),0}") allowscalar() do - print(io, arr[]) + print(io, "NDArray{$(T),0}(", repr(arr[]), ")") end end @@ -253,13 +253,13 @@ function Base.show(io::IO, ::MIME"text/plain", arr::NDArray{T,0}) where {T} end end -function Base.show(io::IO, arr::NDArray{T,D}) where {T,D} - println(io, "NDArray{$(T),$(D)}") - Base.print_array(io, Array(arr)) +function Base.show(io::IO, arr::NDArray{T,N}) where {T,N} + print(io, "NDArray{$(T),$(N)} with size ", size(arr)) end -function Base.show(io::IO, ::MIME"text/plain", arr::NDArray{T}) where {T} - Base.show(io, arr) +function Base.show(io::IO, ::MIME"text/plain", arr::NDArray{T,N}) where {T,N} + println(io, "NDArray{$(T),$(N)} with size ", size(arr)) + Base.print_array(io, Array(arr)) end function Base.print(arr::NDArray{T}) where {T} From 9e50bc62730251ed2a94b0629d1df8896f8e81b7 Mon Sep 17 00:00:00 2001 From: ejmeitz Date: Mon, 27 Apr 2026 14:32:11 -0500 Subject: [PATCH 05/12] mostly there --- src/cuNumeric.jl | 22 ++-- src/cuda/cuda_util.jl | 24 +++- src/fusion/backends.jl | 29 ----- src/fusion/parsing.jl | 219 -------------------------------- src/fusion/ptx_backend.jl | 65 ---------- src/ndarray/broadcast.jl | 43 +------ src/ndarray/broadcast_fusion.jl | 139 ++++++++++++++++++++ 7 files changed, 179 insertions(+), 362 deletions(-) delete mode 100644 src/fusion/backends.jl delete mode 100644 src/fusion/parsing.jl delete mode 100644 src/fusion/ptx_backend.jl create mode 100644 src/ndarray/broadcast_fusion.jl diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index 61a3d086..23389a38 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -31,6 +31,7 @@ using CUDATools: CUDATools using CUDACore: CUDACore import CUDACore: CuArray import KernelAbstractions: @kernel, @index +import KernelAbstractions as KA using cupynumeric_jll using cunumeric_jl_wrapper_jll @@ -143,23 +144,22 @@ include("warnings.jl") # NDArray internal include("ndarray/detail/ndarray.jl") -# NDArray interface +# Utilities +include("cuda/cuda_util.jl") +include("utilities/version.jl") +include("util.jl") + +const FUSE_BROADCAST_EXPRS = load_preference(CNPreferences, "FUSE_BROADCAST_EXPRS", true) + +# Functionality include("ndarray/promotion.jl") +include("cuda/cuda_ptx_task.jl") +include("ndarray/broadcast_fusion.jl") include("ndarray/broadcast.jl") include("ndarray/ndarray.jl") include("ndarray/unary.jl") include("ndarray/binary.jl") - -# special features -const FUSE_BROADCAST_EXPRS = load_preference(CNPreferences, "FUSE_BROADCAST_EXPRS", true) include("scoping.jl") -include("cuda/cuda_util.jl") -include("cuda/cuda_ptx_task.jl") -include("fusion/parsing.jl") - -# Utilities -include("utilities/version.jl") -include("util.jl") # From https://github.com/JuliaGraphics/QML.jl/blob/dca239404135d85fe5d4afe34ed3dc5f61736c63/src/QML.jl#L147 mutable struct ArgcArgv diff --git a/src/cuda/cuda_util.jl b/src/cuda/cuda_util.jl index 5cdc4595..9485b205 100644 --- a/src/cuda/cuda_util.jl +++ b/src/cuda/cuda_util.jl @@ -11,9 +11,31 @@ function _setup_cuda_tasking() end end -ndarray_cuda_type(::Type{<:NDArray{T,N}}) where {T,N} = CuDeviceArray{T,N,1} +# Other memeory types here: https://github.com/JuliaGPU/CUDA.jl/blob/345c1600ebd561135148bb04ee2657f521a40e25/CUDACore/src/device/pointer.jl#L7 +function ndarray_cuda_type(::Type{<:NDArray{T,N}}) where {T,N} + CUDACore.CuDeviceArray{T,N,CUDACore.AS.Global} +end function ndarray_cuda_type(::Type{T}) where {T} Base.isbitstype(T) || throw(ArgumentError("Unsupported argument type: $(T)")) return T end + +""" + map_cuda_type(::Type{T})::Type + +Recursively rewrite cuNumeric broadcast-related types so they can be treated as CUDA-friendly +types for code generation (e.g. mapping `NDArray{...}` to `CuDeviceArray{...}` inside +`Base.Broadcast.Broadcasted{...}` type parameters). +""" +map_cuda_type(::Type{T}) where {T} = T + +map_cuda_type(::Type{<:NDArray{T,N}}) where {T,N} = ndarray_cuda_type(NDArray{T,N}) + +function map_cuda_type(::Type{T}) where {T<:Tuple} + return Tuple{map_cuda_type.(T.parameters)...} +end + +function map_cuda_type(::Type{Base.Broadcast.Broadcasted{S,Ax,F,Args}}) where {S,Ax,F,Args} + return Base.Broadcast.Broadcasted{map_cuda_type(S),Ax,F,map_cuda_type(Args)} +end diff --git a/src/fusion/backends.jl b/src/fusion/backends.jl deleted file mode 100644 index a1791f05..00000000 --- a/src/fusion/backends.jl +++ /dev/null @@ -1,29 +0,0 @@ -# abstract type FusionBackend end - -# """ -# PTXBackend - -# The default backend for broadcasted expressions. This backend -# leverages CUDA.jl to generate PTX code which is executed by the Legate runtime. -# """ -# struct PTXBackend <: FusionBackend end - -# # Backend for expressions that return `Broadcasted` objects. -# # For example, `@fuse z .= x .+ y` uses this backend. -# const DEFAULT_BROADCASTED_FUSION_BACKEND = PTXBackend -# const SUPPORTED_BROADCASTED_FUSION_BACKENDS = (PTXBackend,) - -# #! Call this on init or in global scope -# function set_broadcasted_fusion_backend() -# backend = load_preference( -# CNPreferences, "BCAST_FUSION_BACKEND", DEFAULT_BROADCASTED_FUSION_BACKEND -# ) -# if backend ∉ SUPPORTED_BROADCASTED_FUSION_BACKENDS -# throw( -# ArgumentError( -# "Unsupported broadcasted fusion backend: $backend. Supported backends are: $(SUPPORTED_BROADCASTED_FUSION_BACKENDS)" -# ), -# ) -# end -# return backend -# end diff --git a/src/fusion/parsing.jl b/src/fusion/parsing.jl deleted file mode 100644 index 2347a824..00000000 --- a/src/fusion/parsing.jl +++ /dev/null @@ -1,219 +0,0 @@ -# export @fuse - -# import ExpressionExplorer: compute_symbols_state -# import MacroTools: postwalk - -# struct FusePlan -# outputs::Tuple{Vararg{Symbol}} -# inputs::Tuple{Vararg{Symbol}} -# ops::Tuple{Vararg{Symbol}} -# scalar_literals::Tuple{Vararg{Number}} -# lhs::Any -# rhs::Any -# end - -# # ----------------------------------------------------------------------------- -# # Runtime argument classification hooks -# # ----------------------------------------------------------------------------- - -# is_fusion_scalar(::Number) = true -# is_fusion_scalar(_) = false -# is_fusion_array(::NDArray) = true -# is_fusion_array(_) = false -# is_fusion_arg(x) = is_fusion_array(x) || is_fusion_scalar(x) - -# function parse_fuse_call(args) -# isempty(args) && throw(ArgumentError("@fuse expected at least one argument.")) - -# call = args[end] -# kwargs = Dict{Symbol,Any}() -# for arg in args[1:(end - 1)] -# if !(arg isa Expr) || !Meta.isexpr(arg, :(=)) -# throw( -# ArgumentError( -# "Invalid @fuse argument `$(arg)`. Expected keyword form `name=value`.", -# ), -# ) -# end -# key, value = arg.args -# if !(key isa Symbol) -# throw(ArgumentError("Invalid @fuse keyword `$(arg)`; keyword name must be a symbol.")) -# end -# if key != :blocks && key != :threads -# throw( -# ArgumentError( -# "Invalid @fuse keyword `$(key)`. Supported keywords are `blocks` and `threads`.", -# ), -# ) -# end -# kwargs[key] = value -# end - -# if !(call isa Expr) -# throw( -# ArgumentError( -# "@fuse expected a broadcast assignment expression as its final argument; got `$(typeof(call))`.", -# ), -# ) -# end - -# return kwargs, call -# end - -# function parse_broadcast_assignment(ex::Expr) -# lhs, rhs = if length(ex.args) == 2 -# ex.args[1], ex.args[2] -# else -# throw( -# ArgumentError( -# "@fuse expected a broadcast assignment like `z .= x .+ y .* y`; got malformed expression `$(ex)`.", -# ), -# ) -# end - -# lhs_symbol = lhs isa Symbol ? lhs : nothing -# state = compute_symbols_state(ex) -# is_plain_assignment = lhs_symbol !== nothing && in(lhs_symbol, state.assignments) - -# if is_plain_assignment || Meta.isexpr(ex, :(=)) -# throw( -# ArgumentError("@fuse currently supports broadcast assignment `.=` only; got ordinary `=`."), -# ) -# elseif Meta.isexpr(ex, :(.=)) -# return lhs, rhs -# end - -# throw( -# ArgumentError( -# "@fuse expected a broadcast assignment like `z .= x .+ y .* y`; got expression head `$(ex.head)`.", -# ), -# ) -# end - -# function parse_lhs_outputs(lhs) -# if lhs isa Symbol -# return (lhs,), lhs -# end -# throw( -# ArgumentError( -# "Unsupported @fuse LHS form `$(lhs)`. Supported form for now is a simple symbol target, e.g. `z .= ...`.", -# ), -# ) -# end - -# function collect_rhs_ops(rhs::Expr) -# ops = Symbol[] -# # Lifted from fusion_old's collect_ops idea: post-order walk, collect call heads. -# postwalk(rhs) do node -# if node isa Expr && (node.head === :call || node.head === :.) -# op = node.args[1] -# if op isa Symbol -# push!(ops, op) -# end -# end -# return node -# end -# return Tuple(ops) -# end - -# function collect_rhs_symbol_refs(rhs::Expr) -# symbol_state = compute_symbols_state(rhs) -# referenced_symbols = symbol_state.references -# refs = Symbol[] -# seen = Set{Symbol}() -# postwalk(rhs) do node -# if node isa Symbol && -# node !== :true && -# node !== :false && -# node !== :nothing && -# in(node, referenced_symbols) && -# !in(node, seen) -# push!(refs, node) -# push!(seen, node) -# end -# return node -# end -# return Tuple(refs) -# end - -# function collect_rhs_scalar_literals(rhs::Expr) -# scalars = Number[] -# postwalk(rhs) do node -# if node isa Number -# push!(scalars, node) -# elseif node isa QuoteNode && node.value isa Number -# push!(scalars, node.value) -# end -# return node -# end -# return Tuple(scalars) -# end - -# function build_fuse_plan(ex::Expr) -# lhs, rhs = parse_broadcast_assignment(ex) -# outputs, lhs_expr = parse_lhs_outputs(lhs) -# rhs_expr = rhs isa Expr ? rhs : Expr(:block, rhs) -# lhs_symbol_state = compute_symbols_state(lhs_expr) -# rhs_symbol_state = compute_symbols_state(rhs_expr) -# inputs = collect_rhs_symbol_refs(rhs_expr) -# ops = collect_rhs_ops(rhs_expr) -# scalar_literals = collect_rhs_scalar_literals(rhs_expr) - -# shared_names = intersect(lhs_symbol_state.references, rhs_symbol_state.references) -# if !isempty(shared_names) -# throw( -# ArgumentError( -# "LHS and RHS of @fuse must not share variables. Shared symbol(s): $(collect(shared_names)).", -# ), -# ) -# end - -# return FusePlan(outputs, inputs, ops, scalar_literals, lhs_expr, rhs_expr) -# end - -# function validate_fuse_args(plan::FusePlan, outputs::Tuple, inputs::Tuple) -# if length(outputs) != length(plan.outputs) -# throw( -# ArgumentError("Plan expects $(length(plan.outputs)) output(s), got $(length(outputs)).") -# ) -# end -# if length(inputs) != length(plan.inputs) -# throw(ArgumentError("Plan expects $(length(plan.inputs)) input(s), got $(length(inputs)).")) -# end - -# for (name, output) in zip(plan.outputs, outputs) -# if is_fusion_scalar(output) -# throw( -# ArgumentError( -# "Output `$(name)` is scalar-like. Scalar outputs are not supported; use a 0D store/array." -# ), -# ) -# elseif !is_fusion_array(output) -# throw(ArgumentError("Output `$(name)` has unsupported type `$(typeof(output))`.")) -# end -# end - -# for (name, input) in zip(plan.inputs, inputs) -# if !is_fusion_arg(input) -# throw( -# ArgumentError( -# "Input `$(name)` has unsupported type `$(typeof(input))`; expected array-like or scalar-like value." -# ), -# ) -# end -# end - -# return nothing -# end - -# """ -# @fuse z .= rhs -# @fuse blocks=blocks threads=threads z .= rhs - -# Parse a broadcast assignment into a `FusePlan` that captures output symbols, -# RHS symbol references, operation ordering, and scalar literals. -# """ -# macro fuse(args...) -# _, call = parse_fuse_call(args) -# return esc(:(cuNumeric.build_fuse_plan($(QuoteNode(call))))) -# end diff --git a/src/fusion/ptx_backend.jl b/src/fusion/ptx_backend.jl deleted file mode 100644 index dc6186c5..00000000 --- a/src/fusion/ptx_backend.jl +++ /dev/null @@ -1,65 +0,0 @@ -# const PTX_FUSION_CACHE = Dict{UInt64, CUDATask}() -# const PTX_FUSION_KWARGS = (:blocks, :threads) - -# macro ptx_debug(ex) -# # PRINT GENERATED PTX TO STDOUT -# end - -# function _fuse_broadcast(ex::Expr, ::Type{PTXBackend}) - -# call = ex[end] -# kwargs = map(ex[1:end-1]) do kwarg -# if kwarg in FUSE_KWARGS -# :($kwarg = $kwarg) -# else -# throw(ArgumentError("Invalid keyword argument '$kwarg', expected one of $(FUSE_KWARGS)")) -# end -# end - -# #! TODO SET DEFAULT BLOCKS/THREADS -# #! HOW TO KNOW WHATS THE RIGHT DIMENSION?? - -# blocks = get(kwargs[:blocks], DEFAULT_BLOCKS) -# threads = get(kwargs[:threads], DEFAULT_THREADS) - -# # I think its safe to not put kwargs in the hash -# # as the PTX does not specialize on blocks/threads -# #! NEED TO PUT TYPE INFO OF INPUTS/OUTPUTS HERE -# cache_key = hash(call) - -# return esc(quote - -# local task::CUDATask -# if haskey(PTX_FUSION_CACHE, cache_key) -# task = PTX_FUSION_CACHE[cache_key] -# else -# local _buf = IOBuffer() -# # generate ptx using CUDA.jl -# CUDATools.@device_code_ptx io=_buf $(call) - -# local _ptx = String(take!(_buf)) -# local _func_name = extract_kernel_name(_ptx) - -# # issue ptx_task within legate runtime to register cufunction ptr with cucontext -# ptx_task(_ptx, _func_name) - -# PTX_FUSION_CACHE[cache_key] = CUDATask(_func_name, _types) -# end -# end) - -# end - -# function CUDA.code_ptx( -# io::IO, -# @nospecialize(func), -# @nospecialize(types); -# kernel=false, -# kwargs... -# ) -# compiler_kwargs, kwargs = split_kwargs_runtime(kwargs, CUDACore.COMPILER_KWARGS) -# source = methodinstance(typeof(func), Base.to_tuple_type(types)) -# config = CUDACore.compiler_config(device(); kernel, compiler_kwargs...) -# job = CompilerJob(source, config) -# # use frozen world to avoid recompiling the compiler infrastructure -# CUDACore.invoke_frozen(code_ptx, $(args...); kwargs...) -# end diff --git a/src/ndarray/broadcast.jl b/src/ndarray/broadcast.jl index 8066b3ca..1f0eceeb 100644 --- a/src/ndarray/broadcast.jl +++ b/src/ndarray/broadcast.jl @@ -4,6 +4,11 @@ struct NDArrayStyle{N} <: AbstractArrayStyle{N} end Base.BroadcastStyle(::Type{<:NDArray{<:Any,N}}) where {N} = NDArrayStyle{N}() Base.BroadcastStyle(::NDArrayStyle{N}, ::NDArrayStyle{M}) where {N,M} = NDArrayStyle{max(N, M)}() +# Some other functions in cuda_util.jl +function map_cuda_type(::Type{cuNumeric.NDArrayStyle{N}}) where {N} + CUDACore.CuArrayStyle{N,CUDACore.DeviceMemory} +end # Also can be HostMemory or UnifiedMemory + function _nd_forbid_mix() throw( ArgumentError( @@ -53,24 +58,6 @@ end bcast_depth(bc::Base.Broadcast.Broadcasted) = maximum(bcast_depth, bc.args, init=0) + 1; bcast_depth(::Any) = 0 -# Copied from GPUArrays: https://github.com/JuliaGPU/GPUArrays.jl/blob/a9df2ba41ca2358c1de2f3cc6b020578bf6e39b1/src/host/broadcast.jl#L60-L63 -# Defined with KernelAbstractions.jl. Makes it easier to generate indexing for -# various dimensions of inputs/outputs. Assumes broadcast is `Base.Broadcast.process`ed so that -# dest/bc have singleton dimensions inserted and we can index 1-1 like this. -@kernel function broadcast_kernel_cartesian(dest, bc) - I = @index(Global, Cartesian) - @inbounds dest[I] = bc[I] -end - -@kernel function broadcast_kernel_linear(dest, bc) - I = @index(Global, Linear) - @inbounds dest[I] = bc[I] -end - -# No compilation here, just generating CUDA specific kernel. -const GPU_CARTESIAN_KERNEL = broadcast_kernel_cartesian(CUDACore.CUDAKernels.CUDABackend()) -const GPU_LINEAR_KERNEL = broadcast_kernel_linear(CUDACore.CUDAKernels.CUDABackend()) - struct BrokenBroadcast{T} end Base.convert(::Type{BrokenBroadcast{T}}, x) where {T} = BrokenBroadcast{T}() Base.convert(::Type{BrokenBroadcast{T}}, x::BrokenBroadcast{T}) where {T} = x @@ -112,6 +99,7 @@ function __materialize(bc::Broadcasted{<:NDArrayStyle}) unravel_broadcast_tree(bc) end +# Un-fused implementation of broadcast tree function unravel_broadcast_tree(bc::Broadcasted) # Recursively materialize/unravel any nested broadcasts @@ -137,25 +125,6 @@ function unravel_broadcast_tree(bc::Broadcasted) return __broadcast(bc.f, out, in_args...) end -function fuse_broadcast_tree!(dest::NDArray, bc::Broadcasted) - bc = Base.Broadcast.preprocess(dest, bc) - - # Get proper kernel - broadcast_kernel = - if ndims(dest) == 1 || - (isa(IndexStyle(dest), IndexLinear) && - isa(IndexStyle(bc), IndexLinear)) - GPU_LINEAR_KERNEL - else - GPU_CARTESIAN_KERNEL - end - - #! DO I NEED TO DO TYPE PROMOTION CHECKS?? - # ndims check for 0D support - broadcast_kernel(dest, bc; ndrange=ndims(dest) > 0 ? size(dest) : (1,)) - return dest -end - @inline function _copyto!(dest::NDArray, bc::Broadcasted) axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc)) isempty(dest) && return dest diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl new file mode 100644 index 00000000..4c9db998 --- /dev/null +++ b/src/ndarray/broadcast_fusion.jl @@ -0,0 +1,139 @@ + +# Copied from GPUArrays: https://github.com/JuliaGPU/GPUArrays.jl/blob/a9df2ba41ca2358c1de2f3cc6b020578bf6e39b1/src/host/broadcast.jl#L60-L63 +# Defined with KernelAbstractions.jl. Makes it easier to generate indexing for +# various dimensions of inputs/outputs. Assumes broadcast is `Base.Broadcast.process`ed so that +# dest/bc have singleton dimensions inserted and we can index 1-1 like this. +@kernel function broadcast_kernel_cartesian(dest, bc) + I = @index(Global, Cartesian) + @inbounds dest[I] = bc[I] +end + +@kernel function broadcast_kernel_linear(dest, bc) + I = @index(Global, Linear) + @inbounds dest[I] = bc[I] +end + +# No compilation here, just generating CUDA specific kernel. +const GPU_CARTESIAN_KERNEL = broadcast_kernel_cartesian(CUDACore.CUDAKernels.CUDABackend()) +const GPU_LINEAR_KERNEL = broadcast_kernel_linear(CUDACore.CUDAKernels.CUDABackend()) + +const _BCAST_PTX_CACHE = Dict{Tuple{Any,DataType,DataType,Any},Tuple{CUDATask,Int,Int}}() +const _BCAST_PTX_CACHE_LOCK = ReentrantLock() + +function _default_broadcast_threads(ndrange) + n = prod(ndrange) + return min(256, max(1, n)) +end + +function _cufunction_from_types(f, types_tuple_type; maxthreads=nothing, always_inline=false) + return CUDACore.cufunction( + f, + types_tuple_type; + kernel=true, + maxthreads=maxthreads, + always_inline=always_inline, + ) +end + +""" + get_ptx(obj::KA.Kernel{CUDABackend}, ::Type{DEST_T}, ::Type{BC_T}; + ndrange) -> (ptx::String, threads::Int, blocks::Int) + +Compile a KA CUDA kernel (kernel-body `obj.f(ctx, ...)`) using *types only* for `DEST_T` and `BC_T`, +choose a workgroup size (threads) using CUDA occupancy when possible, and return the generated PTX. +""" +function get_ptx( + obj::KA.Kernel{CUDACore.CUDAKernels.CUDABackend}, + ::Type{DEST_T}, + ::Type{BC_T}; + ndrange, +) where {DEST_T,BC_T} + backend = KA.backend(obj) + + ndrange, workgroupsize, iterspace, dynamic = KA.launch_config(obj, ndrange, nothing) + ctx = KA.mkcontext(obj, ndrange, iterspace) + + # If the kernel is statically sized we can tell the compiler about that + maxthreads = + if KA.workgroupsize(obj) <: KA.StaticSize + prod(KA.get(KA.workgroupsize(obj))) + else + nothing + end + + # Determine threads via occupancy if we can compile from types; otherwise use a broadcast-friendly heuristic. + threads = _default_broadcast_threads(ndrange) + tt = Base.to_tuple_type((typeof(ctx), DEST_T, BC_T)) + host_kernel = CUDACore.cufunction( + obj.f, + tt; + kernel=true, + maxthreads=maxthreads, + always_inline=backend.always_inline, + ) + config = CUDACore.launch_configuration(host_kernel.fun; max_threads=prod(ndrange)) + threads = config.threads + + # if fancy thing doesnt work we can use this + # threads = _default_broadcast_threads(ndrange) + + workgroupsize = CUDACore.CUDAKernels.threads_to_workgroupsize(threads, ndrange) + iterspace, dynamic = KA.partition(obj, ndrange, workgroupsize) + ctx = KA.mkcontext(obj, ndrange, iterspace) + + blocks = length(KA.blocks(iterspace)) + threads = length(KA.workitems(iterspace)) + blocks == 0 && return "", 0, 0 + + buf = IOBuffer() + CUDATools.code_ptx(buf, obj.f, (typeof(ctx), DEST_T, BC_T); raw=false, kernel=true) + return String(take!(buf)), threads, blocks +end + +function get_cuda_task( + obj::KA.Kernel{CUDACore.CUDAKernels.CUDABackend}, + ::Type{DEST_T}, + ::Type{BC_T}; + ndrange, +) where {DEST_T,BC_T} + key = (obj, DEST_T, BC_T, ndrange) + lock(_BCAST_PTX_CACHE_LOCK) do + # Also stores in cache if not found in Dict + return get!(_BCAST_PTX_CACHE, key) do + ptx, blocks, threads = get_ptx(obj, DEST_T, BC_T; ndrange=ndrange) + func_name = extract_kernel_name(_ptx) + ptx_task(ptx, func_name) + (CUDATask(func_name, (DEST_T, BC_T)), blocks, threads) + end + end +end + +function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcast.Broadcasted} + bc = Base.Broadcast.preprocess(dest, bc) + + # Get proper kernel + broadcast_kernel = + if ndims(dest) == 1 || + (isa(IndexStyle(dest), IndexLinear) && + isa(IndexStyle(bc), IndexLinear)) + GPU_LINEAR_KERNEL + else + GPU_CARTESIAN_KERNEL + end + + ndrange = ndims(dest) > 0 ? size(dest) : (1,) + + # Create "fake" type signatures to compile kernel with CUDA.jl infrastructure + DEST_T = map_cuda_type(typeof(dest)) + BC_T = map_cuda_type(B) + + # Lookup in cache, if not found, compile and cache + cuda_task, threads, blocks = get_cuda_task(broadcast_kernel, DEST_T, BC_T; ndrange=ndrange) + # TODO: register with Legate runtime like @cuda_task does + + # https://github.com/JuliaGPU/CUDA.jl/blob/345c1600ebd561135148bb04ee2657f521a40e25/CUDACore/src/CUDAKernels.jl#L111 + + #! DO I NEED TO DO TYPE PROMOTION CHECKS?? + + return dest +end From 1daf211439f30c336ec9e45dc7e012ab08f55af3 Mon Sep 17 00:00:00 2001 From: ejmeitz Date: Tue, 28 Apr 2026 15:01:59 -0500 Subject: [PATCH 06/12] add PTX broadcast task --- lib/cunumeric_jl_wrapper/include/ufi.h | 9 ++ lib/cunumeric_jl_wrapper/src/cuda.cpp | 141 +++++++++++++++++++++++ lib/cunumeric_jl_wrapper/src/wrapper.cpp | 1 + src/cuda/cuda_ptx_task.jl | 68 ++++++++++- src/cuda/cuda_util.jl | 4 + src/ndarray/broadcast_fusion.jl | 86 +++++++++++--- 6 files changed, 289 insertions(+), 20 deletions(-) diff --git a/lib/cunumeric_jl_wrapper/include/ufi.h b/lib/cunumeric_jl_wrapper/include/ufi.h index f5558c16..6a2f04f1 100644 --- a/lib/cunumeric_jl_wrapper/include/ufi.h +++ b/lib/cunumeric_jl_wrapper/include/ufi.h @@ -28,6 +28,7 @@ namespace ufi { enum TaskIDs { LOAD_PTX_TASK = 143432, RUN_PTX_TASK = 143433, + RUN_PTX_BROADCAST_TASK = 143434, }; class LoadPTXTask : public legate::LegateTask { @@ -46,6 +47,14 @@ class RunPTXTask : public legate::LegateTask { static void gpu_variant(legate::TaskContext context); }; +class RunPTXBroadcastTask : public legate::LegateTask { + public: + static inline const auto TASK_CONFIG = + legate::TaskConfig{legate::LocalTaskID{ufi::RUN_PTX_BROADCAST_TASK}}; + + static void gpu_variant(legate::TaskContext context); +}; + } // namespace ufi void wrap_cuda_methods(jlcxx::Module& mod); diff --git a/lib/cunumeric_jl_wrapper/src/cuda.cpp b/lib/cunumeric_jl_wrapper/src/cuda.cpp index d5e779d0..67baf4c4 100644 --- a/lib/cunumeric_jl_wrapper/src/cuda.cpp +++ b/lib/cunumeric_jl_wrapper/src/cuda.cpp @@ -34,6 +34,11 @@ #define BLOCK_START 1 #define THREAD_START 4 #define ARG_OFFSET 7 +#define PREFIX_SCALAR_COUNT_INDEX 7 +#define KERNEL_INPUT_ARG_COUNT_INDEX 8 +#define KERNEL_OUTPUT_ARG_COUNT_INDEX 9 +#define BC_NDARRAY_PATCH_COUNT_INDEX 10 +#define BCAST_ARG_OFFSET 11 // global padding for CUDA.jl kernel state std::size_t padded_bytes_kernel_state = 16; @@ -175,6 +180,28 @@ struct ufiFunctor { } }; +inline void append_array_args(char *&p, const legate::TaskContext &context, + std::size_t count, bool is_input) { + for (std::size_t i = 0; i < count; ++i) { + auto ps = is_input ? context.input(i) : context.output(i); + auto code = ps.type().code(); + auto dim = ps.dim(); + legate::double_dispatch( + dim, code, ufiFunctor{}, + is_input ? ufi::AccessMode::READ : ufi::AccessMode::WRITE, p, ps); + } +} + +inline void write_array_descriptor(char *dst, const legate::PhysicalArray &ps, + bool is_input) { + char *p = dst; + auto code = ps.type().code(); + auto dim = ps.dim(); + legate::double_dispatch( + dim, code, ufiFunctor{}, + is_input ? ufi::AccessMode::READ : ufi::AccessMode::WRITE, p, ps); +} + // https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/eval_udf_gpu.cc /*static*/ void RunPTXTask::gpu_variant(legate::TaskContext context) { cudaStream_t stream_ = context.get_task_stream(); @@ -290,6 +317,118 @@ struct ufiFunctor { // DRIVER_ERROR_CHECK(cuStreamSynchronize(stream_)); } +/*static*/ void RunPTXBroadcastTask::gpu_variant(legate::TaskContext context) { + cudaStream_t stream_ = context.get_task_stream(); + std::string kernel_name = context.scalar(0).value(); // 0 + + std::uint32_t bx = + context.scalar(BLOCK_START + 0).value(); // 1 + std::uint32_t by = + context.scalar(BLOCK_START + 1).value(); // 2 + std::uint32_t bz = + context.scalar(BLOCK_START + 2).value(); // 3 + + std::uint32_t tx = + context.scalar(THREAD_START + 0).value(); // 4 + std::uint32_t ty = + context.scalar(THREAD_START + 1).value(); // 5 + std::uint32_t tz = + context.scalar(THREAD_START + 2).value(); // 6 + + std::uint32_t prefix_scalar_count = + context.scalar(PREFIX_SCALAR_COUNT_INDEX).value(); // 7 + std::uint32_t kernel_input_args_count = + context.scalar(KERNEL_INPUT_ARG_COUNT_INDEX).value(); // 8 + std::uint32_t kernel_output_args_count = + context.scalar(KERNEL_OUTPUT_ARG_COUNT_INDEX) + .value(); // 9 + std::uint32_t bc_ndarray_patch_count = + context.scalar(BC_NDARRAY_PATCH_COUNT_INDEX) + .value(); // 10 + + const std::size_t num_inputs = context.num_inputs(); + const std::size_t num_outputs = context.num_outputs(); + const std::size_t num_scalars = context.num_scalars(); + + assert(kernel_input_args_count <= num_inputs); + assert(kernel_output_args_count <= num_outputs); + assert(bc_ndarray_patch_count <= num_inputs); + assert(BCAST_ARG_OFFSET + bc_ndarray_patch_count + prefix_scalar_count <= + num_scalars); + + CUcontext ctx; + cuStreamGetCtx(stream_, &ctx); + + FunctionKey key = {ctx, kernel_name}; + assert(cufunction_ptr.has_value()); + FunctionMap &fmap = cufunction_ptr.get(); + auto it = fmap.find(key); + assert(it != fmap.end()); + CUfunction func = it->second; + + std::size_t max_buffer_size = + padded_bytes_kernel_state + + (kernel_input_args_count + kernel_output_args_count) * + sizeof(CuDeviceArray); + for (std::size_t i = BCAST_ARG_OFFSET + bc_ndarray_patch_count; + i < num_scalars; ++i) { + max_buffer_size += context.scalar(i).size(); + } + + std::vector arg_buffer(max_buffer_size); + char *p = arg_buffer.data() + padded_bytes_kernel_state; + + const std::size_t patch_offsets_start = BCAST_ARG_OFFSET; + const std::size_t prefix_start = patch_offsets_start + bc_ndarray_patch_count; + const std::size_t prefix_end = prefix_start + prefix_scalar_count; + + for (std::size_t i = prefix_start; i < prefix_end; ++i) { + const auto &scalar = context.scalar(i); + memcpy(p, scalar.ptr(), scalar.size()); + p += scalar.size(); + } + + append_array_args(p, context, kernel_input_args_count, true); + append_array_args(p, context, kernel_output_args_count, false); + + for (std::size_t i = prefix_end; i < num_scalars; ++i) { + const auto &scalar = context.scalar(i); + if (i == prefix_end && bc_ndarray_patch_count > 0) { + std::vector patched_bc(scalar.size()); + memcpy(patched_bc.data(), scalar.ptr(), scalar.size()); + for (std::size_t j = 0; j < bc_ndarray_patch_count; ++j) { + std::uint32_t offset = + context.scalar(patch_offsets_start + j).value(); + assert(offset <= scalar.size()); + char *begin = patched_bc.data() + offset; + char *end = begin; + write_array_descriptor(end, context.input(j), true); + assert(static_cast(end - patched_bc.data()) <= + scalar.size()); + } + memcpy(p, patched_bc.data(), patched_bc.size()); + p += patched_bc.size(); + } else { + memcpy(p, scalar.ptr(), scalar.size()); + p += scalar.size(); + } + } + + std::size_t buffer_size = p - arg_buffer.data(); + + void *config[] = { + CU_LAUNCH_PARAM_BUFFER_POINTER, + static_cast(arg_buffer.data()), + CU_LAUNCH_PARAM_BUFFER_SIZE, + &buffer_size, + CU_LAUNCH_PARAM_END, + }; + + CUstream custream_ = reinterpret_cast(stream_); + DRIVER_ERROR_CHECK(cuLaunchKernel(func, bx, by, bz, tx, ty, tz, 0, custream_, + nullptr, config)); +} + // https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/load_ptx.cc /*static*/ void LoadPTXTask::gpu_variant(legate::TaskContext context) { std::string ptx = context.scalar(0).value(); @@ -417,4 +556,6 @@ void wrap_cuda_methods(jlcxx::Module &mod) { mod.method("extract_kernel_name", &extract_kernel_name); mod.set_const("LOAD_PTX", legate::LocalTaskID{ufi::TaskIDs::LOAD_PTX_TASK}); mod.set_const("RUN_PTX", legate::LocalTaskID{ufi::TaskIDs::RUN_PTX_TASK}); + mod.set_const("RUN_PTX_BROADCAST", + legate::LocalTaskID{ufi::TaskIDs::RUN_PTX_BROADCAST_TASK}); } diff --git a/lib/cunumeric_jl_wrapper/src/wrapper.cpp b/lib/cunumeric_jl_wrapper/src/wrapper.cpp index 29334333..1c5ddb96 100644 --- a/lib/cunumeric_jl_wrapper/src/wrapper.cpp +++ b/lib/cunumeric_jl_wrapper/src/wrapper.cpp @@ -61,6 +61,7 @@ void register_tasks() { auto library = get_lib(); ufi::LoadPTXTask::register_variants(library); ufi::RunPTXTask::register_variants(library); + ufi::RunPTXBroadcastTask::register_variants(library); } #endif diff --git a/src/cuda/cuda_ptx_task.jl b/src/cuda/cuda_ptx_task.jl index 73231a1a..78a3ea7f 100644 --- a/src/cuda/cuda_ptx_task.jl +++ b/src/cuda/cuda_ptx_task.jl @@ -2,7 +2,7 @@ export @cuda_task, @launch, CUDATask struct CUDATask func::String - argtypes::NTuple{N,Type} where {N} + argtypes::NTuple{N,Type} where {N} #! THIS IS TYPE UNSTABLE end #! JUST PASS TYPES HERE INSTEAD OF CALLING typeof() @@ -118,7 +118,8 @@ function Launch(kernel::CUDATask, inputs::Tuple{Vararg{NDArray}}, Legate.submit_auto_task(rt, task) end -function launch(kernel::CUDATask, inputs, outputs, scalars; blocks, threads) +function launch(kernel::CUDATask, inputs, outputs, scalars; + blocks, threads) Launch(kernel, isa(inputs, Tuple) ? inputs : (inputs,), isa(outputs, Tuple) ? outputs : (outputs,), @@ -128,6 +129,66 @@ function launch(kernel::CUDATask, inputs, outputs, scalars; blocks, threads) ) end +function launch_broadcast(kernel::CUDATask, inputs, outputs, scalars; + blocks, threads, prefix_scalars=(), + kernel_input_args_count=0, + kernel_output_args_count=length(isa(outputs, Tuple) ? outputs : (outputs,)), + bc_ndarray_offsets=()) + input_tup = isa(inputs, Tuple) ? inputs : (inputs,) + output_tup = isa(outputs, Tuple) ? outputs : (outputs,) + scalar_tup = isa(scalars, Tuple) ? scalars : (scalars,) + prefix_tup = isa(prefix_scalars, Tuple) ? prefix_scalars : (prefix_scalars,) + offsets_tup = isa(bc_ndarray_offsets, Tuple) ? bc_ndarray_offsets : (bc_ndarray_offsets,) + + ndarrays = vcat(input_tup..., output_tup...) + mx = findmax(arr -> arr.nbytes, ndarrays) + max_shape = size(ndarrays[mx[2]]) + @assert !isnothing(max_shape) + + rt = Legate.get_runtime() + lib = cuNumeric.get_lib() + taskid = cuNumeric.RUN_PTX_BROADCAST + task = Legate.create_auto_task(rt, lib, taskid) + + input_vars = Vector{Legate.Variable}() + for arr in input_tup + check_sz!(arr, max_shape; copy=true) + la = nda_to_logical_array(arr) + p = Legate.add_input(task, la) + push!(input_vars, p) + end + + output_vars = Vector{Legate.Variable}() + for arr in output_tup + check_sz!(arr, max_shape; copy=false) + la = nda_to_logical_array(arr) + p = Legate.add_output(task, la) + push!(output_vars, p) + end + + Legate.add_scalar(task, Legate.string_to_scalar(kernel.func)) # 0 + cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, blocks)) # 1,2,3 + cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, threads)) # 4,5,6 + Legate.add_scalar(task, Legate.Scalar(UInt32(length(prefix_tup)))) # 7 + Legate.add_scalar(task, Legate.Scalar(UInt32(kernel_input_args_count))) # 8 + Legate.add_scalar(task, Legate.Scalar(UInt32(kernel_output_args_count))) # 9 + Legate.add_scalar(task, Legate.Scalar(UInt32(length(offsets_tup)))) # 10 + + for off in offsets_tup + Legate.add_scalar(task, Legate.Scalar(UInt32(off))) + end + + for s in prefix_tup + Legate.add_scalar(task, Legate.Scalar(s)) + end + for s in scalar_tup + Legate.add_scalar(task, Legate.Scalar(s)) + end + + Legate.default_alignment(task, input_vars, output_vars) + Legate.submit_auto_task(rt, task) +end + function ptx_task(ptx::String, kernel_name) rt = Legate.get_runtime() lib = cuNumeric.get_lib() # grab lib of legate app @@ -269,7 +330,8 @@ macro launch(args...) esc( quote cuNumeric.launch( - $task, $inputs, $outputs, $scalars; blocks=($blocks), threads=($threads) + $task, $inputs, $outputs, $scalars; + blocks=($blocks), threads=($threads), ) end, ) diff --git a/src/cuda/cuda_util.jl b/src/cuda/cuda_util.jl index 9485b205..936a4c40 100644 --- a/src/cuda/cuda_util.jl +++ b/src/cuda/cuda_util.jl @@ -39,3 +39,7 @@ end function map_cuda_type(::Type{Base.Broadcast.Broadcasted{S,Ax,F,Args}}) where {S,Ax,F,Args} return Base.Broadcast.Broadcasted{map_cuda_type(S),Ax,F,map_cuda_type(Args)} end + +function map_cuda_type(::Type{Base.Broadcast.Extruded{X,K,D}}) where {X,K,D} + return Base.Broadcast.Extruded{map_cuda_type(X),K,D} +end diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl index 4c9db998..94895d71 100644 --- a/src/ndarray/broadcast_fusion.jl +++ b/src/ndarray/broadcast_fusion.jl @@ -17,9 +17,36 @@ end const GPU_CARTESIAN_KERNEL = broadcast_kernel_cartesian(CUDACore.CUDAKernels.CUDABackend()) const GPU_LINEAR_KERNEL = broadcast_kernel_linear(CUDACore.CUDAKernels.CUDABackend()) -const _BCAST_PTX_CACHE = Dict{Tuple{Any,DataType,DataType,Any},Tuple{CUDATask,Int,Int}}() +struct FusedBroadcastMetadata + ctx::Any # KA.CompilerMetadata + threads::Int + blocks::Int + cuda_task::CUDATask +end + +const _BCAST_PTX_CACHE = Dict{Tuple{Any,DataType,DataType,Any},FusedBroadcastMetadata}() const _BCAST_PTX_CACHE_LOCK = ReentrantLock() +_isbits_size_str(::Type{T}) where {T} = isbitstype(T) ? string(sizeof(T)) : "n/a" + +function _collect_cudevicearray_offsets!(offsets::Vector{Int}, ::Type{T}, base::Int=0) where {T} + if T <: CUDACore.CuDeviceArray + push!(offsets, base) + return offsets + end + if isbitstype(T) && fieldcount(T) > 0 + for i in 1:fieldcount(T) + FT = fieldtype(T, i) + _collect_cudevicearray_offsets!(offsets, FT, base + Int(fieldoffset(T, i))) + end + end + return offsets +end + +function _collect_cudevicearray_offsets(::Type{T}) where {T} + return Tuple(_collect_cudevicearray_offsets!(Int[], T, 0)) +end + function _default_broadcast_threads(ndrange) n = prod(ndrange) return min(256, max(1, n)) @@ -83,33 +110,43 @@ function get_ptx( blocks = length(KA.blocks(iterspace)) threads = length(KA.workitems(iterspace)) - blocks == 0 && return "", 0, 0 + blocks == 0 && return "", 0, 0, ctx buf = IOBuffer() CUDATools.code_ptx(buf, obj.f, (typeof(ctx), DEST_T, BC_T); raw=false, kernel=true) - return String(take!(buf)), threads, blocks + return String(take!(buf)), threads, blocks, ctx end function get_cuda_task( obj::KA.Kernel{CUDACore.CUDAKernels.CUDABackend}, - ::Type{DEST_T}, - ::Type{BC_T}; + dest::D, + bc::B, ndrange, -) where {DEST_T,BC_T} - key = (obj, DEST_T, BC_T, ndrange) +) where {D<:NDArray,B<:Base.Broadcast.Broadcasted} + DEST_T = map_cuda_type(D) + BC_T = map_cuda_type(B) + + key = (obj, D, B, ndrange) lock(_BCAST_PTX_CACHE_LOCK) do # Also stores in cache if not found in Dict return get!(_BCAST_PTX_CACHE, key) do - ptx, blocks, threads = get_ptx(obj, DEST_T, BC_T; ndrange=ndrange) - func_name = extract_kernel_name(_ptx) + ptx, threads, blocks, ctx = get_ptx(obj, DEST_T, BC_T; ndrange=ndrange) + func_name = extract_kernel_name(ptx) + # println(ptx) ptx_task(ptx, func_name) - (CUDATask(func_name, (DEST_T, BC_T)), blocks, threads) + cuda_task = CUDATask(func_name, (DEST_T, BC_T)) + FusedBroadcastMetadata(ctx, threads, blocks, cuda_task) end end end function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcast.Broadcasted} + + #! HOW DOES THIS BEHAVE WHEN BC HAS 2 RESULT ARRAYS? + bc = Base.Broadcast.preprocess(dest, bc) + bc = Base.Broadcast.instantiate(bc) + bc = Base.Broadcast.flatten(bc) # Get proper kernel broadcast_kernel = @@ -123,15 +160,30 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas ndrange = ndims(dest) > 0 ? size(dest) : (1,) - # Create "fake" type signatures to compile kernel with CUDA.jl infrastructure - DEST_T = map_cuda_type(typeof(dest)) - BC_T = map_cuda_type(B) - # Lookup in cache, if not found, compile and cache - cuda_task, threads, blocks = get_cuda_task(broadcast_kernel, DEST_T, BC_T; ndrange=ndrange) - # TODO: register with Legate runtime like @cuda_task does + fused_kernel_metadata = get_cuda_task(broadcast_kernel, dest, bc, ndrange) + input_deps = Tuple(arg for arg in bc.args if arg isa NDArray) + bc_gpu = CUDACore.cudaconvert(bc) + bc_ndarray_offsets = _collect_cudevicearray_offsets(typeof(bc_gpu)) + @assert length(bc_ndarray_offsets) == length(input_deps) + + launch_broadcast( + fused_kernel_metadata.cuda_task, + input_deps, + (dest,), + (bc_gpu,); + blocks=(fused_kernel_metadata.blocks,), + threads=(fused_kernel_metadata.threads,), + prefix_scalars=(fused_kernel_metadata.ctx,), + kernel_input_args_count=0, + kernel_output_args_count=1, + bc_ndarray_offsets=bc_ndarray_offsets, + ) - # https://github.com/JuliaGPU/CUDA.jl/blob/345c1600ebd561135148bb04ee2657f521a40e25/CUDACore/src/CUDAKernels.jl#L111 + #! DOUBLE CHECK bc.args ACTAULLY GIVES BACK ARRAYS + #! HOW TO GET SCALARS??? + # scalars = ???? + # launch(cuda_task, bc.args, (dest,), scalars; blocks, threads) #! DO I NEED TO DO TYPE PROMOTION CHECKS?? From 6bdec7acbfd52d5c08bf08fd193d0e419251df34 Mon Sep 17 00:00:00 2001 From: ejmeitz Date: Tue, 28 Apr 2026 15:13:04 -0500 Subject: [PATCH 07/12] clean up cuda code a bit --- .../include/cuda_macros.h | 74 +++++ lib/cunumeric_jl_wrapper/src/cuda.cpp | 271 +++++------------- 2 files changed, 152 insertions(+), 193 deletions(-) create mode 100644 lib/cunumeric_jl_wrapper/include/cuda_macros.h diff --git a/lib/cunumeric_jl_wrapper/include/cuda_macros.h b/lib/cunumeric_jl_wrapper/include/cuda_macros.h new file mode 100644 index 00000000..de9a9f2f --- /dev/null +++ b/lib/cunumeric_jl_wrapper/include/cuda_macros.h @@ -0,0 +1,74 @@ +#pragma once + +#define ERROR_CHECK(x) \ + { \ + cudaError_t status = x; \ + if (status != cudaSuccess) { \ + fprintf(stderr, "CUDA Error at %s:%d: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(status)); \ + if (stream_) cudaStreamDestroy(stream_); \ + exit(-1); \ + } \ + } + +#define DRIVER_ERROR_CHECK(x) \ + { \ + CUresult status = x; \ + if (status != CUDA_SUCCESS) { \ + const char *err_str = nullptr; \ + cuGetErrorString(status, &err_str); \ + fprintf(stderr, "CUDA Driver Error at %s:%d: %s\n", __FILE__, __LINE__, \ + err_str); \ + if (stream_) cudaStreamDestroy(stream_); \ + exit(-1); \ + } \ + } + +#define TEST_PRINT_DEBUG(dev_ptr, N, T, format, stream, message) \ + { \ + std::vector host_arr(N); \ + ERROR_CHECK(cudaMemcpy(host_arr.data(), \ + reinterpret_cast(dev_ptr), \ + sizeof(T) * N, cudaMemcpyDeviceToHost)); \ + ERROR_CHECK(cudaStreamSynchronize(stream)); \ + fprintf(stderr, "[TEST_PRINT] %s: " format "\n", message, host_arr[0]); \ + } + +#ifdef CUDA_DEBUG +#define CUDA_DEBUG_PRINT(x) \ + do { \ + x; \ + } while (0) +#else +#define CUDA_DEBUG_PRINT(x) \ + do { \ + } while (0) +#endif + +#define CUDA_DEVICE_ARRAY_ARG(MODE, ACCESSOR_CALL) \ + template < \ + typename T, int D, \ + typename std::enable_if<(D >= 1 && D <= REALM_MAX_DIM), int>::type = 0> \ + void cuda_device_array_arg_##MODE(char *&p, \ + const legate::PhysicalArray &rf) { \ + auto shp = rf.shape(); \ + auto acc = rf.data().ACCESSOR_CALL(); \ + CUDA_DEBUG_PRINT(std::cerr << "[RunPTXTask] " #MODE " accessor shape: " \ + << shp.lo << " - " << shp.hi << ", dim: " << D \ + << std::endl; \ + std::cerr << "[RunPTXTask] " #MODE " accessor strides: " \ + << acc.accessor.strides << std::endl;); \ + void *dev_ptr = const_cast(/*.lo to ensure multiple GPU support*/ \ + static_cast( \ + acc.ptr(Realm::Point(shp.lo)))); \ + auto extents = shp.hi - shp.lo + legate::Point::ONES(); \ + CuDeviceArray desc; \ + desc.ptr = dev_ptr; \ + desc.maxsize = shp.volume() * sizeof(T); \ + for (size_t i = 0; i < D; ++i) { \ + desc.dims[i] = extents[i]; \ + } \ + desc.length = shp.volume(); \ + memcpy(p, &desc, sizeof(CuDeviceArray)); \ + p += sizeof(CuDeviceArray); \ + } diff --git a/lib/cunumeric_jl_wrapper/src/cuda.cpp b/lib/cunumeric_jl_wrapper/src/cuda.cpp index 67baf4c4..9ff15dc8 100644 --- a/lib/cunumeric_jl_wrapper/src/cuda.cpp +++ b/lib/cunumeric_jl_wrapper/src/cuda.cpp @@ -23,6 +23,7 @@ #include #include +#include "cuda_macros.h" #include "legate.h" #include "legate/utilities/proc_local_storage.h" #include "legion.h" @@ -43,51 +44,6 @@ // global padding for CUDA.jl kernel state std::size_t padded_bytes_kernel_state = 16; -#define ERROR_CHECK(x) \ - { \ - cudaError_t status = x; \ - if (status != cudaSuccess) { \ - fprintf(stderr, "CUDA Error at %s:%d: %s\n", __FILE__, __LINE__, \ - cudaGetErrorString(status)); \ - if (stream_) cudaStreamDestroy(stream_); \ - exit(-1); \ - } \ - } - -#define DRIVER_ERROR_CHECK(x) \ - { \ - CUresult status = x; \ - if (status != CUDA_SUCCESS) { \ - const char *err_str = nullptr; \ - cuGetErrorString(status, &err_str); \ - fprintf(stderr, "CUDA Driver Error at %s:%d: %s\n", __FILE__, __LINE__, \ - err_str); \ - if (stream_) cudaStreamDestroy(stream_); \ - exit(-1); \ - } \ - } - -#define TEST_PRINT_DEBUG(dev_ptr, N, T, format, stream, message) \ - { \ - std::vector host_arr(N); \ - ERROR_CHECK(cudaMemcpy(host_arr.data(), \ - reinterpret_cast(dev_ptr), \ - sizeof(T) * N, cudaMemcpyDeviceToHost)); \ - ERROR_CHECK(cudaStreamSynchronize(stream)); \ - fprintf(stderr, "[TEST_PRINT] %s: " format "\n", message, host_arr[0]); \ - } - -#ifdef CUDA_DEBUG -#define CUDA_DEBUG_PRINT(x) \ - do { \ - x; \ - } while (0) -#else -#define CUDA_DEBUG_PRINT(x) \ - do { \ - } while (0) -#endif - namespace ufi { using namespace Legion; // TODO CUcontext key hashing is redundant. ProcLocalStorage is local to the @@ -138,34 +94,6 @@ struct CuDeviceArray { uint64_t length; // Number of elements (at the end) }; -#define CUDA_DEVICE_ARRAY_ARG(MODE, ACCESSOR_CALL) \ - template < \ - typename T, int D, \ - typename std::enable_if<(D >= 1 && D <= REALM_MAX_DIM), int>::type = 0> \ - void cuda_device_array_arg_##MODE(char *&p, \ - const legate::PhysicalArray &rf) { \ - auto shp = rf.shape(); \ - auto acc = rf.data().ACCESSOR_CALL(); \ - CUDA_DEBUG_PRINT(std::cerr << "[RunPTXTask] " #MODE " accessor shape: " \ - << shp.lo << " - " << shp.hi << ", dim: " << D \ - << std::endl; \ - std::cerr << "[RunPTXTask] " #MODE " accessor strides: " \ - << acc.accessor.strides << std::endl;); \ - void *dev_ptr = const_cast(/*.lo to ensure multiple GPU support*/ \ - static_cast( \ - acc.ptr(Realm::Point(shp.lo)))); \ - auto extents = shp.hi - shp.lo + legate::Point::ONES(); \ - CuDeviceArray desc; \ - desc.ptr = dev_ptr; \ - desc.maxsize = shp.volume() * sizeof(T); \ - for (size_t i = 0; i < D; ++i) { \ - desc.dims[i] = extents[i]; \ - } \ - desc.length = shp.volume(); \ - memcpy(p, &desc, sizeof(CuDeviceArray)); \ - p += sizeof(CuDeviceArray); \ - } - CUDA_DEVICE_ARRAY_ARG(read, read_accessor); // cuda_device_array_arg_read CUDA_DEVICE_ARRAY_ARG(write, write_accessor); // cuda_device_array_arg_write @@ -180,6 +108,45 @@ struct ufiFunctor { } }; +struct LaunchDims { + std::uint32_t bx, by, bz; + std::uint32_t tx, ty, tz; +}; + +inline LaunchDims read_launch_dims(const legate::TaskContext &context) { + return LaunchDims{ + context.scalar(BLOCK_START + 0).value(), + context.scalar(BLOCK_START + 1).value(), + context.scalar(BLOCK_START + 2).value(), + context.scalar(THREAD_START + 0).value(), + context.scalar(THREAD_START + 1).value(), + context.scalar(THREAD_START + 2).value(), + }; +} + +inline std::pair lookup_kernel_function( + const std::string &kernel_name, cudaStream_t stream_) { + CUcontext ctx; + cuStreamGetCtx(stream_, &ctx); + + FunctionKey key = {ctx, kernel_name}; + assert(cufunction_ptr.has_value()); + FunctionMap &fmap = cufunction_ptr.get(); + auto it = fmap.find(key); +#ifdef CUDA_DEBUG + if (it == fmap.end()) { + std::cerr << "[RunPTXTask] Could not find key: " << key_to_string(key) + << std::endl; + for (const auto &[k, v] : fmap) { + std::cerr << "[RunPTXTask] Map key: " << key_to_string(k) << std::endl; + } + assert(0 && "[RunPTXTask] key is not found in hashmap"); + } +#endif + assert(it != fmap.end()); + return {ctx, it->second}; +} + inline void append_array_args(char *&p, const legate::TaskContext &context, std::size_t count, bool is_input) { for (std::size_t i = 0; i < count; ++i) { @@ -202,58 +169,36 @@ inline void write_array_descriptor(char *dst, const legate::PhysicalArray &ps, is_input ? ufi::AccessMode::READ : ufi::AccessMode::WRITE, p, ps); } +inline void launch_with_buffer(CUfunction func, const LaunchDims &dims, + cudaStream_t stream_, + std::vector &arg_buffer, + std::size_t used_buffer_size) { + void *config[] = { + CU_LAUNCH_PARAM_BUFFER_POINTER, + static_cast(arg_buffer.data()), + CU_LAUNCH_PARAM_BUFFER_SIZE, + &used_buffer_size, + CU_LAUNCH_PARAM_END, + }; + + CUstream custream_ = reinterpret_cast(stream_); + DRIVER_ERROR_CHECK(cuLaunchKernel(func, dims.bx, dims.by, dims.bz, dims.tx, + dims.ty, dims.tz, 0, custream_, nullptr, + config)); +} + // https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/eval_udf_gpu.cc /*static*/ void RunPTXTask::gpu_variant(legate::TaskContext context) { cudaStream_t stream_ = context.get_task_stream(); std::string kernel_name = context.scalar(0).value(); // 0 - - std::uint32_t bx = - context.scalar(BLOCK_START + 0).value(); // 1 - std::uint32_t by = - context.scalar(BLOCK_START + 1).value(); // 2 - std::uint32_t bz = - context.scalar(BLOCK_START + 2).value(); // 3 - - std::uint32_t tx = - context.scalar(THREAD_START + 0).value(); // 4 - std::uint32_t ty = - context.scalar(THREAD_START + 1).value(); // 5 - std::uint32_t tz = - context.scalar(THREAD_START + 2).value(); // 6 - - CUcontext ctx; - cuStreamGetCtx(stream_, &ctx); - - FunctionKey key = {ctx, kernel_name}; - assert(cufunction_ptr.has_value()); - FunctionMap &fmap = cufunction_ptr.get(); - - auto it = fmap.find(key); - -#ifdef CUDA_DEBUG - if (it == fmap.end()) { - // for DEBUG output - std::cerr << "[RunPTXTask] Could not find key: " << key_to_string(key) - << std::endl; - for (const auto &[k, v] : fmap) { - std::cerr << "[RunPTXTask] Map key: " << key_to_string(k) << std::endl; - } - assert(0 && "[RunPTXTask] key is not found in hashmap"); - } -#endif - - assert(it != fmap.end()); - CUfunction func = it->second; + const LaunchDims dims = read_launch_dims(context); + auto kernel = lookup_kernel_function(kernel_name, stream_); + CUfunction func = kernel.second; const std::size_t num_inputs = context.num_inputs(); const std::size_t num_outputs = context.num_outputs(); const std::size_t num_scalars = context.num_scalars(); - const std::size_t num_reductions = - context.num_reductions(); // unused for now - - // compute total size: all device arrays + all scalars - // skip scalar 0-2 (kernel_name, threads, blocks) - // we allocate extra to decrease looping and dynamic dispatching on dim + // Layout: [kernel-state][inputs][outputs][user-scalars] std::size_t max_buffer_size = padded_bytes_kernel_state + (num_inputs + num_outputs) * sizeof(CuDeviceArray); @@ -263,29 +208,8 @@ inline void write_array_descriptor(char *dst, const legate::PhysicalArray &ps, std::vector arg_buffer(max_buffer_size); char *p = arg_buffer.data() + padded_bytes_kernel_state; - for (std::size_t i = 0; i < num_inputs; ++i) { - auto ps = context.input(i); - auto code = ps.type().code(); - auto dim = ps.dim(); -#ifdef CUDA_DEBUG - std::cerr << "[RunPTXTask] Input " << i << " type: " << code - << ", dim: " << dim << std::endl; -#endif - // dispatch on dim and code with ufiFunctor operator() - legate::double_dispatch(dim, code, ufiFunctor{}, ufi::AccessMode::READ, p, - ps); - } - for (std::size_t i = 0; i < num_outputs; ++i) { - auto ps = context.output(i); - auto code = ps.type().code(); - auto dim = ps.dim(); -#ifdef CUDA_DEBUG - std::cerr << "[RunPTXTask] Output " << i << " type: " << code - << ", dim: " << dim << std::endl; -#endif - legate::double_dispatch(dim, code, ufiFunctor{}, ufi::AccessMode::WRITE, p, - ps); - } + append_array_args(p, context, num_inputs, true); + append_array_args(p, context, num_outputs, false); for (std::size_t i = ARG_OFFSET; i < num_scalars; ++i) { const auto &scalar = context.scalar(i); memcpy(p, scalar.ptr(), scalar.size()); @@ -294,46 +218,24 @@ inline void write_array_descriptor(char *dst, const legate::PhysicalArray &ps, std::size_t buffer_size = p - arg_buffer.data(); // calc used buffer - void *config[] = { - CU_LAUNCH_PARAM_BUFFER_POINTER, - static_cast(arg_buffer.data()), - CU_LAUNCH_PARAM_BUFFER_SIZE, - &buffer_size, - CU_LAUNCH_PARAM_END, - }; - - CUstream custream_ = reinterpret_cast(stream_); - -#ifdef CUDA_DEBUG - std::cerr << "[RunPTXTask] Launching kernel " << kernel_name - << " with blocks (" << bx << "," << by << "," << bz - << ") and threads (" << tx << "," << ty << "," << tz << ")" - << " on CUcontext " << context_to_string(ctx) << std::endl; -#endif - // Launch the kernel - DRIVER_ERROR_CHECK(cuLaunchKernel(func, bx, by, bz, tx, ty, tz, 0, custream_, - nullptr, config)); + launch_with_buffer(func, dims, stream_, arg_buffer, buffer_size); // DRIVER_ERROR_CHECK(cuStreamSynchronize(stream_)); } +/* +This code is specifically for the fusion of broadcast operations. +This code uses the passed offsets to patch together the Broadcasted type +that the CUDA kernel expects. The Broadcasted type has all inputs in order, +but mixed scalars and arrays. e.g., [Array, Scalar, Array, Array]. +We must use the offsets to figure out which arguments should be marked as +scalars and which should be marked as inputs. +*/ + /*static*/ void RunPTXBroadcastTask::gpu_variant(legate::TaskContext context) { cudaStream_t stream_ = context.get_task_stream(); std::string kernel_name = context.scalar(0).value(); // 0 - - std::uint32_t bx = - context.scalar(BLOCK_START + 0).value(); // 1 - std::uint32_t by = - context.scalar(BLOCK_START + 1).value(); // 2 - std::uint32_t bz = - context.scalar(BLOCK_START + 2).value(); // 3 - - std::uint32_t tx = - context.scalar(THREAD_START + 0).value(); // 4 - std::uint32_t ty = - context.scalar(THREAD_START + 1).value(); // 5 - std::uint32_t tz = - context.scalar(THREAD_START + 2).value(); // 6 + const LaunchDims dims = read_launch_dims(context); std::uint32_t prefix_scalar_count = context.scalar(PREFIX_SCALAR_COUNT_INDEX).value(); // 7 @@ -356,15 +258,8 @@ inline void write_array_descriptor(char *dst, const legate::PhysicalArray &ps, assert(BCAST_ARG_OFFSET + bc_ndarray_patch_count + prefix_scalar_count <= num_scalars); - CUcontext ctx; - cuStreamGetCtx(stream_, &ctx); - - FunctionKey key = {ctx, kernel_name}; - assert(cufunction_ptr.has_value()); - FunctionMap &fmap = cufunction_ptr.get(); - auto it = fmap.find(key); - assert(it != fmap.end()); - CUfunction func = it->second; + auto kernel = lookup_kernel_function(kernel_name, stream_); + CUfunction func = kernel.second; std::size_t max_buffer_size = padded_bytes_kernel_state + @@ -416,17 +311,7 @@ inline void write_array_descriptor(char *dst, const legate::PhysicalArray &ps, std::size_t buffer_size = p - arg_buffer.data(); - void *config[] = { - CU_LAUNCH_PARAM_BUFFER_POINTER, - static_cast(arg_buffer.data()), - CU_LAUNCH_PARAM_BUFFER_SIZE, - &buffer_size, - CU_LAUNCH_PARAM_END, - }; - - CUstream custream_ = reinterpret_cast(stream_); - DRIVER_ERROR_CHECK(cuLaunchKernel(func, bx, by, bz, tx, ty, tz, 0, custream_, - nullptr, config)); + launch_with_buffer(func, dims, stream_, arg_buffer, buffer_size); } // https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/load_ptx.cc From c794e4ea6e3a367193d848a67a924c4f09de389a Mon Sep 17 00:00:00 2001 From: ejmeitz Date: Sat, 2 May 2026 08:11:36 -0500 Subject: [PATCH 08/12] dev build works --- Project.toml | 2 +- deps/build.jl | 95 +++++++++++++---- lib/cunumeric_jl_wrapper/CMakeLists.txt | 58 +++++++---- lib/cunumeric_jl_wrapper/src/cuda.cpp | 75 +++++++------ scripts/build_cpp_wrapper.sh | 85 +++++++++------ src/cuda/cuda_ptx_task.jl | 42 +++++--- src/ndarray/broadcast_fusion.jl | 133 +++++++++++++++--------- 7 files changed, 320 insertions(+), 170 deletions(-) diff --git a/Project.toml b/Project.toml index 458ff8da..fb5c3074 100644 --- a/Project.toml +++ b/Project.toml @@ -41,5 +41,5 @@ Preferences = "1" Random = "1" StatsBase = "0.34" cunumeric_jl_wrapper_jll = "25.10.3" -cupynumeric_jll = "25.10.2" +cupynumeric_jll = "25.10.3" julia = "1.10" diff --git a/deps/build.jl b/deps/build.jl index d53e576f..ae5e07a6 100644 --- a/deps/build.jl +++ b/deps/build.jl @@ -20,6 +20,11 @@ using Preferences using Legate using CNPreferences +using CUDACore: CUDACore + +# Maybe needed as build deps +using cupynumeric_jll: cupynumeric_jll +using OpenBLAS32_jll: OpenBLAS32_jll include("version.jl") @@ -84,7 +89,7 @@ function build_jlcxxwrap(repo_root, cupynumeric_root) end function build_cpp_wrapper( - repo_root, cupynumeric_loc, legate_loc, blas_loc, install_root + repo_root, cupynumeric_loc, legate_loc, blas_loc, install_root, cuda_header_loc, cuda_driver_loc ) @info "libcunumeric_jl_wrapper: Building C++ Wrapper Library" if isdir(install_root) @@ -95,7 +100,7 @@ function build_cpp_wrapper( build_cpp_wrapper = joinpath(repo_root, "scripts/build_cpp_wrapper.sh") nthreads = Threads.nthreads() - bld_command = `$build_cpp_wrapper $repo_root $cupynumeric_loc $legate_loc $blas_loc $install_root $nthreads` + bld_command = `$build_cpp_wrapper $repo_root $cupynumeric_loc $legate_loc $blas_loc $cuda_header_loc $cuda_driver_loc $install_root $nthreads` # write out a bash script for debugging cmd_str = join(bld_command.exec, " ") @@ -111,13 +116,6 @@ function build_cpp_wrapper( run_sh(`bash $bld_command`, "cpp_wrapper") end -function _find_jll_artifact_dir(jll) - eval(:(using $(jll))) - jll_mod = getfield(Main, jll) - root = jll_mod.artifact_dir - return root -end - function _start_build() pkg_root = up_dir(@__DIR__) deps_dir = joinpath(@__DIR__) @@ -131,10 +129,39 @@ function _start_build() return pkg_root end +function _cuda_include_dir_from_toolkit(cuda_toolkit_root) + cuda_include_dir = joinpath(cuda_toolkit_root, "include") + + isfile(joinpath(cuda_include_dir, "cuda.h")) || + error("Could not find cuda.h at $(joinpath(cuda_include_dir, "cuda.h"))") + + return cuda_include_dir +end + +function _cuda_driver_lib_from_toolkit(cuda_toolkit_root) + candidates = [ + joinpath(cuda_toolkit_root, "lib64", "stubs", "libcuda.so"), + joinpath(cuda_toolkit_root, "lib", "stubs", "libcuda.so"), + joinpath(cuda_toolkit_root, "lib64", "libcuda.so"), + joinpath(cuda_toolkit_root, "lib", "libcuda.so"), + ] + + for path in candidates + isfile(path) && return path + end + + error(""" + Could not find libcuda.so under CUDA_TOOLKIT_ROOT=$cuda_toolkit_root. + + Tried: + $(join(candidates, "\n")) + """) +end + """ build CxxWrap and cunumeric_jl_wrapper """ -function build_deps(pkg_root, cupynumeric_root, blas_root) +function build_deps(pkg_root, cupynumeric_root, blas_root, cuda_header_path, cuda_driver_path) legate_lib = Legate.get_install_liblegate() install_lib = joinpath(pkg_root, "lib", "cunumeric_jl_wrapper", "build") if !cupynumeric_valid(cupynumeric_root) @@ -147,7 +174,7 @@ function build_deps(pkg_root, cupynumeric_root, blas_root) build_jlcxxwrap(pkg_root, cupynumeric_root) build_cpp_wrapper( pkg_root, cupynumeric_root, up_dir(legate_lib), blas_root, - install_lib, + install_lib, cuda_header_path, cuda_driver_path, ) # $pkg_root/lib/cunumeric_jl_wrapper end @@ -161,33 +188,61 @@ function build(::CNPreferences.Conda) pkg_root = _start_build() cupynumeric_root = load_preference(CNPreferences, "cunumeric_conda_env", nothing) + cuda_toolkit_root = load_preference(CNPreferences, "CUDA_TOOLKIT_ROOT", nothing) if isnothing(cupynumeric_root) error("This shouldn't happen. cunumeric_conda_env = nothing?") end + if isnothing(cuda_toolkit_root) + error( + "CUDA_TOOLKIT_ROOT must be set by CNPreferences to point to the CUDA linked in your cupynumeric build." + ) + end + + #!TODO SET LocalPreferences.toml to use local CUDA libraries is_cupynumeric_installed(cupynumeric_root; throw_errors=true) - build_deps(pkg_root, cupynumeric_root, cupynumeric_root) # blas is same root as cupynumeric + build_deps(pkg_root, cupynumeric_root, cupynumeric_root, cuda_toolkit_root) # blas is same root as cupynumeric end function build(::CNPreferences.Developer) pkg_root = _start_build() - # can be nothing so this errors if not set cupynumeric_root = load_preference(CNPreferences, "cunumeric_path", nothing) - blas_lib = load_preference(CNPreferences, "BLAS_LIB", nothing) + blas_lib_dir = load_preference(CNPreferences, "BLAS_LIB_DIR", nothing) + cuda_toolkit_root = load_preference(CNPreferences, "CUDA_TOOLKIT_ROOT", nothing) + + cuda_header_path = nothing + cuda_driver_path = nothing if isnothing(cupynumeric_root) - # we are using cupynumeric_jll - cupynumeric_root = _find_jll_artifact_dir(:cupynumeric_jll) + # We are using cupynumeric_jll. + cupynumeric_root = cupynumeric_jll.artifact_dir + cuda_header_path = dirname(cupynumeric_jll.cuda_header) + cuda_driver_path = joinpath(CUDACore.CUDA_Driver_jll.artifact_dir, "lib", "libcuda.so") else - # this means we have a custom path set + # User provided a custom cupynumeric install. is_cupynumeric_installed(cupynumeric_root; throw_errors=true) + + isnothing(cuda_toolkit_root) && + error( + "CUDA_TOOLKIT_ROOT must be set by CNPreferences when not using cupynumeric_jll in developer mode" + ) + + # Local full CUDA toolkit path. + cuda_header_path = _cuda_include_dir_from_toolkit(cuda_toolkit_root) + cuda_driver_path = _cuda_driver_lib_from_toolkit(cuda_toolkit_root) end - if isnothing(blas_lib) - blas_lib = _find_jll_artifact_dir(:OpenBLAS32_jll) + if isnothing(blas_lib_dir) + blas_lib_dir = joinpath(OpenBLAS32_jll.artifact_dir, "lib") end - build_deps(pkg_root, cupynumeric_root, up_dir(blas_lib)) + build_deps( + pkg_root, + cupynumeric_root, + blas_lib_dir, + cuda_header_path, + cuda_driver_path, + ) end const mode_str = load_preference(CNPreferences, "cunumeric_mode", CNPreferences.MODE_JLL) diff --git a/lib/cunumeric_jl_wrapper/CMakeLists.txt b/lib/cunumeric_jl_wrapper/CMakeLists.txt index 2da543ab..56261c3d 100644 --- a/lib/cunumeric_jl_wrapper/CMakeLists.txt +++ b/lib/cunumeric_jl_wrapper/CMakeLists.txt @@ -3,6 +3,7 @@ project(cuNumericWrapper) set(cuNumericWrapperVersion 0.0.1) message(STATUS "Project version: v${cuNumericWrapperVersion}") + set(CXX_CUNUMERICJL_WRAPPER cunumeric_jl_wrapper) set(C_INTERFACE_LIB cunumeric_c_wrapper) @@ -11,30 +12,30 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib") -# ---- New: NOCUDA option ---- -option(NOCUDA "Build without CUDA support (skip CUDAToolkit and src/cuda.cpp)" OFF) +option(NOCUDA "Build without CUDA support (skip src/cuda.cpp)" OFF) option(BINARYBUILDER "Building with binary builder" ON) - -# Always needed (unless your packages themselves require CUDA at configure time) +# Always needed find_package(legate REQUIRED) find_package(cupynumeric REQUIRED) # CxxWrap Stuff if(NOT BINARYBUILDER) -execute_process( - COMMAND julia -e "println(DEPOT_PATH[1])" - OUTPUT_VARIABLE JULIA_DEP_PATH - OUTPUT_STRIP_TRAILING_WHITESPACE -) - -set(JlCxx_DIR "${JULIA_DEP_PATH}/dev/libcxxwrap_julia_jll/override") -message(STATUS "Setting JlCxx_DIR to ${JlCxx_DIR} (not using BinaryBuilder)") + execute_process( + COMMAND julia -e "println(DEPOT_PATH[1])" + OUTPUT_VARIABLE JULIA_DEP_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + + set(JlCxx_DIR "${JULIA_DEP_PATH}/dev/libcxxwrap_julia_jll/override") + message(STATUS "Setting JlCxx_DIR to ${JlCxx_DIR} (not using BinaryBuilder)") endif() find_package(JlCxx REQUIRED) + get_target_property(JlCxx_location JlCxx::cxxwrap_julia LOCATION) get_filename_component(JlCxx_location ${JlCxx_location} DIRECTORY) + set(CMAKE_INSTALL_RPATH "${CMAKE_INSTALL_PREFIX}/lib;${JlCxx_location}") message(STATUS "Found JlCxx at ${JlCxx_location}") @@ -44,21 +45,34 @@ set(SOURCES src/types.cpp ) -# Conditionally add CUDA bits +# CUDA Driver API only: requires cuda.h and libcuda.so, not a full CUDA Toolkit. if(NOT NOCUDA) - find_package(CUDAToolkit 13.0 REQUIRED) + set(CUDA_DRIVER_INCLUDE_DIR "" CACHE PATH "Directory containing cuda.h") + set(CUDA_DRIVER_LIBRARY "" CACHE FILEPATH "Full path to libcuda.so") + + add_library(CUDA_driver_manual UNKNOWN IMPORTED) + set_target_properties(CUDA_driver_manual PROPERTIES + IMPORTED_LOCATION "${CUDA_DRIVER_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${CUDA_DRIVER_INCLUDE_DIR}" + ) + list(APPEND SOURCES src/cuda.cpp) add_compile_definitions(HAVE_CUDA) set(HAVE_CUDA TRUE) + message(STATUS "CUDA enabled: adding src/cuda.cpp") + message(STATUS "CUDA driver include dir: ${CUDA_DRIVER_INCLUDE_DIR}") + message(STATUS "CUDA driver library: ${CUDA_DRIVER_LIBRARY}") else() set(HAVE_CUDA FALSE) - message(STATUS "NOCUDA=ON: skipping CUDAToolkit and src/cuda.cpp") + message(STATUS "NOCUDA=ON: skipping src/cuda.cpp") endif() # Library: C++ wrapper add_library(${CXX_CUNUMERICJL_WRAPPER} SHARED ${SOURCES}) -set_target_properties(${CXX_CUNUMERICJL_WRAPPER} PROPERTIES VERSION ${cuNumericWrapperVersion}) +set_target_properties(${CXX_CUNUMERICJL_WRAPPER} PROPERTIES + VERSION ${cuNumericWrapperVersion} +) target_link_libraries(${CXX_CUNUMERICJL_WRAPPER} PRIVATE cupynumeric::cupynumeric @@ -67,13 +81,12 @@ target_link_libraries(${CXX_CUNUMERICJL_WRAPPER} PRIVATE JlCxx::cxxwrap_julia_stl ) -# Include dirs (conditionally add CUDA include path) target_include_directories(${CXX_CUNUMERICJL_WRAPPER} PRIVATE include) + if(HAVE_CUDA) - target_include_directories(${CXX_CUNUMERICJL_WRAPPER} PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + target_link_libraries(${CXX_CUNUMERICJL_WRAPPER} PRIVATE CUDA_driver_manual) endif() - install(TARGETS ${CXX_CUNUMERICJL_WRAPPER} DESTINATION lib) # ---- C API ---- @@ -83,7 +96,9 @@ set(C_SOURCES ) add_library(${C_INTERFACE_LIB} SHARED ${C_SOURCES}) -set_target_properties(${C_INTERFACE_LIB} PROPERTIES VERSION ${cuNumericWrapperVersion}) +set_target_properties(${C_INTERFACE_LIB} PROPERTIES + VERSION ${cuNumericWrapperVersion} +) target_link_libraries(${C_INTERFACE_LIB} PRIVATE cupynumeric::cupynumeric @@ -91,8 +106,9 @@ target_link_libraries(${C_INTERFACE_LIB} PRIVATE ) target_include_directories(${C_INTERFACE_LIB} PRIVATE include) + if(HAVE_CUDA) - target_include_directories(${C_INTERFACE_LIB} PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + target_link_libraries(${C_INTERFACE_LIB} PRIVATE CUDA_driver_manual) endif() install(TARGETS ${C_INTERFACE_LIB} DESTINATION lib) diff --git a/lib/cunumeric_jl_wrapper/src/cuda.cpp b/lib/cunumeric_jl_wrapper/src/cuda.cpp index 9ff15dc8..0350cb28 100644 --- a/lib/cunumeric_jl_wrapper/src/cuda.cpp +++ b/lib/cunumeric_jl_wrapper/src/cuda.cpp @@ -39,7 +39,9 @@ #define KERNEL_INPUT_ARG_COUNT_INDEX 8 #define KERNEL_OUTPUT_ARG_COUNT_INDEX 9 #define BC_NDARRAY_PATCH_COUNT_INDEX 10 -#define BCAST_ARG_OFFSET 11 +#define BC_SCALAR_PATCH_COUNT_INDEX 11 +#define BC_BLOB_SIZE_INDEX 12 +#define BCAST_ARG_OFFSET 13 // global padding for CUDA.jl kernel state std::size_t padded_bytes_kernel_state = 16; @@ -247,6 +249,10 @@ scalars and which should be marked as inputs. std::uint32_t bc_ndarray_patch_count = context.scalar(BC_NDARRAY_PATCH_COUNT_INDEX) .value(); // 10 + std::uint32_t bc_scalar_patch_count = + context.scalar(BC_SCALAR_PATCH_COUNT_INDEX).value(); // 11 + std::uint32_t bc_blob_size = + context.scalar(BC_BLOB_SIZE_INDEX).value(); // 12 const std::size_t num_inputs = context.num_inputs(); const std::size_t num_outputs = context.num_outputs(); @@ -254,9 +260,14 @@ scalars and which should be marked as inputs. assert(kernel_input_args_count <= num_inputs); assert(kernel_output_args_count <= num_outputs); - assert(bc_ndarray_patch_count <= num_inputs); - assert(BCAST_ARG_OFFSET + bc_ndarray_patch_count + prefix_scalar_count <= - num_scalars); + const std::size_t array_patch_start = BCAST_ARG_OFFSET; + const std::size_t scalar_patch_start = + array_patch_start + 2 * bc_ndarray_patch_count; + const std::size_t prefix_start = scalar_patch_start + bc_scalar_patch_count; + const std::size_t prefix_end = prefix_start + prefix_scalar_count; + const std::size_t scalar_values_start = prefix_end; + + assert(scalar_values_start + bc_scalar_patch_count <= num_scalars); auto kernel = lookup_kernel_function(kernel_name, stream_); CUfunction func = kernel.second; @@ -264,19 +275,15 @@ scalars and which should be marked as inputs. std::size_t max_buffer_size = padded_bytes_kernel_state + (kernel_input_args_count + kernel_output_args_count) * - sizeof(CuDeviceArray); - for (std::size_t i = BCAST_ARG_OFFSET + bc_ndarray_patch_count; - i < num_scalars; ++i) { + sizeof(CuDeviceArray) + + bc_blob_size; + for (std::size_t i = prefix_start; i < prefix_end; ++i) { max_buffer_size += context.scalar(i).size(); } std::vector arg_buffer(max_buffer_size); char *p = arg_buffer.data() + padded_bytes_kernel_state; - const std::size_t patch_offsets_start = BCAST_ARG_OFFSET; - const std::size_t prefix_start = patch_offsets_start + bc_ndarray_patch_count; - const std::size_t prefix_end = prefix_start + prefix_scalar_count; - for (std::size_t i = prefix_start; i < prefix_end; ++i) { const auto &scalar = context.scalar(i); memcpy(p, scalar.ptr(), scalar.size()); @@ -286,29 +293,33 @@ scalars and which should be marked as inputs. append_array_args(p, context, kernel_input_args_count, true); append_array_args(p, context, kernel_output_args_count, false); - for (std::size_t i = prefix_end; i < num_scalars; ++i) { - const auto &scalar = context.scalar(i); - if (i == prefix_end && bc_ndarray_patch_count > 0) { - std::vector patched_bc(scalar.size()); - memcpy(patched_bc.data(), scalar.ptr(), scalar.size()); - for (std::size_t j = 0; j < bc_ndarray_patch_count; ++j) { - std::uint32_t offset = - context.scalar(patch_offsets_start + j).value(); - assert(offset <= scalar.size()); - char *begin = patched_bc.data() + offset; - char *end = begin; - write_array_descriptor(end, context.input(j), true); - assert(static_cast(end - patched_bc.data()) <= - scalar.size()); - } - memcpy(p, patched_bc.data(), patched_bc.size()); - p += patched_bc.size(); - } else { - memcpy(p, scalar.ptr(), scalar.size()); - p += scalar.size(); - } + std::vector patched_bc(bc_blob_size, 0); + + for (std::size_t j = 0; j < bc_ndarray_patch_count; ++j) { + std::uint32_t offset = + context.scalar(array_patch_start + 2 * j).value(); + std::uint32_t input_index = + context.scalar(array_patch_start + 2 * j + 1).value(); + assert(input_index < num_inputs); + assert(offset <= patched_bc.size()); + char *begin = patched_bc.data() + offset; + char *end = begin; + write_array_descriptor(end, context.input(input_index), true); + assert(static_cast(end - patched_bc.data()) <= + patched_bc.size()); } + for (std::size_t j = 0; j < bc_scalar_patch_count; ++j) { + std::uint32_t offset = + context.scalar(scalar_patch_start + j).value(); + const auto &scalar_value = context.scalar(scalar_values_start + j); + assert(offset + scalar_value.size() <= patched_bc.size()); + memcpy(patched_bc.data() + offset, scalar_value.ptr(), scalar_value.size()); + } + + memcpy(p, patched_bc.data(), patched_bc.size()); + p += patched_bc.size(); + std::size_t buffer_size = p - arg_buffer.data(); launch_with_buffer(func, dims, stream_, arg_buffer, buffer_size); diff --git a/scripts/build_cpp_wrapper.sh b/scripts/build_cpp_wrapper.sh index aae6f432..720779f2 100755 --- a/scripts/build_cpp_wrapper.sh +++ b/scripts/build_cpp_wrapper.sh @@ -1,18 +1,22 @@ -set -e +set -euo pipefail -# Check if exactly one argument is provided -if [[ $# -ne 6 ]]; then - echo "Usage: $0 " +# Default to OFF: CUDA support enabled. +# Set NO_CUDA=ON to skip CUDA. +NO_CUDA=${NO_CUDA:-OFF} + +if [[ $# -ne 8 ]]; then + echo "Usage: $0 " exit 1 fi -CUNUMERICJL_ROOT_DIR=$1 # this is the repo root of cunumeric.jl + +CUNUMERICJL_ROOT_DIR=$1 CUPYNUMERIC_ROOT_DIR=$2 LEGATE_ROOT_DIR=$3 BLAS_LIB_DIR=$4 -INSTALL_DIR=$5 -NTHREADS=$6 - -# Check if the provided argument is a valid directory +CUDA_DRIVER_INCLUDE_DIR=$5 +CUDA_DRIVER_LIBRARY=$6 +INSTALL_DIR=$7 +NTHREADS=$8 if [[ ! -d "$CUNUMERICJL_ROOT_DIR" ]]; then echo "Error: '$CUNUMERICJL_ROOT_DIR' is not a valid directory." @@ -29,35 +33,54 @@ if [[ ! -d "$LEGATE_ROOT_DIR" ]]; then exit 1 fi -CUNUMERIC_WRAPPER_SOURCE=$CUNUMERICJL_ROOT_DIR/lib/cunumeric_jl_wrapper -BUILD_DIR=$CUNUMERIC_WRAPPER_SOURCE/build +if [[ ! -d "$BLAS_LIB_DIR" ]]; then + echo "Error: '$BLAS_LIB_DIR' is not a valid directory." + exit 1 +fi -if [[ ! -d "$BUILD_DIR" ]]; then - mkdir -p $BUILD_DIR +if [[ ! -f "$BLAS_LIB_DIR/libopenblas.so" ]]; then + echo "Error: '$BLAS_LIB_DIR/libopenblas.so' does not exist." + exit 1 fi -if [[ ! -d "$INSTALL_DIR" ]]; then - mkdir -p $INSTALL_DIR +if [[ "$NO_CUDA" != "ON" ]]; then + if [[ ! -f "$CUDA_DRIVER_INCLUDE_DIR/cuda.h" ]]; then + echo "Error: '$CUDA_DRIVER_INCLUDE_DIR/cuda.h' does not exist." + exit 1 + fi + + if [[ ! -f "$CUDA_DRIVER_LIBRARY" ]]; then + echo "Error: '$CUDA_DRIVER_LIBRARY' does not exist." + exit 1 + fi fi -echo $LEGATE_ROOT_DIR +CUNUMERIC_WRAPPER_SOURCE="$CUNUMERICJL_ROOT_DIR/lib/cunumeric_jl_wrapper" +BUILD_DIR="$CUNUMERIC_WRAPPER_SOURCE/build" -# Default to OFF (CUDA support enabled), but allow override via environment variable -NO_CUDA=${NO_CUDA:-OFF} +mkdir -p "$BUILD_DIR" +mkdir -p "$INSTALL_DIR" -if [[ ! -f "$BUILD_DIR/CMakeCache.txt" ]]; then - echo "Configuring project..." - cmake -S "$CUNUMERIC_WRAPPER_SOURCE" -B "$BUILD_DIR" \ - -D BINARYBUILDER=OFF \ - -D NOCUDA=$NO_CUDA \ - -D CMAKE_INSTALL_PREFIX="$INSTALL_DIR" \ - -D CMAKE_PREFIX_PATH="$CUPYNUMERIC_ROOT_DIR;$LEGATE_ROOT_DIR;" \ - -D CUPYNUMERIC_PATH="$CUPYNUMERIC_ROOT_DIR" \ - -D BLAS_LIBRARIES="$BLAS_LIB_DIR/libopenblas.so" \ - -D PROJECT_INSTALL_PATH="$INSTALL_DIR" \ - -D CMAKE_BUILD_TYPE=Releases -else - echo "Skipping configure (already done in $BUILD_DIR)" +echo "LEGATE_ROOT_DIR: $LEGATE_ROOT_DIR" +echo "NO_CUDA: $NO_CUDA" + +if [[ "$NO_CUDA" != "ON" ]]; then + echo "CUDA driver include dir: $CUDA_DRIVER_INCLUDE_DIR" + echo "CUDA driver library: $CUDA_DRIVER_LIBRARY" fi +echo "Configuring project..." + +cmake -S "$CUNUMERIC_WRAPPER_SOURCE" -B "$BUILD_DIR" \ + -D BINARYBUILDER=OFF \ + -D NOCUDA="$NO_CUDA" \ + -D CUDA_DRIVER_INCLUDE_DIR="$CUDA_DRIVER_INCLUDE_DIR" \ + -D CUDA_DRIVER_LIBRARY="$CUDA_DRIVER_LIBRARY" \ + -D CMAKE_INSTALL_PREFIX="$INSTALL_DIR" \ + -D CMAKE_PREFIX_PATH="$CUPYNUMERIC_ROOT_DIR;$LEGATE_ROOT_DIR;" \ + -D CUPYNUMERIC_PATH="$CUPYNUMERIC_ROOT_DIR" \ + -D BLAS_LIBRARIES="$BLAS_LIB_DIR/libopenblas.so" \ + -D PROJECT_INSTALL_PATH="$INSTALL_DIR" \ + -D CMAKE_BUILD_TYPE=Release + cmake --build "$BUILD_DIR" --parallel "$NTHREADS" --verbose diff --git a/src/cuda/cuda_ptx_task.jl b/src/cuda/cuda_ptx_task.jl index 78a3ea7f..226e5233 100644 --- a/src/cuda/cuda_ptx_task.jl +++ b/src/cuda/cuda_ptx_task.jl @@ -129,20 +129,29 @@ function launch(kernel::CUDATask, inputs, outputs, scalars; ) end -function launch_broadcast(kernel::CUDATask, inputs, outputs, scalars; +struct BroadcastPatchInfo + broadcast_size::Int + inputs::Tuple + array_offsets::Tuple{Vararg{Int}} + array_input_indices::Tuple{Vararg{Int}} + scalar_offsets::Tuple{Vararg{Int}} + scalar_values::Tuple +end + +function launch_broadcast(kernel::CUDATask, outputs, patch_info::BroadcastPatchInfo; blocks, threads, prefix_scalars=(), kernel_input_args_count=0, - kernel_output_args_count=length(isa(outputs, Tuple) ? outputs : (outputs,)), - bc_ndarray_offsets=()) - input_tup = isa(inputs, Tuple) ? inputs : (inputs,) + kernel_output_args_count=length(isa(outputs, Tuple) ? outputs : (outputs,))) + input_tup = patch_info.inputs output_tup = isa(outputs, Tuple) ? outputs : (outputs,) - scalar_tup = isa(scalars, Tuple) ? scalars : (scalars,) + scalar_tup = patch_info.scalar_values prefix_tup = isa(prefix_scalars, Tuple) ? prefix_scalars : (prefix_scalars,) - offsets_tup = isa(bc_ndarray_offsets, Tuple) ? bc_ndarray_offsets : (bc_ndarray_offsets,) - ndarrays = vcat(input_tup..., output_tup...) - mx = findmax(arr -> arr.nbytes, ndarrays) - max_shape = size(ndarrays[mx[2]]) + length(patch_info.array_offsets) == length(patch_info.array_input_indices) || + throw(ArgumentError("Broadcast array patch offsets and input indices must match")) + + isempty(output_tup) && throw(ArgumentError("Broadcast launch requires an output array")) + max_shape = size(first(output_tup)) @assert !isnothing(max_shape) rt = Legate.get_runtime() @@ -152,7 +161,7 @@ function launch_broadcast(kernel::CUDATask, inputs, outputs, scalars; input_vars = Vector{Legate.Variable}() for arr in input_tup - check_sz!(arr, max_shape; copy=true) + check_sz(arr, max_shape) la = nda_to_logical_array(arr) p = Legate.add_input(task, la) push!(input_vars, p) @@ -160,7 +169,7 @@ function launch_broadcast(kernel::CUDATask, inputs, outputs, scalars; output_vars = Vector{Legate.Variable}() for arr in output_tup - check_sz!(arr, max_shape; copy=false) + check_sz(arr, max_shape) la = nda_to_logical_array(arr) p = Legate.add_output(task, la) push!(output_vars, p) @@ -172,9 +181,16 @@ function launch_broadcast(kernel::CUDATask, inputs, outputs, scalars; Legate.add_scalar(task, Legate.Scalar(UInt32(length(prefix_tup)))) # 7 Legate.add_scalar(task, Legate.Scalar(UInt32(kernel_input_args_count))) # 8 Legate.add_scalar(task, Legate.Scalar(UInt32(kernel_output_args_count))) # 9 - Legate.add_scalar(task, Legate.Scalar(UInt32(length(offsets_tup)))) # 10 + Legate.add_scalar(task, Legate.Scalar(UInt32(length(patch_info.array_offsets)))) # 10 + Legate.add_scalar(task, Legate.Scalar(UInt32(length(patch_info.scalar_offsets)))) # 11 + Legate.add_scalar(task, Legate.Scalar(UInt32(patch_info.broadcast_size))) # 12 + + for (off, input_index) in zip(patch_info.array_offsets, patch_info.array_input_indices) + Legate.add_scalar(task, Legate.Scalar(UInt32(off))) + Legate.add_scalar(task, Legate.Scalar(UInt32(input_index))) + end - for off in offsets_tup + for off in patch_info.scalar_offsets Legate.add_scalar(task, Legate.Scalar(UInt32(off))) end diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl index 94895d71..08da76fa 100644 --- a/src/ndarray/broadcast_fusion.jl +++ b/src/ndarray/broadcast_fusion.jl @@ -27,40 +27,45 @@ end const _BCAST_PTX_CACHE = Dict{Tuple{Any,DataType,DataType,Any},FusedBroadcastMetadata}() const _BCAST_PTX_CACHE_LOCK = ReentrantLock() -_isbits_size_str(::Type{T}) where {T} = isbitstype(T) ? string(sizeof(T)) : "n/a" +cudevice_array_offset(::Type{T}) where {T<:CUDACore.CuDeviceArray} = 0 +cudevice_array_offset(::Type{T}) where {T<:Base.Broadcast.Extruded} = Int(fieldoffset(T, 1)) + +stores_cudevicearray(::Type{T}) where {T<:CUDACore.CuDeviceArray} = true +stores_cudevicearray(::Type{T}) where {T<:Base.Broadcast.Extruded} = true +stores_cudevicearray(::Type{T}) where {T<:Number} = false +function stores_cudevicearray(::Type{T}) where {T} + throw(error("Broadcast fusion. Don't know what to do with type: $T")) +end -function _collect_cudevicearray_offsets!(offsets::Vector{Int}, ::Type{T}, base::Int=0) where {T} - if T <: CUDACore.CuDeviceArray - push!(offsets, base) - return offsets - end - if isbitstype(T) && fieldcount(T) > 0 - for i in 1:fieldcount(T) - FT = fieldtype(T, i) - _collect_cudevicearray_offsets!(offsets, FT, base + Int(fieldoffset(T, i))) +function find_cudevicearray_offsets_and_indices(::Type{BC_ARGS}) where {BC_ARGS} + offsets = Vector{Int}() + indices = Vector{Int}() + for (i, T) in enumerate(fieldtypes(BC_ARGS)) + if stores_cudevicearray(T) + offset = fieldoffset(BC_ARGS, i) + cudevice_array_offset(T) + push!(offsets, offset) + push!(indices, i) end end - return offsets + return tuple(offsets...), tuple(indices...) end -function _collect_cudevicearray_offsets(::Type{T}) where {T} - return Tuple(_collect_cudevicearray_offsets!(Int[], T, 0)) -end - -function _default_broadcast_threads(ndrange) - n = prod(ndrange) - return min(256, max(1, n)) +function find_scalar_offsets_and_indices(::Type{BC_ARGS}) where {BC_ARGS} + offsets = Vector{Int}() + indices = Vector{Int}() + for (i, T) in enumerate(fieldtypes(BC_ARGS)) + if (T <: Number) + offset = fieldoffset(BC_ARGS, i) + push!(offsets, offset) + push!(indices, i) + end + end + return tuple(offsets...), tuple(indices...) end -function _cufunction_from_types(f, types_tuple_type; maxthreads=nothing, always_inline=false) - return CUDACore.cufunction( - f, - types_tuple_type; - kernel=true, - maxthreads=maxthreads, - always_inline=always_inline, - ) -end +get_ndarray(x::T) where {T<:NDArray} = x +get_ndarray(x::T) where {T<:Base.Broadcast.Extruded} = x.x +get_ndarray(x) = throw(error("Broadcast fusion. Don't know what to do with type: $T")) """ get_ptx(obj::KA.Kernel{CUDABackend}, ::Type{DEST_T}, ::Type{BC_T}; @@ -101,9 +106,6 @@ function get_ptx( config = CUDACore.launch_configuration(host_kernel.fun; max_threads=prod(ndrange)) threads = config.threads - # if fancy thing doesnt work we can use this - # threads = _default_broadcast_threads(ndrange) - workgroupsize = CUDACore.CUDAKernels.threads_to_workgroupsize(threads, ndrange) iterspace, dynamic = KA.partition(obj, ndrange, workgroupsize) ctx = KA.mkcontext(obj, ndrange, iterspace) @@ -113,7 +115,7 @@ function get_ptx( blocks == 0 && return "", 0, 0, ctx buf = IOBuffer() - CUDATools.code_ptx(buf, obj.f, (typeof(ctx), DEST_T, BC_T); raw=false, kernel=true) + CUDATools.code_ptx(buf, obj.f, (typeof(ctx), DEST_T, BC_T); raw=false, kernel=true, ptx=v"7.8") return String(take!(buf)), threads, blocks, ctx end @@ -162,30 +164,57 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas # Lookup in cache, if not found, compile and cache fused_kernel_metadata = get_cuda_task(broadcast_kernel, dest, bc, ndrange) - input_deps = Tuple(arg for arg in bc.args if arg isa NDArray) - bc_gpu = CUDACore.cudaconvert(bc) - bc_ndarray_offsets = _collect_cudevicearray_offsets(typeof(bc_gpu)) - @assert length(bc_ndarray_offsets) == length(input_deps) - - launch_broadcast( - fused_kernel_metadata.cuda_task, - input_deps, - (dest,), - (bc_gpu,); - blocks=(fused_kernel_metadata.blocks,), - threads=(fused_kernel_metadata.threads,), - prefix_scalars=(fused_kernel_metadata.ctx,), - kernel_input_args_count=0, - kernel_output_args_count=1, - bc_ndarray_offsets=bc_ndarray_offsets, + + # Replace NDArrays with CuDeviceArrays in the Broadcasted type so we can figure out bit-offsets + spoofed_bc_type = map_cuda_type(typeof(bc)) + fieldname(spoofed_bc_type, 3) == :args || + throw(ArgumentError("Broadcasted field 3 is not args. Failed to fuse broadcast.")) + args_offset = Int(fieldoffset(spoofed_bc_type, 3)) + + # Replace NDArrays with CuDeviceArrays in the Broadcasted type so we can figure out bit-offsets + spoofed_bc_args_type = map_cuda_type(typeof(bc.args)) + + #!TODO FIGURE OUT HOW TO HANDLE AXES + #! TODO FIGURE OUT IF ITS SAFE TO IGNORE FIRST TWO FIELDS OF BROADCASTED TYPE + + # STEP 1: Figure out bit-offsets for CuDeviceArrays and scalars in args of spoofed type. + # The spoofed type has the same fields and alignment that the PTX kernel expects. + cudevicearray_offsets, cudevicearray_indices = find_cudevicearray_offsets_and_indices( + spoofed_bc_args_type + ) + scalar_offsets, scalar_indices = find_scalar_offsets_and_indices(spoofed_bc_args_type) + cudevicearray_offsets = args_offset .+ cudevicearray_offsets + scalar_offsets = args_offset .+ scalar_offsets + # STEP 2: Get NDarrays corresponding to the offsets in the spoofed type. + input_ndarrays = ntuple( + i -> get_ndarray(bc.args[cudevicearray_indices[i]]), length(cudevicearray_indices) + ) + input_scalars = ntuple(i -> bc.args[scalar_indices[i]], length(scalar_indices)) + patch_info = BroadcastPatchInfo( + sizeof(spoofed_bc_type), + input_ndarrays, + cudevicearray_offsets, + ntuple(i -> i - 1, length(input_ndarrays)), + scalar_offsets, + input_scalars, ) - #! DOUBLE CHECK bc.args ACTAULLY GIVES BACK ARRAYS - #! HOW TO GET SCALARS??? - # scalars = ???? - # launch(cuda_task, bc.args, (dest,), scalars; blocks, threads) + println("Array Offsets: ", cudevicearray_offsets) + println("Scalar Offsets: ", scalar_offsets) + println("Input NDArrays: ", input_ndarrays) + println("Input Scalars: ", input_scalars) + + # launch_broadcast( + # fused_kernel_metadata.cuda_task, + # (dest,), + # patch_info; + # blocks=(fused_kernel_metadata.blocks,), + # threads=(fused_kernel_metadata.threads,), + # prefix_scalars=(fused_kernel_metadata.ctx,), + # kernel_input_args_count=0, + # kernel_output_args_count=1, + # ) #! DO I NEED TO DO TYPE PROMOTION CHECKS?? - return dest end From 680c13fd916b4cf94a825c137bb1f4bfb91ca005 Mon Sep 17 00:00:00 2001 From: ejmeitz Date: Tue, 12 May 2026 12:11:57 -0500 Subject: [PATCH 09/12] change kernel to not take Broadcast object --- lib/cunumeric_jl_wrapper/src/cuda.cpp | 321 ++++++++++++-------------- src/cuda/cuda_ptx_task.jl | 76 ------ src/cuda/cuda_util.jl | 2 +- src/ndarray/broadcast_fusion.jl | 181 +++++++++------ 4 files changed, 256 insertions(+), 324 deletions(-) diff --git a/lib/cunumeric_jl_wrapper/src/cuda.cpp b/lib/cunumeric_jl_wrapper/src/cuda.cpp index 0350cb28..d5e779d0 100644 --- a/lib/cunumeric_jl_wrapper/src/cuda.cpp +++ b/lib/cunumeric_jl_wrapper/src/cuda.cpp @@ -23,7 +23,6 @@ #include #include -#include "cuda_macros.h" #include "legate.h" #include "legate/utilities/proc_local_storage.h" #include "legion.h" @@ -35,17 +34,55 @@ #define BLOCK_START 1 #define THREAD_START 4 #define ARG_OFFSET 7 -#define PREFIX_SCALAR_COUNT_INDEX 7 -#define KERNEL_INPUT_ARG_COUNT_INDEX 8 -#define KERNEL_OUTPUT_ARG_COUNT_INDEX 9 -#define BC_NDARRAY_PATCH_COUNT_INDEX 10 -#define BC_SCALAR_PATCH_COUNT_INDEX 11 -#define BC_BLOB_SIZE_INDEX 12 -#define BCAST_ARG_OFFSET 13 // global padding for CUDA.jl kernel state std::size_t padded_bytes_kernel_state = 16; +#define ERROR_CHECK(x) \ + { \ + cudaError_t status = x; \ + if (status != cudaSuccess) { \ + fprintf(stderr, "CUDA Error at %s:%d: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(status)); \ + if (stream_) cudaStreamDestroy(stream_); \ + exit(-1); \ + } \ + } + +#define DRIVER_ERROR_CHECK(x) \ + { \ + CUresult status = x; \ + if (status != CUDA_SUCCESS) { \ + const char *err_str = nullptr; \ + cuGetErrorString(status, &err_str); \ + fprintf(stderr, "CUDA Driver Error at %s:%d: %s\n", __FILE__, __LINE__, \ + err_str); \ + if (stream_) cudaStreamDestroy(stream_); \ + exit(-1); \ + } \ + } + +#define TEST_PRINT_DEBUG(dev_ptr, N, T, format, stream, message) \ + { \ + std::vector host_arr(N); \ + ERROR_CHECK(cudaMemcpy(host_arr.data(), \ + reinterpret_cast(dev_ptr), \ + sizeof(T) * N, cudaMemcpyDeviceToHost)); \ + ERROR_CHECK(cudaStreamSynchronize(stream)); \ + fprintf(stderr, "[TEST_PRINT] %s: " format "\n", message, host_arr[0]); \ + } + +#ifdef CUDA_DEBUG +#define CUDA_DEBUG_PRINT(x) \ + do { \ + x; \ + } while (0) +#else +#define CUDA_DEBUG_PRINT(x) \ + do { \ + } while (0) +#endif + namespace ufi { using namespace Legion; // TODO CUcontext key hashing is redundant. ProcLocalStorage is local to the @@ -96,6 +133,34 @@ struct CuDeviceArray { uint64_t length; // Number of elements (at the end) }; +#define CUDA_DEVICE_ARRAY_ARG(MODE, ACCESSOR_CALL) \ + template < \ + typename T, int D, \ + typename std::enable_if<(D >= 1 && D <= REALM_MAX_DIM), int>::type = 0> \ + void cuda_device_array_arg_##MODE(char *&p, \ + const legate::PhysicalArray &rf) { \ + auto shp = rf.shape(); \ + auto acc = rf.data().ACCESSOR_CALL(); \ + CUDA_DEBUG_PRINT(std::cerr << "[RunPTXTask] " #MODE " accessor shape: " \ + << shp.lo << " - " << shp.hi << ", dim: " << D \ + << std::endl; \ + std::cerr << "[RunPTXTask] " #MODE " accessor strides: " \ + << acc.accessor.strides << std::endl;); \ + void *dev_ptr = const_cast(/*.lo to ensure multiple GPU support*/ \ + static_cast( \ + acc.ptr(Realm::Point(shp.lo)))); \ + auto extents = shp.hi - shp.lo + legate::Point::ONES(); \ + CuDeviceArray desc; \ + desc.ptr = dev_ptr; \ + desc.maxsize = shp.volume() * sizeof(T); \ + for (size_t i = 0; i < D; ++i) { \ + desc.dims[i] = extents[i]; \ + } \ + desc.length = shp.volume(); \ + memcpy(p, &desc, sizeof(CuDeviceArray)); \ + p += sizeof(CuDeviceArray); \ + } + CUDA_DEVICE_ARRAY_ARG(read, read_accessor); // cuda_device_array_arg_read CUDA_DEVICE_ARRAY_ARG(write, write_accessor); // cuda_device_array_arg_write @@ -110,33 +175,37 @@ struct ufiFunctor { } }; -struct LaunchDims { - std::uint32_t bx, by, bz; - std::uint32_t tx, ty, tz; -}; +// https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/eval_udf_gpu.cc +/*static*/ void RunPTXTask::gpu_variant(legate::TaskContext context) { + cudaStream_t stream_ = context.get_task_stream(); + std::string kernel_name = context.scalar(0).value(); // 0 -inline LaunchDims read_launch_dims(const legate::TaskContext &context) { - return LaunchDims{ - context.scalar(BLOCK_START + 0).value(), - context.scalar(BLOCK_START + 1).value(), - context.scalar(BLOCK_START + 2).value(), - context.scalar(THREAD_START + 0).value(), - context.scalar(THREAD_START + 1).value(), - context.scalar(THREAD_START + 2).value(), - }; -} + std::uint32_t bx = + context.scalar(BLOCK_START + 0).value(); // 1 + std::uint32_t by = + context.scalar(BLOCK_START + 1).value(); // 2 + std::uint32_t bz = + context.scalar(BLOCK_START + 2).value(); // 3 + + std::uint32_t tx = + context.scalar(THREAD_START + 0).value(); // 4 + std::uint32_t ty = + context.scalar(THREAD_START + 1).value(); // 5 + std::uint32_t tz = + context.scalar(THREAD_START + 2).value(); // 6 -inline std::pair lookup_kernel_function( - const std::string &kernel_name, cudaStream_t stream_) { CUcontext ctx; cuStreamGetCtx(stream_, &ctx); FunctionKey key = {ctx, kernel_name}; assert(cufunction_ptr.has_value()); FunctionMap &fmap = cufunction_ptr.get(); + auto it = fmap.find(key); + #ifdef CUDA_DEBUG if (it == fmap.end()) { + // for DEBUG output std::cerr << "[RunPTXTask] Could not find key: " << key_to_string(key) << std::endl; for (const auto &[k, v] : fmap) { @@ -145,62 +214,19 @@ inline std::pair lookup_kernel_function( assert(0 && "[RunPTXTask] key is not found in hashmap"); } #endif - assert(it != fmap.end()); - return {ctx, it->second}; -} -inline void append_array_args(char *&p, const legate::TaskContext &context, - std::size_t count, bool is_input) { - for (std::size_t i = 0; i < count; ++i) { - auto ps = is_input ? context.input(i) : context.output(i); - auto code = ps.type().code(); - auto dim = ps.dim(); - legate::double_dispatch( - dim, code, ufiFunctor{}, - is_input ? ufi::AccessMode::READ : ufi::AccessMode::WRITE, p, ps); - } -} - -inline void write_array_descriptor(char *dst, const legate::PhysicalArray &ps, - bool is_input) { - char *p = dst; - auto code = ps.type().code(); - auto dim = ps.dim(); - legate::double_dispatch( - dim, code, ufiFunctor{}, - is_input ? ufi::AccessMode::READ : ufi::AccessMode::WRITE, p, ps); -} - -inline void launch_with_buffer(CUfunction func, const LaunchDims &dims, - cudaStream_t stream_, - std::vector &arg_buffer, - std::size_t used_buffer_size) { - void *config[] = { - CU_LAUNCH_PARAM_BUFFER_POINTER, - static_cast(arg_buffer.data()), - CU_LAUNCH_PARAM_BUFFER_SIZE, - &used_buffer_size, - CU_LAUNCH_PARAM_END, - }; - - CUstream custream_ = reinterpret_cast(stream_); - DRIVER_ERROR_CHECK(cuLaunchKernel(func, dims.bx, dims.by, dims.bz, dims.tx, - dims.ty, dims.tz, 0, custream_, nullptr, - config)); -} - -// https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/eval_udf_gpu.cc -/*static*/ void RunPTXTask::gpu_variant(legate::TaskContext context) { - cudaStream_t stream_ = context.get_task_stream(); - std::string kernel_name = context.scalar(0).value(); // 0 - const LaunchDims dims = read_launch_dims(context); - auto kernel = lookup_kernel_function(kernel_name, stream_); - CUfunction func = kernel.second; + assert(it != fmap.end()); + CUfunction func = it->second; const std::size_t num_inputs = context.num_inputs(); const std::size_t num_outputs = context.num_outputs(); const std::size_t num_scalars = context.num_scalars(); - // Layout: [kernel-state][inputs][outputs][user-scalars] + const std::size_t num_reductions = + context.num_reductions(); // unused for now + + // compute total size: all device arrays + all scalars + // skip scalar 0-2 (kernel_name, threads, blocks) + // we allocate extra to decrease looping and dynamic dispatching on dim std::size_t max_buffer_size = padded_bytes_kernel_state + (num_inputs + num_outputs) * sizeof(CuDeviceArray); @@ -210,119 +236,58 @@ inline void launch_with_buffer(CUfunction func, const LaunchDims &dims, std::vector arg_buffer(max_buffer_size); char *p = arg_buffer.data() + padded_bytes_kernel_state; - append_array_args(p, context, num_inputs, true); - append_array_args(p, context, num_outputs, false); - for (std::size_t i = ARG_OFFSET; i < num_scalars; ++i) { - const auto &scalar = context.scalar(i); - memcpy(p, scalar.ptr(), scalar.size()); - p += scalar.size(); + for (std::size_t i = 0; i < num_inputs; ++i) { + auto ps = context.input(i); + auto code = ps.type().code(); + auto dim = ps.dim(); +#ifdef CUDA_DEBUG + std::cerr << "[RunPTXTask] Input " << i << " type: " << code + << ", dim: " << dim << std::endl; +#endif + // dispatch on dim and code with ufiFunctor operator() + legate::double_dispatch(dim, code, ufiFunctor{}, ufi::AccessMode::READ, p, + ps); } - - std::size_t buffer_size = p - arg_buffer.data(); // calc used buffer - - launch_with_buffer(func, dims, stream_, arg_buffer, buffer_size); - - // DRIVER_ERROR_CHECK(cuStreamSynchronize(stream_)); -} - -/* -This code is specifically for the fusion of broadcast operations. -This code uses the passed offsets to patch together the Broadcasted type -that the CUDA kernel expects. The Broadcasted type has all inputs in order, -but mixed scalars and arrays. e.g., [Array, Scalar, Array, Array]. -We must use the offsets to figure out which arguments should be marked as -scalars and which should be marked as inputs. -*/ - -/*static*/ void RunPTXBroadcastTask::gpu_variant(legate::TaskContext context) { - cudaStream_t stream_ = context.get_task_stream(); - std::string kernel_name = context.scalar(0).value(); // 0 - const LaunchDims dims = read_launch_dims(context); - - std::uint32_t prefix_scalar_count = - context.scalar(PREFIX_SCALAR_COUNT_INDEX).value(); // 7 - std::uint32_t kernel_input_args_count = - context.scalar(KERNEL_INPUT_ARG_COUNT_INDEX).value(); // 8 - std::uint32_t kernel_output_args_count = - context.scalar(KERNEL_OUTPUT_ARG_COUNT_INDEX) - .value(); // 9 - std::uint32_t bc_ndarray_patch_count = - context.scalar(BC_NDARRAY_PATCH_COUNT_INDEX) - .value(); // 10 - std::uint32_t bc_scalar_patch_count = - context.scalar(BC_SCALAR_PATCH_COUNT_INDEX).value(); // 11 - std::uint32_t bc_blob_size = - context.scalar(BC_BLOB_SIZE_INDEX).value(); // 12 - - const std::size_t num_inputs = context.num_inputs(); - const std::size_t num_outputs = context.num_outputs(); - const std::size_t num_scalars = context.num_scalars(); - - assert(kernel_input_args_count <= num_inputs); - assert(kernel_output_args_count <= num_outputs); - const std::size_t array_patch_start = BCAST_ARG_OFFSET; - const std::size_t scalar_patch_start = - array_patch_start + 2 * bc_ndarray_patch_count; - const std::size_t prefix_start = scalar_patch_start + bc_scalar_patch_count; - const std::size_t prefix_end = prefix_start + prefix_scalar_count; - const std::size_t scalar_values_start = prefix_end; - - assert(scalar_values_start + bc_scalar_patch_count <= num_scalars); - - auto kernel = lookup_kernel_function(kernel_name, stream_); - CUfunction func = kernel.second; - - std::size_t max_buffer_size = - padded_bytes_kernel_state + - (kernel_input_args_count + kernel_output_args_count) * - sizeof(CuDeviceArray) + - bc_blob_size; - for (std::size_t i = prefix_start; i < prefix_end; ++i) { - max_buffer_size += context.scalar(i).size(); + for (std::size_t i = 0; i < num_outputs; ++i) { + auto ps = context.output(i); + auto code = ps.type().code(); + auto dim = ps.dim(); +#ifdef CUDA_DEBUG + std::cerr << "[RunPTXTask] Output " << i << " type: " << code + << ", dim: " << dim << std::endl; +#endif + legate::double_dispatch(dim, code, ufiFunctor{}, ufi::AccessMode::WRITE, p, + ps); } - - std::vector arg_buffer(max_buffer_size); - char *p = arg_buffer.data() + padded_bytes_kernel_state; - - for (std::size_t i = prefix_start; i < prefix_end; ++i) { + for (std::size_t i = ARG_OFFSET; i < num_scalars; ++i) { const auto &scalar = context.scalar(i); memcpy(p, scalar.ptr(), scalar.size()); p += scalar.size(); } - append_array_args(p, context, kernel_input_args_count, true); - append_array_args(p, context, kernel_output_args_count, false); - - std::vector patched_bc(bc_blob_size, 0); - - for (std::size_t j = 0; j < bc_ndarray_patch_count; ++j) { - std::uint32_t offset = - context.scalar(array_patch_start + 2 * j).value(); - std::uint32_t input_index = - context.scalar(array_patch_start + 2 * j + 1).value(); - assert(input_index < num_inputs); - assert(offset <= patched_bc.size()); - char *begin = patched_bc.data() + offset; - char *end = begin; - write_array_descriptor(end, context.input(input_index), true); - assert(static_cast(end - patched_bc.data()) <= - patched_bc.size()); - } + std::size_t buffer_size = p - arg_buffer.data(); // calc used buffer - for (std::size_t j = 0; j < bc_scalar_patch_count; ++j) { - std::uint32_t offset = - context.scalar(scalar_patch_start + j).value(); - const auto &scalar_value = context.scalar(scalar_values_start + j); - assert(offset + scalar_value.size() <= patched_bc.size()); - memcpy(patched_bc.data() + offset, scalar_value.ptr(), scalar_value.size()); - } + void *config[] = { + CU_LAUNCH_PARAM_BUFFER_POINTER, + static_cast(arg_buffer.data()), + CU_LAUNCH_PARAM_BUFFER_SIZE, + &buffer_size, + CU_LAUNCH_PARAM_END, + }; - memcpy(p, patched_bc.data(), patched_bc.size()); - p += patched_bc.size(); + CUstream custream_ = reinterpret_cast(stream_); - std::size_t buffer_size = p - arg_buffer.data(); +#ifdef CUDA_DEBUG + std::cerr << "[RunPTXTask] Launching kernel " << kernel_name + << " with blocks (" << bx << "," << by << "," << bz + << ") and threads (" << tx << "," << ty << "," << tz << ")" + << " on CUcontext " << context_to_string(ctx) << std::endl; +#endif + // Launch the kernel + DRIVER_ERROR_CHECK(cuLaunchKernel(func, bx, by, bz, tx, ty, tz, 0, custream_, + nullptr, config)); - launch_with_buffer(func, dims, stream_, arg_buffer, buffer_size); + // DRIVER_ERROR_CHECK(cuStreamSynchronize(stream_)); } // https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/load_ptx.cc @@ -452,6 +417,4 @@ void wrap_cuda_methods(jlcxx::Module &mod) { mod.method("extract_kernel_name", &extract_kernel_name); mod.set_const("LOAD_PTX", legate::LocalTaskID{ufi::TaskIDs::LOAD_PTX_TASK}); mod.set_const("RUN_PTX", legate::LocalTaskID{ufi::TaskIDs::RUN_PTX_TASK}); - mod.set_const("RUN_PTX_BROADCAST", - legate::LocalTaskID{ufi::TaskIDs::RUN_PTX_BROADCAST_TASK}); } diff --git a/src/cuda/cuda_ptx_task.jl b/src/cuda/cuda_ptx_task.jl index 226e5233..7cbb3e07 100644 --- a/src/cuda/cuda_ptx_task.jl +++ b/src/cuda/cuda_ptx_task.jl @@ -129,82 +129,6 @@ function launch(kernel::CUDATask, inputs, outputs, scalars; ) end -struct BroadcastPatchInfo - broadcast_size::Int - inputs::Tuple - array_offsets::Tuple{Vararg{Int}} - array_input_indices::Tuple{Vararg{Int}} - scalar_offsets::Tuple{Vararg{Int}} - scalar_values::Tuple -end - -function launch_broadcast(kernel::CUDATask, outputs, patch_info::BroadcastPatchInfo; - blocks, threads, prefix_scalars=(), - kernel_input_args_count=0, - kernel_output_args_count=length(isa(outputs, Tuple) ? outputs : (outputs,))) - input_tup = patch_info.inputs - output_tup = isa(outputs, Tuple) ? outputs : (outputs,) - scalar_tup = patch_info.scalar_values - prefix_tup = isa(prefix_scalars, Tuple) ? prefix_scalars : (prefix_scalars,) - - length(patch_info.array_offsets) == length(patch_info.array_input_indices) || - throw(ArgumentError("Broadcast array patch offsets and input indices must match")) - - isempty(output_tup) && throw(ArgumentError("Broadcast launch requires an output array")) - max_shape = size(first(output_tup)) - @assert !isnothing(max_shape) - - rt = Legate.get_runtime() - lib = cuNumeric.get_lib() - taskid = cuNumeric.RUN_PTX_BROADCAST - task = Legate.create_auto_task(rt, lib, taskid) - - input_vars = Vector{Legate.Variable}() - for arr in input_tup - check_sz(arr, max_shape) - la = nda_to_logical_array(arr) - p = Legate.add_input(task, la) - push!(input_vars, p) - end - - output_vars = Vector{Legate.Variable}() - for arr in output_tup - check_sz(arr, max_shape) - la = nda_to_logical_array(arr) - p = Legate.add_output(task, la) - push!(output_vars, p) - end - - Legate.add_scalar(task, Legate.string_to_scalar(kernel.func)) # 0 - cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, blocks)) # 1,2,3 - cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, threads)) # 4,5,6 - Legate.add_scalar(task, Legate.Scalar(UInt32(length(prefix_tup)))) # 7 - Legate.add_scalar(task, Legate.Scalar(UInt32(kernel_input_args_count))) # 8 - Legate.add_scalar(task, Legate.Scalar(UInt32(kernel_output_args_count))) # 9 - Legate.add_scalar(task, Legate.Scalar(UInt32(length(patch_info.array_offsets)))) # 10 - Legate.add_scalar(task, Legate.Scalar(UInt32(length(patch_info.scalar_offsets)))) # 11 - Legate.add_scalar(task, Legate.Scalar(UInt32(patch_info.broadcast_size))) # 12 - - for (off, input_index) in zip(patch_info.array_offsets, patch_info.array_input_indices) - Legate.add_scalar(task, Legate.Scalar(UInt32(off))) - Legate.add_scalar(task, Legate.Scalar(UInt32(input_index))) - end - - for off in patch_info.scalar_offsets - Legate.add_scalar(task, Legate.Scalar(UInt32(off))) - end - - for s in prefix_tup - Legate.add_scalar(task, Legate.Scalar(s)) - end - for s in scalar_tup - Legate.add_scalar(task, Legate.Scalar(s)) - end - - Legate.default_alignment(task, input_vars, output_vars) - Legate.submit_auto_task(rt, task) -end - function ptx_task(ptx::String, kernel_name) rt = Legate.get_runtime() lib = cuNumeric.get_lib() # grab lib of legate app diff --git a/src/cuda/cuda_util.jl b/src/cuda/cuda_util.jl index 936a4c40..70760bcb 100644 --- a/src/cuda/cuda_util.jl +++ b/src/cuda/cuda_util.jl @@ -30,7 +30,7 @@ types for code generation (e.g. mapping `NDArray{...}` to `CuDeviceArray{...}` i """ map_cuda_type(::Type{T}) where {T} = T -map_cuda_type(::Type{<:NDArray{T,N}}) where {T,N} = ndarray_cuda_type(NDArray{T,N}) +map_cuda_type(::Type{<:NDArray{T,N}}) where {T,N} = CUDACore.CuDeviceArray{T,N,CUDACore.AS.Global} function map_cuda_type(::Type{T}) where {T<:Tuple} return Tuple{map_cuda_type.(T.parameters)...} diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl index 08da76fa..078de5c8 100644 --- a/src/ndarray/broadcast_fusion.jl +++ b/src/ndarray/broadcast_fusion.jl @@ -1,21 +1,34 @@ -# Copied from GPUArrays: https://github.com/JuliaGPU/GPUArrays.jl/blob/a9df2ba41ca2358c1de2f3cc6b020578bf6e39b1/src/host/broadcast.jl#L60-L63 -# Defined with KernelAbstractions.jl. Makes it easier to generate indexing for -# various dimensions of inputs/outputs. Assumes broadcast is `Base.Broadcast.process`ed so that -# dest/bc have singleton dimensions inserted and we can index 1-1 like this. -@kernel function broadcast_kernel_cartesian(dest, bc) - I = @index(Global, Cartesian) - @inbounds dest[I] = bc[I] -end +# The indexing in these kernels is based off this internal function from julia/broadcast.jl + +# Base.@propagate_inbounds function _broadcast_getindex(bc::Broadcasted{<:Any,<:Any,<:Any,<:Any}, I) +# args = _getindex(bc.args, I) +# return _broadcast_getindex_evalf(bc.f, args...) +# end + +function make_linear_kernel(dest, bc::Base.Broadcast.Broadcasted) + f = bc.f -@kernel function broadcast_kernel_linear(dest, bc) - I = @index(Global, Linear) - @inbounds dest[I] = bc[I] + @kernel function broadcast_kernel_linear_splat(dest, args...) + I = @index(Global, Linear) + @inbounds args_modified = Base.Broadcast._getindex(args, I) + @inbounds dest[I] = Base.Broadcast._broadcast_getindex_evalf(f, args_modified...) + end + + return broadcast_kernel_linear_splat end -# No compilation here, just generating CUDA specific kernel. -const GPU_CARTESIAN_KERNEL = broadcast_kernel_cartesian(CUDACore.CUDAKernels.CUDABackend()) -const GPU_LINEAR_KERNEL = broadcast_kernel_linear(CUDACore.CUDAKernels.CUDABackend()) +function make_cartesian_kernel(dest, bc::Base.Broadcast.Broadcasted) + f = bc.f + + @kernel function broadcast_kernel_cartesian_splat(dest, args...) + I = @index(Global, Cartesian) + @inbounds args_modified = Base.Broadcast._getindex(args, I) + @inbounds dest[I] = Base.Broadcast._broadcast_getindex_evalf(f, args_modified...) + end + + return broadcast_kernel_cartesian_splat +end struct FusedBroadcastMetadata ctx::Any # KA.CompilerMetadata @@ -67,19 +80,12 @@ get_ndarray(x::T) where {T<:NDArray} = x get_ndarray(x::T) where {T<:Base.Broadcast.Extruded} = x.x get_ndarray(x) = throw(error("Broadcast fusion. Don't know what to do with type: $T")) -""" - get_ptx(obj::KA.Kernel{CUDABackend}, ::Type{DEST_T}, ::Type{BC_T}; - ndrange) -> (ptx::String, threads::Int, blocks::Int) - -Compile a KA CUDA kernel (kernel-body `obj.f(ctx, ...)`) using *types only* for `DEST_T` and `BC_T`, -choose a workgroup size (threads) using CUDA occupancy when possible, and return the generated PTX. -""" -function get_ptx( +function _threads_from_occupancy( obj::KA.Kernel{CUDACore.CUDAKernels.CUDABackend}, ::Type{DEST_T}, - ::Type{BC_T}; + ARG_TYPES...; ndrange, -) where {DEST_T,BC_T} +) where {DEST_T} backend = KA.backend(obj) ndrange, workgroupsize, iterspace, dynamic = KA.launch_config(obj, ndrange, nothing) @@ -93,9 +99,8 @@ function get_ptx( nothing end - # Determine threads via occupancy if we can compile from types; otherwise use a broadcast-friendly heuristic. - threads = _default_broadcast_threads(ndrange) - tt = Base.to_tuple_type((typeof(ctx), DEST_T, BC_T)) + # Determine threads via occupancy + tt = Base.to_tuple_type((typeof(ctx), DEST_T, ARG_TYPES...)) host_kernel = CUDACore.cufunction( obj.f, tt; @@ -112,10 +117,33 @@ function get_ptx( blocks = length(KA.blocks(iterspace)) threads = length(KA.workitems(iterspace)) - blocks == 0 && return "", 0, 0, ctx + return threads, blocks, ctx +end + +""" + get_ptx(obj::KA.Kernel{CUDABackend}, ::Type{DEST_T}, ::Type{BC_T}; + ndrange) -> (ptx::String, threads::Int, blocks::Int) + +Compile a KA CUDA kernel (kernel-body `obj.f(ctx, ...)`) using *types only* for `DEST_T` and `BC_T`, +choose a workgroup size (threads) using CUDA occupancy when possible, and return the generated PTX. +""" +function get_ptx( + obj::KA.Kernel{CUDACore.CUDAKernels.CUDABackend}, + ::Type{DEST_T}, + arg_types...; + ndrange, +) where {DEST_T,BC_T} + println(Base.isbitstype.(arg_types)) + threads, blocks, ctx = _threads_from_occupancy(obj, DEST_T, arg_types...; ndrange=ndrange) + blocks == 0 && return "", 0, 0, ctx #! MAYBE ERROR HERE? + + # Generate PTX buf = IOBuffer() - CUDATools.code_ptx(buf, obj.f, (typeof(ctx), DEST_T, BC_T); raw=false, kernel=true, ptx=v"7.8") + #!TODO REMOVE THE MANUAL PTX VERSION HERE! + CUDATools.code_ptx(buf, obj.f, (typeof(ctx), DEST_T, arg_types...); + raw=false, kernel=true, ptx=v"7.8") + return String(take!(buf)), threads, blocks, ctx end @@ -126,17 +154,17 @@ function get_cuda_task( ndrange, ) where {D<:NDArray,B<:Base.Broadcast.Broadcasted} DEST_T = map_cuda_type(D) - BC_T = map_cuda_type(B) + ARG_TYPES = map_cuda_type.(typeof.(bc.args)) key = (obj, D, B, ndrange) lock(_BCAST_PTX_CACHE_LOCK) do # Also stores in cache if not found in Dict return get!(_BCAST_PTX_CACHE, key) do - ptx, threads, blocks, ctx = get_ptx(obj, DEST_T, BC_T; ndrange=ndrange) + ptx, threads, blocks, ctx = get_ptx(obj, DEST_T, ARG_TYPES...; ndrange=ndrange) func_name = extract_kernel_name(ptx) - # println(ptx) + println(ptx) ptx_task(ptx, func_name) - cuda_task = CUDATask(func_name, (DEST_T, BC_T)) + cuda_task = CUDATask(func_name, (DEST_T, ARG_TYPES...)) FusedBroadcastMetadata(ctx, threads, blocks, cuda_task) end end @@ -146,6 +174,7 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas #! HOW DOES THIS BEHAVE WHEN BC HAS 2 RESULT ARRAYS? + # Normalize the Braodcasted type bc = Base.Broadcast.preprocess(dest, bc) bc = Base.Broadcast.instantiate(bc) bc = Base.Broadcast.flatten(bc) @@ -155,66 +184,82 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas if ndims(dest) == 1 || (isa(IndexStyle(dest), IndexLinear) && isa(IndexStyle(bc), IndexLinear)) - GPU_LINEAR_KERNEL + make_linear_kernel(dest, bc) else - GPU_CARTESIAN_KERNEL + make_cartesian_kernel(dest, bc) end ndrange = ndims(dest) > 0 ? size(dest) : (1,) + # Tell KernelAsbtractions.jl we are using CUDA.jl + bck_cuda = broadcast_kernel(CUDACore.CUDAKernels.CUDABackend()) + # Lookup in cache, if not found, compile and cache - fused_kernel_metadata = get_cuda_task(broadcast_kernel, dest, bc, ndrange) + fkm = get_cuda_task(bck_cuda, dest, bc, ndrange) # Replace NDArrays with CuDeviceArrays in the Broadcasted type so we can figure out bit-offsets spoofed_bc_type = map_cuda_type(typeof(bc)) fieldname(spoofed_bc_type, 3) == :args || throw(ArgumentError("Broadcasted field 3 is not args. Failed to fuse broadcast.")) - args_offset = Int(fieldoffset(spoofed_bc_type, 3)) - # Replace NDArrays with CuDeviceArrays in the Broadcasted type so we can figure out bit-offsets + args_offset = Int(fieldoffset(spoofed_bc_type, 3)) + args_offset == 0 || + throw( + ArgumentError( + "Broadcast fusion only supports Broadcasted layouts where args starts at offset 0; got offset $args_offset. This is likely a bug in the compiler." + ), + ) + + # Spoof NDArrays with CuDeviceArrays in the Broadcasted type spoofed_bc_args_type = map_cuda_type(typeof(bc.args)) - #!TODO FIGURE OUT HOW TO HANDLE AXES - #! TODO FIGURE OUT IF ITS SAFE TO IGNORE FIRST TWO FIELDS OF BROADCASTED TYPE - - # STEP 1: Figure out bit-offsets for CuDeviceArrays and scalars in args of spoofed type. - # The spoofed type has the same fields and alignment that the PTX kernel expects. + # Figure out where input arrays are cudevicearray_offsets, cudevicearray_indices = find_cudevicearray_offsets_and_indices( spoofed_bc_args_type ) + + # Figure out where scalars are scalar_offsets, scalar_indices = find_scalar_offsets_and_indices(spoofed_bc_args_type) - cudevicearray_offsets = args_offset .+ cudevicearray_offsets - scalar_offsets = args_offset .+ scalar_offsets - # STEP 2: Get NDarrays corresponding to the offsets in the spoofed type. + input_ndarrays = ntuple( i -> get_ndarray(bc.args[cudevicearray_indices[i]]), length(cudevicearray_indices) ) input_scalars = ntuple(i -> bc.args[scalar_indices[i]], length(scalar_indices)) - patch_info = BroadcastPatchInfo( - sizeof(spoofed_bc_type), + + #! I AM NOT SURE IT IS CORRECT TO ASSUME SCALARS ARE ALWAYS AFTER NDARARYS + #! THE PTX KERNEL SIGNATURE BEFORE MODIFICATION BY CUDA.jl: + #! (fkm.ctx, dest, bc.args...) + + launch( + fkm.cuda_task, input_ndarrays, - cudevicearray_offsets, - ntuple(i -> i - 1, length(input_ndarrays)), - scalar_offsets, - input_scalars, + (dest,), + input_scalars; + blocks=fkm.blocks, + threads=fkm.threads, ) - println("Array Offsets: ", cudevicearray_offsets) - println("Scalar Offsets: ", scalar_offsets) - println("Input NDArrays: ", input_ndarrays) - println("Input Scalars: ", input_scalars) - - # launch_broadcast( - # fused_kernel_metadata.cuda_task, - # (dest,), - # patch_info; - # blocks=(fused_kernel_metadata.blocks,), - # threads=(fused_kernel_metadata.threads,), - # prefix_scalars=(fused_kernel_metadata.ctx,), - # kernel_input_args_count=0, - # kernel_output_args_count=1, - # ) - #! DO I NEED TO DO TYPE PROMOTION CHECKS?? return dest end + +# STEP 1: Figure out bit-offsets for CuDeviceArrays and scalars in args of spoofed type. +# The spoofed type has the same fields and alignment that the PTX kernel expects. +# cudevicearray_offsets, cudevicearray_indices = +# find_cudevicearray_offsets_and_indices(spoofed_bc_args_type) + +# scalar_offsets, scalar_indices = find_scalar_offsets_and_indices(spoofed_bc_args_type) + +# # STEP 2: Get NDarrays corresponding to the offsets in the spoofed type. +# input_ndarrays = ntuple( +# i -> get_ndarray(bc.args[cudevicearray_indices[i]]), length(cudevicearray_indices) +# ) +# input_scalars = ntuple(i -> bc.args[scalar_indices[i]], length(scalar_indices)) +# patch_info = BroadcastPatchInfo( +# sizeof(spoofed_bc_type), +# input_ndarrays, +# cudevicearray_offsets, +# ntuple(i -> i - 1, length(input_ndarrays)), +# scalar_offsets, +# input_scalars, +# ) From 6ae32dcc63153fec6346720ecf0a713201b8734e Mon Sep 17 00:00:00 2001 From: krasow Date: Thu, 14 May 2026 16:54:26 -0500 Subject: [PATCH 10/12] broadcast fusion support. --- .githash | 2 +- lib/cunumeric_jl_wrapper/src/cuda.cpp | 213 ++++++++++++++++++-------- src/cuda/cuda_ptx_task.jl | 26 +++- src/ndarray/broadcast_fusion.jl | 72 ++++++--- test/runtests.jl | 9 +- test/tests/broadcast_fusion_tests.jl | 185 ++++++++++++++++++++++ 6 files changed, 411 insertions(+), 96 deletions(-) create mode 100644 test/tests/broadcast_fusion_tests.jl diff --git a/.githash b/.githash index 6cdfeea5..2c5d87fb 100644 --- a/.githash +++ b/.githash @@ -1 +1 @@ -dec047f1bd1c8287513c6c437f946982e516ccd4 +e4cc4b6a778aa828e5647803f594222053797f54 diff --git a/lib/cunumeric_jl_wrapper/src/cuda.cpp b/lib/cunumeric_jl_wrapper/src/cuda.cpp index d5e779d0..3582b8a6 100644 --- a/lib/cunumeric_jl_wrapper/src/cuda.cpp +++ b/lib/cunumeric_jl_wrapper/src/cuda.cpp @@ -29,7 +29,7 @@ #include "types.h" #include "ufi.h" -#define CUDA_DEBUG 0 +// #define CUDA_DEBUG #define BLOCK_START 1 #define THREAD_START 4 @@ -175,37 +175,40 @@ struct ufiFunctor { } }; -// https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/eval_udf_gpu.cc -/*static*/ void RunPTXTask::gpu_variant(legate::TaskContext context) { - cudaStream_t stream_ = context.get_task_stream(); - std::string kernel_name = context.scalar(0).value(); // 0 - - std::uint32_t bx = - context.scalar(BLOCK_START + 0).value(); // 1 - std::uint32_t by = - context.scalar(BLOCK_START + 1).value(); // 2 - std::uint32_t bz = - context.scalar(BLOCK_START + 2).value(); // 3 - - std::uint32_t tx = - context.scalar(THREAD_START + 0).value(); // 4 - std::uint32_t ty = - context.scalar(THREAD_START + 1).value(); // 5 - std::uint32_t tz = - context.scalar(THREAD_START + 2).value(); // 6 +struct PTXLaunchParams { + cudaStream_t stream; + CUstream custream; + CUfunction func; + std::string kernel_name; + std::uint32_t bx, by, bz; + std::uint32_t tx, ty, tz; +}; + +// Reads common scalars (kernel_name, blocks, threads) and looks up the +// compiled CUfunction. Shared by RunPTXTask and RunPTXBroadcastTask. +static PTXLaunchParams read_launch_params(legate::TaskContext &context) { + PTXLaunchParams p; + p.stream = context.get_task_stream(); + p.kernel_name = context.scalar(0).value(); + + p.bx = context.scalar(BLOCK_START + 0).value(); + p.by = context.scalar(BLOCK_START + 1).value(); + p.bz = context.scalar(BLOCK_START + 2).value(); + + p.tx = context.scalar(THREAD_START + 0).value(); + p.ty = context.scalar(THREAD_START + 1).value(); + p.tz = context.scalar(THREAD_START + 2).value(); CUcontext ctx; - cuStreamGetCtx(stream_, &ctx); + cuStreamGetCtx(p.stream, &ctx); - FunctionKey key = {ctx, kernel_name}; + FunctionKey key = {ctx, p.kernel_name}; assert(cufunction_ptr.has_value()); FunctionMap &fmap = cufunction_ptr.get(); - auto it = fmap.find(key); #ifdef CUDA_DEBUG if (it == fmap.end()) { - // for DEBUG output std::cerr << "[RunPTXTask] Could not find key: " << key_to_string(key) << std::endl; for (const auto &[k, v] : fmap) { @@ -216,17 +219,51 @@ struct ufiFunctor { #endif assert(it != fmap.end()); - CUfunction func = it->second; + p.func = it->second; + p.custream = reinterpret_cast(p.stream); + return p; +} + +// Launch the kernel with the filled arg_buffer. +static void launch_kernel(const PTXLaunchParams &lp, + std::vector &arg_buffer, + std::size_t buffer_size) { + cudaStream_t stream_ = lp.stream; // alias for DRIVER_ERROR_CHECK macro + void *config[] = { + CU_LAUNCH_PARAM_BUFFER_POINTER, + static_cast(arg_buffer.data()), + CU_LAUNCH_PARAM_BUFFER_SIZE, + &buffer_size, + CU_LAUNCH_PARAM_END, + }; + +#ifdef CUDA_DEBUG + std::cerr << "[RunPTXTask] Launching kernel " << lp.kernel_name + << " with blocks (" << lp.bx << "," << lp.by << "," << lp.bz + << ") and threads (" << lp.tx << "," << lp.ty << "," << lp.tz << ")" + << std::endl; +#endif + + DRIVER_ERROR_CHECK(cuLaunchKernel(lp.func, lp.bx, lp.by, lp.bz, lp.tx, lp.ty, + lp.tz, 0, lp.custream, nullptr, config)); +} + +// Helper: align pointer to 8-byte boundary. +static inline void align8(char *&ptr) { + std::uintptr_t addr = reinterpret_cast(ptr); + ptr = reinterpret_cast((addr + 7) & ~std::uintptr_t(7)); +} + +// RunPTXTask: user-defined @cuda_task kernels +// Arg buffer: [kernel_state | inputs... | outputs... | scalars...] +// https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/eval_udf_gpu.cc +/*static*/ void RunPTXTask::gpu_variant(legate::TaskContext context) { + auto lp = read_launch_params(context); const std::size_t num_inputs = context.num_inputs(); const std::size_t num_outputs = context.num_outputs(); const std::size_t num_scalars = context.num_scalars(); - const std::size_t num_reductions = - context.num_reductions(); // unused for now - // compute total size: all device arrays + all scalars - // skip scalar 0-2 (kernel_name, threads, blocks) - // we allocate extra to decrease looping and dynamic dispatching on dim std::size_t max_buffer_size = padded_bytes_kernel_state + (num_inputs + num_outputs) * sizeof(CuDeviceArray); @@ -238,26 +275,13 @@ struct ufiFunctor { for (std::size_t i = 0; i < num_inputs; ++i) { auto ps = context.input(i); - auto code = ps.type().code(); - auto dim = ps.dim(); -#ifdef CUDA_DEBUG - std::cerr << "[RunPTXTask] Input " << i << " type: " << code - << ", dim: " << dim << std::endl; -#endif - // dispatch on dim and code with ufiFunctor operator() - legate::double_dispatch(dim, code, ufiFunctor{}, ufi::AccessMode::READ, p, - ps); + legate::double_dispatch(ps.dim(), ps.type().code(), ufiFunctor{}, + ufi::AccessMode::READ, p, ps); } for (std::size_t i = 0; i < num_outputs; ++i) { auto ps = context.output(i); - auto code = ps.type().code(); - auto dim = ps.dim(); -#ifdef CUDA_DEBUG - std::cerr << "[RunPTXTask] Output " << i << " type: " << code - << ", dim: " << dim << std::endl; -#endif - legate::double_dispatch(dim, code, ufiFunctor{}, ufi::AccessMode::WRITE, p, - ps); + legate::double_dispatch(ps.dim(), ps.type().code(), ufiFunctor{}, + ufi::AccessMode::WRITE, p, ps); } for (std::size_t i = ARG_OFFSET; i < num_scalars; ++i) { const auto &scalar = context.scalar(i); @@ -265,29 +289,78 @@ struct ufiFunctor { p += scalar.size(); } - std::size_t buffer_size = p - arg_buffer.data(); // calc used buffer + launch_kernel(lp, arg_buffer, p - arg_buffer.data()); +} - void *config[] = { - CU_LAUNCH_PARAM_BUFFER_POINTER, - static_cast(arg_buffer.data()), - CU_LAUNCH_PARAM_BUFFER_SIZE, - &buffer_size, - CU_LAUNCH_PARAM_END, - }; +// RunPTXBroadcastTask: broadcast fusion kernels +// Arg buffer: [kernel_state | ctx | arg_map-driven args...] +// +// Scalars after ARG_OFFSET: +// [7] = ctx (CompilerMetadata, raw bytes) +// [8] = num_kernel_args (Int32) +// [9..8+N] = arg_map entries (Int32 each) +// [9+N..] = actual scalar values +// +// arg_map encoding: +// val >= 0, val < num_outputs → output[val] (write CuDeviceArray) +// val >= num_outputs → input[val - num_outputs] (read +// CuDeviceArray) val < 0 → scalar at index -(val + 1) in +// trailing scalars +/*static*/ void RunPTXBroadcastTask::gpu_variant(legate::TaskContext context) { + auto lp = read_launch_params(context); - CUstream custream_ = reinterpret_cast(stream_); + const std::size_t num_inputs = context.num_inputs(); + const std::size_t num_outputs = context.num_outputs(); + const std::size_t num_scalars = context.num_scalars(); + // Read num_kernel_args first so we can size the buffer precisely + std::int32_t num_kernel_args = + context.scalar(ARG_OFFSET + 1).value(); + std::size_t map_start = ARG_OFFSET + 2; + std::size_t scalar_values_start = map_start + num_kernel_args; -#ifdef CUDA_DEBUG - std::cerr << "[RunPTXTask] Launching kernel " << kernel_name - << " with blocks (" << bx << "," << by << "," << bz - << ") and threads (" << tx << "," << ty << "," << tz << ")" - << " on CUcontext " << context_to_string(ctx) << std::endl; -#endif - // Launch the kernel - DRIVER_ERROR_CHECK(cuLaunchKernel(func, bx, by, bz, tx, ty, tz, 0, custream_, - nullptr, config)); + std::size_t max_buffer_size = + padded_bytes_kernel_state + + context.scalar(ARG_OFFSET).size() + // ctx (CompilerMetadata) + num_kernel_args * + (sizeof(CuDeviceArray) + + 8); // worst case: all args are CuDeviceArrays + alignment + for (std::size_t i = scalar_values_start; i < num_scalars; ++i) { + max_buffer_size += context.scalar(i).size(); + } - // DRIVER_ERROR_CHECK(cuStreamSynchronize(stream_)); + std::vector arg_buffer(max_buffer_size); + char *p = arg_buffer.data() + padded_bytes_kernel_state; + + // 1. Write ctx (CompilerMetadata) + if (num_scalars > ARG_OFFSET) { + const auto &ctx_scalar = context.scalar(ARG_OFFSET); + memcpy(p, ctx_scalar.ptr(), ctx_scalar.size()); + p += ctx_scalar.size(); + } + + // 2. Read arg_map and reconstruct arg buffer + for (std::int32_t i = 0; i < num_kernel_args; ++i) { + std::int32_t val = context.scalar(map_start + i).value(); + + if (val >= 0 && val < static_cast(num_outputs)) { + align8(p); + auto ps = context.output(val); + legate::double_dispatch(ps.dim(), ps.type().code(), ufiFunctor{}, + ufi::AccessMode::WRITE, p, ps); + } else if (val >= static_cast(num_outputs)) { + align8(p); + auto ps = context.input(val - num_outputs); + legate::double_dispatch(ps.dim(), ps.type().code(), ufiFunctor{}, + ufi::AccessMode::READ, p, ps); + } else { + std::size_t scalar_idx = static_cast(-(val + 1)); + const auto &scalar = context.scalar(scalar_values_start + scalar_idx); + memcpy(p, scalar.ptr(), scalar.size()); + p += scalar.size(); + } + } + + launch_kernel(lp, arg_buffer, p - arg_buffer.data()); } // https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/load_ptx.cc @@ -389,6 +462,13 @@ inline void add_xyz_scalars(legate::AutoTask &task, task.add_scalar_arg(legate::Scalar(xyz[2])); } +inline void add_scalar_from_ptr(legate::AutoTask &task, void *ptr, + size_t size) { + uint8_t *byte_ptr = static_cast(ptr); + std::vector vec(byte_ptr, byte_ptr + size); + task.add_scalar_arg(legate::Scalar(vec)); +} + void gpu_sync() { cudaStream_t stream_ = nullptr; ERROR_CHECK(cudaDeviceSynchronize()); @@ -412,9 +492,12 @@ void register_kernel_state_size(uint64_t st_size) { void wrap_cuda_methods(jlcxx::Module &mod) { mod.method("add_xyz_scalars", &add_xyz_scalars); + mod.method("add_scalar_from_ptr", &add_scalar_from_ptr); mod.method("register_kernel_state_size", ®ister_kernel_state_size); mod.method("gpu_sync", &gpu_sync); mod.method("extract_kernel_name", &extract_kernel_name); mod.set_const("LOAD_PTX", legate::LocalTaskID{ufi::TaskIDs::LOAD_PTX_TASK}); mod.set_const("RUN_PTX", legate::LocalTaskID{ufi::TaskIDs::RUN_PTX_TASK}); + mod.set_const("RUN_PTX_BROADCAST", + legate::LocalTaskID{ufi::TaskIDs::RUN_PTX_BROADCAST_TASK}); } diff --git a/src/cuda/cuda_ptx_task.jl b/src/cuda/cuda_ptx_task.jl index 7cbb3e07..3c2594b8 100644 --- a/src/cuda/cuda_ptx_task.jl +++ b/src/cuda/cuda_ptx_task.jl @@ -73,7 +73,8 @@ function nda_to_logical_array(arr::NDArray{T,N}) where {T,N} end function Launch(kernel::CUDATask, inputs::Tuple{Vararg{NDArray}}, - outputs::Tuple{Vararg{NDArray}}, scalars::Tuple{Vararg{Any}}; blocks, threads) + outputs::Tuple{Vararg{NDArray}}, scalars::Tuple{Vararg{Any}}; + blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing) # we find the largest input/output. ndarrays = vcat(inputs..., outputs...) @@ -84,7 +85,6 @@ function Launch(kernel::CUDATask, inputs::Tuple{Vararg{NDArray}}, rt = Legate.get_runtime() lib = cuNumeric.get_lib() - taskid = cuNumeric.RUN_PTX task = Legate.create_auto_task(rt, lib, taskid) input_vars = Vector{Legate.Variable}() @@ -103,14 +103,22 @@ function Launch(kernel::CUDATask, inputs::Tuple{Vararg{NDArray}}, push!(output_vars, p) end - # next 3 lines are reserved scalars in the RUN_PTX task + # Reserved scalars: kernel_name (0), blocks (1,2,3), threads (4,5,6) Legate.add_scalar(task, Legate.string_to_scalar(kernel.func)) # 0 - cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, blocks)) # bx,by,bz 1,2,3 - cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, threads)) # tx,ty,tz 4,5,6 + cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, blocks)) # 1,2,3 + cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, threads)) # 4,5,6 + + # CompilerMetadata ctx for broadcast tasks (scalar 7, raw bytes) + if !isnothing(ctx) + ref = Ref(ctx) + GC.@preserve ref begin + cuNumeric.add_scalar_from_ptr(task, Base.unsafe_convert(Ptr{Cvoid}, ref), sizeof(ctx)) + end + end - # any user defined scalars in the launch macro + # User-defined scalars for s in scalars - Legate.add_scalar(task, Legate.Scalar(s)) # 7+ -> ARG_OFFSET + Legate.add_scalar(task, Legate.Scalar(s)) end # all inputs are aligned with all outputs @@ -119,13 +127,15 @@ function Launch(kernel::CUDATask, inputs::Tuple{Vararg{NDArray}}, end function launch(kernel::CUDATask, inputs, outputs, scalars; - blocks, threads) + blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing) Launch(kernel, isa(inputs, Tuple) ? inputs : (inputs,), isa(outputs, Tuple) ? outputs : (outputs,), isa(scalars, Tuple) ? scalars : (scalars,); blocks=isa(blocks, Tuple) ? blocks : (blocks,), threads=isa(threads, Tuple) ? threads : (threads,), + taskid=taskid, + ctx=ctx, ) end diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl index 078de5c8..b7ef5f9c 100644 --- a/src/ndarray/broadcast_fusion.jl +++ b/src/ndarray/broadcast_fusion.jl @@ -134,7 +134,7 @@ function get_ptx( arg_types...; ndrange, ) where {DEST_T,BC_T} - println(Base.isbitstype.(arg_types)) + # println(Base.isbitstype.(arg_types)) threads, blocks, ctx = _threads_from_occupancy(obj, DEST_T, arg_types...; ndrange=ndrange) blocks == 0 && return "", 0, 0, ctx #! MAYBE ERROR HERE? @@ -161,10 +161,13 @@ function get_cuda_task( # Also stores in cache if not found in Dict return get!(_BCAST_PTX_CACHE, key) do ptx, threads, blocks, ctx = get_ptx(obj, DEST_T, ARG_TYPES...; ndrange=ndrange) - func_name = extract_kernel_name(ptx) - println(ptx) - ptx_task(ptx, func_name) - cuda_task = CUDATask(func_name, (DEST_T, ARG_TYPES...)) + orig_name = extract_kernel_name(ptx) + # Append a PTX content hash to make each name unique. + unique_name = orig_name * "_" * string(hash(ptx); base=16) + ptx = replace(ptx, orig_name => unique_name) + # println(ptx) + ptx_task(ptx, unique_name) + cuda_task = CUDATask(unique_name, (DEST_T, ARG_TYPES...)) FusedBroadcastMetadata(ctx, threads, blocks, cuda_task) end end @@ -213,30 +216,59 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas # Spoof NDArrays with CuDeviceArrays in the Broadcasted type spoofed_bc_args_type = map_cuda_type(typeof(bc.args)) - # Figure out where input arrays are - cudevicearray_offsets, cudevicearray_indices = find_cudevicearray_offsets_and_indices( - spoofed_bc_args_type - ) - - # Figure out where scalars are - scalar_offsets, scalar_indices = find_scalar_offsets_and_indices(spoofed_bc_args_type) + # Build the arg_map: for each kernel arg (after ctx), record its source. + # Convention: index into combined [outputs..., inputs...] for NDArrays. + # idx < num_outputs → output[idx] + # idx >= num_outputs → input[idx - num_outputs] + # Scalar values are passed separately after the mapping. + # + # PTX kernel signature: f(kernel_state, ctx, dest, bc.args...) + # So the arg_map covers: [dest, bc.args...] in order. + + num_outputs = 1 # dest is always the single output + + # Deduplicate NDArray inputs: map each unique NDArray to a unique input index. + # The arg_map can reference the same input index multiple times (e.g. a .+ a). + unique_ndarrays = NDArray[] + ndarray_to_input_idx = Dict{UInt,Int}() # objectid → input index + + arg_map = Int32[] + actual_scalars = Any[] + + # First arg in PTX (after ctx) is always dest = output[0] + push!(arg_map, Int32(0)) # output[0] + + # Then bc.args... in order + for (i, arg) in enumerate(bc.args) + if stores_cudevicearray(map_cuda_type(typeof(arg))) + nda = get_ndarray(arg) + oid = objectid(nda) + if !haskey(ndarray_to_input_idx, oid) + push!(unique_ndarrays, nda) + ndarray_to_input_idx[oid] = length(unique_ndarrays) - 1 # 0-based + end + input_idx = ndarray_to_input_idx[oid] + push!(arg_map, Int32(num_outputs + input_idx)) # offset by num_outputs + else + # Scalar — record its position and save the value + push!(arg_map, Int32(-1 - length(actual_scalars))) # negative = scalar + push!(actual_scalars, arg) + end + end - input_ndarrays = ntuple( - i -> get_ndarray(bc.args[cudevicearray_indices[i]]), length(cudevicearray_indices) - ) - input_scalars = ntuple(i -> bc.args[scalar_indices[i]], length(scalar_indices)) + input_ndarrays = tuple(unique_ndarrays...) - #! I AM NOT SURE IT IS CORRECT TO ASSUME SCALARS ARE ALWAYS AFTER NDARARYS - #! THE PTX KERNEL SIGNATURE BEFORE MODIFICATION BY CUDA.jl: - #! (fkm.ctx, dest, bc.args...) + # @show arg_map actual_scalars length(unique_ndarrays) launch( fkm.cuda_task, input_ndarrays, (dest,), - input_scalars; + (Int32(length(arg_map)), arg_map..., actual_scalars...); blocks=fkm.blocks, threads=fkm.threads, + taskid=cuNumeric.RUN_PTX_BROADCAST, + ctx=fkm.ctx, ) #! DO I NEED TO DO TYPE PROMOTION CHECKS?? diff --git a/test/runtests.jl b/test/runtests.jl index d2396b3f..7aa99150 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,4 @@ -#= Copyright 2026 Northwestern University, +#= Copyright 2026 Northwestern University, * Carnegie Mellon University University * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -57,6 +57,7 @@ include("tests/unary_tests.jl") include("tests/binary_tests.jl") include("tests/scoping.jl") include("tests/scoping-advanced.jl") +include("tests/broadcast_fusion_tests.jl") @testset verbose = true "AXPY" begin N = 100 @@ -430,11 +431,15 @@ end end if run_gpu_tests + @testset verbose = true "Broadcast Fusion" begin + test_broadcast_fusion() + end + # @testset verbose = true "CUDA Tests" begin # cuda_unaryop(rtol(Float32)) # cuda_binaryop(rtol(Float32)) # end - @warn "CUDA tests are turned off inside Pkg.test for now. --check-bounds=yes causes issues." + @warn "CUDA @cuda_task tests are turned off inside Pkg.test for now. --check-bounds=yes causes issues." else @warn "The CUDA tests will not be run as a CUDA-enabled device is not available" end diff --git a/test/tests/broadcast_fusion_tests.jl b/test/tests/broadcast_fusion_tests.jl new file mode 100644 index 00000000..9c5be8a2 --- /dev/null +++ b/test/tests/broadcast_fusion_tests.jl @@ -0,0 +1,185 @@ +#= Copyright 2026 Northwestern University, + * Carnegie Mellon University University + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * Author(s): David Krasowska + * Ethan Meitz +=# + +#= Broadcast fusion arg buffer tests. + * Tests every combination of array and scalar argument positions + * to verify the arg_map protocol correctly reconstructs the CUDA + * kernel arg buffer. + * + * Pattern naming convention: + * A = NDArray input + * S = scalar input + * Position = left-to-right in the broadcast expression + * "same" = same array appears multiple times (deduplication test) +=# + +function test_broadcast_fusion(; T=Float32, N=100, atol=1e-5, rtol=1e-5) + # Create test arrays with known non-zero values + julia_a = rand(T, N) + julia_b = rand(T, N) + julia_c = rand(T, N) + + a = @allowscalar NDArray(julia_a) + b = @allowscalar NDArray(julia_b) + c = @allowscalar NDArray(julia_c) + + s1 = T(2.5) + s2 = T(1.0) + s3 = T(0.5) + + # two different arrays + @testset "A + B (two different arrays)" begin + expected = julia_a .+ julia_b + result = a .+ b + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # same array, deduplication + @testset "A + A (same array twice)" begin + expected = julia_a .+ julia_a + result = a .+ a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # array then scalar + @testset "A + scalar (array first)" begin + expected = julia_a .+ s1 + result = a .+ s1 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # scalar then array + @testset "scalar + A (scalar first)" begin + expected = s1 .+ julia_a + result = s1 .+ a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # array, two scalars, fused + @testset "A * scalar - scalar (fused)" begin + expected = julia_a .* s1 .- s2 + result = a .* s1 .- s2 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # scalar-array-scalar + @testset "scalar * A + scalar" begin + expected = s1 .* julia_a .+ s2 + result = s1 .* a .+ s2 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # two arrays then scalar + @testset "A + B + scalar" begin + expected = julia_a .+ julia_b .+ s1 + result = a .+ b .+ s1 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # scalar then two arrays + @testset "scalar + A + B" begin + expected = s1 .+ julia_a .+ julia_b + result = s1 .+ a .+ b + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # same array twice + scalar (dedup) + @testset "A + A + scalar (dedup + scalar)" begin + expected = julia_a .+ julia_a .+ s1 + result = a .+ a .+ s1 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # three different arrays + @testset "A + B + C (three arrays)" begin + expected = julia_a .+ julia_b .+ julia_c + result = a .+ b .+ c + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # same array three times (triple dedup) + @testset "A + A + A (triple dedup)" begin + expected = julia_a .+ julia_a .+ julia_a + result = a .+ a .+ a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # two scalars then array + @testset "scalar * scalar + A" begin + expected = s1 .* s2 .+ julia_a + result = s1 .* s2 .+ a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # multiply, two arrays (different PTX kernel name collision test) + @testset "A * B (multiply)" begin + expected = julia_a .* julia_b + result = a .* b + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # same array squared, dedup + @testset "A * A (self multiply, dedup)" begin + expected = julia_a .* julia_a + result = a .* a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # subtraction, order matters + @testset "A - B (subtraction)" begin + expected = julia_a .- julia_b + result = a .- b + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # scalar minus array + @testset "scalar - A" begin + expected = s1 .- julia_a + result = s1 .- a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # three arrays, mixed ops + @testset "A * B + C (three arrays, mixed ops)" begin + expected = julia_a .* julia_b .+ julia_c + result = a .* b .+ c + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # two arrays fused, then scaled + @testset "(A + B) * scalar" begin + expected = (julia_a .+ julia_b) .* s1 + result = (a .+ b) .* s1 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # scalar-array pairs + @testset "scalar * A + scalar * B" begin + expected = s1 .* julia_a .+ s2 .* julia_b + result = s1 .* a .+ s2 .* b + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # same array subtracted, should be all zeros + @testset "A - A (same array, expect zeros)" begin + expected = julia_a .- julia_a + result = a .- a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end +end From 50b08a2896e6c3b1243eccb6ac05ccfc09b90685 Mon Sep 17 00:00:00 2001 From: Ethan Meitz Date: Fri, 22 May 2026 16:25:48 +0000 Subject: [PATCH 11/12] gray scott benchmark Co-authored-by: Copilot --- benchmark/gray_scott_fusion.jl | 159 ++++++++++++++++++++++++ deps/build.jl | 43 ++++--- lib/cunumeric_jl_wrapper/CMakeLists.txt | 60 ++++----- scripts/build_cpp_wrapper.sh | 91 +++++--------- scripts/install_cxxwrap.sh | 2 +- 5 files changed, 239 insertions(+), 116 deletions(-) create mode 100644 benchmark/gray_scott_fusion.jl diff --git a/benchmark/gray_scott_fusion.jl b/benchmark/gray_scott_fusion.jl new file mode 100644 index 00000000..3ce91151 --- /dev/null +++ b/benchmark/gray_scott_fusion.jl @@ -0,0 +1,159 @@ +using cuNumeric +using Random +using StatsBase + +cuNumeric.versioninfo() + +Random.seed!(1234) + +struct Params{T} + dx::T + dt::T + c_u::T + c_v::T + f::T + k::T + + function Params(dx=1.0f0, c_u=1.0f0, c_v=0.3f0, f=0.03f0, k=0.06f0) + new{Float32}(dx, dx/5, c_u, c_v, f, k) + end +end + +function bc!(u_new, v_new, u, v) + u_new[:, 1] = u[:, end - 1] + u_new[:, end] = u[:, 2] + u_new[1, :] = u[end - 1, :] + u_new[end, :] = u[2, :] + v_new[:, 1] = v[:, end - 1] + v_new[:, end] = v[:, 2] + v_new[1, :] = v[end - 1, :] + v_new[end, :] = v[2, :] +end + +function step!(u, v, u_new, v_new, args::Params) + + dx_sq = args.dx^2 + + # AVOID SOME ISSUES WITH ^ operator + v_sq = v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)] + + # calculate F_u and F_v functions + F_u = ( + (-u[2:(end - 1), 2:(end - 1)] .* (v_sq)) .+ + args.f*(1.0f0 .- u[2:(end - 1), 2:(end - 1)]) + ) + F_v = ( + (u[2:(end - 1), 2:(end - 1)] .* (v_sq)) .- + (args.f+args.k) .* v[2:(end - 1), 2:(end - 1)] + ) + # 2-D Laplacian of f using array slicing, excluding boundaries + # For an N x N array f, f_lap is the Nend x Nend array in the "middle" + u_lap = ( + (u[3:end, 2:(end - 1)] - 2*u[2:(end - 1), 2:(end - 1)] + u[1:(end - 2), 2:(end - 1)]) ./ + dx_sq + + + (u[2:(end - 1), 3:end] - 2*u[2:(end - 1), 2:(end - 1)] + u[2:(end - 1), 1:(end - 2)]) ./ + dx_sq + ) + v_lap = ( + (v[3:end, 2:(end - 1)] - 2*v[2:(end - 1), 2:(end - 1)] + v[1:(end - 2), 2:(end - 1)]) ./ + dx_sq + + + (v[2:(end - 1), 3:end] - 2*v[2:(end - 1), 2:(end - 1)] + v[2:(end - 1), 1:(end - 2)]) ./ + dx_sq + ) + + # Forward-Euler time step for all points except the boundaries + u_new[2:(end - 1), 2:(end - 1)] = + ((args.c_u * u_lap) + F_u) * args.dt + u[2:(end - 1), 2:(end - 1)] + v_new[2:(end - 1), 2:(end - 1)] = + ((args.c_v * v_lap) + F_v) * args.dt + v[2:(end - 1), 2:(end - 1)] + + # Apply periodic boundary conditions + bc!(u_new, v_new, u, v) +end + +function initialize(N) + u = cuNumeric.ones(N, N) + v = cuNumeric.zeros(N, N) + u_new = cuNumeric.zeros(N, N) + v_new = cuNumeric.zeros(N, N) + + rand_frac = round(Int, 0.15 * N) + + # Make julia arrays first for reproducability, I don't think + # we have a way to set cuNumeric's random seed. + u_tmp = rand(Float32, rand_frac, rand_frac) + v_tmp = rand(Float32, rand_frac, rand_frac) + + # Copy over values to cuNumeric arrays + allowscalar() do + for i in 1:rand_frac, j in 1:rand_frac + u[i,j] = u_tmp[i,j] + v[i,j] = v_tmp[i,j] + end + end + + GC.gc() # remove any intermediates + + return u, v, u_new, v_new +end + +# Re-iplement loop without plotting, and add timing +function gray_scott(N::Int, n_steps::Int, n_warmup::Int) + + args = Params() + u, v, u_new, v_new = initialize(N) + @info "Initialized Arrays" + + for n in 1:n_warmup + @info "Warmup Step $n / $n_warmup" + step!(u, v, u_new, v_new, args) + u, u_new = u_new, u + v, v_new = v_new, v + end + + step_times_us = Vector{Float64}(undef, n_steps) + for n in 1:n_steps + @info "Step $n / $n_steps" + start_time = get_time_microseconds() + step!(u, v, u_new, v_new, args) + u, u_new = u_new, u + v, v_new = v_new, v + step_times_us[n] = get_time_microseconds() - start_time + + if n%10 == 0 + GC.gc() + end + end + + return u, v, step_times_us +end + +gpus = parse(Int, ARGS[1]) +N = parse(Int, ARGS[2]) +n_steps = parse(Int, ARGS[3]) +n_warmup = parse(Int, ARGS[4]) + +println( + "[cuNumeric] Gray-Scott benchmark on $(N)x$(N) grid for $(n_steps) iterations, $(n_warmup) warmups" +) + +u, v, step_times_us = gray_scott(N, n_steps, n_warmup) + +step_times_ms = step_times_us ./ 1000.0 +mean_time_ms = mean(step_times_ms) +std_time_ms = std(step_times_ms) + +total_time_s = sum(step_times_ms) / 1000.0 +iter_per_second = n_steps / total_time_s + +println("Mean Step Time: $(mean_time_ms) ms") +println("StdDev Step Time: $(std_time_ms) ms") +println("Min Step Time: $(minimum(step_times_ms)) ms") +println("Max Step Time: $(maximum(step_times_ms)) ms") +println("Throughput: $(iter_per_second) ") + +# open("./gray_scott.csv", "a") do io +# @printf(io, "%s,%d,%d,%d,%.6f,%.6f\n", "cunumeric", gpus, N, M, mean_time_ms, gflops) +# end \ No newline at end of file diff --git a/deps/build.jl b/deps/build.jl index ae5e07a6..7899822c 100644 --- a/deps/build.jl +++ b/deps/build.jl @@ -67,17 +67,22 @@ end function build_jlcxxwrap(repo_root, cupynumeric_root) build_libcxxwrap = joinpath(repo_root, "scripts/install_cxxwrap.sh") version_path = joinpath(DEPOT_PATH[1], "dev/libcxxwrap_julia_jll/override/LEGATE_INSTALL.txt") - if isfile(version_path) - version = VersionNumber(strip(read(version_path, String))) - @info "libcxxwrap: Found cuNumeric $version" - if is_supported_version(version) - @info "libcxxwrap: Found supported version built with cuNumeric.jl: $version" - return nothing + jlcxxwrap_libpath = joinpath(DEPOT_PATH[1], "dev/libcxxwrap_julia_jll/override/lib/libjlcxxwrap_julia.so") + if isfile(jlcxxwrap_libpath) + if isfile(version_path) + version = VersionNumber(strip(read(version_path, String))) + @info "libcxxwrap: Found cuNumeric $version" + if is_supported_version(version) + @info "libcxxwrap: Found supported version built with cuNumeric.jl: $version" + return nothing + else + @info "libcxxwrap: Unsupported version found: $version. Rebuilding..." + end else - @info "libcxxwrap: Unsupported version found: $version. Rebuilding..." + @info "libcxxwrap: No version file found. Starting build..." end else - @info "libcxxwrap: No version file found. Starting build..." + @info "libcxxwrap: No jlcxxwrap_julia library found. Starting build..." end @info "libcxxwrap: Running build script: $build_libcxxwrap" @@ -89,7 +94,7 @@ function build_jlcxxwrap(repo_root, cupynumeric_root) end function build_cpp_wrapper( - repo_root, cupynumeric_loc, legate_loc, blas_loc, install_root, cuda_header_loc, cuda_driver_loc + repo_root, cupynumeric_loc, legate_loc, blas_loc, install_root ) @info "libcunumeric_jl_wrapper: Building C++ Wrapper Library" if isdir(install_root) @@ -100,7 +105,7 @@ function build_cpp_wrapper( build_cpp_wrapper = joinpath(repo_root, "scripts/build_cpp_wrapper.sh") nthreads = Threads.nthreads() - bld_command = `$build_cpp_wrapper $repo_root $cupynumeric_loc $legate_loc $blas_loc $cuda_header_loc $cuda_driver_loc $install_root $nthreads` + bld_command = `$build_cpp_wrapper $repo_root $cupynumeric_loc $legate_loc $blas_loc $install_root $nthreads` # write out a bash script for debugging cmd_str = join(bld_command.exec, " ") @@ -161,7 +166,7 @@ end """ build CxxWrap and cunumeric_jl_wrapper """ -function build_deps(pkg_root, cupynumeric_root, blas_root, cuda_header_path, cuda_driver_path) +function build_deps(pkg_root, cupynumeric_root, blas_root) legate_lib = Legate.get_install_liblegate() install_lib = joinpath(pkg_root, "lib", "cunumeric_jl_wrapper", "build") if !cupynumeric_valid(cupynumeric_root) @@ -174,7 +179,7 @@ function build_deps(pkg_root, cupynumeric_root, blas_root, cuda_header_path, cud build_jlcxxwrap(pkg_root, cupynumeric_root) build_cpp_wrapper( pkg_root, cupynumeric_root, up_dir(legate_lib), blas_root, - install_lib, cuda_header_path, cuda_driver_path, + install_lib ) # $pkg_root/lib/cunumeric_jl_wrapper end @@ -201,7 +206,7 @@ function build(::CNPreferences.Conda) #!TODO SET LocalPreferences.toml to use local CUDA libraries is_cupynumeric_installed(cupynumeric_root; throw_errors=true) - build_deps(pkg_root, cupynumeric_root, cupynumeric_root, cuda_toolkit_root) # blas is same root as cupynumeric + build_deps(pkg_root, cupynumeric_root, cupynumeric_root) # blas is same root as cupynumeric end function build(::CNPreferences.Developer) @@ -216,8 +221,8 @@ function build(::CNPreferences.Developer) if isnothing(cupynumeric_root) # We are using cupynumeric_jll. cupynumeric_root = cupynumeric_jll.artifact_dir - cuda_header_path = dirname(cupynumeric_jll.cuda_header) - cuda_driver_path = joinpath(CUDACore.CUDA_Driver_jll.artifact_dir, "lib", "libcuda.so") + # cuda_header_path = dirname(cupynumeric_jll.cuda_header) + # cuda_driver_path = joinpath(CUDACore.CUDA_Driver_jll.artifact_dir, "lib", "libcuda.so") else # User provided a custom cupynumeric install. is_cupynumeric_installed(cupynumeric_root; throw_errors=true) @@ -228,8 +233,8 @@ function build(::CNPreferences.Developer) ) # Local full CUDA toolkit path. - cuda_header_path = _cuda_include_dir_from_toolkit(cuda_toolkit_root) - cuda_driver_path = _cuda_driver_lib_from_toolkit(cuda_toolkit_root) + # cuda_header_path = _cuda_include_dir_from_toolkit(cuda_toolkit_root) + # cuda_driver_path = _cuda_driver_lib_from_toolkit(cuda_toolkit_root) end if isnothing(blas_lib_dir) @@ -239,9 +244,7 @@ function build(::CNPreferences.Developer) build_deps( pkg_root, cupynumeric_root, - blas_lib_dir, - cuda_header_path, - cuda_driver_path, + blas_lib_dir ) end diff --git a/lib/cunumeric_jl_wrapper/CMakeLists.txt b/lib/cunumeric_jl_wrapper/CMakeLists.txt index 56261c3d..87848a77 100644 --- a/lib/cunumeric_jl_wrapper/CMakeLists.txt +++ b/lib/cunumeric_jl_wrapper/CMakeLists.txt @@ -3,7 +3,6 @@ project(cuNumericWrapper) set(cuNumericWrapperVersion 0.0.1) message(STATUS "Project version: v${cuNumericWrapperVersion}") - set(CXX_CUNUMERICJL_WRAPPER cunumeric_jl_wrapper) set(C_INTERFACE_LIB cunumeric_c_wrapper) @@ -12,30 +11,30 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib") -option(NOCUDA "Build without CUDA support (skip src/cuda.cpp)" OFF) +# ---- New: NOCUDA option ---- +option(NOCUDA "Build without CUDA support (skip CUDAToolkit and src/cuda.cpp)" OFF) option(BINARYBUILDER "Building with binary builder" ON) -# Always needed + +# Always needed (unless your packages themselves require CUDA at configure time) find_package(legate REQUIRED) find_package(cupynumeric REQUIRED) # CxxWrap Stuff if(NOT BINARYBUILDER) - execute_process( - COMMAND julia -e "println(DEPOT_PATH[1])" - OUTPUT_VARIABLE JULIA_DEP_PATH - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - - set(JlCxx_DIR "${JULIA_DEP_PATH}/dev/libcxxwrap_julia_jll/override") - message(STATUS "Setting JlCxx_DIR to ${JlCxx_DIR} (not using BinaryBuilder)") +execute_process( + COMMAND julia -e "println(DEPOT_PATH[1])" + OUTPUT_VARIABLE JULIA_DEP_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE +) + +set(JlCxx_DIR "${JULIA_DEP_PATH}/dev/libcxxwrap_julia_jll/override") +message(STATUS "Setting JlCxx_DIR to ${JlCxx_DIR} (not using BinaryBuilder)") endif() find_package(JlCxx REQUIRED) - get_target_property(JlCxx_location JlCxx::cxxwrap_julia LOCATION) get_filename_component(JlCxx_location ${JlCxx_location} DIRECTORY) - set(CMAKE_INSTALL_RPATH "${CMAKE_INSTALL_PREFIX}/lib;${JlCxx_location}") message(STATUS "Found JlCxx at ${JlCxx_location}") @@ -45,34 +44,21 @@ set(SOURCES src/types.cpp ) -# CUDA Driver API only: requires cuda.h and libcuda.so, not a full CUDA Toolkit. +# Conditionally add CUDA bits if(NOT NOCUDA) - set(CUDA_DRIVER_INCLUDE_DIR "" CACHE PATH "Directory containing cuda.h") - set(CUDA_DRIVER_LIBRARY "" CACHE FILEPATH "Full path to libcuda.so") - - add_library(CUDA_driver_manual UNKNOWN IMPORTED) - set_target_properties(CUDA_driver_manual PROPERTIES - IMPORTED_LOCATION "${CUDA_DRIVER_LIBRARY}" - INTERFACE_INCLUDE_DIRECTORIES "${CUDA_DRIVER_INCLUDE_DIR}" - ) - + find_package(CUDAToolkit 13.0 REQUIRED) list(APPEND SOURCES src/cuda.cpp) add_compile_definitions(HAVE_CUDA) set(HAVE_CUDA TRUE) - message(STATUS "CUDA enabled: adding src/cuda.cpp") - message(STATUS "CUDA driver include dir: ${CUDA_DRIVER_INCLUDE_DIR}") - message(STATUS "CUDA driver library: ${CUDA_DRIVER_LIBRARY}") else() set(HAVE_CUDA FALSE) - message(STATUS "NOCUDA=ON: skipping src/cuda.cpp") + message(STATUS "NOCUDA=ON: skipping CUDAToolkit and src/cuda.cpp") endif() # Library: C++ wrapper add_library(${CXX_CUNUMERICJL_WRAPPER} SHARED ${SOURCES}) -set_target_properties(${CXX_CUNUMERICJL_WRAPPER} PROPERTIES - VERSION ${cuNumericWrapperVersion} -) +set_target_properties(${CXX_CUNUMERICJL_WRAPPER} PROPERTIES VERSION ${cuNumericWrapperVersion}) target_link_libraries(${CXX_CUNUMERICJL_WRAPPER} PRIVATE cupynumeric::cupynumeric @@ -81,12 +67,13 @@ target_link_libraries(${CXX_CUNUMERICJL_WRAPPER} PRIVATE JlCxx::cxxwrap_julia_stl ) +# Include dirs (conditionally add CUDA include path) target_include_directories(${CXX_CUNUMERICJL_WRAPPER} PRIVATE include) - if(HAVE_CUDA) - target_link_libraries(${CXX_CUNUMERICJL_WRAPPER} PRIVATE CUDA_driver_manual) + target_include_directories(${CXX_CUNUMERICJL_WRAPPER} PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) endif() + install(TARGETS ${CXX_CUNUMERICJL_WRAPPER} DESTINATION lib) # ---- C API ---- @@ -96,9 +83,7 @@ set(C_SOURCES ) add_library(${C_INTERFACE_LIB} SHARED ${C_SOURCES}) -set_target_properties(${C_INTERFACE_LIB} PROPERTIES - VERSION ${cuNumericWrapperVersion} -) +set_target_properties(${C_INTERFACE_LIB} PROPERTIES VERSION ${cuNumericWrapperVersion}) target_link_libraries(${C_INTERFACE_LIB} PRIVATE cupynumeric::cupynumeric @@ -106,9 +91,8 @@ target_link_libraries(${C_INTERFACE_LIB} PRIVATE ) target_include_directories(${C_INTERFACE_LIB} PRIVATE include) - if(HAVE_CUDA) - target_link_libraries(${C_INTERFACE_LIB} PRIVATE CUDA_driver_manual) + target_include_directories(${C_INTERFACE_LIB} PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) endif() -install(TARGETS ${C_INTERFACE_LIB} DESTINATION lib) +install(TARGETS ${C_INTERFACE_LIB} DESTINATION lib) \ No newline at end of file diff --git a/scripts/build_cpp_wrapper.sh b/scripts/build_cpp_wrapper.sh index 720779f2..6bd7e73d 100755 --- a/scripts/build_cpp_wrapper.sh +++ b/scripts/build_cpp_wrapper.sh @@ -1,22 +1,18 @@ -set -euo pipefail +set -e -# Default to OFF: CUDA support enabled. -# Set NO_CUDA=ON to skip CUDA. -NO_CUDA=${NO_CUDA:-OFF} - -if [[ $# -ne 8 ]]; then - echo "Usage: $0 " +# Check if exactly one argument is provided +if [[ $# -ne 6 ]]; then + echo "Usage: $0 " exit 1 fi - -CUNUMERICJL_ROOT_DIR=$1 +CUNUMERICJL_ROOT_DIR=$1 # this is the repo root of cunumeric.jl CUPYNUMERIC_ROOT_DIR=$2 LEGATE_ROOT_DIR=$3 BLAS_LIB_DIR=$4 -CUDA_DRIVER_INCLUDE_DIR=$5 -CUDA_DRIVER_LIBRARY=$6 -INSTALL_DIR=$7 -NTHREADS=$8 +INSTALL_DIR=$5 +NTHREADS=$6 + +# Check if the provided argument is a valid directory if [[ ! -d "$CUNUMERICJL_ROOT_DIR" ]]; then echo "Error: '$CUNUMERICJL_ROOT_DIR' is not a valid directory." @@ -33,54 +29,35 @@ if [[ ! -d "$LEGATE_ROOT_DIR" ]]; then exit 1 fi -if [[ ! -d "$BLAS_LIB_DIR" ]]; then - echo "Error: '$BLAS_LIB_DIR' is not a valid directory." - exit 1 -fi - -if [[ ! -f "$BLAS_LIB_DIR/libopenblas.so" ]]; then - echo "Error: '$BLAS_LIB_DIR/libopenblas.so' does not exist." - exit 1 -fi +CUNUMERIC_WRAPPER_SOURCE=$CUNUMERICJL_ROOT_DIR/lib/cunumeric_jl_wrapper +BUILD_DIR=$CUNUMERIC_WRAPPER_SOURCE/build -if [[ "$NO_CUDA" != "ON" ]]; then - if [[ ! -f "$CUDA_DRIVER_INCLUDE_DIR/cuda.h" ]]; then - echo "Error: '$CUDA_DRIVER_INCLUDE_DIR/cuda.h' does not exist." - exit 1 - fi - - if [[ ! -f "$CUDA_DRIVER_LIBRARY" ]]; then - echo "Error: '$CUDA_DRIVER_LIBRARY' does not exist." - exit 1 - fi +if [[ ! -d "$BUILD_DIR" ]]; then + mkdir -p $BUILD_DIR fi -CUNUMERIC_WRAPPER_SOURCE="$CUNUMERICJL_ROOT_DIR/lib/cunumeric_jl_wrapper" -BUILD_DIR="$CUNUMERIC_WRAPPER_SOURCE/build" - -mkdir -p "$BUILD_DIR" -mkdir -p "$INSTALL_DIR" - -echo "LEGATE_ROOT_DIR: $LEGATE_ROOT_DIR" -echo "NO_CUDA: $NO_CUDA" - -if [[ "$NO_CUDA" != "ON" ]]; then - echo "CUDA driver include dir: $CUDA_DRIVER_INCLUDE_DIR" - echo "CUDA driver library: $CUDA_DRIVER_LIBRARY" +if [[ ! -d "$INSTALL_DIR" ]]; then + mkdir -p $INSTALL_DIR fi -echo "Configuring project..." +echo $LEGATE_ROOT_DIR -cmake -S "$CUNUMERIC_WRAPPER_SOURCE" -B "$BUILD_DIR" \ - -D BINARYBUILDER=OFF \ - -D NOCUDA="$NO_CUDA" \ - -D CUDA_DRIVER_INCLUDE_DIR="$CUDA_DRIVER_INCLUDE_DIR" \ - -D CUDA_DRIVER_LIBRARY="$CUDA_DRIVER_LIBRARY" \ - -D CMAKE_INSTALL_PREFIX="$INSTALL_DIR" \ - -D CMAKE_PREFIX_PATH="$CUPYNUMERIC_ROOT_DIR;$LEGATE_ROOT_DIR;" \ - -D CUPYNUMERIC_PATH="$CUPYNUMERIC_ROOT_DIR" \ - -D BLAS_LIBRARIES="$BLAS_LIB_DIR/libopenblas.so" \ - -D PROJECT_INSTALL_PATH="$INSTALL_DIR" \ - -D CMAKE_BUILD_TYPE=Release +# Default to OFF (CUDA support enabled), but allow override via environment variable +NO_CUDA=${NO_CUDA:-OFF} -cmake --build "$BUILD_DIR" --parallel "$NTHREADS" --verbose +if [[ ! -f "$BUILD_DIR/CMakeCache.txt" ]]; then + echo "Configuring project..." + cmake -S "$CUNUMERIC_WRAPPER_SOURCE" -B "$BUILD_DIR" \ + -D BINARYBUILDER=OFF \ + -D NOCUDA=$NO_CUDA \ + -D CMAKE_INSTALL_PREFIX="$INSTALL_DIR" \ + -D CMAKE_PREFIX_PATH="$CUPYNUMERIC_ROOT_DIR;$LEGATE_ROOT_DIR;" \ + -D CUPYNUMERIC_PATH="$CUPYNUMERIC_ROOT_DIR" \ + -D BLAS_LIBRARIES="$BLAS_LIB_DIR/libopenblas.so" \ + -D PROJECT_INSTALL_PATH="$INSTALL_DIR" \ + -D CMAKE_BUILD_TYPE=Releases +else + echo "Skipping configure (already done in $BUILD_DIR)" +fi + +cmake --build "$BUILD_DIR" --parallel "$NTHREADS" --verbose \ No newline at end of file diff --git a/scripts/install_cxxwrap.sh b/scripts/install_cxxwrap.sh index ba430d50..8335858c 100755 --- a/scripts/install_cxxwrap.sh +++ b/scripts/install_cxxwrap.sh @@ -69,7 +69,7 @@ JULIA_CXXWRAP=$JULIA_CXXWRAP_DEV/override cd $CUNUMERIC_ROOT_DIR [ -f Manifest.toml ] && rm Manifest.toml rm -rf $JULIA_CXXWRAP_DEV -julia -e 'using Pkg; Pkg.activate("."); Pkg.add(url="https://github.com/JuliaLegate/Legate.jl")' +julia -e 'using Pkg; Pkg.activate("."); Pkg.add("Legate")' julia -e 'using Pkg; Pkg.activate("."); Pkg.precompile(["CxxWrap"])' # https://github.com/JuliaInterop/libcxxwrap-julia/tree/v0.13.3?tab=readme-ov-file#preparing-the-install-location From 11a1ee6ee20758e74d065b7994c902d95b278bb1 Mon Sep 17 00:00:00 2001 From: ejmeitz Date: Mon, 1 Jun 2026 12:14:40 -0500 Subject: [PATCH 12/12] handle exponentiation --- src/ndarray/broadcast_fusion.jl | 245 +++++++++++++++++++++----------- 1 file changed, 160 insertions(+), 85 deletions(-) diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl index b7ef5f9c..6961d6a9 100644 --- a/src/ndarray/broadcast_fusion.jl +++ b/src/ndarray/broadcast_fusion.jl @@ -1,4 +1,122 @@ +struct RuntimeBroadcastArg{J} end +struct StaticBroadcastArg{J} end + +Base.@propagate_inbounds @inline function _gpu_broadcast_getindex(x, I) + return @inbounds Base.Broadcast._broadcast_getindex(x, I) +end + +Base.@propagate_inbounds @inline function _gpu_broadcast_getindex(x::Number, I) + return x +end + +Base.@propagate_inbounds @inline function _materialize_broadcast_arg( + ::RuntimeBroadcastArg{J}, + runtime_args, + static_args, + I, +) where {J} + arg = getfield(runtime_args, J) + return @inbounds _gpu_broadcast_getindex(arg, I) +end + +Base.@propagate_inbounds @inline function _materialize_broadcast_arg( + ::StaticBroadcastArg{J}, + runtime_args, + static_args, + I, +) where {J} + return getfield(static_args, J) +end + +Base.@propagate_inbounds @inline function _materialize_broadcast_args( + ::Tuple{}, + runtime_args, + static_args, + I, +) + return () +end + +Base.@propagate_inbounds @inline function _materialize_broadcast_args( + plan::Tuple, + runtime_args, + static_args, + I, +) + head = getfield(plan, 1) + tail = Base.tail(plan) + + return ( + @inbounds(_materialize_broadcast_arg(head, runtime_args, static_args, I)), + @inbounds(_materialize_broadcast_args(tail, runtime_args, static_args, I))..., + ) +end + +_is_runtime_broadcast_arg(x::NDArray) = true +_is_runtime_broadcast_arg(x::Base.Broadcast.Extruded) = true +_is_runtime_broadcast_arg(x::Number) = true +_is_runtime_broadcast_arg(x) = false + +function _push_runtime_arg!(runtime_args, arg_plan, x) + push!(runtime_args, x) + push!(arg_plan, RuntimeBroadcastArg{length(runtime_args)}()) + return nothing +end + +function _push_static_arg!(static_args, arg_plan, x) + isbits(x) || throw( + ArgumentError( + "Broadcast fusion cannot statically capture non-isbits broadcast leaf " * + "$(repr(x)) of type $(typeof(x))", + ), + ) + + push!(static_args, x) + push!(arg_plan, StaticBroadcastArg{length(static_args)}()) + return nothing +end + +function split_broadcast_args_for_kernel(args::Tuple) + runtime_args = Any[] + static_args = Any[] + arg_plan = Any[] + + for arg in args + if arg isa Base.RefValue + value = arg[] + + # Keep ordinary numeric broadcast scalars dynamic so scalar values + # do not cause recompilation. + if value isa Number + _push_runtime_arg!(runtime_args, arg_plan, value) + else + # Function singletons, Val{N}(), etc. + _push_static_arg!(static_args, arg_plan, value) + end + + elseif _is_runtime_broadcast_arg(arg) + _push_runtime_arg!(runtime_args, arg_plan, arg) + + elseif isbits(arg) + # Conservative fallback for singleton/isbits scalar broadcast leaves. + _push_static_arg!(static_args, arg_plan, arg) + + else + throw( + ArgumentError( + "Broadcast fusion does not know how to lower argument " * + "$(repr(arg)) of type $(typeof(arg))", + ), + ) + end + end + + return tuple(runtime_args...), tuple(static_args...), tuple(arg_plan...) +end + +############## + # The indexing in these kernels is based off this internal function from julia/broadcast.jl # Base.@propagate_inbounds function _broadcast_getindex(bc::Broadcasted{<:Any,<:Any,<:Any,<:Any}, I) @@ -6,24 +124,28 @@ # return _broadcast_getindex_evalf(bc.f, args...) # end -function make_linear_kernel(dest, bc::Base.Broadcast.Broadcasted) +function make_linear_kernel(dest, bc::Base.Broadcast.Broadcasted, arg_plan, static_args) f = bc.f - @kernel function broadcast_kernel_linear_splat(dest, args...) + @kernel function broadcast_kernel_linear_splat(dest, runtime_args...) I = @index(Global, Linear) - @inbounds args_modified = Base.Broadcast._getindex(args, I) + @inbounds args_modified = _materialize_broadcast_args( + arg_plan, runtime_args, static_args, I + ) @inbounds dest[I] = Base.Broadcast._broadcast_getindex_evalf(f, args_modified...) end return broadcast_kernel_linear_splat end -function make_cartesian_kernel(dest, bc::Base.Broadcast.Broadcasted) +function make_cartesian_kernel(dest, bc::Base.Broadcast.Broadcasted, arg_plan, static_args) f = bc.f - @kernel function broadcast_kernel_cartesian_splat(dest, args...) + @kernel function broadcast_kernel_cartesian_splat(dest, runtime_args...) I = @index(Global, Cartesian) - @inbounds args_modified = Base.Broadcast._getindex(args, I) + @inbounds args_modified = _materialize_broadcast_args( + arg_plan, runtime_args, static_args, I + ) @inbounds dest[I] = Base.Broadcast._broadcast_getindex_evalf(f, args_modified...) end @@ -46,6 +168,8 @@ cudevice_array_offset(::Type{T}) where {T<:Base.Broadcast.Extruded} = Int(fieldo stores_cudevicearray(::Type{T}) where {T<:CUDACore.CuDeviceArray} = true stores_cudevicearray(::Type{T}) where {T<:Base.Broadcast.Extruded} = true stores_cudevicearray(::Type{T}) where {T<:Number} = false +stores_cudevicearray(::Type{T}) where {T<:Bool} = false + function stores_cudevicearray(::Type{T}) where {T} throw(error("Broadcast fusion. Don't know what to do with type: $T")) end @@ -78,7 +202,7 @@ end get_ndarray(x::T) where {T<:NDArray} = x get_ndarray(x::T) where {T<:Base.Broadcast.Extruded} = x.x -get_ndarray(x) = throw(error("Broadcast fusion. Don't know what to do with type: $T")) +get_ndarray(x::T) where {T} = throw(error("Broadcast fusion. Don't know what to do with type: $T")) function _threads_from_occupancy( obj::KA.Kernel{CUDACore.CUDAKernels.CUDABackend}, @@ -150,116 +274,88 @@ end function get_cuda_task( obj::KA.Kernel{CUDACore.CUDAKernels.CUDABackend}, dest::D, - bc::B, + runtime_args::RT, ndrange, -) where {D<:NDArray,B<:Base.Broadcast.Broadcasted} +) where {D<:NDArray,RT<:Tuple} DEST_T = map_cuda_type(D) - ARG_TYPES = map_cuda_type.(typeof.(bc.args)) + ARG_TYPES = map_cuda_type.(typeof.(runtime_args)) + + key = (obj, D, RT, ndrange) - key = (obj, D, B, ndrange) lock(_BCAST_PTX_CACHE_LOCK) do - # Also stores in cache if not found in Dict return get!(_BCAST_PTX_CACHE, key) do ptx, threads, blocks, ctx = get_ptx(obj, DEST_T, ARG_TYPES...; ndrange=ndrange) + orig_name = extract_kernel_name(ptx) - # Append a PTX content hash to make each name unique. unique_name = orig_name * "_" * string(hash(ptx); base=16) ptx = replace(ptx, orig_name => unique_name) - # println(ptx) + ptx_task(ptx, unique_name) cuda_task = CUDATask(unique_name, (DEST_T, ARG_TYPES...)) + FusedBroadcastMetadata(ctx, threads, blocks, cuda_task) end end end function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcast.Broadcasted} - - #! HOW DOES THIS BEHAVE WHEN BC HAS 2 RESULT ARRAYS? - - # Normalize the Braodcasted type bc = Base.Broadcast.preprocess(dest, bc) bc = Base.Broadcast.instantiate(bc) bc = Base.Broadcast.flatten(bc) - # Get proper kernel + # Things like exponentiation generate arguments like Base.RefValue + # which do not work with our pattern for making CUDA kernels as they are + # not is-bits types. We split these out manually into static args and handle + # them separately from runtime args (i.e., arrays, scalars) + runtime_args, static_args, arg_plan = split_broadcast_args_for_kernel(bc.args) + broadcast_kernel = if ndims(dest) == 1 || (isa(IndexStyle(dest), IndexLinear) && isa(IndexStyle(bc), IndexLinear)) - make_linear_kernel(dest, bc) + make_linear_kernel(dest, bc, arg_plan, static_args) else - make_cartesian_kernel(dest, bc) + make_cartesian_kernel(dest, bc, arg_plan, static_args) end ndrange = ndims(dest) > 0 ? size(dest) : (1,) - # Tell KernelAsbtractions.jl we are using CUDA.jl bck_cuda = broadcast_kernel(CUDACore.CUDAKernels.CUDABackend()) - # Lookup in cache, if not found, compile and cache - fkm = get_cuda_task(bck_cuda, dest, bc, ndrange) - - # Replace NDArrays with CuDeviceArrays in the Broadcasted type so we can figure out bit-offsets - spoofed_bc_type = map_cuda_type(typeof(bc)) - fieldname(spoofed_bc_type, 3) == :args || - throw(ArgumentError("Broadcasted field 3 is not args. Failed to fuse broadcast.")) - - args_offset = Int(fieldoffset(spoofed_bc_type, 3)) - args_offset == 0 || - throw( - ArgumentError( - "Broadcast fusion only supports Broadcasted layouts where args starts at offset 0; got offset $args_offset. This is likely a bug in the compiler." - ), - ) - - # Spoof NDArrays with CuDeviceArrays in the Broadcasted type - spoofed_bc_args_type = map_cuda_type(typeof(bc.args)) - - # Build the arg_map: for each kernel arg (after ctx), record its source. - # Convention: index into combined [outputs..., inputs...] for NDArrays. - # idx < num_outputs → output[idx] - # idx >= num_outputs → input[idx - num_outputs] - # Scalar values are passed separately after the mapping. - # - # PTX kernel signature: f(kernel_state, ctx, dest, bc.args...) - # So the arg_map covers: [dest, bc.args...] in order. + fkm = get_cuda_task(bck_cuda, dest, runtime_args, ndrange) - num_outputs = 1 # dest is always the single output + num_outputs = 1 - # Deduplicate NDArray inputs: map each unique NDArray to a unique input index. - # The arg_map can reference the same input index multiple times (e.g. a .+ a). unique_ndarrays = NDArray[] - ndarray_to_input_idx = Dict{UInt,Int}() # objectid → input index + ndarray_to_input_idx = Dict{UInt,Int}() arg_map = Int32[] actual_scalars = Any[] - # First arg in PTX (after ctx) is always dest = output[0] - push!(arg_map, Int32(0)) # output[0] + # First PTX argument after ctx is dest. + push!(arg_map, Int32(0)) - # Then bc.args... in order - for (i, arg) in enumerate(bc.args) + # Now map only runtime args, not original bc.args. + for arg in runtime_args if stores_cudevicearray(map_cuda_type(typeof(arg))) nda = get_ndarray(arg) oid = objectid(nda) + if !haskey(ndarray_to_input_idx, oid) push!(unique_ndarrays, nda) - ndarray_to_input_idx[oid] = length(unique_ndarrays) - 1 # 0-based + ndarray_to_input_idx[oid] = length(unique_ndarrays) - 1 end + input_idx = ndarray_to_input_idx[oid] - push!(arg_map, Int32(num_outputs + input_idx)) # offset by num_outputs + push!(arg_map, Int32(num_outputs + input_idx)) else - # Scalar — record its position and save the value - push!(arg_map, Int32(-1 - length(actual_scalars))) # negative = scalar + push!(arg_map, Int32(-1 - length(actual_scalars))) push!(actual_scalars, arg) end end input_ndarrays = tuple(unique_ndarrays...) - # @show arg_map actual_scalars length(unique_ndarrays) - launch( fkm.cuda_task, input_ndarrays, @@ -271,27 +367,6 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas ctx=fkm.ctx, ) - #! DO I NEED TO DO TYPE PROMOTION CHECKS?? + #! PROMOTION CHECKS? return dest end - -# STEP 1: Figure out bit-offsets for CuDeviceArrays and scalars in args of spoofed type. -# The spoofed type has the same fields and alignment that the PTX kernel expects. -# cudevicearray_offsets, cudevicearray_indices = -# find_cudevicearray_offsets_and_indices(spoofed_bc_args_type) - -# scalar_offsets, scalar_indices = find_scalar_offsets_and_indices(spoofed_bc_args_type) - -# # STEP 2: Get NDarrays corresponding to the offsets in the spoofed type. -# input_ndarrays = ntuple( -# i -> get_ndarray(bc.args[cudevicearray_indices[i]]), length(cudevicearray_indices) -# ) -# input_scalars = ntuple(i -> bc.args[scalar_indices[i]], length(scalar_indices)) -# patch_info = BroadcastPatchInfo( -# sizeof(spoofed_bc_type), -# input_ndarrays, -# cudevicearray_offsets, -# ntuple(i -> i - 1, length(input_ndarrays)), -# scalar_offsets, -# input_scalars, -# )