From 0f3f92612f15b3495d95a2a70077627f303530ed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 20 Oct 2026 12:11:00 +0000 Subject: [PATCH 1/2] Add torch-stable-abi-shim This is a small shim that vendors the header-only API from the latest Torch. We always build stable ABI kernels with the latest Torch version that supports a given backend version. However, in some cases this leads to the use of older Torch versions (e.g. Torch 2.11 for CUDA 12.8), which do not have the latest header-only API. If a kernel needs a newer API, depend on the shim to get these newer headers, even on older Torch versions. We might also use the shim in the future for our own functions that are compatible with the stable ABI, but are more efficient than the upstream counterparts (e.g. because they go through the dispatcher). --- examples/kernels/flake.nix | 6 + .../index-select-torch-stable-abi/build.toml | 23 + .../index-select-torch-stable-abi/flake.nix | 17 + .../index_select_cuda/index_select.cu | 9 + .../tests/test_index_select.py | 25 + .../index_select_torch_stable_abi/__init__.py | 10 + .../torch-ext/torch_binding.cpp | 14 + .../torch-ext/torch_binding.h | 9 + kernels-common/src/config/deps.rs | 15 + nix-builder/lib/cpp-deps.nix | 3 + nix-builder/overlay.nix | 2 + .../pkgs/torch-stable-abi-shim/default.nix | 27 + torch-stable-abi-shim/LICENSE | 84 + torch-stable-abi-shim/README.md | 15 + .../csrc/stable/macros.h | 172 ++ .../torch_stable_abi_shim/csrc/stable/ops.h | 1628 +++++++++++++++++ 16 files changed, 2059 insertions(+) create mode 100644 examples/kernels/index-select-torch-stable-abi/build.toml create mode 100644 examples/kernels/index-select-torch-stable-abi/flake.nix create mode 100644 examples/kernels/index-select-torch-stable-abi/index_select_cuda/index_select.cu create mode 100644 examples/kernels/index-select-torch-stable-abi/tests/test_index_select.py create mode 100644 examples/kernels/index-select-torch-stable-abi/torch-ext/index_select_torch_stable_abi/__init__.py create mode 100644 examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.cpp create mode 100644 examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.h create mode 100644 nix-builder/pkgs/torch-stable-abi-shim/default.nix create mode 100644 torch-stable-abi-shim/LICENSE create mode 100644 torch-stable-abi-shim/README.md create mode 100644 torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/macros.h create mode 100644 torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index 28498e74d..8a1364687 100644 --- a/examples/kernels/flake.nix +++ b/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"; diff --git a/examples/kernels/index-select-torch-stable-abi/build.toml b/examples/kernels/index-select-torch-stable-abi/build.toml new file mode 100644 index 000000000..2c29baeb7 --- /dev/null +++ b/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"] diff --git a/examples/kernels/index-select-torch-stable-abi/flake.nix b/examples/kernels/index-select-torch-stable-abi/flake.nix new file mode 100644 index 000000000..db0fcdfaf --- /dev/null +++ b/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 = ./.; + }; +} diff --git a/examples/kernels/index-select-torch-stable-abi/index_select_cuda/index_select.cu b/examples/kernels/index-select-torch-stable-abi/index_select_cuda/index_select.cu new file mode 100644 index 000000000..23328ea97 --- /dev/null +++ b/examples/kernels/index-select-torch-stable-abi/index_select_cuda/index_select.cu @@ -0,0 +1,9 @@ +#include +#include + +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); +} diff --git a/examples/kernels/index-select-torch-stable-abi/tests/test_index_select.py b/examples/kernels/index-select-torch-stable-abi/tests/test_index_select.py new file mode 100644 index 000000000..1bdf2e83e --- /dev/null +++ b/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) diff --git a/examples/kernels/index-select-torch-stable-abi/torch-ext/index_select_torch_stable_abi/__init__.py b/examples/kernels/index-select-torch-stable-abi/torch-ext/index_select_torch_stable_abi/__init__.py new file mode 100644 index 000000000..384cb2fbc --- /dev/null +++ b/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"] diff --git a/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.cpp b/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.cpp new file mode 100644 index 000000000..864f78a43 --- /dev/null +++ b/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.cpp @@ -0,0 +1,14 @@ +#include + +#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) diff --git a/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.h b/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.h new file mode 100644 index 000000000..d1bb39412 --- /dev/null +++ b/examples/kernels/index-select-torch-stable-abi/torch-ext/torch_binding.h @@ -0,0 +1,9 @@ +#pragma once + +#include +#include + +torch::stable::Tensor index_select_impl( + const torch::stable::Tensor& input, + int64_t dim, + const torch::stable::Tensor& index); diff --git a/kernels-common/src/config/deps.rs b/kernels-common/src/config/deps.rs index aea871df8..dfd3fa6b2 100644 --- a/kernels-common/src/config/deps.rs +++ b/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); + } +} diff --git a/nix-builder/lib/cpp-deps.nix b/nix-builder/lib/cpp-deps.nix index f5c38ec0e..416fbc23e 100644 --- a/nix-builder/lib/cpp-deps.nix +++ b/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; }) ]; diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix index b41df3f24..905b5d8f3 100644 --- a/nix-builder/overlay.nix +++ b/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"; diff --git a/nix-builder/pkgs/torch-stable-abi-shim/default.nix b/nix-builder/pkgs/torch-stable-abi-shim/default.nix new file mode 100644 index 000000000..c85fe00c5 --- /dev/null +++ b/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; + }; +} diff --git a/torch-stable-abi-shim/LICENSE b/torch-stable-abi-shim/LICENSE new file mode 100644 index 000000000..c23172f7a --- /dev/null +++ b/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. diff --git a/torch-stable-abi-shim/README.md b/torch-stable-abi-shim/README.md new file mode 100644 index 000000000..d088238d0 --- /dev/null +++ b/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. diff --git a/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/macros.h b/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/macros.h new file mode 100644 index 000000000..99acd15f5 --- /dev/null +++ b/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/macros.h @@ -0,0 +1,172 @@ +#pragma once + +// Vendored from PyTorch commit d3c397fef74cebde85dde5ca5e6638a2cb2da9c5. +// Helpers and macros are namespaced to coexist with PyTorch headers. +#include +#include + +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#endif + +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + +// Users of this macro are expected to include cuda_runtime.h +#define TORCH_STABLE_ABI_SHIM_CUDA_CHECK(EXPR) \ + do { \ + const cudaError_t __err = EXPR; \ + char* __error_msg = nullptr; \ + torch_c10_cuda_check_msg( \ + static_cast(__err), \ + __FILE__, \ + __func__, \ + static_cast(__LINE__), \ + true, \ + &__error_msg); \ + if (__error_msg != nullptr) { \ + std::string __msg(__error_msg); \ + torch_c10_cuda_free_error_msg(__error_msg); \ + /* @allow-raw-throw: the message already names the failing call */ \ + throw std::runtime_error(__msg); \ + } \ + } while (0) + +// Users of this macro are expected to include cuda_runtime.h +#define TORCH_STABLE_ABI_SHIM_CUDA_KERNEL_LAUNCH_CHECK() TORCH_STABLE_ABI_SHIM_CUDA_CHECK(cudaGetLastError()) + +#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + +namespace torch_stable_abi_shim::detail { +// Look up a symbol already loaded in the current process. Used to reach +// stable-ABI shims that were added after the extension's TORCH_TARGET_VERSION +// when the running libtorch is new enough to provide them. Returns nullptr if +// the symbol is absent or the platform has no in-process dynamic lookup. This +// is the single place that holds the per-platform lookup; the macros below are +// platform agnostic. +[[maybe_unused]] static void* lookup_stable_symbol(const char* name) { +#if defined(C10_MOBILE) + (void)name; + return nullptr; +#elif defined(_WIN32) + // Windows has no RTLD_DEFAULT equivalent, so resolve against the modules that + // export the stable shims. Today these live in torch_cpu.dll and + // torch_cuda.dll, so try each in turn. GetModuleHandleW does not take a + // reference on the returned handle. + const wchar_t* dlls[] = {L"torch_cpu.dll", L"torch_cuda.dll"}; + for (const wchar_t* dll : dlls) { + HMODULE handle = GetModuleHandleW(dll); + if (handle != nullptr) { + void* sym = reinterpret_cast(GetProcAddress(handle, name)); + if (sym != nullptr) { + return sym; + } + } + } + return nullptr; +#else + return dlsym(RTLD_DEFAULT, name); +#endif +} + +// Fallback for the const char* exception getters, returns a nullptr. A fallback +// MUST share the exact signature (return type and arguments) of the shim it +// backs: the looked-up pointer is called through decltype(&FALLBACK_FUNCTION). +[[maybe_unused]] static const char* torch_shim_bc_const_char_ptr() { + return nullptr; +} +} // namespace torch_stable_abi_shim::detail + +// IMPORTANT: this macro is intended for RARE, EXCEPTIONAL use only. It performs +// an at-runtime symbol lookup (dlsym / GetProcAddress), which is relatively +// expensive and only works on platforms that support dynamic symbol lookup. +// Reach for it only when an older-target extension genuinely needs a newer shim +// that would otherwise be unavailable (e.g. richer error messages); prefer +// normal version-guarded shim calls everywhere else. +// +// Calls SHIM_FUNCTION via the runtime lookup when the running libtorch is at +// least VERSION_SHIM_PRESENT, otherwise calls FALLBACK_FUNCTION. +// +// We structure this macro as an immediately-invoked lambda so it can 'return' a +// value from the macro call. The called pointer type is derived from +// FALLBACK_FUNCTION, which must therefore share SHIM_FUNCTION's exact signature +// (including arguments). +#define TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL( \ + VERSION_SHIM_PRESENT, SHIM_FUNCTION, FALLBACK_FUNCTION, ...) \ + ([&]() { \ + if (aoti_torch_abi_version() >= (VERSION_SHIM_PRESENT)) { \ + void* fn_ptr = \ + torch_stable_abi_shim::detail::lookup_stable_symbol(#SHIM_FUNCTION); \ + if (fn_ptr != nullptr) { \ + using FunctionType = decltype(&FALLBACK_FUNCTION); \ + return reinterpret_cast(fn_ptr)(__VA_ARGS__); \ + } \ + } \ + return FALLBACK_FUNCTION(__VA_ARGS__); \ + })() + +// Entry point for dynamic version calls for shims added in 2.13.0. Putting the +// version expectation in the macro lets us bypass the symbol lookup entirely +// when the target version already includes the shim and call it directly; +// otherwise we fall through to the runtime lookup. +#if TORCH_FEATURE_VERSION >= (((0ULL + 2) << 56) | ((0ULL + 13) << 48)) +// Target version already includes the shim, call it directly. +#define TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL_2_13_0( \ + SHIM_FUNCTION, FALLBACK_FUNCTION, ...) \ + ([&]() { return SHIM_FUNCTION(__VA_ARGS__); })() +#else +// Target version predates the shim, try a dynamic lookup. +#define TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL_2_13_0( \ + SHIM_FUNCTION, FALLBACK_FUNCTION, ...) \ + TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL( \ + (((0ULL + 2) << 56) | ((0ULL + 13) << 48)), SHIM_FUNCTION, FALLBACK_FUNCTION, __VA_ARGS__) +#endif // TORCH_FEATURE_VERSION >= (((0ULL + 2) << 56) | ((0ULL + 13) << 48)) + +namespace torch_stable_abi_shim::detail { +[[maybe_unused]] C10_NOINLINE static void throw_exception( + const char* call, + const char* file, + int64_t line) { + std::stringstream ss; + const auto& error_msg_without = TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL_2_13_0( + torch_exception_get_what_without_backtrace, torch_shim_bc_const_char_ptr); + if (error_msg_without) { + ss << error_msg_without; + ss << " (originally from " << call << " API call failed at " << file + << ", line " << line << ")"; + } else { + ss << call << " API call failed at " << file << ", line " << line; + } + + const auto& error_msg = TORCH_STABLE_ABI_SHIM_DYNAMIC_VERSION_CALL_2_13_0( + torch_exception_get_what, torch_shim_bc_const_char_ptr); + if (error_msg) { + std::time_t t = std::time(nullptr); + std::tm tm = *std::localtime(&t); + std::cerr << "[" << std::put_time(&tm, "%H:%M:%S") << " " << file << ":" + << line + << "] Exception across libtorch C API boundary: " << error_msg; + } + + // STD_TORCH_CHECK would stamp this helper's file and line over the + // caller's, which the message above already carries. + // @allow-raw-throw: keeps the caller's location in the message + throw std::runtime_error(ss.str()); +} +} // namespace torch_stable_abi_shim::detail + +// This macro is similar to the header-only macro TORCH_ERROR_CODE_CHECK, +// but this macro is NOT header-only! It depends on the stable ABI but +// provides more info in the exception, including the error message as +// retrieved through the c shims from the original error message. +#define TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(call) \ + if ((call) != TORCH_SUCCESS) { \ + torch_stable_abi_shim::detail::throw_exception(#call, __FILE__, __LINE__); \ + } diff --git a/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h b/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h new file mode 100644 index 000000000..7b7bc6619 --- /dev/null +++ b/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h @@ -0,0 +1,1628 @@ +#pragma once + +// Vendored from PyTorch commit d3c397fef74cebde85dde5ca5e6638a2cb2da9c5. +// APIs are gated by the targeted PyTorch stable ABI version. +// Encoded version checks also work with older headers lacking newer constants. +// Upstream header availability records when PyTorch shipped each wrapper; +// it is not a build-time requirement for this shim. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace torch_stable_abi_shim { + +/// A function pointer type for data deleters used with from_blob. +/// The deleter is called with the data pointer when the tensor's storage +/// is deallocated. +using DeleterFnPtr = void (*)(void*); + +/// Stable version of the empty_like op. +/// +/// Creates a new uninitialized tensor with the same size, dtype, layout, and +/// device as the input tensor. This version does not support kwargs (device, +/// dtype, layout, memory_format) - kwargs support may be added in the future. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor whose properties will be used for the new +/// tensor. +/// @return A new uninitialized tensor with the same properties as self. +inline torch::stable::Tensor empty_like(const torch::stable::Tensor& self) { + const auto num_args = 6; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(std::nullopt), + torch::stable::detail::from(std::nullopt), + torch::stable::detail::from(std::nullopt), + torch::stable::detail::from(std::nullopt), + torch::stable::detail::from(std::nullopt)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::empty_like", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::empty_like", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the fill_.Scalar op. +/// +/// Fills the input tensor with the specified scalar value in-place and returns +/// it. This has identical semantics to the existing fill_.Scalar op. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note The value parameter is typed as double +/// This is because Scalar.h is currently not header-only. +/// +/// @param self The tensor to fill. +/// @param value The scalar value to fill the tensor with. +/// @return The input tensor, now filled with the specified value. +inline torch::stable::Tensor fill_( + const torch::stable::Tensor& self, + double value) { + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_aten_fill__Scalar(self.get(), value)); + return self; +} + +/// Stable version of the narrow.default op. +/// +/// Returns a new tensor that is a narrowed version of the input tensor. The +/// dimension dim is narrowed from start to start + length. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note The start and length parameters +/// is not yet header-only. +/// +/// @param self The input tensor to narrow. +/// @param dim The dimension along which to narrow. +/// @param start The starting index for the narrowed dimension. +/// @param length The length of the narrowed dimension. +/// @return A new tensor that is a narrowed view of the input. +inline torch::stable::Tensor narrow( + torch::stable::Tensor& self, + int64_t dim, + int64_t start, + int64_t length) { + AtenTensorHandle ret0 = nullptr; + + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_aten_narrow(self.get(), dim, start, length, &ret0)); + return torch::stable::Tensor(ret0); +} + +#if TORCH_FEATURE_VERSION < TORCH_VERSION_2_10_0 +/// Stable version of the new_empty op (2.9 version). +/// +/// Creates a new uninitialized tensor with the specified size, inheriting +/// device and layout from the input tensor. This version only supports the +/// dtype kwarg. For the full kwargs version, use PyTorch 2.10+. +/// +/// Minimum compatible version: PyTorch 2.9. For full kwargs support, use +/// PyTorch 2.10+. +/// +/// @param self The input tensor whose device +/// @param size The desired size of the output tensor. +/// @param dtype Optional scalar type for the tensor elements. If not provided, +/// inherits from self. +/// @return A new uninitialized tensor with the specified properties. +inline torch::stable::Tensor new_empty( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef size, + std::optional dtype = std::nullopt) { + int32_t device_type; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_device_type(self.get(), &device_type)); + + int32_t device_index; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_device_index(self.get(), &device_index)); + + int32_t target_dtype; + if (dtype.has_value()) { + target_dtype = torch::stable::detail::to( + torch::stable::detail::from(dtype.value())); + } else { + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_dtype(self.get(), &target_dtype)); + } + + int32_t layout; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_get_layout(self.get(), &layout)); + + AtenTensorHandle ret0; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_aten_new_empty( + self.get(), + size.data(), + static_cast(size.size()), + &target_dtype, + &layout, + &device_type, + device_index, + nullptr, // pin_memory (nullptr for default) + &ret0)); + + return torch::stable::Tensor(ret0); +} + +/// Stable version of the new_zeros op (2.9 version). +/// +/// Creates a new tensor filled with zeros with the specified size, inheriting +/// device and layout from the input tensor. This version only supports the +/// dtype kwarg. For the full kwargs version, use PyTorch 2.10+. +/// +/// Minimum compatible version: PyTorch 2.9. For full kwargs support, use +/// PyTorch 2.10+. +/// +/// @param self The input tensor whose device and layout will be inherited. +/// @param size The desired size of the output tensor. +/// @param dtype Optional scalar type for the tensor elements. If not provided, +/// inherits from self. +/// @return A new zero-filled tensor with the specified properties. +inline torch::stable::Tensor new_zeros( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef size, + std::optional dtype = std::nullopt) { + int32_t device_type; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_device_type(self.get(), &device_type)); + + int32_t device_index; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_device_index(self.get(), &device_index)); + + int32_t target_dtype; + if (dtype.has_value()) { + target_dtype = torch::stable::detail::to( + torch::stable::detail::from(dtype.value())); + } else { + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_get_dtype(self.get(), &target_dtype)); + } + + int32_t layout; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_get_layout(self.get(), &layout)); + + AtenTensorHandle ath; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_aten_new_zeros( + self.get(), + size.data(), + static_cast(size.size()), + &target_dtype, + &layout, + &device_type, + device_index, + nullptr, // pin_memory (nullptr for default) + &ath)); + + return torch::stable::Tensor(ath); +} +#endif // TORCH_FEATURE_VERSION < TORCH_VERSION_2_10_0 + +/// Stable version of the pad.default op. +/// +/// Pads the input tensor according to the specified padding sizes. The padding +/// is applied symmetrically to each dimension, with the padding sizes specified +/// in reverse order (last dimension first). +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note The pad parameter is typed +/// not yet header-only. +/// +/// @param self The input tensor to pad. +/// @param pad The padding sizes for each dimension (in pairs, starting from +/// the last dimension). +/// @param mode The padding mode: "constant", "reflect", "replicate", or +/// "circular". Defaults to "constant". +/// @param value The fill value for constant padding. Defaults to 0.0. +/// @return A new padded tensor. +inline torch::stable::Tensor pad( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef pad, + const std::string& mode = "constant", + double value = 0.0) { + AtenTensorHandle ret0 = nullptr; + + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_aten_pad( + self.get(), pad.data(), pad.size(), mode.c_str(), &value, &ret0)); + return torch::stable::Tensor(ret0); +} + +/// Stable version of the amax.default op (single dimension). +/// +/// Computes the maximum value along the specified dimension. If keepdim is +/// true, the output tensor has the same number of dimensions as the input, +/// with the reduced dimension having size 1. Otherwise, the reduced dimension +/// is removed. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor. +/// @param dim The dimension along which to compute the maximum. +/// @param keepdim Whether to retain +/// @return A tensor containing the maximum values along the specified +/// dimension. +inline torch::stable::Tensor amax( + const torch::stable::Tensor& self, + int64_t dim, + bool keepdim = false) { + AtenTensorHandle ret = nullptr; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_aten_amax(self.get(), &dim, 1, keepdim, &ret)); + return torch::stable::Tensor(ret); +} + +/// Stable version of the amax.default op (multiple dimensions). +/// +/// Computes the maximum value reducing over all the specified dimensions. If +/// keepdim is true, the output tensor has the same number of dimensions as the +/// input, with the reduced dimensions having size 1. Otherwise, the reduced +/// dimensions are removed. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note The dims parameter is typed +/// is not yet header-only. +/// +/// @param self The input tensor. +/// @param dims The dimensions along which to compute the maximum. +/// @param keepdim Whether to retain the reduced dimensions. Defaults to false. +/// @return A tensor containing the maximum values. +inline torch::stable::Tensor amax( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef dims, + bool keepdim = false) { + AtenTensorHandle ret = nullptr; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_aten_amax( + self.get(), + dims.data(), + static_cast(dims.size()), + keepdim, + &ret)); + return torch::stable::Tensor(ret); +} + +/// Stable version of the transpose.int op. +/// +/// Returns a tensor that is a transposed version of the input, with dimensions +/// dim0 and dim1 swapped. The returned tensor shares storage with the input. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor. +/// @param dim0 The first dimension to transpose. +/// @param dim1 The second dimension to transpose. +/// @return A transposed view of the input tensor. +inline torch::stable::Tensor transpose( + const torch::stable::Tensor& self, + int64_t dim0, + int64_t dim1) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim0), + torch::stable::detail::from(dim1)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::transpose", "int", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::transpose", "int", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the zero_ op. +/// +/// Fills the input tensor with zeros in-place and returns it. Unlike the +/// tensor method version (t.zero_()), this is called as a function: zero_(t). +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The tensor to fill with zeros. +/// @return The input tensor, now filled with zeros. +inline torch::stable::Tensor zero_(torch::stable::Tensor& self) { + const auto num_args = 1; + std::array stack{torch::stable::detail::from(self)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::zero_", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::zero_", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the copy_ op. +/// +/// Copies the elements from the source tensor into the destination tensor +/// in-place and returns the destination tensor. The tensors must be +/// broadcastable. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The destination tensor (modified in-place). +/// @param src The source tensor to copy from. +/// @param non_blocking If true, the copy may occur asynchronously with respect +/// to the host. Defaults to false. +/// @return The destination tensor with copied values. +inline torch::stable::Tensor copy_( + torch::stable::Tensor& self, + const torch::stable::Tensor& src, + std::optional non_blocking = std::nullopt) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(src), + torch::stable::detail::from(non_blocking.value_or(false))}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::copy_", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::copy_", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the clone op. +/// +/// Returns a copy of the input tensor. The returned tensor has the same data +/// and type as the input, but is stored in a new memory location. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note Optional memory_format kwarg support +/// +/// @param self The input tensor to clone. +/// @return A new tensor with copied data. +inline torch::stable::Tensor clone(const torch::stable::Tensor& self) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(std::nullopt)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::clone", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::clone", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the flatten.using_ints op. +/// +/// Flattens the input tensor by reshaping it into a one-dimensional tensor. +/// If start_dim or end_dim are specified, only dimensions starting from +/// start_dim to end_dim are flattened. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor to flatten. +/// @param start_dim The first dimension to flatten. Defaults to 0. +/// @param end_dim The last dimension to flatten. Defaults to -1 (last dim). +/// @return A flattened tensor. +inline torch::stable::Tensor flatten( + const torch::stable::Tensor& self, + int64_t start_dim = 0, + int64_t end_dim = -1) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(start_dim), + torch::stable::detail::from(end_dim)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::flatten", "using_ints", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::flatten", "using_ints", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the unsqueeze op. +/// +/// Returns a new tensor with a dimension of size one inserted at the specified +/// position. The returned tensor shares the same underlying data with the input +/// tensor. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor. +/// @param dim The index at which to insert +/// values are supported. +/// @return A tensor with an additional dimension. +inline torch::stable::Tensor unsqueeze( + const torch::stable::Tensor& self, + int64_t dim) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(dim)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::unsqueeze", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::unsqueeze", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the squeeze.dim op. +/// +/// Returns a tensor with the dimension of size one at the specified position +/// removed. The returned tensor shares the same underlying data with the input +/// tensor. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The input tensor. +/// @param dim The dimension to squeeze. +/// the tensor is returned unchanged. +/// @return A tensor with the specified dimension removed (if size was 1). +inline torch::stable::Tensor squeeze( + const torch::stable::Tensor& self, + int64_t dim) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(dim)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::squeeze", "dim", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::squeeze", "dim", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the select.int op. +/// +/// Slices the input tensor along the specified dimension at the given index. +/// This function returns a view of the original tensor with the given dimension +/// removed. +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @note The index parameter is typed +/// header-only. +/// +/// @param self The input tensor. +/// @param dim The dimension to slice. +/// @param index The index to select along the dimension. +/// @return A tensor with one fewer dimension. +inline torch::stable::Tensor select( + const torch::stable::Tensor& self, + int64_t dim, + int64_t index) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim), + torch::stable::detail::from(index)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::select", "int", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::select", "int", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the matmul op. +/// +/// Performs matrix multiplication between two tensors. The behavior depends on +/// the dimensionality of the tensors (see PyTorch documentation for details on +/// broadcasting rules for matmul). +/// +/// Minimum compatible version: PyTorch 2.9. +/// +/// @param self The first input tensor. +/// @param other The second input tensor. +/// @return The result of matrix multiplication. +inline torch::stable::Tensor matmul( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::matmul", "", stack.data(), TORCH_ABI_VERSION)); +#else + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_call_dispatcher("aten::matmul", "", stack.data())); +#endif + return torch::stable::detail::to(stack[0]); +} + +#if TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + +/// Stable parallel_for utility. +/// +/// Provides a stable interface to at::parallel_for for parallel execution. +/// The function f will be called with (begin, end) ranges to process in +/// parallel. grain_size controls the minimum work size per thread for efficient +/// parallelization. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @tparam F The callable type +/// @param begin The start of the iteration range. +/// @param end The end of the iteration range (exclusive). +/// @param grain_size The minimum number of iterations per thread. +/// @param f The function to execute in parallel. +template +inline void parallel_for( + const int64_t begin, + const int64_t end, + const int64_t grain_size, + const F& f) { + auto callback = [](int64_t cb_begin, int64_t cb_end, void* ctx) { + const F* func = static_cast(ctx); + (*func)(cb_begin, cb_end); + }; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_parallel_for( + begin, + end, + grain_size, + callback, + const_cast(static_cast(&f)))); +} + +/// Gets the number of threads for the parallel backend. +/// +/// Provides a stable interface to at::get_num_threads. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @return The number of threads +inline uint32_t get_num_threads() { + uint32_t num_threads; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_get_num_threads(&num_threads)); + return num_threads; +} + +/// Stable version of the empty.memory_format op. +/// +/// Creates a new uninitialized tensor with the specified size and options. +/// This function supports full tensor creation options including device, +/// dtype, layout, and memory format. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param size The desired size of the output tensor. +/// @param dtype Optional scalar type for the tensor elements. +/// @param layout Optional memory layout (e.g., strided, sparse). +/// @param device Optional device to place the tensor on. +/// @param pin_memory Optional flag to use pinned memory (for CUDA tensors). +/// @param memory_format Optional memory format for the tensor. +/// @return A new uninitialized tensor with the specified properties. +inline torch::stable::Tensor empty( + torch::headeronly::IntHeaderOnlyArrayRef size, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt, + std::optional memory_format = + std::nullopt) { + const auto num_args = 6; + std::array stack{ + torch::stable::detail::from(size), + torch::stable::detail::from(dtype), + torch::stable::detail::from(layout), + torch::stable::detail::from(device), + torch::stable::detail::from(pin_memory), + torch::stable::detail::from(memory_format)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::empty", "memory_format", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the reshape op. +/// +/// Returns a tensor with the same data and number of elements as the input, +/// but with the specified shape. When possible, the returned tensor will be +/// a view of the input. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param shape The desired output shape. +/// @return A tensor with the specified shape. +inline torch::stable::Tensor reshape( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef shape) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(shape)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::reshape", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the view op. +/// +/// Returns a new tensor with the same data as the input tensor but with a +/// different shape. The returned tensor shares the same data and must have +/// the same number of elements. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param size The desired output shape. +/// @return A view tensor with the specified shape. +inline torch::stable::Tensor view( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef size) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(size)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + torch_call_dispatcher("aten::view", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the view.dtype op. +/// +/// Returns a view of the input tensor with the same data reinterpreted as the +/// given dtype. This is a different overload than the shape-based view op. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor. +/// @param dtype The dtype to reinterpret the data as. +/// @return A view of the input tensor with the specified dtype. +inline torch::stable::Tensor view( + const torch::stable::Tensor& self, + torch::headeronly::ScalarType dtype) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(dtype)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::view", "dtype", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Creates a tensor from an existing data blob. +/// +/// Creates a tensor that uses the provided data pointer as its storage. +/// The tensor does not own the data, so the caller must ensure the data +/// remains valid for the lifetime of the tensor. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param data Pointer to the data buffer. +/// @param sizes The size of each dimension of the tensor. +/// @param strides The stride for each dimension. +/// @param device The device where the data resides. +/// @param dtype The scalar type of the data. +/// @param storage_offset The offset into the data buffer. Defaults to 0. +/// @param layout The memory layout. Defaults to Strided. +/// @return A tensor backed by the provided data. +inline torch::stable::Tensor from_blob( + void* data, + torch::headeronly::IntHeaderOnlyArrayRef sizes, + torch::headeronly::IntHeaderOnlyArrayRef strides, + torch::stable::Device device, + torch::headeronly::ScalarType dtype, + int64_t storage_offset = 0, + torch::headeronly::Layout layout = torch::headeronly::Layout::Strided) { + auto shim_dtype = + torch::stable::detail::to(torch::stable::detail::from(dtype)); + auto shim_device_type = torch::stable::detail::to( + torch::stable::detail::from(device.type())); + auto shim_layout = + torch::stable::detail::to(torch::stable::detail::from(layout)); + AtenTensorHandle ath; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob_v2( + data, + sizes.size(), + sizes.data(), + strides.data(), + storage_offset, + shim_dtype, + shim_device_type, + device.index(), + &ath, + shim_layout, + nullptr, + 0)); + return torch::stable::Tensor(ath); +} + +#if TORCH_FEATURE_VERSION >= (((0ULL + 2) << 56) | ((0ULL + 11) << 48)) +/// Creates a tensor from an existing data blob with a custom deleter. +/// +/// This is the same as the from_blob function above, but allows specifying a +/// custom deleter function that will be called when the tensor's storage is +/// deallocated. Accepts both plain function pointers and capturing lambdas. +/// +/// Minimum compatible version: PyTorch 2.11. +/// +/// @tparam F The callable type. Must be invocable with (void*). +/// @param data Pointer to the data buffer. +/// @param sizes The size of each dimension of the tensor. +/// @param strides The stride for each dimension. +/// @param device The device where the data resides. +/// @param dtype The scalar type of the data. +/// @param deleter Callable to invoke when the tensor is deallocated. +/// @param storage_offset The offset into the data buffer. Defaults to 0. +/// @param layout The memory layout. Defaults to Strided. +/// @return A tensor backed by the provided data. +template , int> = 0> +inline torch::stable::Tensor from_blob( + void* data, + torch::headeronly::IntHeaderOnlyArrayRef sizes, + torch::headeronly::IntHeaderOnlyArrayRef strides, + torch::stable::Device device, + torch::headeronly::ScalarType dtype, + F deleter, + int64_t storage_offset = 0, + torch::headeronly::Layout layout = torch::headeronly::Layout::Strided) { + auto shim_dtype = + torch::stable::detail::to(torch::stable::detail::from(dtype)); + auto shim_device_type = torch::stable::detail::to( + torch::stable::detail::from(device.type())); + auto shim_layout = + torch::stable::detail::to(torch::stable::detail::from(layout)); + + AtenTensorHandle ath; + if constexpr (std::is_convertible_v) { + // Simple function pointer: pass it as ctx, no heap allocation. + auto deleter_callback = [](void* data, void* ctx) { + auto fn = reinterpret_cast(ctx); + fn(data); + }; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_from_blob( + data, + sizes.size(), + sizes.data(), + strides.data(), + storage_offset, + shim_dtype, + shim_device_type, + device.index(), + &ath, + shim_layout, + nullptr, + 0, + deleter_callback, + reinterpret_cast(static_cast(deleter)))); + } else { + // Capturing lambda: heap-allocate and type-erase. + F* heap_allocated_deleter = new F(std::move(deleter)); + auto deleter_callback = [](void* data, void* ctx) { + F* func = static_cast(ctx); + (*func)(data); + delete func; + }; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_from_blob( + data, + sizes.size(), + sizes.data(), + strides.data(), + storage_offset, + shim_dtype, + shim_device_type, + device.index(), + &ath, + shim_layout, + nullptr, + 0, + deleter_callback, + static_cast(heap_allocated_deleter))); + } + return torch::stable::Tensor(ath); +} +#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_11_0 + +/// Stable version of the to.dtype_layout op. +/// +/// Converts a tensor to the specified dtype, layout, device, and/or memory +/// format. Returns a new tensor with the specified properties. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param dtype Optional target scalar type. +/// @param layout Optional target memory layout. +/// @param device Optional target device. +/// @param pin_memory Optional flag to use pinned memory. +/// @param non_blocking If true, the operation may be asynchronous. Defaults to +/// false. +/// @param copy If true, always create a copy. Defaults to false. +/// @param memory_format Optional target memory format. +/// @return A tensor with the specified properties. +inline torch::stable::Tensor to( + const torch::stable::Tensor& self, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt, + bool non_blocking = false, + bool copy = false, + std::optional memory_format = + std::nullopt) { + const auto num_args = 8; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dtype), + torch::stable::detail::from(layout), + torch::stable::detail::from(device), + torch::stable::detail::from(pin_memory), + torch::stable::detail::from(non_blocking), + torch::stable::detail::from(copy), + torch::stable::detail::from(memory_format)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::to", "dtype_layout", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Convenience overload for moving a tensor to a device. +/// +/// Moves the tensor to the specified device. This is a convenience wrapper +/// around the full to() function. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param device The target device. +/// @param non_blocking If true, the operation may be asynchronous. Defaults to +/// false. +/// @param copy If true, always create a copy. Defaults to false. +/// @return A tensor on the specified device. +inline torch::stable::Tensor to( + const torch::stable::Tensor& self, + torch::stable::Device device, + bool non_blocking = false, + bool copy = false) { + return to( + self, + std::nullopt, + std::nullopt, + device, + std::nullopt, + non_blocking, + copy, + std::nullopt); +} + +/// Stable version of the contiguous op. +/// +/// Returns a contiguous in memory tensor containing the same data as the input +/// tensor. If the input tensor is already contiguous in the specified memory +/// format, the input tensor is returned. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param memory_format The desired memory format. +/// @return A contiguous tensor. +inline torch::stable::Tensor contiguous( + const torch::stable::Tensor& self, + torch::headeronly::MemoryFormat memory_format = + torch::headeronly::MemoryFormat::Contiguous) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(memory_format)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::contiguous", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the new_empty op (2.10 version with full kwargs). +/// +/// Creates a new uninitialized tensor with the specified size and options. +/// This version supports all tensor creation kwargs. For versions < 2.10, +/// a simpler overload that only takes dtype is available. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor whose properties may be inherited if kwargs are +/// not provided. +/// @param size The desired size of the output tensor. +/// @param dtype Optional scalar type for the tensor elements. +/// @param layout Optional memory layout (e.g., strided, sparse). +/// @param device Optional device to place the tensor on. +/// @param pin_memory Optional flag to use pinned memory (for CUDA tensors). +/// @return A new uninitialized tensor with the specified properties. +inline torch::stable::Tensor new_empty( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef size, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt) { + const auto num_args = 6; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(size), + torch::stable::detail::from(dtype), + torch::stable::detail::from(layout), + torch::stable::detail::from(device), + torch::stable::detail::from(pin_memory)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::new_empty", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the new_zeros op (2.10 version with full kwargs). +/// +/// Creates a new zero-filled tensor with the specified size and options. +/// This version supports all tensor creation kwargs. For versions < 2.10, +/// a simpler overload that only takes dtype is available. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor whose properties may be inherited if kwargs are +/// not provided. +/// @param size The desired size of the output tensor. +/// @param dtype Optional scalar type for the tensor elements. +/// @param layout Optional memory layout (e.g., strided, sparse). +/// @param device Optional device to place the tensor on. +/// @param pin_memory Optional flag to use pinned memory (for CUDA tensors). +/// @return A new zero-filled tensor with the specified properties. +inline torch::stable::Tensor new_zeros( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef size, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt) { + const auto num_args = 6; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(size), + torch::stable::detail::from(dtype), + torch::stable::detail::from(layout), + torch::stable::detail::from(device), + torch::stable::detail::from(pin_memory)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::new_zeros", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the sum.dim_IntList op. +/// +/// Computes the sum of the input tensor along the specified dimensions. +/// If dim is not provided, sums over all dimensions. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param self The input tensor. +/// @param dim Optional dimensions to reduce. If not provided, reduces all +/// dimensions. +/// @param keepdim Whether to retain the reduced dimensions. Defaults to false. +/// @param dtype Optional output dtype. If not provided, uses the input dtype. +/// @return A tensor containing the sum. +inline torch::stable::Tensor sum( + const torch::stable::Tensor& self, + std::optional dim = std::nullopt, + bool keepdim = false, + std::optional dtype = std::nullopt) { + const auto num_args = 4; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim), + torch::stable::detail::from(keepdim), + torch::stable::detail::from(dtype)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::sum", "dim_IntList", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the sum.IntList_out op. +/// +/// Computes the sum of the input tensor along the specified dimensions, +/// storing the result in the provided output tensor. Following C++ convention, +/// the out parameter comes first. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @param out The output tensor (modified in-place). +/// @param self The input tensor. +/// @param dim Optional dimensions to reduce. +/// @param keepdim Whether to retain the reduced dimensions. Defaults to false. +/// @param dtype Optional output dtype. +/// @return Reference to the output tensor. +inline torch::stable::Tensor& sum_out( + torch::stable::Tensor& out, + const torch::stable::Tensor& self, + std::optional dim = std::nullopt, + bool keepdim = false, + std::optional dtype = std::nullopt) { + const auto num_args = 5; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim), + torch::stable::detail::from(keepdim), + torch::stable::detail::from(dtype), + torch::stable::detail::from(out)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::sum", "IntList_out", stack.data(), TORCH_ABI_VERSION)); + // Clean up the handle in stack[0], discard the temporary + (void)torch::stable::detail::to(stack[0]); + return out; +} + +/// Stable version of the subtract.Tensor op. +/// +/// Subtracts the other tensor from self, with an optional scaling factor alpha. +/// Computes: self - alpha * other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @note The alpha parameter is typed as double +/// API uses double for the Scalar parameter. +/// +/// @param self The input tensor. +/// @param other The tensor to subtract. +/// @param alpha The scaling factor for other. Defaults to 1.0. +/// @return The result of self - alpha * other. +inline torch::stable::Tensor subtract( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other, + double alpha = 1.0) { + AtenTensorHandle ret0; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_aten_subtract_Tensor(self.get(), other.get(), alpha, &ret0)); + return torch::stable::Tensor(ret0); +} + +/// Stable version of the full.default op. +/// +/// Creates a tensor of the specified size filled with the given value. +/// +/// Minimum compatible version: PyTorch 2.10. +/// +/// @note The fill_value parameter is typed +/// C shim API uses double for the Scalar parameter. +/// +/// @param size The desired size of the output tensor. +/// @param fill_value The value to fill the tensor with. +/// @param dtype Optional scalar type for the tensor elements. +/// @param layout Optional memory layout. +/// @param device Optional device to place the tensor on. +/// @param pin_memory Optional flag to use pinned memory. +/// @return A new tensor filled with the specified value. +inline torch::stable::Tensor full( + torch::headeronly::IntHeaderOnlyArrayRef size, + double fill_value, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt) { + int32_t* dtype_ptr = nullptr; + int32_t dtype_val; + if (dtype.has_value()) { + dtype_val = torch::stable::detail::to( + torch::stable::detail::from(dtype.value())); + dtype_ptr = &dtype_val; + } + + int32_t* layout_ptr = nullptr; + int32_t layout_val; + if (layout.has_value()) { + layout_val = torch::stable::detail::to( + torch::stable::detail::from(layout.value())); + layout_ptr = &layout_val; + } + + int32_t* device_type_ptr = nullptr; + int32_t device_type_val; + int32_t device_index = 0; + if (device.has_value()) { + device_type_val = torch::stable::detail::to( + torch::stable::detail::from(device.value().type())); + device_type_ptr = &device_type_val; + device_index = device.value().index(); + } + + int32_t* pin_memory_ptr = nullptr; + int32_t pin_memory_val; + if (pin_memory.has_value()) { + pin_memory_val = pin_memory.value() ? 1 : 0; + pin_memory_ptr = &pin_memory_val; + } + + AtenTensorHandle ret0; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(aoti_torch_aten_full( + size.data(), + static_cast(size.size()), + fill_value, + dtype_ptr, + layout_ptr, + device_type_ptr, + device_index, + pin_memory_ptr, + &ret0)); + + return torch::stable::Tensor(ret0); +} + +/// Stable version of the bitwise_and.Tensor op. +/// +/// Computes the element-wise bitwise AND of self and other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The first input tensor. +/// @param other The second input tensor. +/// @return The result of the bitwise AND. +inline torch::stable::Tensor bitwise_and( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bitwise_and", "Tensor", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the bitwise_or.Tensor op. +/// +/// Computes the element-wise bitwise OR of self and other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The first input tensor. +/// @param other The second input tensor. +/// @return The result of the bitwise OR. +inline torch::stable::Tensor bitwise_or( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bitwise_or", "Tensor", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the bitwise_left_shift.Tensor op. +/// +/// Computes the element-wise left bitwise shift of self by other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor to shift. +/// @param other The tensor of shift amounts. +/// @return The result of the left bitwise shift. +inline torch::stable::Tensor bitwise_left_shift( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bitwise_left_shift", "Tensor", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the bitwise_right_shift.Tensor op. +/// +/// Computes the element-wise right bitwise shift of self by other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor to shift. +/// @param other The tensor of shift amounts. +/// @return The result of the right bitwise shift. +inline torch::stable::Tensor bitwise_right_shift( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bitwise_right_shift", "Tensor", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the permute op. +/// +/// Returns a view of the input tensor with its dimensions rearranged according +/// to dims. The returned tensor shares the same underlying data with the input. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor. +/// @param dims The desired ordering of dimensions. +/// @return A permuted view of the input tensor. +inline torch::stable::Tensor permute( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef dims) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(dims)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::permute", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the index_select op. +/// +/// Returns a new tensor which indexes the input tensor along dimension dim +/// using the entries in index. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor. +/// @param dim The dimension to index along. +/// @param index A 1-D tensor of indices to select. +/// @return A tensor with values selected along dim. +inline torch::stable::Tensor index_select( + const torch::stable::Tensor& self, + int64_t dim, + const torch::stable::Tensor& index) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim), + torch::stable::detail::from(index)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::index_select", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the floor_divide op. +/// +/// Computes the element-wise floor division of self by other. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The dividend tensor. +/// @param other The divisor tensor. +/// @return The result of floor division. +inline torch::stable::Tensor floor_divide( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(other)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::floor_divide", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the is_pinned op. +/// +/// Returns true if the tensor's storage is allocated in pinned memory. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.14. +/// +/// @param self The input tensor. +/// @return true if the tensor is pinned, false otherwise. +inline bool is_pinned(const torch::stable::Tensor& self) { + const auto num_args = 2; + // Pass nullopt for aten::is_pinned's deprecated optional Device arg. + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(std::nullopt)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::is_pinned", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the sort op. +/// +/// Sorts the input tensor along dim, returning the values and their indices. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param self The input tensor. +/// @param dim The dimension to sort along. +/// @param descending Whether to sort in descending order. +/// @return A tuple of the sorted values and their indices. +inline std::tuple sort( + const torch::stable::Tensor& self, + int64_t dim = -1, + bool descending = false) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(dim), + torch::stable::detail::from(descending)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + torch_call_dispatcher("aten::sort", "", stack.data(), TORCH_ABI_VERSION)); + return std::make_tuple( + torch::stable::detail::to(stack[0]), + torch::stable::detail::to(stack[1])); +} + +/// Stable version of the sort.stable op. +/// +/// Sorts the input tensor along dim, returning the values and their indices. +/// This is a different overload than the sort op without the stable argument. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param self The input tensor. +/// @param stable Whether equal elements keep their original order. +/// @param dim The dimension to sort along. +/// @param descending Whether to sort in descending order. +/// @return A tuple of the sorted values and their indices. +inline std::tuple sort( + const torch::stable::Tensor& self, + std::optional stable, + int64_t dim = -1, + bool descending = false) { + const auto num_args = 4; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(stable), + torch::stable::detail::from(dim), + torch::stable::detail::from(descending)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::sort", "stable", stack.data(), TORCH_ABI_VERSION)); + return std::make_tuple( + torch::stable::detail::to(stack[0]), + torch::stable::detail::to(stack[1])); +} + +/// Stable version of the masked_select op. +/// +/// Returns a 1-D tensor of the elements of self where mask is true. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param self The input tensor. +/// @param mask A boolean tensor broadcastable with self. +/// @return A 1-D tensor of the selected elements. +inline torch::stable::Tensor masked_select( + const torch::stable::Tensor& self, + const torch::stable::Tensor& mask) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(mask)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::masked_select", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the mm.out op. +/// +/// Computes the matrix product of the 2-D tensors self and mat2, storing the +/// result in out. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param out The output tensor (modified in-place). +/// @param self The first 2-D input tensor. +/// @param mat2 The second 2-D input tensor. +/// @return Reference to the output tensor. +inline torch::stable::Tensor& mm_out( + torch::stable::Tensor& out, + const torch::stable::Tensor& self, + const torch::stable::Tensor& mat2) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(mat2), + torch::stable::detail::from(out)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::mm", "out", stack.data(), TORCH_ABI_VERSION)); + // Clean up the handle in stack[0], discard the temporary + (void)torch::stable::detail::to(stack[0]); + return out; +} + +/// Stable version of the mm.dtype_out op. +/// +/// Computes the matrix product of self and mat2 in out_dtype. This is a +/// different overload than mm_out without the out_dtype argument. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param out The output tensor (modified in-place). +/// @param self The first 2-D input tensor. +/// @param mat2 The second 2-D input tensor. +/// @param out_dtype The dtype of the result. +/// @return Reference to the output tensor. +inline torch::stable::Tensor& mm_out( + torch::stable::Tensor& out, + const torch::stable::Tensor& self, + const torch::stable::Tensor& mat2, + torch::headeronly::ScalarType out_dtype) { + const auto num_args = 4; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(mat2), + torch::stable::detail::from(out_dtype), + torch::stable::detail::from(out)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::mm", "dtype_out", stack.data(), TORCH_ABI_VERSION)); + // Clean up the handle in stack[0], discard the temporary + (void)torch::stable::detail::to(stack[0]); + return out; +} + +/// Stable version of the bmm.out op. +/// +/// Computes the batch matrix product of the 3-D tensors self and mat2, +/// storing the result in out. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param out The output tensor (modified in-place). +/// @param self The first 3-D input tensor of shape (b, n, m). +/// @param mat2 The second 3-D input tensor of shape (b, m, p). +/// @return Reference to the output tensor. +inline torch::stable::Tensor& bmm_out( + torch::stable::Tensor& out, + const torch::stable::Tensor& self, + const torch::stable::Tensor& mat2) { + const auto num_args = 3; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(mat2), + torch::stable::detail::from(out)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bmm", "out", stack.data(), TORCH_ABI_VERSION)); + // Clean up the handle in stack[0], discard the temporary + (void)torch::stable::detail::to(stack[0]); + return out; +} + +/// Stable version of the bmm.dtype_out op. +/// +/// Computes the batch matrix product of self and mat2 in out_dtype. This is a +/// different overload than bmm_out without the out_dtype argument. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param out The output tensor (modified in-place). +/// @param self The first 3-D input tensor of shape (b, n, m). +/// @param mat2 The second 3-D input tensor of shape (b, m, p). +/// @param out_dtype The dtype of the result. +/// @return Reference to the output tensor. +inline torch::stable::Tensor& bmm_out( + torch::stable::Tensor& out, + const torch::stable::Tensor& self, + const torch::stable::Tensor& mat2, + torch::headeronly::ScalarType out_dtype) { + const auto num_args = 4; + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(mat2), + torch::stable::detail::from(out_dtype), + torch::stable::detail::from(out)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bmm", "dtype_out", stack.data(), TORCH_ABI_VERSION)); + // Clean up the handle in stack[0], discard the temporary + (void)torch::stable::detail::to(stack[0]); + return out; +} + +/// Stable version of the flip op. +/// +/// Reverses the order of elements along the given dimensions and returns +/// the result as a new tensor. Negative dimensions are wrapped. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param self The input tensor. +/// @param dims The dimensions to flip. +/// @return A new tensor with the given dimensions reversed. +inline torch::stable::Tensor flip( + const torch::stable::Tensor& self, + torch::headeronly::IntHeaderOnlyArrayRef dims) { + const auto num_args = 2; + std::array stack{ + torch::stable::detail::from(self), torch::stable::detail::from(dims)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + torch_call_dispatcher("aten::flip", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +/// Stable version of the from_file op. +/// +/// Creates a CPU tensor backed by a memory-map of the file at filename. +/// Unlike at::from_file, filename is a const std::string& rather than a +/// std::string_view. +/// +/// Minimum compatible version: PyTorch 2.10. +/// Upstream header availability: PyTorch 2.15. +/// +/// @param filename Path of the file to read. +/// @param shared Optional flag to write tensor changes back to the file. +/// @param size Optional number of elements to read from the file. +/// @param dtype Optional scalar type for the tensor elements. +/// @param layout Optional memory layout, which must be strided for this op. +/// @param device Optional device, which must be CPU for this op. +/// @param pin_memory Optional pinning flag, which must be false for this op. +/// @return A new tensor holding the file contents. +inline torch::stable::Tensor from_file( + const std::string& filename, + std::optional shared = std::nullopt, + std::optional size = 0, + std::optional dtype = std::nullopt, + std::optional layout = std::nullopt, + std::optional device = std::nullopt, + std::optional pin_memory = std::nullopt) { + const auto num_args = 7; + std::array stack{ + torch::stable::detail::from(filename), + torch::stable::detail::from(shared), + torch::stable::detail::from(size), + torch::stable::detail::from(dtype), + torch::stable::detail::from(layout), + torch::stable::detail::from(device), + torch::stable::detail::from(pin_memory)}; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::from_file", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); +} + +#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_10_0 + +#if TORCH_FEATURE_VERSION >= (((0ULL + 2) << 56) | ((0ULL + 16) << 48)) + +/// Stable version of the add.Tensor op. +/// +/// Adds the other tensor to self, scaled by alpha: self + alpha * other. +/// Unlike at::add, alpha is a double rather than a Scalar. Adding two +/// integral or boolean tensors is not supported. +/// +/// Minimum compatible version: PyTorch 2.16. +/// Upstream header availability: PyTorch 2.16. +/// +/// @param self The input tensor. +/// @param other The tensor to add. +/// @param alpha The scaling factor for other. Defaults to 1.0. +/// @return The result of self + alpha * other. +inline torch::stable::Tensor add( + const torch::stable::Tensor& self, + const torch::stable::Tensor& other, + double alpha = 1.0) { + AtenTensorHandle ret0; + TORCH_STABLE_ABI_SHIM_TORCH_ERROR_CODE_CHECK( + aoti_torch_aten_add_Tensor(self.get(), other.get(), alpha, &ret0)); + return torch::stable::Tensor(ret0); +} + +#endif // TORCH_FEATURE_VERSION >= TORCH_VERSION_2_16_0 + +} // namespace torch_stable_abi_shim From e36846758de2ad62e1456ab3e5479d1ab34529d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 20 Oct 2026 12:26:04 +0000 Subject: [PATCH 2/2] Qualify `to` call. --- .../include/torch_stable_abi_shim/csrc/stable/ops.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h b/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h index 7b7bc6619..8b1a2fa11 100644 --- a/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h +++ b/torch-stable-abi-shim/include/torch_stable_abi_shim/csrc/stable/ops.h @@ -900,7 +900,7 @@ inline torch::stable::Tensor to( torch::stable::Device device, bool non_blocking = false, bool copy = false) { - return to( + return torch_stable_abi_shim::to( self, std::nullopt, std::nullopt,