allow for unary reductions with arbitrary dims

.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: ""

.githash

@@ -1 +1 @@
-dec047f1bd1c8287513c6c437f946982e516ccd4
+3713a2acb6579891706b08f970e99f9ae5e6472f

.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

.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

.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

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<int32_t> axis_vec(axes, axes + num_axes);
+    NDArray result = input->obj._perform_unary_reduction(
+        static_cast<int32_t>(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<legate::Slice> slices) {
   switch (slices.size()) {
     case 1: {

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),

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

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"

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
 

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