CUDA

.github/workflows/cmake-build-cuda.yml

@@ -0,0 +1,55 @@
+name: CMake on CUDA Docker
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    container:
+      image: nvidia/cuda:12.6.3-devel-ubuntu22.04
+    steps:
+    - name: Install Git 2.18 or higher
+      run: |
+        apt-get update
+        apt-get install -y git cmake python3-dev ninja-build python3-pip
+        git --version
+        python3 --version
+        ln -s /usr/bin/python3 /usr/bin/python
+
+    - name: Install cmake
+      run: |
+        python -m pip install cmake
+        cmake --version
+
+    - uses: actions/checkout@v4
+      with:
+        submodules: true
+
+    - name: Set reusable strings
+      # Turn repeated input strings (such as the build output directory) into step outputs. These step outputs can be used throughout the workflow file.
+      id: strings
+      shell: bash
+      run: |
+        echo "build-output-dir=$(pwd)/build" >> "$GITHUB_OUTPUT"
+    - name: Check CUDA installation
+      run: |
+        nvcc -V
+
+    - name: Configure CMake for python part
+      # Configure CMake in a 'build' subdirectory. `CMAKE_BUILD_TYPE` is only required if you are using a single-configuration generator such as make.
+      # See https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html?highlight=cmake_build_type
+      run: >
+        cmake -B ${{ steps.strings.outputs.build-output-dir }}
+        -S $(pwd)
+        -DBUILD_TESTS=ON  # Enable testing
+
+    - name: Build
+      run: cmake --build ${{ steps.strings.outputs.build-output-dir }}
+
+    - name: Test
+      working-directory: ${{ steps.strings.outputs.build-output-dir }}
+      run: ctest

.github/workflows/cmake-build.yml

@@ -65,6 +65,7 @@ jobs:
         -DCMAKE_C_COMPILER=${{ matrix.c_compiler }}
         -DCMAKE_BUILD_TYPE=${{ matrix.build_type }}
         -S ${{ github.workspace }}
+        -DBUILD_TESTS=ON  # Enable testing
 
     - name: Build
       # Build your program with the given configuration. Note that --config is needed because the default Windows generator is a multi-config generator (Visual Studio generator).
@@ -74,4 +75,4 @@ jobs:
       working-directory: ${{ steps.strings.outputs.build-output-dir }}
       # Execute tests defined by the CMake configuration. Note that --build-config is needed because the default Windows generator is a multi-config generator (Visual Studio generator).
       # See https://cmake.org/cmake/help/latest/manual/ctest.1.html for more detail
-      run: ctest --build-config ${{ matrix.build_type }}
+      run: ctest --build-config ${{ matrix.build_type }}

.github/workflows/gradle.yml

@@ -15,8 +15,38 @@ on:
 
 jobs:
   build:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      # Set fail-fast to false to ensure that feedback is delivered for all matrix combinations. Consider changing this to true when your workflow is stable.
+      fail-fast: false
 
-    runs-on: ubuntu-latest
+      # Set up a matrix to run the following 3 configurations:
+      # 1. <Windows, Release, latest MSVC compiler toolchain on the default runner image, default generator>
+      # 2. <Linux, Release, latest GCC compiler toolchain on the default runner image, default generator>
+      # 3. <Linux, Release, latest Clang compiler toolchain on the default runner image, default generator>
+      #
+      # To add more build types (Release, Debug, RelWithDebInfo, etc.) customize the build_type list.
+      matrix:
+        os: [ubuntu-latest, windows-latest]
+        build_type: [Release]
+        c_compiler: [gcc, clang, cl]
+        include:
+          - os: windows-latest
+            c_compiler: cl
+            cpp_compiler: cl
+          - os: ubuntu-latest
+            c_compiler: gcc
+            cpp_compiler: g++
+          - os: ubuntu-latest
+            c_compiler: clang
+            cpp_compiler: clang++
+        exclude:
+          - os: windows-latest
+            c_compiler: gcc
+          - os: windows-latest
+            c_compiler: clang
+          - os: ubuntu-latest
+            c_compiler: cl
     permissions:
       contents: read
 

.github/workflows/publish-package.yml

@@ -64,6 +64,20 @@ jobs:
 
     - uses: astral-sh/setup-uv@v4
 
+    - name: Install CUDA Toolkit (Windows/Linux only)
+      if: matrix.os == 'windows-latest' || matrix.os == 'ubuntu-latest'
+      uses: Jimver/[email protected]
+      id: cuda-toolkit
+      with:
+        cuda: '12.5.0'
+
+    - name: Check CUDA installation
+      if: matrix.os == 'windows-latest' || matrix.os == 'ubuntu-latest'
+      run: |
+        echo "Installed cuda version is: ${{steps.cuda-toolkit.outputs.cuda}}"
+        echo "Cuda install location: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}"
+        nvcc -V
+
     - name: Build wheels
       uses: pypa/[email protected]
       env:

.github/workflows/python-ci.yml

@@ -38,7 +38,7 @@ jobs:
         flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Test with pytest
       run: |
-        pytest
+        pytest tests/python/unit
 
   build-package:
     runs-on: ubuntu-latest

CMakeLists.txt

@@ -1,5 +1,6 @@
 cmake_minimum_required(VERSION 3.15...3.29)
 project(craftground LANGUAGES CXX)
+include(GNUInstallDirs)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 set(CMAKE_CXX_STANDARD 17)
@@ -29,7 +30,7 @@ endif()
 # Collect source files for the module
 set(CRAFTGROUND_PY_SOURCES src/cpp/ipc.cpp)
 
-if (APPLE)
+if(APPLE)
     # Add Apple-specific source files
     list(APPEND CRAFTGROUND_PY_SOURCES src/cpp/ipc_apple.mm src/cpp/ipc_apple_torch.cpp)
 
@@ -54,25 +55,17 @@ if(APPLE)
         target_link_libraries(craftground_native PRIVATE "${TORCH_LIBRARIES}")
         set(LIBTORCH_DIR ${CMAKE_SOURCE_DIR}/src/cpp/libtorch)
         target_link_libraries(craftground_native PRIVATE
-           ${LIBTORCH_DIR}/libtorch_cpu_minimal.a
+            ${LIBTORCH_DIR}/libtorch_cpu_minimal.a
             ${LIBTORCH_DIR}/libc10.a
         )
         target_include_directories(craftground_native PRIVATE "${TORCH_INCLUDE_DIRS}")
         target_compile_definitions(craftground_native PRIVATE HAS_TORCH)
         target_compile_options(craftground_native PRIVATE -ffunction-sections -fdata-sections)
         target_link_options(craftground_native PRIVATE -Wl,-dead_strip)
-        # add_custom_command(
-        # TARGET craftground_native POST_BUILD
-        # COMMAND install_name_tool -change @rpath/libc10.dylib @loader_path/libc10.dylib $<TARGET_FILE:craftground_native>
-        # COMMAND install_name_tool -change @rpath/libtorch.dylib @loader_path/libtorch.dylib $<TARGET_FILE:craftground_native>
-        # COMMAND install_name_tool -change @rpath/libtorch_cpu.dylib @loader_path/libtorch_cpu.dylib $<TARGET_FILE:craftground_native>
-        # COMMENT "Updating RPATH to use relative paths"
-        # )
-
     endif()
 elseif(CUDAToolkit_FOUND)
     target_include_directories(craftground_native PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
-    target_link_libraries(craftground_native PRIVATE ${CUDAToolkit_LIBRARIES})
+    target_link_libraries(craftground_native PRIVATE CUDA::cudart CUDA::cudart_static)
 endif()
 
 
@@ -85,13 +78,44 @@ target_compile_definitions(
     PRIVATE VERSION_INFO=${PRIVATE_VERSION_INFO}
 )
 
-if(APPLE AND Torch_FOUND)
-    # install(FILES
-    #     /usr/local/libtorch/lib/libc10.dylib
-    #     /usr/local/libtorch/lib/libtorch.dylib
-    #     /usr/local/libtorch/lib/libtorch_cpu.dylib
-    #     DESTINATION craftground
-    # )
-endif()
+install(TARGETS craftground_native LIBRARY DESTINATION craftground)
+
+option(BUILD_TESTS "Build tests" ON)
+if(BUILD_TESTS)
+    enable_testing()
+    add_library(craftground STATIC ${CRAFTGROUND_PY_SOURCES})
+    find_package(Python3 REQUIRED COMPONENTS Development)
+    target_include_directories(
+        craftground PRIVATE
+        ${CMAKE_CURRENT_SOURCE_DIR}/pybind11/include
+        ${Python3_INCLUDE_DIRS}
+    )
+    if(APPLE)
+        if(Torch_FOUND)
+            target_link_libraries(craftground PUBLIC "${TORCH_LIBRARIES}")
+            set(LIBTORCH_DIR ${CMAKE_SOURCE_DIR}/src/cpp/libtorch)
+            target_link_libraries(craftground PUBLIC
+                ${LIBTORCH_DIR}/libtorch_cpu_minimal.a
+                ${LIBTORCH_DIR}/libc10.a
+            )
+            target_include_directories(craftground PUBLIC "${TORCH_INCLUDE_DIRS}")
+            target_compile_definitions(craftground PUBLIC HAS_TORCH)
+            target_compile_options(craftground PUBLIC -ffunction-sections -fdata-sections)
+            # target_link_options(craftground PUBLIC -Wl,-dead_strip)
+        endif()
+    elseif(CUDAToolkit_FOUND)
+        message(STATUS "CUDA is available in tests")
+        target_include_directories(craftground PUBLIC ${CUDAToolkit_INCLUDE_DIRS})
+        target_link_libraries(craftground PUBLIC CUDA::cudart CUDA::cudart_static)
+    endif()
+    target_link_options(craftground PRIVATE "-Wl,--whole-archive" "-Wl,--no-whole-archive")
+    target_include_directories(craftground PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/src/cpp)
+
+    target_compile_options(craftground PUBLIC ${CRAFTGROUND_PY_COMPILE_OPTIONS})
 
-install(TARGETS craftground_native LIBRARY DESTINATION craftground)
+    target_compile_definitions(
+        craftground
+        PUBLIC VERSION_INFO=${PRIVATE_VERSION_INFO}
+    )
+    add_subdirectory(tests)
+endif()

README.md

@@ -260,7 +260,15 @@ https://dejavu-fonts.github.io/License.html
 ```zsh
 brew install ktlint clang-format google-java-format
 ```
+```bash
+wget https://apt.llvm.org/llvm.sh
+sudo ./llvm.sh 19
+sudo apt install clang-format-19
+sudo ln -s /usr/bin/clang-format-19 /usr/bin/clang-format
+```
+
 ## Run formatters
+
 ```bash
 find . \( -iname '*.h' -o -iname '*.cpp' -o -iname '*.mm' \) | xargs clang-format -i
 ktlint '!src/craftground/MinecraftEnv/src/main/java/com/kyhsgeekcode/minecraftenv/proto/**'
@@ -285,4 +293,21 @@ protoc proto/observation_space.proto --java_out=craftground/MinecraftEnv/src/mai
 ### Solution
 ```bash
 pip install --upgrade protobuf
+```
+
+
+# Dev setup & build (conda, linux)
+```
+conda create --name craftground python=3.11
+conda activate craftground
+conda install gymnasium Pillow numpy protobuf typing_extensions psutil pytorch ninja build cmake
+conda install -c conda-forge openjdk=21 libgl-devel
+conda install glew
+python -m build
+```
+
+## Build jvm c++ part
+```bash
+ cmake src/main/cpp -DCMAKE_PREFIX_PATH=$CONDA_PREFIX
+ cmake --build .
 ```

src/cpp/ipc.cpp

@@ -9,33 +9,40 @@ py::object
 initialize_from_mach_port(unsigned int machPort, int width, int height) {
     return mtl_tensor_from_mach_port(machPort, width, height);
 }
-py::object mtl_tensor_from_cuda_mem_handle(
+py::capsule mtl_tensor_from_cuda_mem_handle(
     const char *cuda_ipc_handle, int width, int height
 ) {
     return py::none();
 }
 
-#elif __CUDA__
+#elif HAS_CUDA
 #include "ipc_cuda.h"
-py::object initialize_from_mach_port(int machPort, int width, int height) {
+py::object
+initialize_from_mach_port(unsigned int machPort, int width, int height) {
     return py::none();
 }
 
-py::object mtl_tensor_from_cuda_mem_handle(
+py::capsule mtl_tensor_from_cuda_mem_handle(
     const char *cuda_ipc_handle, int width, int height
 ) {
     DLManagedTensor *tensor = mtl_tensor_from_cuda_ipc_handle(
-        const_cast<void *>(cuda_ipc_handle), width, height
+        reinterpret_cast<void *>(const_cast<char *>(cuda_ipc_handle)),
+        width,
+        height
     );
-    return py::reinterpret_steal<py::object>(PyCapsule_New(
-        tensor,
-        "dltensor",
-        [](PyObject *capsule) {
-            DLManagedTensor *tensor =
-                (DLManagedTensor *)PyCapsule_GetPointer(capsule, "dltensor");
-            tensor->deleter(tensor);
+
+    if (!tensor) {
+        throw std::runtime_error(
+            "Failed to create DLManagedTensor from CUDA IPC handle"
+        );
+    }
+
+    return py::capsule(tensor, "dltensor", [](void *ptr) {
+        DLManagedTensor *managed_tensor = static_cast<DLManagedTensor *>(ptr);
+        if (managed_tensor && managed_tensor->deleter) {
+            managed_tensor->deleter(managed_tensor);
         }
-    ));
+    });
 }
 
 #else
@@ -44,7 +51,7 @@ initialize_from_mach_port(unsigned int machPort, int width, int height) {
     return py::none();
 }
 
-py::object mtl_tensor_from_cuda_mem_handle(
+py::capsule mtl_tensor_from_cuda_mem_handle(
     const char *cuda_ipc_handle, int width, int height
 ) {
     return py::none();

src/cpp/ipc.h

@@ -4,7 +4,7 @@
 namespace py = pybind11;
 py::object
 initialize_from_mach_port(unsigned int machPort, int width, int height);
-py::object mtl_tensor_from_cuda_mem_handle(
+py::capsule mtl_tensor_from_cuda_mem_handle(
     const char *cuda_ipc_handle, int width, int height
 );
 

src/cpp/ipc_apple.mm

@@ -66,7 +66,7 @@ static void deleteDLManagedTensor(DLManagedTensor *self) {
     DLManagedTensor *tensor =
         (DLManagedTensor *)malloc(sizeof(DLManagedTensor));
 
-    tensor->dl_tensor.data = mtlBuffer;
+    tensor->dl_tensor.data = (void *)mtlBuffer;
     tensor->dl_tensor.ndim = 3; // H x W x C
     tensor->dl_tensor.shape = (int64_t *)malloc(3 * sizeof(int64_t));
     tensor->dl_tensor.shape[0] = height;