From 4cd05ad11e205874eebdfe02467d58a538356c8d Mon Sep 17 00:00:00 2001
From: jamesETsmith <james.smith9113@gmail.com>
Date: Fri, 20 Mar 2026 14:49:05 -0400
Subject: [PATCH 1/8] Making the Dockerfile.TheRock_ci configurable, some of
 the versions were hardcoded

---
 docker/rocm_ci/Dockerfile.TheRock_ci | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/docker/rocm_ci/Dockerfile.TheRock_ci b/docker/rocm_ci/Dockerfile.TheRock_ci
index b2339332a6..472922be70 100644
--- a/docker/rocm_ci/Dockerfile.TheRock_ci
+++ b/docker/rocm_ci/Dockerfile.TheRock_ci
@@ -56,8 +56,9 @@ RUN apt update \
 
 
 # Note: the index page is at https://therock-nightly-tarball.s3.amazonaws.com/index.html
-# This will directly download the TheROCK nightly ROCm SDK tarball
-ENV THEROCK_TARBALL=https://rocm.nightlies.amd.com/tarball/therock-dist-linux-gfx94X-dcgpu-7.9.0rc20250930.tar.gz
+# This directly downloads the TheRock nightly ROCm SDK tarball.
+# Override THE_ROCK_VERSION/GFX_FAMILY at build time as needed.
+ENV THEROCK_TARBALL=https://rocm.nightlies.amd.com/tarball/therock-dist-linux-${GFX_FAMILY}-${THE_ROCK_VERSION}.tar.gz
 
 RUN mkdir -p /opt/rocm \
     && cd /opt/rocm \

From 28b86dc8b6a6d26e77cecbe320f94900bc006fe6 Mon Sep 17 00:00:00 2001
From: jamesETsmith <james.smith9113@gmail.com>
Date: Wed, 25 Mar 2026 09:25:24 -0400
Subject: [PATCH 2/8] making some changes to the docker build file to install
 theRock with pip and address some of the issues with the build system that
 expected things in /opt/rocm

---
 docker/rocm_ci/Dockerfile.TheRock_ci | 33 ++++++------------
 warp/_src/build_dll.py               | 50 ++++++++++++++++++++++++----
 2 files changed, 55 insertions(+), 28 deletions(-)

diff --git a/docker/rocm_ci/Dockerfile.TheRock_ci b/docker/rocm_ci/Dockerfile.TheRock_ci
index 472922be70..4327f44717 100644
--- a/docker/rocm_ci/Dockerfile.TheRock_ci
+++ b/docker/rocm_ci/Dockerfile.TheRock_ci
@@ -6,8 +6,7 @@
 # context should contain.
 #
 #   docker buildx (local):
-#     docker buildx build --build-context warp_src=/path/to/warp \
-#         -f docker/Dockerfile.rocm_ci --target warp_base .
+#      docker buildx build --build-context warp_src=/path/to/warp -f docker/rocm_ci/Dockerfile.TheRock_ci --target warp_base .
 #
 #   buildctl (CI entrypoint.sh):
 #     buildctl build --local warp_src="${SRC_FOLDER}/warp" \
@@ -27,9 +26,9 @@ ENV PYTORCH_ROCM_ARCH="gfx942"
 ENV ROCM_AMDGPU_TARGETS="gfx942"
 
 
-ARG THE_ROCK_VERSION=7.9.0rc20250930
+ARG THE_ROCK_VERSION=7.10.0
 ARG GFX_FAMILY="gfx94X-dcgpu"
-ARG PY_VERSION=3.12
+ARG PY_VERSION=3.12 
 ARG TORCH_VERSION=2.9.1
 ARG MAMBA_ENV_NAME=base
 
@@ -55,19 +54,6 @@ RUN apt update \
     && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-14 100
 
 
-# Note: the index page is at https://therock-nightly-tarball.s3.amazonaws.com/index.html
-# This directly downloads the TheRock nightly ROCm SDK tarball.
-# Override THE_ROCK_VERSION/GFX_FAMILY at build time as needed.
-ENV THEROCK_TARBALL=https://rocm.nightlies.amd.com/tarball/therock-dist-linux-${GFX_FAMILY}-${THE_ROCK_VERSION}.tar.gz
-
-RUN mkdir -p /opt/rocm \
-    && cd /opt/rocm \
-    && wget ${THEROCK_TARBALL} \
-    && tar -xvf *.tar.gz -C /opt/rocm \
-    && rm *.tar.gz \
-    # && rm /opt/rocm/bin/flatc \
-    && rm -rf /opt/rocm/include/flatbuffers \
-    && rm -rf /opt/rocm/lib/libflatbuffers*
 
 # warp_src is a named build context — see file header for invocation details.
 COPY --from=warp_src . /root/warp/
@@ -81,21 +67,24 @@ RUN git config --global --add safe.directory /root/warp
 
 # Setup ROCm stack related envs
 # Note: /opt/rocm/lib is intentionally NOT in PATH (it is a library dir, not a bin dir).
-ENV PATH="/opt/rocm/bin:$PATH"
-ENV LD_LIBRARY_PATH=/opt/rocm/lib
+# ENV PATH="/opt/rocm/bin:$PATH"
+# ENV LD_LIBRARY_PATH=/opt/rocm/lib
 ENV ROCM_PATH=/opt/rocm
 ENV HIP_PLATFORM=amd
-ENV HIP_DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode
+# ENV HIP_DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode
 # Workaround for the HAS_STATUS_ERROR_OUT_OF_RESOURCES issue (runtime)
 ENV HAS_ENABLE_SCRATCH_ASYNC_RECLAIM=0
 ENV HAS_NO_SCRATCH_RECLAIM=1
-
+# https://repo.amd.com/rocm/whl/gfx94X-dcgpu/torch/torch-2.9.1+rocm7.10.0-cp312-cp312-linux_x86_64.whl
+# https://repo.amd.com/rocm/whl/gfx94X-dcgpu/torch-2.9.1%2Brocm7.10.0-cp312-cp312-linux_x86_64.whl
 RUN curl -L micro.mamba.pm/install.sh | bash && \
     ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} ~/.local/bin/micromamba install python=${PY_VERSION} -c conda-forge --override-channels -y && \
     ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install cmake pybind11 build ninja scikit-build-core setuptools-scm numpy pytest pytest-xdist Pillow && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://rocm.nightlies.amd.com/v2-staging/${GFX_FAMILY}/ --pre torch && \
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://repo.amd.com/rocm/whl/${GFX_FAMILY}/ "rocm[libraries,devel]==${THE_ROCK_VERSION}" && \
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://repo.amd.com/rocm/whl/${GFX_FAMILY}/ torch=="${TORCH_VERSION}" && \
     /bin/bash -c "eval \"\$(~/.local/bin/micromamba shell hook --shell bash)\" && \
     micromamba activate ${MAMBA_ENV_NAME} && \
+    rocm-sdk init && export ROCM_PATH=\$(rocm-sdk path --root) &&\
     cd /root/warp/ && \
     python build_lib.py --jobs \$(nproc) && \
     python -m build --wheel -C--build-option=-Plinux-x86_64 && \
diff --git a/warp/_src/build_dll.py b/warp/_src/build_dll.py
index cc87bafc72..c3b12b36bb 100644
--- a/warp/_src/build_dll.py
+++ b/warp/_src/build_dll.py
@@ -53,6 +53,37 @@ def find_rocm_sdk() -> str | None:
     return None
 
 
+def find_rocm_device_lib_path(rocm_path: str | None) -> str | None:
+    """Find the ROCm device library bitcode directory.
+
+    Clang needs bitcode files (e.g. ``ocml.bc``) for GPU code generation.
+    In a standard ``/opt/rocm`` install they live under ``amdgcn/bitcode/``,
+    but pip-installed TheRock SDKs may place them in a separate
+    ``_rocm_sdk_core`` site-packages tree.
+    """
+    if rocm_path is None:
+        return None
+
+    # Standard locations relative to the ROCm root
+    candidates = [
+        os.path.join(rocm_path, "amdgcn", "bitcode"),
+        os.path.join(rocm_path, "lib", "llvm", "amdgcn", "bitcode"),
+    ]
+
+    # TheRock pip-installed SDK: the core package lives next to _rocm_sdk_devel
+    # in site-packages and holds the bitcode files.
+    if "_rocm_sdk_devel" in rocm_path:
+        core_root = rocm_path.replace("_rocm_sdk_devel", "_rocm_sdk_core")
+        candidates.append(os.path.join(core_root, "amdgcn", "bitcode"))
+        candidates.append(os.path.join(core_root, "lib", "llvm", "amdgcn", "bitcode"))
+
+    for path in candidates:
+        if os.path.isdir(path):
+            return path
+
+    return None
+
+
 def find_hipcc_executable(rocm_path: str | None) -> str:
     hipcc_name = "hipcc.exe" if os.name == "nt" else "hipcc"
     if rocm_path:
@@ -823,17 +854,24 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_paths, arch, libs: list[str
                         hip_arches = _parse_hip_arches(args)
                         hip_arch_flags = " ".join([f"--offload-arch={arch}" for arch in hip_arches])
                         # Match nvcc/NVRTC default: strict IEEE 754 FP semantics
-                        hip_fp_flags = "-fno-finite-math-only -fno-associative-math -fno-reciprocal-math -fno-strict-aliasing"
+                        hip_fp_flags = (
+                            "-fno-finite-math-only -fno-associative-math -fno-reciprocal-math -fno-strict-aliasing"
+                        )
+                        hip_includes = f'-I"{native_dir}" -I"{args.rocm_path}/include"'
+                        hip_rocm_path_flag = f'--rocm-path="{args.rocm_path}"'
+                        device_lib_path = find_rocm_device_lib_path(args.rocm_path)
+                        if device_lib_path:
+                            hip_rocm_path_flag += f' --rocm-device-lib-path="{device_lib_path}"'
                         if mode == "debug":
                             cuda_cmd = (
-                                f'{hipcc_cmd} -x hip -std=c++17 -g -O0 -fPIC -fvisibility=hidden '
-                                f'-D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 {hip_fp_flags} {hip_arch_flags} -DWP_ENABLE_CUDA=1 '
-                                f'-I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
+                                f"{hipcc_cmd} -x hip -std=c++17 -g -O0 -fPIC -fvisibility=hidden "
+                                f"-D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 {hip_fp_flags} {hip_arch_flags} {hip_rocm_path_flag} -DWP_ENABLE_CUDA=1 "
+                                f'{hip_includes} -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
                             )
                         elif mode == "release":
                             cuda_cmd = (
-                                f'{hipcc_cmd} -x hip -std=c++17 -O3 -fPIC -fvisibility=hidden -DNDEBUG '
-                                f'{hip_fp_flags} {hip_arch_flags} -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} '
+                                f"{hipcc_cmd} -x hip -std=c++17 -O3 -fPIC -fvisibility=hidden -DNDEBUG "
+                                f"{hip_fp_flags} {hip_arch_flags} {hip_rocm_path_flag} -DWP_ENABLE_CUDA=1 {hip_includes} -D{mathdx_enabled} "
                                 f'{libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
                             )
                     elif cuda_compiler == "nvcc":

From ad12acd8d2f72b553a7f8efddc45d5e128920840 Mon Sep 17 00:00:00 2001
From: jamesETsmith <james.smith9113@gmail.com>
Date: Wed, 25 Mar 2026 14:35:13 -0400
Subject: [PATCH 3/8] reverting some of my changes. I've kept the splitting of
 the docker stages for installing prereqs and building warp to speed up
 iterative development on the image

---
 docker/rocm_ci/Dockerfile.TheRock_ci | 51 ++++++++++++++++------------
 1 file changed, 29 insertions(+), 22 deletions(-)

diff --git a/docker/rocm_ci/Dockerfile.TheRock_ci b/docker/rocm_ci/Dockerfile.TheRock_ci
index 4327f44717..6f6d8a847d 100644
--- a/docker/rocm_ci/Dockerfile.TheRock_ci
+++ b/docker/rocm_ci/Dockerfile.TheRock_ci
@@ -6,7 +6,7 @@
 # context should contain.
 #
 #   docker buildx (local):
-#      docker buildx build --build-context warp_src=/path/to/warp -f docker/rocm_ci/Dockerfile.TheRock_ci --target warp_base .
+#      docker buildx build --build-context warp_src=/path/to/warp -f docker/rocm_ci/Dockerfile.TheRock_ci --target warp_base -t warp-base .
 #
 #   buildctl (CI entrypoint.sh):
 #     buildctl build --local warp_src="${SRC_FOLDER}/warp" \
@@ -26,10 +26,9 @@ ENV PYTORCH_ROCM_ARCH="gfx942"
 ENV ROCM_AMDGPU_TARGETS="gfx942"
 
 
-ARG THE_ROCK_VERSION=7.10.0
+ARG THE_ROCK_VERSION=7.9.0rc20250930
 ARG GFX_FAMILY="gfx94X-dcgpu"
 ARG PY_VERSION=3.12 
-ARG TORCH_VERSION=2.9.1
 ARG MAMBA_ENV_NAME=base
 
 # Install Ubuntu dependencies
@@ -55,6 +54,32 @@ RUN apt update \
 
 
 
+# Setup ROCm stack related envs
+# Note: the index page is at https://therock-nightly-tarball.s3.amazonaws.com/index.html
+# This will directly download the TheROCK nightly ROCm SDK tarball
+ENV THEROCK_TARBALL=https://rocm.nightlies.amd.com/tarball/therock-dist-linux-${GFX_FAMILY}-${THE_ROCK_VERSION}.tar.gz
+
+RUN mkdir -p /opt/rocm \
+    && cd /opt/rocm \
+    && wget ${THEROCK_TARBALL} \
+    && tar -xvf *.tar.gz -C /opt/rocm \
+    && rm *.tar.gz \
+    && rm -rf /opt/rocm/include/flatbuffers \
+    && rm -rf /opt/rocm/lib/libflatbuffers*
+
+ENV PATH="/opt/rocm/bin:$PATH"
+ENV LD_LIBRARY_PATH=/opt/rocm/lib
+ENV ROCM_PATH=/opt/rocm
+ENV HIP_PLATFORM=amd
+ENV HIP_DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode
+# Workaround for the HAS_STATUS_ERROR_OUT_OF_RESOURCES issue (runtime)
+ENV HAS_ENABLE_SCRATCH_ASYNC_RECLAIM=0
+ENV HAS_NO_SCRATCH_RECLAIM=1
+RUN curl -L micro.mamba.pm/install.sh | bash && \
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} ~/.local/bin/micromamba install python=${PY_VERSION} -c conda-forge --override-channels -y && \
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install cmake pybind11 build ninja scikit-build-core setuptools-scm numpy pytest pytest-xdist Pillow && \
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://rocm.nightlies.amd.com/v2-staging/${GFX_FAMILY}/ --pre torch
+
 # warp_src is a named build context — see file header for invocation details.
 COPY --from=warp_src . /root/warp/
 RUN test -f /root/warp/pyproject.toml && \
@@ -65,26 +90,8 @@ RUN test -f /root/warp/pyproject.toml && \
     }
 RUN git config --global --add safe.directory /root/warp
 
-# Setup ROCm stack related envs
-# Note: /opt/rocm/lib is intentionally NOT in PATH (it is a library dir, not a bin dir).
-# ENV PATH="/opt/rocm/bin:$PATH"
-# ENV LD_LIBRARY_PATH=/opt/rocm/lib
-ENV ROCM_PATH=/opt/rocm
-ENV HIP_PLATFORM=amd
-# ENV HIP_DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode
-# Workaround for the HAS_STATUS_ERROR_OUT_OF_RESOURCES issue (runtime)
-ENV HAS_ENABLE_SCRATCH_ASYNC_RECLAIM=0
-ENV HAS_NO_SCRATCH_RECLAIM=1
-# https://repo.amd.com/rocm/whl/gfx94X-dcgpu/torch/torch-2.9.1+rocm7.10.0-cp312-cp312-linux_x86_64.whl
-# https://repo.amd.com/rocm/whl/gfx94X-dcgpu/torch-2.9.1%2Brocm7.10.0-cp312-cp312-linux_x86_64.whl
-RUN curl -L micro.mamba.pm/install.sh | bash && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} ~/.local/bin/micromamba install python=${PY_VERSION} -c conda-forge --override-channels -y && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install cmake pybind11 build ninja scikit-build-core setuptools-scm numpy pytest pytest-xdist Pillow && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://repo.amd.com/rocm/whl/${GFX_FAMILY}/ "rocm[libraries,devel]==${THE_ROCK_VERSION}" && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://repo.amd.com/rocm/whl/${GFX_FAMILY}/ torch=="${TORCH_VERSION}" && \
-    /bin/bash -c "eval \"\$(~/.local/bin/micromamba shell hook --shell bash)\" && \
+RUN /bin/bash -c "eval \"\$(~/.local/bin/micromamba shell hook --shell bash)\" && \
     micromamba activate ${MAMBA_ENV_NAME} && \
-    rocm-sdk init && export ROCM_PATH=\$(rocm-sdk path --root) &&\
     cd /root/warp/ && \
     python build_lib.py --jobs \$(nproc) && \
     python -m build --wheel -C--build-option=-Plinux-x86_64 && \

From f141b73e466a1a20eeee2c3e286354ca63b2ffc8 Mon Sep 17 00:00:00 2001
From: jamesETsmith <james.smith9113@gmail.com>
Date: Wed, 25 Mar 2026 15:56:04 -0400
Subject: [PATCH 4/8] Adding rocprof-compute to docker and bumping theRock
 version so we can get it to match the rockprof compute

---
 docker/rocm_ci/Dockerfile.TheRock_ci | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/docker/rocm_ci/Dockerfile.TheRock_ci b/docker/rocm_ci/Dockerfile.TheRock_ci
index 6f6d8a847d..9d6c002c6f 100644
--- a/docker/rocm_ci/Dockerfile.TheRock_ci
+++ b/docker/rocm_ci/Dockerfile.TheRock_ci
@@ -25,8 +25,7 @@ ENV MAX_JOBS=128
 ENV PYTORCH_ROCM_ARCH="gfx942"
 ENV ROCM_AMDGPU_TARGETS="gfx942"
 
-
-ARG THE_ROCK_VERSION=7.9.0rc20250930
+ARG THE_ROCK_VERSION=7.10.0a20251029
 ARG GFX_FAMILY="gfx94X-dcgpu"
 ARG PY_VERSION=3.12 
 ARG MAMBA_ENV_NAME=base
@@ -80,6 +79,15 @@ RUN curl -L micro.mamba.pm/install.sh | bash && \
     ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install cmake pybind11 build ninja scikit-build-core setuptools-scm numpy pytest pytest-xdist Pillow && \
     ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://rocm.nightlies.amd.com/v2-staging/${GFX_FAMILY}/ --pre torch
 
+# Install rocprof-compute
+RUN wget https://github.com/ROCm/rocm-systems/releases/download/therock-7.10/rocprofiler-compute.tar.gz && \
+    tar -xvf rocprofiler-compute.tar.gz -C /tmp && \
+    cd /tmp/rocprofiler-compute && \
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install -r requirements.txt && \
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} cmake -B build -G Ninja && \
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} cmake --build build --target install && \
+    which rocprof-compute
+
 # warp_src is a named build context — see file header for invocation details.
 COPY --from=warp_src . /root/warp/
 RUN test -f /root/warp/pyproject.toml && \

From 00ddd60b1e2b9ed03853872df1951b62d7bc2448 Mon Sep 17 00:00:00 2001
From: jamesETsmith <james.smith9113@gmail.com>
Date: Wed, 25 Mar 2026 16:08:46 -0400
Subject: [PATCH 5/8] Reverting change to build_dll.py

---
 warp/_src/build_dll.py | 50 +++++-------------------------------------
 1 file changed, 6 insertions(+), 44 deletions(-)

diff --git a/warp/_src/build_dll.py b/warp/_src/build_dll.py
index c3b12b36bb..cc87bafc72 100644
--- a/warp/_src/build_dll.py
+++ b/warp/_src/build_dll.py
@@ -53,37 +53,6 @@ def find_rocm_sdk() -> str | None:
     return None
 
 
-def find_rocm_device_lib_path(rocm_path: str | None) -> str | None:
-    """Find the ROCm device library bitcode directory.
-
-    Clang needs bitcode files (e.g. ``ocml.bc``) for GPU code generation.
-    In a standard ``/opt/rocm`` install they live under ``amdgcn/bitcode/``,
-    but pip-installed TheRock SDKs may place them in a separate
-    ``_rocm_sdk_core`` site-packages tree.
-    """
-    if rocm_path is None:
-        return None
-
-    # Standard locations relative to the ROCm root
-    candidates = [
-        os.path.join(rocm_path, "amdgcn", "bitcode"),
-        os.path.join(rocm_path, "lib", "llvm", "amdgcn", "bitcode"),
-    ]
-
-    # TheRock pip-installed SDK: the core package lives next to _rocm_sdk_devel
-    # in site-packages and holds the bitcode files.
-    if "_rocm_sdk_devel" in rocm_path:
-        core_root = rocm_path.replace("_rocm_sdk_devel", "_rocm_sdk_core")
-        candidates.append(os.path.join(core_root, "amdgcn", "bitcode"))
-        candidates.append(os.path.join(core_root, "lib", "llvm", "amdgcn", "bitcode"))
-
-    for path in candidates:
-        if os.path.isdir(path):
-            return path
-
-    return None
-
-
 def find_hipcc_executable(rocm_path: str | None) -> str:
     hipcc_name = "hipcc.exe" if os.name == "nt" else "hipcc"
     if rocm_path:
@@ -854,24 +823,17 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_paths, arch, libs: list[str
                         hip_arches = _parse_hip_arches(args)
                         hip_arch_flags = " ".join([f"--offload-arch={arch}" for arch in hip_arches])
                         # Match nvcc/NVRTC default: strict IEEE 754 FP semantics
-                        hip_fp_flags = (
-                            "-fno-finite-math-only -fno-associative-math -fno-reciprocal-math -fno-strict-aliasing"
-                        )
-                        hip_includes = f'-I"{native_dir}" -I"{args.rocm_path}/include"'
-                        hip_rocm_path_flag = f'--rocm-path="{args.rocm_path}"'
-                        device_lib_path = find_rocm_device_lib_path(args.rocm_path)
-                        if device_lib_path:
-                            hip_rocm_path_flag += f' --rocm-device-lib-path="{device_lib_path}"'
+                        hip_fp_flags = "-fno-finite-math-only -fno-associative-math -fno-reciprocal-math -fno-strict-aliasing"
                         if mode == "debug":
                             cuda_cmd = (
-                                f"{hipcc_cmd} -x hip -std=c++17 -g -O0 -fPIC -fvisibility=hidden "
-                                f"-D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 {hip_fp_flags} {hip_arch_flags} {hip_rocm_path_flag} -DWP_ENABLE_CUDA=1 "
-                                f'{hip_includes} -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
+                                f'{hipcc_cmd} -x hip -std=c++17 -g -O0 -fPIC -fvisibility=hidden '
+                                f'-D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 {hip_fp_flags} {hip_arch_flags} -DWP_ENABLE_CUDA=1 '
+                                f'-I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
                             )
                         elif mode == "release":
                             cuda_cmd = (
-                                f"{hipcc_cmd} -x hip -std=c++17 -O3 -fPIC -fvisibility=hidden -DNDEBUG "
-                                f"{hip_fp_flags} {hip_arch_flags} {hip_rocm_path_flag} -DWP_ENABLE_CUDA=1 {hip_includes} -D{mathdx_enabled} "
+                                f'{hipcc_cmd} -x hip -std=c++17 -O3 -fPIC -fvisibility=hidden -DNDEBUG '
+                                f'{hip_fp_flags} {hip_arch_flags} -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} '
                                 f'{libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
                             )
                     elif cuda_compiler == "nvcc":

From 04ad69b066e8fb00d1e680f5f541542f70969087 Mon Sep 17 00:00:00 2001
From: jamesETsmith <james.smith9113@gmail.com>
Date: Thu, 26 Mar 2026 16:09:36 -0400
Subject: [PATCH 6/8] Locking the torch and rocm versions together

---
 docker/rocm_ci/Dockerfile.TheRock_ci | 31 ++++++++++------------------
 1 file changed, 11 insertions(+), 20 deletions(-)

diff --git a/docker/rocm_ci/Dockerfile.TheRock_ci b/docker/rocm_ci/Dockerfile.TheRock_ci
index 9d6c002c6f..4762026e64 100644
--- a/docker/rocm_ci/Dockerfile.TheRock_ci
+++ b/docker/rocm_ci/Dockerfile.TheRock_ci
@@ -25,9 +25,14 @@ ENV MAX_JOBS=128
 ENV PYTORCH_ROCM_ARCH="gfx942"
 ENV ROCM_AMDGPU_TARGETS="gfx942"
 
-ARG THE_ROCK_VERSION=7.10.0a20251029
+
+# 7.10.0a20251121
+# 7.11.0a20251201
+# 7.12.0a20260311
+ARG THE_ROCK_VERSION=7.10.0a20251121
 ARG GFX_FAMILY="gfx94X-dcgpu"
-ARG PY_VERSION=3.12 
+ARG TORCH_VERSION=2.9.1
+ARG PY_VERSION=3.12
 ARG MAMBA_ENV_NAME=base
 
 # Install Ubuntu dependencies
@@ -51,21 +56,6 @@ RUN apt update \
     && update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-14 100 \
     && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-14 100
 
-
-
-# Setup ROCm stack related envs
-# Note: the index page is at https://therock-nightly-tarball.s3.amazonaws.com/index.html
-# This will directly download the TheROCK nightly ROCm SDK tarball
-ENV THEROCK_TARBALL=https://rocm.nightlies.amd.com/tarball/therock-dist-linux-${GFX_FAMILY}-${THE_ROCK_VERSION}.tar.gz
-
-RUN mkdir -p /opt/rocm \
-    && cd /opt/rocm \
-    && wget ${THEROCK_TARBALL} \
-    && tar -xvf *.tar.gz -C /opt/rocm \
-    && rm *.tar.gz \
-    && rm -rf /opt/rocm/include/flatbuffers \
-    && rm -rf /opt/rocm/lib/libflatbuffers*
-
 ENV PATH="/opt/rocm/bin:$PATH"
 ENV LD_LIBRARY_PATH=/opt/rocm/lib
 ENV ROCM_PATH=/opt/rocm
@@ -77,8 +67,9 @@ ENV HAS_NO_SCRATCH_RECLAIM=1
 RUN curl -L micro.mamba.pm/install.sh | bash && \
     ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} ~/.local/bin/micromamba install python=${PY_VERSION} -c conda-forge --override-channels -y && \
     ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install cmake pybind11 build ninja scikit-build-core setuptools-scm numpy pytest pytest-xdist Pillow && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://rocm.nightlies.amd.com/v2-staging/${GFX_FAMILY}/ --pre torch
-
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://rocm.nightlies.amd.com/v2/${GFX_FAMILY}/ "rocm[libraries,devel]"==${THE_ROCK_VERSION} && \
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} rocm-sdk init && \
+    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://rocm.nightlies.amd.com/v2/${GFX_FAMILY}/ torch==${TORCH_VERSION}
 # Install rocprof-compute
 RUN wget https://github.com/ROCm/rocm-systems/releases/download/therock-7.10/rocprofiler-compute.tar.gz && \
     tar -xvf rocprofiler-compute.tar.gz -C /tmp && \
@@ -101,7 +92,7 @@ RUN git config --global --add safe.directory /root/warp
 RUN /bin/bash -c "eval \"\$(~/.local/bin/micromamba shell hook --shell bash)\" && \
     micromamba activate ${MAMBA_ENV_NAME} && \
     cd /root/warp/ && \
-    python build_lib.py --jobs \$(nproc) && \
+    python build_lib.py --jobs \$(nproc) 2>&1 | tee build_lib_{THE_ROCK_VERSION}.log && \
     python -m build --wheel -C--build-option=-Plinux-x86_64 && \
     pip install dist/warp_lang-*.whl"
 

From a45ad3382b6b73f38c94c92fdc02127e38fb73f0 Mon Sep 17 00:00:00 2001
From: jamesETsmith <james.smith9113@gmail.com>
Date: Thu, 26 Mar 2026 22:32:32 -0400
Subject: [PATCH 7/8] Simplifying the theRock_ci dockerfile by switching from
 mamba -> uv and installing rocm via uv pip (not working 100% yet)

---
 docker/rocm_ci/Dockerfile.TheRock_ci | 67 ++++++++++------------------
 1 file changed, 24 insertions(+), 43 deletions(-)

diff --git a/docker/rocm_ci/Dockerfile.TheRock_ci b/docker/rocm_ci/Dockerfile.TheRock_ci
index 4762026e64..e244a25b9d 100644
--- a/docker/rocm_ci/Dockerfile.TheRock_ci
+++ b/docker/rocm_ci/Dockerfile.TheRock_ci
@@ -33,7 +33,6 @@ ARG THE_ROCK_VERSION=7.10.0a20251121
 ARG GFX_FAMILY="gfx94X-dcgpu"
 ARG TORCH_VERSION=2.9.1
 ARG PY_VERSION=3.12
-ARG MAMBA_ENV_NAME=base
 
 # Install Ubuntu dependencies
 ENV DEBIAN_FRONTEND=noninteractive
@@ -56,28 +55,17 @@ RUN apt update \
     && update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-14 100 \
     && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-14 100
 
-ENV PATH="/opt/rocm/bin:$PATH"
-ENV LD_LIBRARY_PATH=/opt/rocm/lib
-ENV ROCM_PATH=/opt/rocm
-ENV HIP_PLATFORM=amd
-ENV HIP_DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode
-# Workaround for the HAS_STATUS_ERROR_OUT_OF_RESOURCES issue (runtime)
-ENV HAS_ENABLE_SCRATCH_ASYNC_RECLAIM=0
-ENV HAS_NO_SCRATCH_RECLAIM=1
-RUN curl -L micro.mamba.pm/install.sh | bash && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} ~/.local/bin/micromamba install python=${PY_VERSION} -c conda-forge --override-channels -y && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install cmake pybind11 build ninja scikit-build-core setuptools-scm numpy pytest pytest-xdist Pillow && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://rocm.nightlies.amd.com/v2/${GFX_FAMILY}/ "rocm[libraries,devel]"==${THE_ROCK_VERSION} && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} rocm-sdk init && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --index-url https://rocm.nightlies.amd.com/v2/${GFX_FAMILY}/ torch==${TORCH_VERSION}
-# Install rocprof-compute
-RUN wget https://github.com/ROCm/rocm-systems/releases/download/therock-7.10/rocprofiler-compute.tar.gz && \
-    tar -xvf rocprofiler-compute.tar.gz -C /tmp && \
-    cd /tmp/rocprofiler-compute && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} pip install -r requirements.txt && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} cmake -B build -G Ninja && \
-    ~/.local/bin/micromamba run -n ${MAMBA_ENV_NAME} cmake --build build --target install && \
-    which rocprof-compute
+# Setup uv
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+RUN uv venv -p ${PY_VERSION} /opt/venv
+ENV VIRTUAL_ENV=/opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+
+# rocm and torch MUST be installed together or torch nightlies will 
+# mess up the rocm version
+RUN uv pip install cmake pybind11 build ninja scikit-build-core setuptools-scm numpy pytest pytest-xdist Pillow && \
+    uv pip install --index-url https://rocm.nightlies.amd.com/v2/${GFX_FAMILY}/ "rocm[libraries,devel]"==${THE_ROCK_VERSION} torch==${TORCH_VERSION} && \
+    rocm-sdk init
 
 # warp_src is a named build context — see file header for invocation details.
 COPY --from=warp_src . /root/warp/
@@ -89,12 +77,19 @@ RUN test -f /root/warp/pyproject.toml && \
     }
 RUN git config --global --add safe.directory /root/warp
 
-RUN /bin/bash -c "eval \"\$(~/.local/bin/micromamba shell hook --shell bash)\" && \
-    micromamba activate ${MAMBA_ENV_NAME} && \
-    cd /root/warp/ && \
-    python build_lib.py --jobs \$(nproc) 2>&1 | tee build_lib_{THE_ROCK_VERSION}.log && \
-    python -m build --wheel -C--build-option=-Plinux-x86_64 && \
-    pip install dist/warp_lang-*.whl"
+# Build warp_lang wheel
+WORKDIR /root/warp
+RUN cat <<EOF > build.sh
+set -exuo pipefail
+export ROCM_PATH=$(rocm-sdk path --root)
+export LD_LIBRARY_PATH=$(rocm-sdk path --root)/lib:"${LD_LIBRARY_PATH}"
+# necessary for rocm < 7.12 https://github.com/ROCm/TheRock/issues/3213
+export HIP_DEVICE_LIB_PATH=$(rocm-sdk path --root)/lib/llvm/amdgcn/bitcode
+python build_lib.py --jobs \$(nproc) 2>&1 | tee build_lib_{THE_ROCK_VERSION}.log
+python -m build --wheel -C--build-option=-Plinux-x86_64
+uv pip install dist/warp_lang-*.whl
+EOF
+RUN /bin/bash build.sh
 
 #-----------------------------------------------------------------------------
 # Stage 2 (warp_base): clean release image shipped to users.
@@ -103,22 +98,8 @@ RUN /bin/bash -c "eval \"\$(~/.local/bin/micromamba shell hook --shell bash)\" &
 # -----------------------------------------------------------------------------
 FROM warp_build AS warp_base
 
-# ARG must be redeclared in each stage that uses it (Docker multi-stage scoping).
-ARG MAMBA_ENV_NAME=base
-
 RUN rm -rf /root/warp/
 
-# Activate the micromamba base env for all shell types:
-#   - ENV PATH: covers non-interactive shells (docker exec, entrypoints, scripts)
-#   - shell init + bashrc append: covers interactive bash sessions
-ENV MAMBA_ROOT_PREFIX=/root/micromamba
-ENV CONDA_PREFIX=/root/micromamba
-ENV CONDA_DEFAULT_ENV=${MAMBA_ENV_NAME}
-ENV PATH=/root/micromamba/bin:${PATH}
-
-RUN /root/.local/bin/micromamba shell init --shell bash --root-prefix /root/micromamba && \
-echo "micromamba activate base" >> /root/.bashrc
-
 # -----------------------------------------------------------------------------
 # Stage 3 (whl_export): minimal scratch image used only to extract the wheel.
 #   Built by CI with --output type=local — never shipped to users.

From 2ea8a4b2a40cb96c13f8fb77d5bb0dedcb916a63 Mon Sep 17 00:00:00 2001
From: jamesETsmith <james.smith9113@gmail.com>
Date: Fri, 27 Mar 2026 12:16:41 -0400
Subject: [PATCH 8/8] Fixing some env problems in the docker image for rocm

---
 docker/rocm_ci/Dockerfile.TheRock_ci | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/docker/rocm_ci/Dockerfile.TheRock_ci b/docker/rocm_ci/Dockerfile.TheRock_ci
index e244a25b9d..18c7efcb22 100644
--- a/docker/rocm_ci/Dockerfile.TheRock_ci
+++ b/docker/rocm_ci/Dockerfile.TheRock_ci
@@ -77,15 +77,21 @@ RUN test -f /root/warp/pyproject.toml && \
     }
 RUN git config --global --add safe.directory /root/warp
 
-# Build warp_lang wheel
-WORKDIR /root/warp
-RUN cat <<EOF > build.sh
-set -exuo pipefail
+# Persist ROCm environment for interactive containers and build scripts
+RUN cat <<'EOF' > ~/.rocm-env.sh
 export ROCM_PATH=$(rocm-sdk path --root)
-export LD_LIBRARY_PATH=$(rocm-sdk path --root)/lib:"${LD_LIBRARY_PATH}"
+export LD_LIBRARY_PATH=$(rocm-sdk path --root)/lib
 # necessary for rocm < 7.12 https://github.com/ROCm/TheRock/issues/3213
 export HIP_DEVICE_LIB_PATH=$(rocm-sdk path --root)/lib/llvm/amdgcn/bitcode
-python build_lib.py --jobs \$(nproc) 2>&1 | tee build_lib_{THE_ROCK_VERSION}.log
+EOF
+RUN echo 'source ~/.rocm-env.sh' >> ~/.bashrc
+
+# Build warp_lang wheel
+WORKDIR /root/warp
+RUN cat <<'EOF' > build.sh
+set -euxo pipefail
+source ~/.rocm-env.sh
+python build_lib.py --jobs $(nproc) 2>&1 | tee build_lib_${THE_ROCK_VERSION}.log
 python -m build --wheel -C--build-option=-Plinux-x86_64
 uv pip install dist/warp_lang-*.whl
 EOF