Improvements to large buffer tests

.buildkite/pipeline.yml

@@ -116,6 +116,34 @@ steps:
             build.message !~ /\[skip validation\]/ &&
             build.message !~ /\[skip special\]/
         timeout_in_minutes: 60
+      - label: ":ram: Large buffers"
+        plugins:
+          - JuliaCI/julia#v1:
+              version: "1.12"
+          - JuliaCI/julia-test#v1:
+              test_args: "--jobs=1 --quickfail largecopy largebroadcast"
+          - JuliaCI/julia-coverage#v1:
+              codecov: true
+              dirs:
+                - src
+                - lib
+        env:
+          MTL_DEBUG_LAYER: '1'
+          MTL_SHADER_VALIDATION: '1'
+        agents:
+          queue: "juliaecosystem"
+          os: "macos"
+          arch: "aarch64"
+          macos_version: "15.0"
+        if: |
+          build.message =~ /\[only tests\]/ ||
+          build.message =~ /\[only large\]/ ||
+          build.message =~ /\[only special\]/ ||
+          build.message !~ /\[only/ && !build.pull_request.draft &&
+            build.message !~ /\[skip tests\]/ &&
+            build.message !~ /\[skip large\]/ &&
+            build.message !~ /\[skip special\]/
+        timeout_in_minutes: 15
 
   - wait: ~
     continue_on_failure: true

test/largebroadcast.jl

@@ -0,0 +1,35 @@
+const N = Int(typemax(UInt32)) + 1
+const T = Int8
+
+@testset "len = $n" for n in ((N÷2) - 4, (N÷2), (N÷2) + 4, N - 1024, N - 3, N - 1, N, N + 4)
+    A = MtlArray{T}(undef, n)
+    # Known working method to zero out array
+    Metal.unsafe_fill!(device(A), pointer(A), T(0), n * sizeof(T); async = false)
+
+    _dims = [(n,), (n, 1), (1, n), (n, 1, 1), (1, n, 1), (1, 1, n),
+             (n, 1, 1, 1), (1, n, 1, 1), (1, 1, n, 1), (1, 1, 1, n),]
+    if n == 2^32
+        push!(_dims, (2^16, 2^16))
+        push!(_dims, (2^16, 2^8, 2^8))
+    end
+
+    @testset "$dims" for (i, dims) in enumerate(_dims)
+        # These must be run first to ensure we test
+        # the unspecialized broadcast kernels
+        Metal._broadcast_shapes[CartesianIndices(dims)] = Metal.BROADCAST_SPECIALIZATION_THRESHOLD - 1
+        unspec_val = T((i-1) * 2 + 1)
+        arr = reshape(A, dims)
+        arr .= unspec_val
+        @test all(==(unspec_val), arr)
+
+        # Test specialized compiled kernel
+        # Not necessary but just in case
+        Metal._broadcast_shapes[CartesianIndices(dims)] = Metal.BROADCAST_SPECIALIZATION_THRESHOLD
+        spec_val = T(unspec_val + 1)
+        arr .= spec_val
+        @test all(==(spec_val), arr)
+    end
+    A = nothing
+    GC.gc()
+    GC.gc()
+end

test/runtests.jl

@@ -85,9 +85,14 @@ if filter_tests!(testsuite, args)
     end
 
     # only run large copy test on machines with >12GiB memory
-    if Sys.total_memory() < 12 * 2^30
+    if parse(Bool, get(ENV, "CI", "false")) || Sys.total_memory() < 12 * 2^30
         delete!(testsuite, "largecopy")
     end
+
+    # only run large broadcast test on machines with >12GiB memory
+    if parse(Bool, get(ENV, "CI", "false")) || Sys.total_memory() < 12 * 2^30
+        delete!(testsuite, "largebroadcast")
+    end
 end
 
 # workers to run tests on