Broadcast Kernel Fusion

.githash

@@ -1 +1 @@
-dec047f1bd1c8287513c6c437f946982e516ccd4
+e4cc4b6a778aa828e5647803f594222053797f54

Project.toml

@@ -4,8 +4,12 @@ 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"
+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"
@@ -20,17 +24,14 @@ 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"
+ExpressionExplorer = "1.1.4"
 JuliaFormatter = "2.3.0"
+KernelAbstractions = "0.9.41"
 Legate = "0.1.2"
 LegatePreferences = "0.1.6"
 MacroTools = "0.5.16"

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

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")
 
@@ -62,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"
@@ -111,13 +121,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,6 +134,35 @@ 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
 """
@@ -147,7 +179,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
     ) # $pkg_root/lib/cunumeric_jl_wrapper
 end
 
@@ -161,9 +193,17 @@ 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
@@ -172,22 +212,40 @@ 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
+    )
 end
 
 const mode_str = load_preference(CNPreferences, "cunumeric_mode", CNPreferences.MODE_JLL)