diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index f7bcfd4f8..e9abf561e 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -28,7 +28,7 @@ steps: agents: queue: "juliagpu" cuda: "*" - if: build.message !~ /\[skip tests\]/ + if: build.message !~ /\[skip tests\]/ && (build.branch == "main" || build.pull_request.base_branch == "main") timeout_in_minutes: 44 env: LD_LIBRARY_PATH: "" diff --git a/.githash b/.githash index 6cdfeea5d..f2d1e6df2 100644 --- a/.githash +++ b/.githash @@ -1 +1 @@ -dec047f1bd1c8287513c6c437f946982e516ccd4 +3713a2acb6579891706b08f970e99f9ae5e6472f diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0daab8c40..e979e8968 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,8 +19,6 @@ on: - 'deps/build.jl' - 'Project.toml' - 'lib/CNPreferences/src/**' - - '.github/workflows/ci.yml' - - 'Dockerfile' # container.yml depends on this tags: - 'v*' branches: @@ -32,11 +30,25 @@ on: - 'deps/build.jl' - 'Project.toml' - 'lib/CNPreferences/src/**' - - '.github/workflows/ci.yml' - - 'Dockerfile' # container.yml depends on this jobs: + check_changes: + name: Check for wrapper changes + runs-on: ubuntu-latest + outputs: + wrapper_changed: ${{ steps.filter.outputs.wrapper }} + steps: + - uses: actions/checkout@v4 + - uses: dorny/paths-filter@v3 + id: filter + with: + filters: | + wrapper: + - 'lib/cunumeric_jl_wrapper/**' + test: name: Julia ${{ matrix.julia }} - ${{ matrix.os }} + needs: check_changes + if: ${{ needs.check_changes.outputs.wrapper_changed != 'true' }} runs-on: ${{ matrix.os }} strategy: fail-fast: false diff --git a/.github/workflows/container.yml b/.github/workflows/container.yml index 00d5d995e..07a68c927 100644 --- a/.github/workflows/container.yml +++ b/.github/workflows/container.yml @@ -18,6 +18,7 @@ on: - main jobs: push_to_registry: + if: ${{ github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success' }} name: Container for ${{ matrix.platform }} - Julia ${{ matrix.julia }} - CUDA ${{ matrix.cuda }} permissions: contents: read diff --git a/.github/workflows/developer.yml b/.github/workflows/developer.yml index 10c5094bd..f2bef3a89 100644 --- a/.github/workflows/developer.yml +++ b/.github/workflows/developer.yml @@ -95,6 +95,11 @@ jobs: using CNPreferences; CNPreferences.use_developer_mode(); Pkg.develop(PackageSpec(path = ".")) ' + julia --color=yes --project=test -e ' + using Pkg; + Pkg.develop(PackageSpec(path = "lib/CNPreferences")) + using CNPreferences; CNPreferences.use_developer_mode(); + ' julia --color=yes -e 'using Pkg; Pkg.build("cuNumeric")' - name: Perform Test diff --git a/lib/cunumeric_jl_wrapper/src/ndarray.cpp b/lib/cunumeric_jl_wrapper/src/ndarray.cpp index 47cdf3b02..ed157b077 100644 --- a/lib/cunumeric_jl_wrapper/src/ndarray.cpp +++ b/lib/cunumeric_jl_wrapper/src/ndarray.cpp @@ -237,6 +237,23 @@ void nda_unary_reduction(CN_NDArray* out, CuPyNumericUnaryRedCode op_code, out->obj.unary_reduction(op_code, input->obj); } +CN_NDArray* nda_unary_reduction_axes(CuPyNumericUnaryRedCode op_code, CN_NDArray* input, const int32_t* axes, int32_t num_axes, bool keepdims) { + std::vector axis_vec(axes, axes + num_axes); + NDArray result = input->obj._perform_unary_reduction( + static_cast(op_code), + input->obj, + axis_vec, + std::nullopt, // dtype + std::nullopt, // res_dtype + std::nullopt, // out + keepdims, + {}, // args + std::nullopt, // initial + std::nullopt // where + ); + return new CN_NDArray{NDArray(std::move(result))}; +} + NDArray get_slice(NDArray arr, std::vector slices) { switch (slices.size()) { case 1: { diff --git a/src/ndarray/detail/ndarray.jl b/src/ndarray/detail/ndarray.jl index b11c41015..57160b835 100644 --- a/src/ndarray/detail/ndarray.jl +++ b/src/ndarray/detail/ndarray.jl @@ -219,6 +219,14 @@ function nda_unary_reduction(out::NDArray, op_code::UnaryRedCode, input::NDArray return out end +function nda_unary_reduction_axes(op_code::UnaryRedCode, input::NDArray{T,N}, axes::Vector{Int32}, keepdims::Bool) where {T,N} + axes_c = collect(Int32, axes) + ptr = ccall((:nda_unary_reduction_axes, libnda), + NDArray_t, (UnaryRedCode, NDArray_t, Ptr{Int32}, Int32, Cint), + op_code, input.ptr, axes_c, Int32(length(axes_c)), keepdims) + return NDArray(ptr) +end + function nda_array_equal(rhs1::NDArray{T,N}, rhs2::NDArray{T,N}) where {T,N} ptr = ccall((:nda_array_equal, libnda), NDArray_t, (NDArray_t, NDArray_t), diff --git a/src/ndarray/unary.jl b/src/ndarray/unary.jl index 6cc7b7491..96e69f981 100644 --- a/src/ndarray/unary.jl +++ b/src/ndarray/unary.jl @@ -209,9 +209,11 @@ The following unary reduction operations are supported and can be applied direct • `prod` • `sum` - These operations follow standard Julia semantics. +Reduction over specific dimensions is supported via the `dims` keyword argument, +following the same semantics as Julia's base reduction functions. + Examples -------- @@ -220,6 +222,14 @@ A = cuNumeric.ones(5) maximum(A) sum(A) + +# Reduce over a specific dimension +B = cuNumeric.ones(3, 4) +sum(B, dims=1) # 1×4 result +sum(B, dims=2) # 3×1 result + +# Reduce over multiple dimensions +sum(B, dims=(1,2)) # 1×1 result ``` """ global const unary_reduction_map = Dict{Function,UnaryRedCode}( @@ -242,26 +252,57 @@ global const unary_reduction_map = Dict{Function,UnaryRedCode}( #! IT WOULD BE NICE IF THESE JUST RETURNED SCALARS WHEN APPROPRIATE # #*TODO HOW TO GET THESE ACTING ON CERTAIN DIMS + +function _unary_reduction_impl(base_func, op_code, input::NDArray{T}, ::Colon) where {T} + T_OUT = Base.promote_op(base_func, Vector{T}) + is_wider_type(T_OUT, T) && assertpromotion(base_func, T, T_OUT) + out = cuNumeric.zeros(T_OUT) + return nda_unary_reduction(out, op_code, unchecked_promote_arr(input, T_OUT)) +end + +function _unary_reduction_impl(base_func, op_code, input::NDArray{T,N}, dims::Integer) where {T,N} + T_OUT = Base.promote_op(base_func, Vector{T}) + is_wider_type(T_OUT, T) && assertpromotion(base_func, T, T_OUT) + axes = Int32[dims - 1] + return nda_unary_reduction_axes(op_code, unchecked_promote_arr(input, T_OUT), axes, true) +end + +function _unary_reduction_impl(base_func, op_code, input::NDArray{T,N}, dims::Tuple) where {T,N} + if length(dims) > 1 + error("$(base_func): reducing over multiple dimensions is not yet supported. Got dims=$dims") + end + # single element tuple + T_OUT = Base.promote_op(base_func, Vector{T}) + is_wider_type(T_OUT, T) && assertpromotion(base_func, T, T_OUT) + axes = Int32[dims[1] - 1] + return nda_unary_reduction_axes(op_code, unchecked_promote_arr(input, T_OUT), axes, true) +end + # Generate code for all unary reductions. for (base_func, op_code) in unary_reduction_map @eval begin - function $(Symbol(base_func))(input::NDArray{T}) where {T} - T_OUT = Base.promote_op($base_func, Vector{T}) - is_wider_type(T_OUT, T) && assertpromotion($base_func, T, T_OUT) - out = cuNumeric.zeros(T_OUT) #0D result (not right if reducing along dims) - return nda_unary_reduction(out, $(op_code), unchecked_promote_arr(input, T_OUT)) + function $(Symbol(base_func))(input::NDArray{T,N}; dims=Colon()) where {T,N} + return _unary_reduction_impl($base_func, $(op_code), input, dims) end end end -function Base.all(input::NDArray{Bool}) +function _bool_reduction_impl(op_code, input::NDArray{Bool}, ::Colon) out = cuNumeric.zeros(Bool) - return nda_unary_reduction(out, cuNumeric.ALL, input) + return nda_unary_reduction(out, op_code, input) end -function Base.any(input::NDArray{Bool}) - out = cuNumeric.zeros(Bool) - return nda_unary_reduction(out, cuNumeric.ANY, input) +function _bool_reduction_impl(op_code, input::NDArray{Bool}, dims) + axes = collect(Int32, (d - 1 for d in (dims isa Integer ? (dims,) : dims))) + return nda_unary_reduction_axes(op_code, input, axes, true) +end + +function Base.all(input::NDArray{Bool}; dims=Colon()) + return _bool_reduction_impl(cuNumeric.ALL, input, dims) +end + +function Base.any(input::NDArray{Bool}; dims=Colon()) + return _bool_reduction_impl(cuNumeric.ANY, input, dims) end #! ONLY ADD ONCE REDUCTIONS RETURN A SCALAR diff --git a/test/Project.toml b/test/Project.toml index 7dd0521e1..8e41dd68b 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +CNPreferences = "3e078157-ea10-49d5-bf32-908f777cd46f" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CUDA_Driver_jll = "4ee394cb-3365-5eb0-8335-949819d2adfc" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" diff --git a/test/runtests.jl b/test/runtests.jl index d2396b3f8..7cf26acfb 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"); @@ -167,6 +167,32 @@ end end end +@testset verbose=true "Unary Reductions with Dims" begin + N = 100 + + @testset for T in Base.uniontypes(cuNumeric.SUPPORTED_ARRAY_TYPES) + julia_arr_1D = my_rand(T, N) + julia_arr_2D = my_rand(T, isqrt(N), isqrt(N)) + + cunumeric_arr_1D = @allowscalar NDArray(julia_arr_1D) + cunumeric_arr_2D = @allowscalar NDArray(julia_arr_2D) + + @testset "$(func)" for (func, _) in cuNumeric.unary_reduction_map + # Skip reductions not supported by the cuNumeric backend for complex types + if T <: Complex && ( + func == Base.maximum || + func == Base.minimum || + func == Base.prod + ) + continue + end + + test_unary_reduction_dims(func, julia_arr_1D, cunumeric_arr_1D) + test_unary_reduction_dims(func, julia_arr_2D, cunumeric_arr_2D) + end + end +end + @testset verbose = true "Binary Ops" begin N = 100 diff --git a/test/tests/unary_tests.jl b/test/tests/unary_tests.jl index 88f284ff6..8c1e4f527 100644 --- a/test/tests/unary_tests.jl +++ b/test/tests/unary_tests.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"); @@ -88,3 +88,23 @@ function test_unary_function_set(func_dict, T, N) test_unary_operation(func, julia_arr_2D, cunumeric_arr_2D, T) end end + +function test_unary_reduction_dims( + func, julia_arr::AbstractArray{T,N}, cunumeric_arr::NDArray{T,N} +) where {T,N} + allowpromotion(true) do + for d in 1:N + julia_res = func(julia_arr; dims=d) + cunumeric_res = func(cunumeric_arr; dims=d) + allowscalar() do + @test cuNumeric.compare(julia_res, cunumeric_res, atol(T), rtol(T)) + end + end + + # we are testing a multi axis reduction. This will throw a runtime error. + # https://github.com/nv-legate/cupynumeric/blob/main/src/cupynumeric/ndarray.cc#L1132 + if N >= 2 + @test_throws Exception func(cunumeric_arr, dims=(1, 2)) + end + end +end