CUDA build local (#94)

Co-authored-by: William Moses <[email protected]>

Author: wsmoses

deps/ReactantExtra/.bazelrc

@@ -14,3 +14,17 @@ build --define=allow_oversize_protos=true
 
 build -c opt
 
+
+build:cuda --repo_env TF_NEED_CUDA=1
+build:cuda --repo_env TF_NCCL_USE_STUB=1
+build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.3.2"
+build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.1.1"
+# "sm" means we emit only cubin, which is forward compatible within a GPU generation.
+# "compute" means we emit both cubin and PTX, which is larger but also forward compatible to future GPU generations.
+build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,compute_90"
+build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
+build:cuda --@local_config_cuda//:enable_cuda
+build:cuda --@xla//xla/python:jax_cuda_pip_rpaths=true
+# Default hermetic CUDA and CUDNN versions.
+build:cuda --@local_config_cuda//cuda:include_cuda_libs=true
+build:cuda --@local_config_cuda//:cuda_compiler=nvcc

deps/build_local.jl

@@ -37,9 +37,14 @@ run(Cmd(`$(Base.julia_cmd().exec[1]) --project=. -e "using Pkg; Pkg.instantiate(
 # --@local_config_cuda//:cuda_compiler=nvcc
 # --crosstool_top="@local_config_cuda//crosstool:toolchain"
 
-run(Cmd(`bazel query --experimental_repo_remote_exec --repo_env HERMETIC_PYTHON_VERSION="3.10" "allpaths(:libReactantExtra.so, @llvm-project//mlir:CAPIIR)" --output graph`, dir=source_dir,
-))
-run(Cmd(`bazel build -c dbg --action_env=JULIA=$(Base.julia_cmd().exec[1])
+arg = try
+	run(Cmd(`nvidia-smi`))
+	"--config=cuda"
+catch
+	""
+end
+
+run(Cmd(`bazel build $(arg) -c dbg --action_env=JULIA=$(Base.julia_cmd().exec[1])
 --repo_env HERMETIC_PYTHON_VERSION="3.10"
 --check_visibility=false --verbose_failures :libReactantExtra.so :Builtin.inc.jl :Arith.inc.jl :Affine.inc.jl :Func.inc.jl :Enzyme.inc.jl :StableHLO.inc.jl :CHLO.inc.jl :VHLO.inc.jl`, dir=source_dir,
 ))

src/XLA.jl

@@ -38,6 +38,8 @@ function from_row_major(x::Matrix{T}) where {T}
     return transpose(x)
 end
 
+SetLogLevel(x) = @ccall MLIR.API.mlir_c.SetLogLevel(x::Cint)::Cvoid
+
 const cpuclientcount = Ref(0)
 # TODO synchronization when async is not working because `future` in `ConcreteRArray` is always `nothing`
 function CPUClient(asynchronous=false, node_id=0, num_nodes=1)
@@ -93,6 +95,8 @@ using Scratch, Downloads
 function __init__()
     initLogs = Libdl.dlsym(Reactant_jll.libReactantExtra_handle, "InitializeLogs")
     ccall(initLogs, Cvoid, ())
+    # Add most log level
+    SetLogLevel(0)
     cpu = CPUClient()
     backends["cpu"] = cpu
     default_backend[] = cpu