Add torch backend

.github/workflows/clang-tidy-review.yml

@@ -35,7 +35,7 @@ jobs:
         with:
           build_dir: build
           apt_packages: cmake,libxml2,libxml2-dev,libtinfo-dev,zlib1g-dev,libzstd-dev,libthrust-dev,libbenchmark-dev,libomp-20-dev,libopenblas-dev
-          exclude: "test/*,unittests/*"
+          exclude: "test/*,unittests/*,benchmark/LibTorch/*,include/clad/Differentiator/TorchBuiltins.h,include/clad/Differentiator/TorchBuiltins/*"
           split_workflow: true
           config_file: .clang-tidy
           cmake_command: >

.github/workflows/libtorch.yml

@@ -0,0 +1,109 @@
+name: LibTorch
+
+on:
+  push:
+    branches:
+      - master
+    paths:
+      - '.github/workflows/libtorch.yml'
+      - '.clang-tidy'
+      - '**/CMakeLists.txt'
+      - 'cmake/**'
+      - 'benchmark/LibTorch/**'
+      - 'include/clad/Differentiator/**'
+      - 'lib/Differentiator/**'
+  pull_request:
+    branches:
+      - master
+    paths:
+      - '.github/workflows/libtorch.yml'
+      - '.clang-tidy'
+      - '**/CMakeLists.txt'
+      - 'cmake/**'
+      - 'benchmark/LibTorch/**'
+      - 'include/clad/Differentiator/**'
+      - 'lib/Differentiator/**'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  cpu:
+    name: cpu-clang21
+    runs-on: ubuntu-24.04
+    timeout-minutes: 30
+    env:
+      LLVM_ROOT: /usr/lib/llvm-21
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install LLVM and build dependencies
+        run: |
+          curl --fail --location --silent --show-error \
+            https://apt.llvm.org/llvm-snapshot.gpg.key \
+            | sudo gpg --dearmor --yes \
+              --output /usr/share/keyrings/apt.llvm.org.gpg
+          echo 'deb [signed-by=/usr/share/keyrings/apt.llvm.org.gpg] https://apt.llvm.org/noble/ llvm-toolchain-noble-21 main' \
+            | sudo tee /etc/apt/sources.list.d/llvm-21.list
+          sudo apt-get update
+          sudo apt-get install -y \
+            clang-21 clang-tidy-21 libclang-21-dev llvm-21-dev \
+            ninja-build unzip
+
+      - name: Cache LibTorch
+        id: cache-libtorch
+        uses: actions/cache@v4
+        with:
+          path: ${{ runner.temp }}/libtorch
+          key: libtorch-cpu-2.11.0
+
+      - name: Download LibTorch
+        if: steps.cache-libtorch.outputs.cache-hit != 'true'
+        run: |
+          curl --fail --location --retry 3 \
+            --output "$RUNNER_TEMP/libtorch.zip" \
+            'https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.11.0%2Bcpu.zip'
+          echo 'ee3f453f6c3e3d934339875d76a2b04d6a1731554c016fc82be1f4b8d1d4ae62  '"$RUNNER_TEMP/libtorch.zip" \
+            | sha256sum --check
+          unzip -q "$RUNNER_TEMP/libtorch.zip" -d "$RUNNER_TEMP"
+
+      - name: Configure
+        run: |
+          cmake -S . -B build/libtorch -G Ninja \
+            -DCLAD_DISABLE_TESTS=ON \
+            -DCLAD_ENABLE_BENCHMARKS=OFF \
+            -DCLAD_ENABLE_LIBTORCH_BENCHMARKS=ON \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DCMAKE_C_COMPILER="$LLVM_ROOT/bin/clang" \
+            -DCMAKE_CXX_COMPILER="$LLVM_ROOT/bin/clang++" \
+            -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
+            -DClang_DIR="$LLVM_ROOT/lib/cmake/clang" \
+            -DLLVM_DIR="$LLVM_ROOT/lib/cmake/llvm" \
+            -DTorch_DIR="$RUNNER_TEMP/libtorch/share/cmake/Torch"
+
+      - name: Build
+        run: >-
+          cmake --build build/libtorch
+          --target LibTorchOpsTest LibTorchOpsBenchmark
+          --parallel 2
+
+      - name: Test
+        run: >-
+          ctest --test-dir build/libtorch
+          --output-on-failure
+          -R '^clad-LibTorchOpsTest$'
+
+      - name: Run clang-tidy
+        run: |
+          # clang-tidy interprets compiler plugins as analyzer checker plugins.
+          sed -E -i 's# -fplugin=[^ ]+# -D__CLAD__=1#g' \
+            build/libtorch/compile_commands.json
+          "$LLVM_ROOT/bin/clang-tidy" \
+            -p=build/libtorch \
+            --checks='-misc-include-cleaner' \
+            --warnings-as-errors='*' \
+            --header-filter='^.*/include/clad/Differentiator/TorchBuiltins.*' \
+            benchmark/LibTorch/TorchOpsBenchmark.cpp \
+            benchmark/LibTorch/TorchOpsTest.cpp

CMakeLists.txt

@@ -300,6 +300,8 @@ add_definitions( -D_GNU_SOURCE
   -DCLAD_INSTDIR_INCL="${CLAD_BINARY_DIR}/include" )
 
 option(CLAD_BUILD_STATIC_ONLY "Does not build shared libraries. Useful when we have LLVM_ENABLE_PLUGINS=OFF (eg. Win or CYGWIN)" OFF)
+option(CLAD_ENABLE_LIBTORCH_BENCHMARKS
+  "Build benchmarks for the Clad-backed LibTorch integration." OFF)
 if (NOT CLAD_BUILD_STATIC_ONLY AND NOT LLVM_ENABLE_PLUGINS)
   message(FATAL_ERROR "LLVM_ENABLE_PLUGINS is set to OFF. Please build clad with -DCLAD_BUILD_STATIC_ONLY=ON.")
 endif()
@@ -327,11 +329,12 @@ endif()
 
 
 if (NOT CLAD_BUILD_STATIC_ONLY)
-  if (CLAD_ENABLE_BENCHMARKS)
+  if (CLAD_ENABLE_BENCHMARKS OR CLAD_ENABLE_LIBTORCH_BENCHMARKS)
     include(GoogleBenchmark)
-  endif(CLAD_ENABLE_BENCHMARKS)
+  endif()
 
-  if (NOT CLAD_DISABLE_TESTS OR CLAD_ENABLE_BENCHMARKS)
+  if (NOT CLAD_DISABLE_TESTS OR CLAD_ENABLE_BENCHMARKS OR
+      CLAD_ENABLE_LIBTORCH_BENCHMARKS)
     # Change the default compiler to the clang which we run clad upon. Our unittests
     # need to use a supported by clad compiler. Note that's a huge hack and it is
     # not guaranteed to work with cmake.
@@ -361,6 +364,10 @@ if (NOT CLAD_BUILD_STATIC_ONLY)
   if (CLAD_ENABLE_BENCHMARKS)
     add_subdirectory(benchmark)
   endif(CLAD_ENABLE_BENCHMARKS)
+  if (CLAD_ENABLE_LIBTORCH_BENCHMARKS)
+    enable_testing()
+    add_subdirectory(benchmark/LibTorch)
+  endif()
 
   if (stored_cxx_compiler)
     # Restore the default compilers.

benchmark/LibTorch/CMakeLists.txt

@@ -0,0 +1,11 @@
+find_package(Torch REQUIRED)
+
+CB_ADD_GBENCHMARK(LibTorchOpsBenchmark TorchOpsBenchmark.cpp)
+target_link_libraries(LibTorchOpsBenchmark PRIVATE torch)
+
+ADD_CLAD_EXECUTABLE(LibTorchOpsTest TorchOpsTest.cpp)
+target_link_libraries(LibTorchOpsTest PRIVATE torch)
+add_test(NAME clad-LibTorchOpsTest COMMAND LibTorchOpsTest)
+set_tests_properties(clad-LibTorchOpsTest PROPERTIES
+  LABELS "libtorch;benchmark"
+  TIMEOUT 120)

benchmark/LibTorch/TorchOpsBenchmark.cpp

@@ -0,0 +1,89 @@
+#include "clad/Differentiator/Differentiator.h"
+#include "clad/Differentiator/TorchBuiltins.h" // IWYU pragma: keep
+
+#include "benchmark/benchmark.h"
+
+#include <torch/torch.h> // IWYU pragma: keep
+
+#include <cstdint>
+
+namespace {
+
+float squared_error(const torch::Tensor& prediction,
+                    const torch::Tensor& target) {
+  auto residual = prediction - target;
+  auto loss = residual.dot(residual);
+  return loss.item<float>();
+}
+
+struct Inputs {
+  torch::Tensor prediction{};
+  torch::Tensor target{};
+};
+
+Inputs make_inputs(std::int64_t size, bool requires_grad) {
+  const auto options = torch::TensorOptions()
+                           .dtype(torch::kFloat32)
+                           .requires_grad(requires_grad);
+  return {torch::linspace(-1.0, 1.0, size, options),
+          torch::linspace(0.75, -0.25, size, options)};
+}
+
+void set_items_processed(benchmark::State& state, std::int64_t size) {
+  state.SetItemsProcessed(state.iterations() * size);
+}
+
+void BM_CladGradient(benchmark::State& state) {
+  const auto size = state.range(0);
+  const auto inputs = make_inputs(size, /*requires_grad=*/false);
+  auto prediction_gradient = torch::zeros_like(inputs.prediction);
+  auto target_gradient = torch::zeros_like(inputs.target);
+  auto gradient = clad::gradient(squared_error, "prediction, target");
+
+  for ([[maybe_unused]] auto _ : state) {
+    prediction_gradient.zero_();
+    target_gradient.zero_();
+    gradient.execute(inputs.prediction, inputs.target, &prediction_gradient,
+                     &target_gradient);
+    benchmark::DoNotOptimize(prediction_gradient.data_ptr<float>());
+    benchmark::DoNotOptimize(target_gradient.data_ptr<float>());
+  }
+  set_items_processed(state, size);
+}
+
+void BM_LibTorchAutograd(benchmark::State& state) {
+  const auto size = state.range(0);
+  const auto inputs = make_inputs(size, /*requires_grad=*/true);
+
+  for ([[maybe_unused]] auto _ : state) {
+    auto residual = inputs.prediction - inputs.target;
+    auto loss = residual.dot(residual);
+    auto gradients =
+        torch::autograd::grad({loss}, {inputs.prediction, inputs.target});
+    benchmark::DoNotOptimize(gradients[0].data_ptr<float>());
+    benchmark::DoNotOptimize(gradients[1].data_ptr<float>());
+  }
+  set_items_processed(state, size);
+}
+
+void configure(benchmark::internal::Benchmark* benchmark) {
+  benchmark->RangeMultiplier(4)
+      ->Range(256, 1 << 20)
+      ->Unit(benchmark::kMicrosecond);
+}
+
+BENCHMARK(BM_CladGradient)->Apply(configure);
+BENCHMARK(BM_LibTorchAutograd)->Apply(configure);
+
+} // namespace
+
+int main(int argc, char** argv) {
+  torch::set_num_threads(1);
+  torch::set_num_interop_threads(1);
+  benchmark::Initialize(&argc, argv);
+  if (benchmark::ReportUnrecognizedArguments(argc, argv))
+    return 1;
+  benchmark::RunSpecifiedBenchmarks();
+  benchmark::Shutdown();
+  return 0;
+}