Add torch-stable-abi-shim

examples/kernels/flake.nix

@@ -102,6 +102,12 @@
           path = ./relu-torch-stable-abi;
           drv = sys: out: out.packages.${sys}.redistributable.${"torch-stable-abi211-${cudaVersion}-${sys}"};
         }
+        {
+          # Exercise the shim with Torch 2.11 headers, which lack index_select.
+          name = "index-select-torch-stable-abi-kernel";
+          path = ./index-select-torch-stable-abi;
+          drv = sys: out: out.packages.${sys}.redistributable.${"torch-stable-abi210-cu128-${sys}"};
+        }
         # Check that we can also build the non-stable ABI version preceding the stable ABI version.
         {
           name = "relu-torch-stable-abi-kernel";

examples/kernels/index-select-torch-stable-abi/build.toml

@@ -0,0 +1,23 @@
+[general]
+name = "index-select-torch-stable-abi"
+version = 1
+edition = 6
+license = "Apache-2.0"
+backends = ["cuda"]
+
+[general.hub]
+repo-id = "kernels-test/index-select-torch-stable-abi"
+
+[torch]
+src = [
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
+]
+
+[torch.stable-abi]
+cuda = "2.10"
+
+[kernel.index_select]
+language = "cpp-cuda"
+depends = ["torch", "torch-stable-abi-shim"]
+src = ["index_select_cuda/index_select.cu"]

examples/kernels/index-select-torch-stable-abi/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for the Torch stable ABI index_select example";
+
+  inputs = {
+    kernel-builder.url = "path:../../..";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

examples/kernels/index-select-torch-stable-abi/index_select_cuda/index_select.cu

@@ -0,0 +1,9 @@
+#include <cstdint>
+#include <torch_stable_abi_shim/csrc/stable/ops.h>
+
+torch::stable::Tensor index_select_impl(
+    const torch::stable::Tensor& input,
+    int64_t dim,
+    const torch::stable::Tensor& index) {
+  return torch_stable_abi_shim::index_select(input, dim, index);
+}

examples/kernels/index-select-torch-stable-abi/tests/test_index_select.py

@@ -0,0 +1,25 @@
+import kernels
+import pytest
+import torch
+
+index_select_torch_stable_abi = kernels.get_kernel(
+    "kernels-test/index-select-torch-stable-abi", version=1
+)
+
+
+@pytest.mark.kernels_ci
+@pytest.mark.parametrize("dim", [0, 1, -1])
+@pytest.mark.parametrize("contiguous", [True, False])
+@pytest.mark.parametrize("indices", [[2, 0, 2], []])
+def test_index_select(device, dim, contiguous, indices):
+    if torch.device(device).type != "cuda":
+        pytest.skip("This example requires CUDA")
+
+    input = torch.arange(24, device=device, dtype=torch.float32).reshape(4, 6)
+    if not contiguous:
+        input = input.t()
+    index = torch.tensor(indices, device=device, dtype=torch.int64)
+
+    actual = index_select_torch_stable_abi.index_select(input, dim, index)
+    expected = torch.index_select(input, dim, index)
+    torch.testing.assert_close(actual, expected)

examples/kernels/index-select-torch-stable-abi/torch-ext/index_select_torch_stable_abi/__init__.py

@@ -0,0 +1,10 @@
+import torch
+
+from ._ops import ops
+
+
+def index_select(input: torch.Tensor, dim: int, index: torch.Tensor) -> torch.Tensor:
+    return ops.index_select(input, dim, index)
+
+
+__all__ = ["index_select"]

examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.cpp

@@ -0,0 +1,14 @@
+#include <torch/csrc/stable/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+STABLE_TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("index_select(Tensor input, int dim, Tensor index) -> Tensor");
+}
+
+STABLE_TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, ops) {
+  ops.impl("index_select", TORCH_BOX(&index_select_impl));
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.h

@@ -0,0 +1,9 @@
+#pragma once
+
+#include <cstdint>
+#include <torch/csrc/stable/tensor.h>
+
+torch::stable::Tensor index_select_impl(
+    const torch::stable::Tensor& input,
+    int64_t dim,
+    const torch::stable::Tensor& index);

kernels-common/src/config/deps.rs

@@ -34,6 +34,8 @@ pub enum Dependency {
     #[serde(rename = "metal-cpp")]
     MetalCpp,
     Torch,
+    #[serde(rename = "torch-stable-abi-shim")]
+    TorchStableAbiShim,
 }
 
 #[derive(Debug, Deserialize, Serialize)]
@@ -102,3 +104,16 @@ pub enum DependencyError {
     #[error("Unknown dependency: `{dependency:?}`")]
     GeneralDependency { dependency: String },
 }
+
+#[cfg(test)]
+mod tests {
+    use super::Dependency;
+
+    #[test]
+    fn torch_stable_abi_shim_round_trips() {
+        let json = r#""torch-stable-abi-shim""#;
+        let dependency: Dependency = serde_json::from_str(json).unwrap();
+        assert_eq!(dependency, Dependency::TorchStableAbiShim);
+        assert_eq!(serde_json::to_string(&dependency).unwrap(), json);
+    }
+}

nix-builder/lib/cpp-deps.nix

@@ -34,6 +34,9 @@ let
     "torch" = [
       torch
     ];
+    "torch-stable-abi-shim" = [
+      pkgs.torch-stable-abi-shim
+    ];
     "sycl_tla" = [
       (torch.xpuPackages.sycl-tla.override { inherit stdenv; })
     ];

nix-builder/overlay.nix

@@ -59,6 +59,8 @@ final: prev:
 
   torch-ops-check = final.callPackage ./pkgs/torch-ops-check { };
 
+  torch-stable-abi-shim = final.callPackage ./pkgs/torch-stable-abi-shim { };
+
   ruff = prev.ruff.overrideAttrs (
     finalAttrs: prevAttrs: {
       version = "0.15.10";

nix-builder/pkgs/torch-stable-abi-shim/default.nix

@@ -0,0 +1,27 @@
+{
+  lib,
+  stdenvNoCC,
+}:
+
+stdenvNoCC.mkDerivation {
+  name = "torch-stable-abi-shim";
+
+  src = ../../../torch-stable-abi-shim;
+
+  dontConfigure = true;
+  dontBuild = true;
+
+  installPhase = ''
+    runHook preInstall
+
+    mkdir -p "$out/include"
+    cp -r include/. "$out/include/"
+
+    runHook postInstall
+  '';
+
+  meta = {
+    description = "Header-only shim for the PyTorch stable ABI";
+    platforms = lib.platforms.all;
+  };
+}

torch-stable-abi-shim/LICENSE

@@ -0,0 +1,84 @@
+From PyTorch:
+
+Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
+Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
+Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
+Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
+Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
+Copyright (c) 2011-2013 NYU                      (Clement Farabet)
+Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
+Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
+Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
+
+From Caffe2:
+
+Copyright (c) 2016-present, Facebook Inc. All rights reserved.
+
+All contributions by Facebook:
+Copyright (c) 2016 Facebook Inc.
+
+All contributions by Google:
+Copyright (c) 2015 Google Inc.
+All rights reserved.
+
+All contributions by Yangqing Jia:
+Copyright (c) 2015 Yangqing Jia
+All rights reserved.
+
+All contributions by Kakao Brain:
+Copyright 2019-2020 Kakao Brain
+
+All contributions by Cruise LLC:
+Copyright (c) 2022 Cruise LLC.
+All rights reserved.
+
+All contributions by Tri Dao:
+Copyright (c) 2024 Tri Dao.
+All rights reserved.
+
+All contributions by Arm:
+Copyright (c) 2021, 2023-2025 Arm Limited and/or its affiliates
+
+All contributions from Caffe:
+Copyright(c) 2013, 2014, 2015, the respective contributors
+All rights reserved.
+
+All other contributions:
+Copyright(c) 2015, 2016 the respective contributors
+All rights reserved.
+
+Caffe2 uses a copyright model similar to Caffe: each contributor holds
+copyright over their contributions to Caffe2. The project versioning records
+all such contribution and copyright details. If a contributor wants to further
+mark their specific copyright on a particular contribution, they should
+indicate their copyright solely in the commit message of the change when it is
+committed.
+
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright
+   notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright
+   notice, this list of conditions and the following disclaimer in the
+   documentation and/or other materials provided with the distribution.
+
+3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America
+   and IDIAP Research Institute nor the names of its contributors may be
+   used to endorse or promote products derived from this software without
+   specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+POSSIBILITY OF SUCH DAMAGE.

torch-stable-abi-shim/README.md

@@ -0,0 +1,15 @@
+# Torch stable ABI shim
+
+A header-only shim that makes newer PyTorch stable C++ API wrappers available
+when building against older PyTorch stable headers. Wrappers introduced in later
+headers can use the stable dispatcher API already available in older versions.
+APIs requiring a newer runtime are guarded by `TORCH_FEATURE_VERSION` and are
+only exposed when the targeted stable ABI version supports them.
+
+The headers are vendored from PyTorch, with names adjusted to avoid collisions
+with PyTorch. Tensor and device types remain in `torch::stable`.
+
+## Copyright and license
+
+See [LICENSE](LICENSE) for the unmodified upstream PyTorch license, including
+all copyright notices, conditions, and the disclaimer.