diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 016b7b1..942a191 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -104,7 +104,7 @@ jobs: LEGATE_SHOW_CONFIG: "1" LEGATE_AUTO_CONFIG: "0" GPUTESTS: "0" # parsed by runtests.jl - LEGATE_CONFIG: "--cpus 1 --utility 1 --sysmem 4000" + LEGATE_CONFIG: "--cpus 2 --utility 1 --sysmem 4000" run: | if julia -e 'exit(VERSION >= v"1.12" ? 0 : 1)'; then export JULIA_NUM_THREADS=1,0 diff --git a/.github/workflows/developer.yml b/.github/workflows/developer.yml index fba4476..4a646d4 100644 --- a/.github/workflows/developer.yml +++ b/.github/workflows/developer.yml @@ -79,6 +79,8 @@ jobs: julia --color=yes -e 'using Pkg; Pkg.build("Legate")' - name: Perform Test + env: + LEGATE_CONFIG: "--cpus 2 --gpus 0 --utility 1 --sysmem 4000" run: | julia --color=yes -e 'using Pkg; Pkg.test("Legate"; julia_args=["--threads=4,1"])' diff --git a/lib/legate_jl_wrapper/include/wrapper.inl b/lib/legate_jl_wrapper/include/wrapper.inl index b196cea..cd878d3 100644 --- a/lib/legate_jl_wrapper/include/wrapper.inl +++ b/lib/legate_jl_wrapper/include/wrapper.inl @@ -109,6 +109,20 @@ inline Constraint align(const Variable& a, const Variable& b) { return legate::align(a, b); } +/** + * @ingroup legate_wrapper + * @brief Bloat target partitions around their aligned source partitions. + */ +inline Constraint bloat(const Variable& source, const Variable& target, + const std::vector& low_offsets, + const std::vector& high_offsets) { + return legate::bloat( + source, target, + legate::Span{low_offsets.data(), low_offsets.size()}, + legate::Span{high_offsets.data(), + high_offsets.size()}); +} + /** * @ingroup legate_wrapper * @brief Create an auto task in the runtime. diff --git a/lib/legate_jl_wrapper/src/module.cpp b/lib/legate_jl_wrapper/src/module.cpp index 75cc73c..4566327 100644 --- a/lib/legate_jl_wrapper/src/module.cpp +++ b/lib/legate_jl_wrapper/src/module.cpp @@ -179,6 +179,10 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) { mod.method("slice", [](LogicalStore& s, int32_t dim, legate::Slice sl) { return s.slice(dim, sl); }); + mod.method("slice", + [](LogicalStore& s, int32_t dim, int64_t start, int64_t stop) { + return s.slice(dim, legate::Slice{start, stop}); + }); mod.method( "get_physical_store", [](LogicalStore& s, std::optional target) { @@ -228,6 +232,8 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) { for (int i = 0; i < arr.dim(); i++) result.push_back(s[i]); return result; }); + mod.method("array_from_store", + [](const LogicalStore& store) { return LogicalArray{store}; }); mod.add_type("AutoTask") .method("add_input", static_cast( @@ -294,6 +300,7 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) { mod.method("runtime_sync", &legate_wrapper::runtime::runtime_sync); /* tasking */ mod.method("align", &legate_wrapper::tasking::align); + mod.method("bloat", &legate_wrapper::tasking::bloat); mod.method("domain_from_shape", &legate_wrapper::tasking::domain_from_shape); mod.method("create_manual_task", &legate_wrapper::tasking::create_manual_task); diff --git a/lib/legate_jl_wrapper/src/task.cpp b/lib/legate_jl_wrapper/src/task.cpp index 9bbc304..7550764 100644 --- a/lib/legate_jl_wrapper/src/task.cpp +++ b/lib/legate_jl_wrapper/src/task.cpp @@ -68,16 +68,21 @@ UFI(write, write_accessor); struct ufiFunctor { int* ndim_ptr = nullptr; int64_t* dims_ptr = nullptr; + int64_t* arg_dims_ptr = nullptr; ufiFunctor() = default; - ufiFunctor(int* ndim, int64_t* dims) : ndim_ptr(ndim), dims_ptr(dims) {} + ufiFunctor(int* ndim, int64_t* dims, int64_t* arg_dims) + : ndim_ptr(ndim), dims_ptr(dims), arg_dims_ptr(arg_dims) {} template void operator()(ufi::AccessMode mode, std::uintptr_t& p, int64_t* strides, const legate::PhysicalArray& rf) { + auto shp = rf.shape(); + for (int i = 0; i < DIM && i < REALM_MAX_DIM; ++i) { + arg_dims_ptr[i] = shp.hi[i] - shp.lo[i] + 1; + } if (ndim_ptr && *ndim_ptr == 0) { *ndim_ptr = DIM; - auto shp = rf.shape(); for (int i = 0; i < DIM && i < REALM_MAX_DIM; ++i) { dims_ptr[i] = shp.hi[i] - shp.lo[i] + 1; } @@ -130,6 +135,8 @@ struct UFISlot { char scalar_data[MAX_UFI_ARGS][MAX_SCALAR_SIZE]; int64_t input_strides[MAX_UFI_ARGS][REALM_MAX_DIM]; int64_t output_strides[MAX_UFI_ARGS][REALM_MAX_DIM]; + int64_t input_dims[MAX_UFI_ARGS][REALM_MAX_DIM]; + int64_t output_dims[MAX_UFI_ARGS][REALM_MAX_DIM]; UFISlot() { task_done.store(false); @@ -182,6 +189,16 @@ JULIA_LEGATE_UFI_EXPORT void* legate_get_slot_request_ptr(int slot_id) { return static_cast(&g_ufi_slots[slot_id].request); } +JULIA_LEGATE_UFI_EXPORT int64_t* legate_get_slot_input_dims_ptr(int slot_id) { + if (slot_id < 0 || slot_id >= MAX_UFI_SLOTS) return nullptr; + return &g_ufi_slots[slot_id].input_dims[0][0]; +} + +JULIA_LEGATE_UFI_EXPORT int64_t* legate_get_slot_output_dims_ptr(int slot_id) { + if (slot_id < 0 || slot_id >= MAX_UFI_SLOTS) return nullptr; + return &g_ufi_slots[slot_id].output_dims[0][0]; +} + JULIA_LEGATE_UFI_EXPORT int legate_get_active_call_count() { return g_active_calls.load(); } @@ -273,11 +290,11 @@ inline void JuliaTaskInterface(legate::TaskContext context, bool is_gpu) { slot.request.task_id = task_id; slot.request.ndim = 0; - ufiFunctor functor{&slot.request.ndim, slot.request.dims}; - for (size_t i = 0; i < ni; ++i) { auto ps = context.input(i); std::uintptr_t p = 0; + ufiFunctor functor{no == 0 && i == 0 ? &slot.request.ndim : nullptr, + slot.request.dims, slot.input_dims[i]}; legate::double_dispatch(ps.dim(), ps.type().code(), functor, ufi::AccessMode::READ, p, slot.input_strides[i], ps); @@ -287,6 +304,8 @@ inline void JuliaTaskInterface(legate::TaskContext context, bool is_gpu) { for (size_t i = 0; i < no; ++i) { auto ps = context.output(i); std::uintptr_t p = 0; + ufiFunctor functor{i == 0 ? &slot.request.ndim : nullptr, slot.request.dims, + slot.output_dims[i]}; legate::double_dispatch(ps.dim(), ps.type().code(), functor, ufi::AccessMode::WRITE, p, slot.output_strides[i], ps); @@ -355,6 +374,10 @@ void wrap_ufi(jlcxx::Module& mod) { mod.method("_initialize_async_system", &ufi::initialize_async_system); mod.method("legate_get_max_slots", &ufi::legate_get_max_slots); mod.method("legate_get_slot_request_ptr", &ufi::legate_get_slot_request_ptr); + mod.method("legate_get_slot_input_dims_ptr", + &ufi::legate_get_slot_input_dims_ptr); + mod.method("legate_get_slot_output_dims_ptr", + &ufi::legate_get_slot_output_dims_ptr); mod.method("legate_pop_pending_slot_nonblocking", &ufi::legate_pop_pending_slot_nonblocking); mod.method("legate_get_active_call_count", diff --git a/src/api/data.jl b/src/api/data.jl index f60decb..c0c7653 100644 --- a/src/api/data.jl +++ b/src/api/data.jl @@ -202,6 +202,22 @@ function slice(store::LogicalStore, indices...) return LogicalStore{eltype(store),Int(LegateInternal.dim(impl))}(impl, nothing) end +""" + slice(array::LogicalArray, dim, start, stop) -> LogicalArray + +Return a view of the half-open interval `[start, stop)` in zero-based dimension `dim`. +""" +function slice( + array::LogicalArray{T,N}, dim::Integer, start::Integer, stop::Integer +) where {T,N} + store = LegateInternal.slice( + LegateInternal.data(array.handle), Int32(dim), Int64(start), Int64(stop) + ) + impl = LegateInternal.array_from_store(store) + dims = Tuple(Int.(collect(LegateInternal.shape(impl)))) + return LogicalArray{T,N}(impl, dims, array.order) +end + """ get_physical_store(LogicalStore) -> PhysicalStore get_physical_store(LogicalArray) -> PhysicalStore diff --git a/src/api/runtime.jl b/src/api/runtime.jl index 4496957..83ffaf5 100644 --- a/src/api/runtime.jl +++ b/src/api/runtime.jl @@ -60,7 +60,7 @@ Block until all pending Legate tasks have completed. Useful before reading files written by `h5write` or other async operations. """ -runtime_sync() = LegateInternal.runtime_sync() +runtime_sync() = issue_execution_fence() """ create_library(name::String) -> Library @@ -113,7 +113,11 @@ Issues an execution fence to the runtime. """ function issue_execution_fence(; blocking::Bool=true) if blocking - @threadcall((:legate_issue_execution_fence_blocking, Legate.WRAPPER_LIB_PATH), Cvoid, ()) + @static if VERSION >= v"1.12" + @ccall gc_safe = true WRAPPER_LIB_PATH.legate_issue_execution_fence_blocking()::Cvoid + else + @ccall WRAPPER_LIB_PATH.legate_issue_execution_fence_blocking()::Cvoid + end else LegateInternal.issue_execution_fence(false) end diff --git a/src/api/tasks.jl b/src/api/tasks.jl index 941c164..b0e7918 100644 --- a/src/api/tasks.jl +++ b/src/api/tasks.jl @@ -60,6 +60,39 @@ function align(a::Variable, b::Variable) return LegateInternal.align(a, b) end +""" + bloat(source, target, low_offsets, high_offsets) -> Constraint + +Partition `target` like `source`, expanded by the given halo width in each dimension. +""" +function bloat(source::Variable, target::Variable, low_offsets, high_offsets) + isdefined(LegateInternal, :bloat) || + error("bloat constraints require a newer Legate.jl wrapper") + length(low_offsets) == length(high_offsets) || + throw(DimensionMismatch("low and high bloat offsets must have equal lengths")) + all(offset -> offset >= 0, low_offsets) || + throw(ArgumentError("bloat offsets must be nonnegative")) + all(offset -> offset >= 0, high_offsets) || + throw(ArgumentError("bloat offsets must be nonnegative")) + return LegateInternal.bloat( + source, target, to_cxx_vector(low_offsets), to_cxx_vector(high_offsets) + ) +end + +""" + broadcast(var) -> Constraint + broadcast(var, axes) -> Constraint + +Give every task the whole of `var`, or only the full extent of zero-based `axes`. +""" +broadcast(var::Variable) = LegateInternal.broadcast(var) + +function broadcast(var::Variable, axes) + all(axis -> axis >= 0, axes) || + throw(ArgumentError("broadcast axes must be nonnegative")) + return LegateInternal.broadcast(var, CxxWrap.StdVector([UInt32(a) for a in axes])) +end + function default_alignment( task::LegateTask, inputs::Vector{<:Variable}, outputs::Vector{<:Variable} ) @@ -161,8 +194,16 @@ _gpu_precompile(args...) = nothing function submit_task(rt::CxxPtr{Runtime}, task::LegateTask) drain_pending_frees!() if !isnothing(task.fun) - in_t = Tuple{task.input_types...} - out_t = Tuple{task.output_types...} + n_inputs = length(task.input_types) + in_t = Tuple{ + [Array{T,length(task.arg_dims[i])} for (i, T) in enumerate(task.input_types)]... + } + out_t = Tuple{ + [ + Array{T,length(task.arg_dims[n_inputs + i])} for + (i, T) in enumerate(task.output_types) + ]..., + } sc_t = Tuple{task.scalar_types...} sig = UfiSignature{in_t,out_t,sc_t}() @@ -188,6 +229,8 @@ function submit_task(rt::CxxPtr{Runtime}, task::LegateTask) Ptr{Ptr{Cvoid}}, Ptr{Int64}, Ptr{Int64}, + Ptr{Int64}, + Ptr{Int64}, local_dims_type, typeof(sig), ), @@ -199,11 +242,11 @@ function submit_task(rt::CxxPtr{Runtime}, task::LegateTask) # 2. Precompile the user-provided function with exact types user_arg_types = Any[] - for (T, d) in zip(task.input_types, task.arg_dims) - push!(user_arg_types, Array{T,length(d)}) + for (i, T) in enumerate(task.input_types) + push!(user_arg_types, Array{T,length(task.arg_dims[i])}) end - for (T, d) in zip(task.output_types, task.arg_dims) - push!(user_arg_types, Array{T,length(d)}) + for (i, T) in enumerate(task.output_types) + push!(user_arg_types, Array{T,length(task.arg_dims[n_inputs + i])}) end for T in task.scalar_types push!(user_arg_types, T) @@ -217,14 +260,13 @@ function submit_task(rt::CxxPtr{Runtime}, task::LegateTask) # event loop. Compiling here caches the cubin so the worker's launch is a # cache hit. No-op without CUDA (see CUDAExt). if task.is_gpu - # get_ptr only needs to be GC-safe once a GPU task is in flight; flag it so - # CPU-only programs keep the original accessor path. - LegateInternal.set_gpu_tasking_active(true) _gpu_precompile( task.fun, task.input_types, task.output_types, task.scalar_types, task.arg_dims ) end + # A worker GC would otherwise deadlock against the caller parked in get_ptr. + LegateInternal.set_gpu_tasking_active(true) Threads.atomic_add!(SUBMITTED_COUNT, 1) end return _submit_task(rt, task) diff --git a/src/ufi.jl b/src/ufi.jl index ae3c866..057432b 100644 --- a/src/ufi.jl +++ b/src/ufi.jl @@ -21,6 +21,8 @@ const REALM_MAX_DIM = 6 const MAX_UFI_SLOTS_VAL = 32 const PhysArrPtr = Ptr{Cvoid} const SLOT_REQUEST_PTRS = Vector{Ptr{Cvoid}}(undef, MAX_UFI_SLOTS_VAL) +const SLOT_INPUT_DIMS_PTRS = fill(Ptr{Int64}(C_NULL), MAX_UFI_SLOTS_VAL) +const SLOT_OUTPUT_DIMS_PTRS = fill(Ptr{Int64}(C_NULL), MAX_UFI_SLOTS_VAL) const UFI_INIT_LOCK = ReentrantLock() @@ -76,6 +78,8 @@ struct TaskJob{B<:TaskBackend,M<:UfiMetadata,D<:Tuple} scal_args::Vector{Ptr{Cvoid}} in_strides::Vector{Int64} # flat [arg][REALM_MAX_DIM] element strides out_strides::Vector{Int64} + in_dims::Vector{Int64} # flat [arg][REALM_MAX_DIM] tile dimensions + out_dims::Vector{Int64} local_dims::D meta::M end @@ -161,27 +165,32 @@ end scal_p_ptr::Ptr{Ptr{Cvoid}}, in_str_ptr::Ptr{Int64}, out_str_ptr::Ptr{Int64}, + in_dim_ptr::Ptr{Int64}, + out_dim_ptr::Ptr{Int64}, local_dims::Tuple, ::UfiSignature{InT,OutT,ScT}, ) where {B<:TaskBackend,InT,OutT,ScT} - nd = length(local_dims.parameters) pre = [] # prepare input and output buffers callargs = [] # args passed to the user task post = [] # commit outputs to their destination tiles - # Unroll per-rank stride loads at generation time. - instr(base) = Expr(:tuple, [:(Int(unsafe_load(in_str_ptr, $(base + d)))) for d in 1:nd]...) - outstr(base) = Expr(:tuple, [:(Int(unsafe_load(out_str_ptr, $(base + d)))) for d in 1:nd]...) + metadata(ptr, base, nd) = + Expr(:tuple, [:(Int(unsafe_load($ptr, $(base + d)))) for d in 1:nd]...) for (i, T) in enumerate(InT.parameters) E = eltype(T) + nd = ndims(T) base = (i - 1) * REALM_MAX_DIM sym = gensym(:in) + dsym = gensym(:indims) + ssym = gensym(:instrides) + push!(pre, :($dsym = $(metadata(:in_dim_ptr, base, nd)))) + push!(pre, :($ssym = $(metadata(:in_str_ptr, base, nd)))) push!( pre, :( $sym = _ufi_prepare_input( - backend, $E, unsafe_load(in_p_ptr, $i), local_dims, $(instr(base))) + backend, $E, unsafe_load(in_p_ptr, $i), $dsym, $ssym) ), ) push!(callargs, sym) @@ -189,20 +198,23 @@ end for (i, T) in enumerate(OutT.parameters) E = eltype(T) + nd = ndims(T) base = (i - 1) * REALM_MAX_DIM sym = gensym(:out) state = gensym(:state) + dsym = gensym(:outdims) ssym = gensym(:outs) - push!(pre, :($ssym = $(outstr(base)))) + push!(pre, :($dsym = $(metadata(:out_dim_ptr, base, nd)))) + push!(pre, :($ssym = $(metadata(:out_str_ptr, base, nd)))) push!( pre, :( ($sym, $state) = _ufi_prepare_output( - backend, $E, unsafe_load(out_p_ptr, $i), local_dims, $ssym) + backend, $E, unsafe_load(out_p_ptr, $i), $dsym, $ssym) ), ) push!(callargs, sym) - push!(post, :(_ufi_commit_output(backend, $state, $sym, local_dims, $ssym))) + push!(post, :(_ufi_commit_output(backend, $state, $sym, $dsym, $ssym))) end for (i, T) in enumerate(ScT.parameters) @@ -224,7 +236,9 @@ function _execute_task(job::TaskJob) scalar_args = job.scal_args in_strides = job.in_strides out_strides = job.out_strides - GC.@preserve in_args out_args scalar_args in_strides out_strides begin + in_dims = job.in_dims + out_dims = job.out_dims + GC.@preserve in_args out_args scalar_args in_strides out_strides in_dims out_dims begin _do_call( job.backend, job.meta.fun, @@ -233,6 +247,8 @@ function _execute_task(job::TaskJob) pointer(scalar_args), pointer(in_strides), pointer(out_strides), + pointer(in_dims), + pointer(out_dims), job.local_dims, job.meta.sig, ) @@ -273,6 +289,18 @@ function _copy_strides(ptr::Ptr{Int64}, count::Int) return strides end +function _copy_dims(ptr::Ptr{Int64}, count::Int, common_dims::Tuple) + dims = zeros(Int64, count * REALM_MAX_DIM) + if ptr == C_NULL + for arg in 0:(count - 1), dim in eachindex(common_dims) + dims[arg * REALM_MAX_DIM + dim] = common_dims[dim] + end + elseif count > 0 + unsafe_copyto!(pointer(dims), ptr, length(dims)) + end + return dims +end + # C++ keeps the referenced tile storage valid until the completion callback. function _make_job(slot_id::Int, req::TaskRequest, meta::UfiMetadata) sig_type = typeof(meta.sig) @@ -280,6 +308,8 @@ function _make_job(slot_id::Int, req::TaskRequest, meta::UfiMetadata) out_count = length(sig_type.parameters[2].parameters) scalar_count = length(sig_type.parameters[3].parameters) local_dims = ntuple(i -> Int(max(0, req.dims[i])), Int(req.ndim)) + in_dims_ptr = SLOT_INPUT_DIMS_PTRS[slot_id + 1] + out_dims_ptr = SLOT_OUTPUT_DIMS_PTRS[slot_id + 1] return TaskJob( slot_id, @@ -289,6 +319,8 @@ function _make_job(slot_id::Int, req::TaskRequest, meta::UfiMetadata) _copy_pointer_args(req.scalars_ptr, scalar_count), _copy_strides(req.input_strides_ptr, in_count), _copy_strides(req.output_strides_ptr, out_count), + _copy_dims(in_dims_ptr, in_count, local_dims), + _copy_dims(out_dims_ptr, out_count, local_dims), local_dims, meta, ) @@ -386,6 +418,9 @@ function init_ufi() exit(UFI_ERROR) end + has_arg_dims = + isdefined(LegateInternal, :legate_get_slot_input_dims_ptr) && + isdefined(LegateInternal, :legate_get_slot_output_dims_ptr) for i in 1:max_slots SLOT_REQUEST_PTRS[i] = ccall( (:legate_get_slot_request_ptr, Legate.WRAPPER_LIB_PATH), @@ -393,6 +428,12 @@ function init_ufi() (Cint,), Cint(i-1), ) + if has_arg_dims + SLOT_INPUT_DIMS_PTRS[i] = + LegateInternal.legate_get_slot_input_dims_ptr(Cint(i - 1)).cpp_object + SLOT_OUTPUT_DIMS_PTRS[i] = + LegateInternal.legate_get_slot_output_dims_ptr(Cint(i - 1)).cpp_object + end end LegateInternal._initialize_async_system() diff --git a/test/runtests.jl b/test/runtests.jl index 11e5ccb..812718e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -61,7 +61,9 @@ include("tests/hdf5.jl") include("tests/stability.jl") include("tests/basic.jl") -include("tests/tasking.jl") +include("tests/tasking/setup.jl") +include("tests/tasking/basic.jl") +include("tests/tasking/constraints.jl") if run_gpu_tests - include("tests/tasking_gpu.jl") + include("tests/tasking/gpu.jl") end diff --git a/test/tests/tasking.jl b/test/tests/tasking/basic.jl similarity index 90% rename from test/tests/tasking.jl rename to test/tests/tasking/basic.jl index 38ce05f..9ec5d50 100644 --- a/test/tests/tasking.jl +++ b/test/tests/tasking/basic.jl @@ -69,9 +69,6 @@ expected_a = expected_c .* 2.5f0 @testset verbose=true "CPU Tasking" begin Legate.Experimental(true) # tasking is experimental - rt = Legate.get_runtime() - lib = Legate.create_library("test_comparison") - my_task = Legate.wrap_task(task_test, Legate.CPUBackend) my_4arg_task = Legate.wrap_task(task_4arg, Legate.CPUBackend) my_scalar_task = Legate.wrap_task(task_scalar, Legate.CPUBackend) @@ -94,6 +91,15 @@ expected_a = expected_c .* 2.5f0 @test val_b ≈ base_results.b_init end + @testset "LogicalArray Slice" begin + reference = reshape(Float32.(1:36), 6, 6) + array = Legate.create_array([6, 6], Float32) + copyto!(array, reference) + view = Legate.slice(array, 0, 1, 5) + view = Legate.slice(view, 1, 2, 6) + @test size(view) == (4, 4) + end + a = Legate.create_array([10, 10], Float32) b = Legate.create_array([10, 10], Float32) c = Legate.create_array([10, 10], Float32) @@ -143,4 +149,9 @@ expected_a = expected_c .* 2.5f0 val_a = Array(a) @test val_a ≈ expected_a end + + @testset "Execution Fence" begin + @test Legate.issue_execution_fence() === nothing + @test Legate.runtime_sync() === nothing + end end diff --git a/test/tests/tasking/constraints.jl b/test/tests/tasking/constraints.jl new file mode 100644 index 0000000..e875b61 --- /dev/null +++ b/test/tests/tasking/constraints.jl @@ -0,0 +1,130 @@ +# Operands must be large enough that Legate would otherwise split them. + +function stencil_task(core_indices, halo, output) + offset = Int(core_indices[1] - halo[1]) + @assert length(halo) >= length(output) + @assert 0 <= offset <= 1 + @inbounds for i in eachindex(output) + h = i + offset + center = halo[h] + left = h == 1 ? center : halo[h - 1] + right = h == length(halo) ? center : halo[h + 1] + output[i] = left + 2 * center + right + end + return nothing +end + +function gather_task(indices, table, out) + @inbounds for i in eachindex(out) + out[i] = table[Int(indices[i])] + end + return nothing +end + +function row_normalize_task(a, out) + @inbounds for i in axes(a, 1) + s = zero(eltype(a)) + for j in axes(a, 2) + s += a[i, j] + end + for j in axes(a, 2) + out[i, j] = a[i, j] / s + end + end + return nothing +end + +@testset verbose=true "Tasking Constraints" begin + Legate.Experimental(true) # tasking is experimental + my_stencil_task = Legate.wrap_task(stencil_task, Legate.CPUBackend) + my_gather_task = Legate.wrap_task(gather_task, Legate.CPUBackend) + my_row_normalize_task = Legate.wrap_task(row_normalize_task, Legate.CPUBackend) + + @testset "Bloat Constraint (Radius-One Stencil)" begin + if !isdefined(Legate.LegateInternal, :bloat) + @test_skip false + else + n = 100_000 + reference = Float64.(1:n) + core_indices = Legate.create_array([n], Float64) + stencil_input = Legate.create_array([n], Float64) + stencil_output = Legate.create_array([n], Float64) + copyto!(core_indices, reference) + copyto!(stencil_input, reference) + + task = Legate.create_julia_task(rt, lib, my_stencil_task) + core_var = Legate.add_input(task, core_indices) + halo_var = Legate.add_input(task, stencil_input) + output_var = Legate.add_output(task, stencil_output) + Legate.add_constraint(task, Legate.align(core_var, output_var)) + Legate.add_constraint(task, Legate.bloat(core_var, halo_var, (1,), (1,))) + + started_before = Legate.LegateInternal.legate_get_started_count() + Legate.submit_task(rt, task) + result = Array(stencil_output) + started = + Legate.LegateInternal.legate_get_started_count() - started_before + + expected = 4 .* reference + expected[1] = 5 + expected[end] = 4n - 1 + @test result == expected + if Legate.LegateInternal.num_procs() > 1 + @test started > 1 + end + end + end + + @testset "Broadcast Constraint (Full)" begin + n = 100_000 + indices_host = Float64.(n:-1:1) + table_host = Float64.(1:n) .^ 2 + indices = Legate.create_array([n], Float64) + table = Legate.create_array([n], Float64) + out = Legate.create_array([n], Float64) + copyto!(indices, indices_host) + copyto!(table, table_host) + + task = Legate.create_julia_task(rt, lib, my_gather_task) + indices_var = Legate.add_input(task, indices) + table_var = Legate.add_input(task, table) + out_var = Legate.add_output(task, out) + Legate.add_constraint(task, Legate.align(indices_var, out_var)) + Legate.add_constraint(task, Legate.broadcast(table_var)) + + started_before = Legate.LegateInternal.legate_get_started_count() + Legate.submit_task(rt, task) + result = Array(out) + started = Legate.LegateInternal.legate_get_started_count() - started_before + + @test result == reverse(table_host) + if Legate.LegateInternal.num_procs() > 1 + @test started > 1 + end + end + + @testset "Broadcast Constraint (Axes)" begin + m, k = 1024, 1024 + a_host = reshape(Float64.(1:(m * k)), m, k) + a = Legate.create_array([m, k], Float64) + out = Legate.create_array([m, k], Float64) + copyto!(a, a_host) + + task = Legate.create_julia_task(rt, lib, my_row_normalize_task) + a_var = Legate.add_input(task, a) + out_var = Legate.add_output(task, out) + Legate.add_constraint(task, Legate.align(a_var, out_var)) + # Row sums need every column in each tile. + Legate.add_constraint(task, Legate.broadcast(a_var, (1,))) + + started_before = Legate.LegateInternal.legate_get_started_count() + Legate.submit_task(rt, task) + result = Array(out) + started = Legate.LegateInternal.legate_get_started_count() - started_before + + @test result ≈ a_host ./ sum(a_host; dims=2) + if Legate.LegateInternal.num_procs() > 1 + @test started > 1 + end + end +end diff --git a/test/tests/tasking_gpu.jl b/test/tests/tasking/gpu.jl similarity index 100% rename from test/tests/tasking_gpu.jl rename to test/tests/tasking/gpu.jl diff --git a/test/tests/tasking/setup.jl b/test/tests/tasking/setup.jl new file mode 100644 index 0000000..5bbc079 --- /dev/null +++ b/test/tests/tasking/setup.jl @@ -0,0 +1,3 @@ +# Shared by the tasking test files. +const rt = Legate.get_runtime() +const lib = Legate.create_library("test_tasking")