PyTorch convenience features

.github/workflows/lint.yml

@@ -7,41 +7,38 @@
 
 jobs:
   linters:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-latest
 
     steps:
     - name: Check out Git repository
       uses: actions/checkout@v4
 
     - name: Install ClangFormat
       run: sudo apt-get install -y clang-format
 
-    - name: Run git-clang-format
-      run: git clang-format --style=file --diff
-
     - name: Set up Python
       uses: actions/setup-python@v4
       with:
-        python-version: 3.8
+        python-version: '3.12'
 
     - name: Install Python dependencies
-      run: python3.8 -m pip install black
+      run: pip install black
 
-    - name: Run black
-      run: python3.8 -m black --check --config pyproject.toml .
+    - name: Run lint
+      run: bash tools/lint.sh dry
 
   spelling:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-latest
 
     steps:
     - name: Check out Git repository
       uses: actions/checkout@v4
 
     - name: Download misspell
       run: |
         curl -L https://github.com/client9/misspell/releases/download/v0.3.4/misspell_0.3.4_linux_64bit.tar.gz -o /tmp/misspell_0.3.4_linux_64bit.tar.gz
         tar -xzf /tmp/misspell_0.3.4_linux_64bit.tar.gz -C .
 
     - name: Check spelling
       run: |
         ./misspell -error .github ark examples python scripts

.github/workflows/ut-cuda.yml → .github/workflows/ut.yml

@@ -1,4 +1,4 @@
-name: "Unit Tests (CUDA)"
+name: "Unit Tests"
 
 on:
   push:
@@ -11,38 +11,50 @@ on:
 
 jobs:
   UnitTest:
-    runs-on: [ self-hosted, A100 ]
     defaults:
       run:
         shell: bash
-    timeout-minutes: 30
+    timeout-minutes: 60
+    permissions:
+      actions: read
+      contents: read
+      security-events: write
     strategy:
+      fail-fast: false
       matrix:
-        cuda: [ cuda11.8, cuda12.2 ]
+        include:
+          - platform: cuda
+            runner: [self-hosted, CUDA]
+            container: nvcr.io/nvidia/pytorch:26.03-py3
+            container_options: --privileged --ipc=host --gpus=all --ulimit memlock=-1:-1
+          - platform: rocm
+            runner: [self-hosted, ROCM]
+            container: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.12_pytorch_release_2.9.1
+            container_options: --privileged --ipc=host --security-opt seccomp=unconfined --group-add video --ulimit memlock=-1:-1
+    runs-on: ${{ matrix.runner }}
     concurrency:
-      group: ${{ github.workflow }}-${{ github.ref }}-${{ matrix.cuda }}
+      group: ${{ github.workflow }}-${{ matrix.platform }}-${{ github.ref }}
       cancel-in-progress: true
     container:
-      image: "ghcr.io/microsoft/ark/ark:base-dev-${{ matrix.cuda }}"
-      options: --privileged --ipc=host --gpus=all --ulimit memlock=-1:-1
+      image: ${{ matrix.container }}
+      options: ${{ matrix.container_options }}
 
     steps:
       - name: Checkout
         uses: actions/checkout@v4
 
-      - name: LockGPUClock
-        run: |
-          sudo nvidia-smi -pm 1
-          for i in $(seq 0 $(( $(nvidia-smi -L | wc -l) - 1 ))); do
-            sudo nvidia-smi -ac $(nvidia-smi --query-gpu=clocks.max.memory,clocks.max.sm --format=csv,noheader,nounits -i $i | sed 's/\ //') -i $i
-          done
-
       - name: Dubious ownership exception
         run: |
           git config --global --add safe.directory /__w/ark/ark
 
+      - name: Initialize CodeQL
+        uses: github/codeql-action/init@v3
+        with:
+          languages: cpp
+
       - name: Build
         run: |
+          apt-get update && apt-get install -y lcov
           mkdir build && cd build
           cmake -DCMAKE_BUILD_TYPE=Debug ..
           make -j ut ark_py
@@ -51,6 +63,10 @@ jobs:
         run: |
           cd build
           ARK_ROOT=$PWD ctest --stop-on-failure --verbose --schedule-random
+
+      - name: C++ Coverage
+        run: |
+          cd build
           lcov --capture --directory . --output-file cpp_coverage.info
           lcov --remove cpp_coverage.info \
               '/usr/*' \
@@ -75,7 +91,7 @@ jobs:
               --cov=python/ark \
               --cov-report lcov:py_coverage.info \
               --verbose \
-              ../python/unittest/test.py
+              ../python/unittest/
 
       - name: Report Coverage
         env:
@@ -92,3 +108,8 @@ jobs:
       - name: Run Tutorials
         run: |
           python3 ./examples/tutorial/quickstart_tutorial.py
+
+      - name: Perform CodeQL Analysis
+        uses: github/codeql-action/analyze@v3
+        with:
+          category: "/language:cpp-${{ matrix.platform }}"

ark/api/context.cpp

@@ -29,8 +29,6 @@ void Context::set(const std::string& key, const std::string& value,
     this->impl_->set(key, value_json, type);
 }
 
-std::string Context::dump() const {
-    return this->impl_->dump().dump();
-}
+std::string Context::dump() const { return this->impl_->dump().dump(); }
 
 }  // namespace ark

ark/api/planner.cpp

@@ -211,8 +211,8 @@ std::string Planner::Impl::plan(bool pretty) const {
             Dims tile(trim_leading_ones);
 
             std::stringstream ss;
-            ss << "Result shape is not divided by tile "
-               << tile << ". Op: " << op->serialize().dump();
+            ss << "Result shape is not divided by tile " << tile
+               << ". Op: " << op->serialize().dump();
             auto not_divided_error = ss.str();
 
             auto &result_shape = result_tensors[0]->padded_shape();
@@ -224,11 +224,10 @@ std::string Planner::Impl::plan(bool pretty) const {
             max_num_tasks = 1;
             for (int i = 0; i < tile4.ndims(); i++) {
                 if (tile4[i] == 0) {
-                    ERR(PlanError, "Tile dimension is zero. Op: ",
-                        op->serialize().dump());
+                    ERR(PlanError,
+                        "Tile dimension is zero. Op: ", op->serialize().dump());
                 }
-                max_num_tasks *=
-                    (result_shape4[i] + tile4[i] - 1) / tile4[i];
+                max_num_tasks *= (result_shape4[i] + tile4[i] - 1) / tile4[i];
             }
             if (max_num_tasks == 0) ERR(InternalError, "max_num_tasks == 0");
         }
@@ -328,10 +327,13 @@ std::string Planner::Impl::plan(bool pretty) const {
                 max_processor_id = std::max(max_processor_id, num_processors);
             } else if (processor_group_root == -1) {
                 processor_group_root = ctx_processor_range_list.front()[0];
-                processor_group["ProcessorRange"] = ctx_processor_range_list.front()[1];
-                resource_group["ProcessorRange"] = ctx_processor_range_list.back()[1];
+                processor_group["ProcessorRange"] =
+                    ctx_processor_range_list.front()[1];
+                resource_group["ProcessorRange"] =
+                    ctx_processor_range_list.back()[1];
                 max_processor_id = std::max(
-                    max_processor_id, ctx_processor_range_list.front()[1][1].get<size_t>());
+                    max_processor_id,
+                    ctx_processor_range_list.front()[1][1].get<size_t>());
             } else {
                 new_processor_group = false;
                 resource_group["ProcessorRange"] =