Merge branch 'private/tile' into 'main'

Tile Support

See merge request omniverse/warp!828

.gitlab-ci.yml

@@ -34,6 +34,7 @@ variables:
   PIP_CACHE_DIR: '$CI_PROJECT_DIR/.cache/pip'
   CUDA_BIN: '$CI_PROJECT_DIR/_build/target-deps/cuda/bin'
   CUDA: '$CI_PROJECT_DIR/_build/target-deps/cuda'
+  CUDA_HOME: '$CI_PROJECT_DIR/_build/target-deps/cuda'
   PYTHON: '$CI_PROJECT_DIR/_build/target-deps/python/python'
   LINBUILD: '$CI_PROJECT_DIR/_build/host-deps/linbuild/linbuild.sh'
   WARP_CACHE_ROOT: '$CI_PROJECT_DIR/.cache/warp' # Used by the parallel test runner
@@ -61,15 +62,23 @@ stages:
 
 linux-aarch64 build:
   stage: build
-  image: ubuntu:22.04
+  image: ubuntu:20.04
   extends:
     - .save_warp_bin_artifact
   before_script:
     - echo -e "\\e[0Ksection_start:`date +%s`:install_dependencies[collapsed=true]\\r\\e[0KInstalling dependencies"
     - apt-get update && apt-get install build-essential curl --no-install-recommends -y
+    - >
+      curl -k -H "Authorization: Bearer $ARTIFACTORY_ACCESS_TOKEN"
+      $ARTIFACTORY_BASE_URL/sw-cuda-math-mathdx-generic-local/cicd/libmathdx/main/PostMerge/99/libmathdx_build_aarch64_rockylinux8_cuda12.0.0_release.tar.gz
+      -o libmathdx.tar.gz
+    - mkdir -p _build/target-deps
+    - tar -xzf libmathdx.tar.gz -C _build/target-deps
+    - export LIBMATHDX_HOME="$CI_PROJECT_DIR/_build/target-deps/libmathdx-0.0.1-Linux"
+    - gcc --version
     - echo -e "\\e[0Ksection_end:`date +%s`:install_dependencies\\r\\e[0K"
   script:
-    - ./tools/ci/building/build-linux-aarch64/build.sh --no-docker # We are already using the builder image
+    - ./tools/ci/building/build-linux-aarch64/build.sh --no-docker
     - mkdir -p warp/bin/linux-aarch64
     - mv warp/bin/warp.so warp/bin/linux-aarch64
     - mv warp/bin/warp-clang.so warp/bin/linux-aarch64
@@ -78,12 +87,24 @@ linux-aarch64 build:
 
 linux-x86_64 build:
   stage: build
-  image: urm.nvidia.com/ct-omniverse-docker/centos7-gcc10-builder:3.2.0
+  image: ubuntu:20.04
   extends:
     - .save_warp_bin_artifact
     - .runner-build-linux-x86_64
+  before_script:
+    - echo -e "\\e[0Ksection_start:`date +%s`:install_dependencies[collapsed=true]\\r\\e[0KInstalling dependencies"
+    - apt-get update && apt-get install build-essential curl --no-install-recommends -y
+    - >
+      curl -k -H "Authorization: Bearer $ARTIFACTORY_ACCESS_TOKEN"
+      $ARTIFACTORY_BASE_URL/sw-cuda-math-mathdx-generic-local/cicd/libmathdx/main/PostMerge/99/libmathdx_build_x86_64_rockylinux8_cuda12.0.0_release.tar.gz
+      -o libmathdx.tar.gz
+    - mkdir -p _build/target-deps
+    - tar -xzf libmathdx.tar.gz -C _build/target-deps
+    - export LIBMATHDX_HOME="$CI_PROJECT_DIR/_build/target-deps/libmathdx-0.0.1-Linux"
+    - gcc --version
+    - echo -e "\\e[0Ksection_end:`date +%s`:install_dependencies\\r\\e[0K"
   script:
-    - ./tools/ci/building/build-linux-x86_64/build.sh --no-docker # We are already using the builder image
+    - ./tools/ci/building/build-linux-x86_64/build.sh --no-docker
     - mkdir -p warp/bin/linux-x86_64
     - mv warp/bin/warp.so warp/bin/linux-x86_64
     - mv warp/bin/warp-clang.so warp/bin/linux-x86_64
@@ -216,6 +237,7 @@ linux-aarch64 test jetson:
     - echo -e "\\e[0Ksection_start:`date +%s`:install_dependencies[collapsed=true]\\r\\e[0KInstalling dependencies"
     - !reference [.snippets, install-python+warp-aarch64]
     - python -m pip install coverage[toml]
+    - python -m pip install -U "jax[cuda12]"
     - echo -e "\\e[0Ksection_end:`date +%s`:install_dependencies\\r\\e[0K"
   script:
     - python -m warp.tests --junit-report-xml rspec.xml --coverage --coverage-xml coverage.xml -s autodetect --failfast
@@ -240,7 +262,7 @@ linux-x86_64 test:
     - python -m pip install --upgrade pip
     - python -m pip install --upgrade usd-core coverage[toml]
     - python -m pip install --upgrade torch --extra-index-url https://download.pytorch.org/whl/cu121
-    - python -m pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+    - python -m pip install -U "jax[cuda12]"
     - python -m pip install -e .
     - echo -e "\\e[0Ksection_end:`date +%s`:install_dependencies\\r\\e[0K"
     # HACK: disable P2P tests due to misbehaving agents

.gitlab/ci/additional-tests.yml

@@ -43,7 +43,7 @@ linux-x86_64 test:
     - python -m pip install --upgrade pip
     - python -m pip install --upgrade usd-core
     - python -m pip install --upgrade torch --extra-index-url https://download.pytorch.org/whl/cu121
-    - python -m pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+    - python -m pip install -U "jax[cuda12]"
     - python -m pip install -e .
     - echo -e "\\e[0Ksection_end:`date +%s`:install_dependencies\\r\\e[0K"
     # HACK: disable P2P tests due to misbehaving agents

.gitlab/ci/common.yml

@@ -106,6 +106,7 @@ include:
     PIP_CACHE_DIR: '$PARENT_PROJECT_DIR/.cache/pip'
     CUDA_BIN: '$PARENT_PROJECT_DIR/_build/target-deps/cuda/bin'
     CUDA: '$PARENT_PROJECT_DIR/_build/target-deps/cuda'
+    CUDA_HOME: '$CI_PROJECT_DIR/_build/target-deps/cuda'
     PYTHON: '$PARENT_PROJECT_DIR/_build/target-deps/python/python'
     LINBUILD: '$PARENT_PROJECT_DIR/_build/host-deps/linbuild/linbuild.sh'
     WARP_CACHE_ROOT: '$PARENT_PROJECT_DIR/.cache/warp' # Used by the parallel test runner

.gitlab/ci/cuda-11-build-and-test.yml

@@ -45,7 +45,7 @@ stages:
 
 linux-aarch64 build:
   stage: build
-  image: ubuntu:22.04
+  image: ubuntu:20.04
   extends:
     - .save_warp_bin_artifact
   before_script:
@@ -123,7 +123,7 @@ linux-x86_64 test:
     - python -m pip install --upgrade pip
     - python -m pip install --upgrade usd-core
     - python -m pip install --upgrade torch --extra-index-url https://download.pytorch.org/whl/cu121
-    - python -m pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+    - python -m pip install -U "jax[cuda12]"
     - python -m pip install -e .
     - echo -e "\\e[0Ksection_end:`date +%s`:install_dependencies\\r\\e[0K"
     # HACK: disable P2P tests due to misbehaving agents

.gitlab/ci/debug-build-and-test.yml

@@ -35,7 +35,7 @@ stages:
 # Hide this job for now until debug aarch64 builds work
 .linux-aarch64 build:
   stage: build
-  image: ubuntu:22.04
+  image: ubuntu:20.04
   extends:
     - .save_warp_bin_artifact
   before_script:
@@ -114,7 +114,7 @@ linux-x86_64 test:
     - python -m pip install --upgrade pip
     - python -m pip install --upgrade usd-core
     - python -m pip install --upgrade torch --extra-index-url https://download.pytorch.org/whl/cu121
-    - python -m pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+    - python -m pip install -U "jax[cuda12]"
     - python -m pip install -e .
     - echo -e "\\e[0Ksection_end:`date +%s`:install_dependencies\\r\\e[0K"
     # HACK: disable P2P tests due to misbehaving agents

build_lib.py

@@ -17,6 +17,7 @@
 import argparse
 import glob
 import os
+import platform
 import shutil
 
 from warp.build_dll import build_dll, find_host_compiler, set_msvc_env, verbose_cmd
@@ -26,6 +27,7 @@
 parser.add_argument("--msvc_path", type=str, help="Path to MSVC compiler (optional if already on PATH)")
 parser.add_argument("--sdk_path", type=str, help="Path to WinSDK (optional if already on PATH)")
 parser.add_argument("--cuda_path", type=str, help="Path to CUDA SDK")
+parser.add_argument("--libmathdx_path", type=str, help="Path to libmathdx (optional if LIBMATHDX_HOME is defined)")
 parser.add_argument(
     "--mode",
     type=str,
@@ -52,6 +54,7 @@
 parser.set_defaults(fast_math=False)
 
 parser.add_argument("--quick", action="store_true", help="Only generate PTX code, disable CUTLASS ops")
+parser.set_defaults(quick=False)
 
 parser.add_argument("--build_llvm", action="store_true", help="Build Clang/LLVM compiler from source, default disabled")
 parser.add_argument("--no_build_llvm", dest="build_llvm", action="store_false")
@@ -69,6 +72,7 @@
 parser.add_argument("--no_standalone", dest="standalone", action="store_false")
 parser.set_defaults(standalone=True)
 
+
 args = parser.parse_args()
 
 # set build output path off this file
@@ -96,7 +100,7 @@ def find_cuda_sdk():
         return cuda_sdk
 
     # check default paths
-    if os.name == "nt":
+    if platform.system() == "Windows":
         cuda_paths = glob.glob("C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v*.*")
         if len(cuda_paths) >= 1:
             cuda_sdk = cuda_paths[0]
@@ -114,16 +118,21 @@ def find_cuda_sdk():
 
 
 # setup CUDA Toolkit path
-if sys.platform == "darwin":
+if platform.system() == "Darwin":
     args.cuda_path = None
-
 else:
     if not args.cuda_path:
         args.cuda_path = find_cuda_sdk()
 
+    if not args.libmathdx_path:
+        libmathdx_path = os.environ.get("LIBMATHDX_HOME")
+
+        if libmathdx_path:
+            print(f"Using libmathdx path '{libmathdx_path}' provided through the 'LIBMATHDX_HOME' environment variable")
+            args.libmathdx_path = libmathdx_path
 
 # setup MSVC and WinSDK paths
-if os.name == "nt":
+if platform.system() == "Windows":
     if args.msvc_path or args.sdk_path:
         # user provided MSVC and Windows SDK
         assert args.msvc_path and args.sdk_path, "--msvc_path and --sdk_path must be used together."
@@ -140,9 +149,9 @@ def find_cuda_sdk():
 
 # return platform specific shared library name
 def lib_name(name):
-    if sys.platform == "win32":
+    if platform.system() == "Windows":
         return f"{name}.dll"
-    elif sys.platform == "darwin":
+    elif platform.system() == "Darwin":
         return f"lib{name}.dylib"
     else:
         return f"{name}.so"
@@ -188,6 +197,7 @@ def generate_exports_header_file():
         "native/volume.cpp",
         "native/marching.cpp",
         "native/cutlass_gemm.cpp",
+        "native/mathdx.cpp",
     ]
     warp_cpp_paths = [os.path.join(build_path, cpp) for cpp in cpp_sources]
 
@@ -197,6 +207,9 @@ def generate_exports_header_file():
     else:
         warp_cu_path = os.path.join(build_path, "native/warp.cu")
 
+    if args.libmathdx_path is None:
+        print("Warning: libmathdx not found, building without MathDx support")
+
     warp_dll_path = os.path.join(build_path, f"bin/{lib_name('warp')}")
 
     build_dll(args, dll_path=warp_dll_path, cpp_paths=warp_cpp_paths, cu_path=warp_cu_path)

docs/index.rst

@@ -350,6 +350,7 @@ Full Table of Contents
     modules/devices
     modules/differentiability
     modules/generics
+    modules/tiles
     modules/interoperability
     configuration
     debugging

docs/modules/differentiability.rst

@@ -1,3 +1,5 @@
+.. _differentiability:
+
 Differentiability
 =================