chore: import upstream snapshot with attribution

test/cpp/compat/torch_library_dispatch_test.cc

@@ -0,0 +1,200 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// The file has been adapted from pytorch project
+// Licensed under BSD-style license -
+// https://github.com/pytorch/pytorch/blob/main/LICENSE
+
+// Tests for the compat-layer dispatch key priority selection logic introduced
+// in OperationInvoker::get_op_with_args (torch_compat.h).
+//
+// The lookup order is: CPU → BackendSelect → CatchAll.
+// If none of those keys exist and exactly one implementation is registered it
+// is used directly (deterministic). If multiple unrecognised keys exist the
+// invoker raises an Ambiguous error rather than picking arbitrarily from an
+// unordered_map (which has no stable iteration order).
+// These tests exercise scenarios where the registrant uses BackendSelect
+// (e.g. TORCH_LIBRARY_IMPL(..., BackendSelect, m)) so that the Python-facing
+// invoker can reach it even when no CPU implementation exists.
+
+#include <torch/library.h>
+
+#include <vector>
+
+#include "gtest/gtest.h"
+
+// ---------------------------------------------------------------------------
+// Operator implementations used by the tests below
+// ---------------------------------------------------------------------------
+
+namespace {
+
+int backend_select_probe(int x) { return x + 10; }
+
+int backend_select_and_cpu_cpu_fn(int x) { return x + 1; }
+int backend_select_and_cpu_bs_fn(int x) { return x + 2; }
+
+}  // namespace
+
+int unique_non_preferred_fn(int x) { return x + 7; }
+int ambiguous_cuda_fn(int x) { return x + 100; }
+int ambiguous_xpu_fn(int x) { return x + 200; }
+
+TORCH_LIBRARY(compat_dispatch_test_lib, m) {
+  m.def("backend_select_only(int x) -> int");
+  m.def("backend_select_and_cpu(int x) -> int");
+  m.def("unique_non_preferred(int x) -> int");
+  m.def("ambiguous_multi_key(int x) -> int");
+}
+
+TORCH_LIBRARY_IMPL(compat_dispatch_test_lib, BackendSelect, m) {
+  m.impl("backend_select_only", &backend_select_probe);
+  m.impl("backend_select_and_cpu", &backend_select_and_cpu_bs_fn);
+}
+
+TORCH_LIBRARY_IMPL(compat_dispatch_test_lib, CPU, m) {
+  m.impl("backend_select_and_cpu", &backend_select_and_cpu_cpu_fn);
+}
+
+TORCH_LIBRARY_IMPL(compat_dispatch_test_lib, CUDA, m) {
+  m.impl("unique_non_preferred", &unique_non_preferred_fn);
+  m.impl("ambiguous_multi_key", &ambiguous_cuda_fn);
+}
+
+TORCH_LIBRARY_IMPL(compat_dispatch_test_lib, XPU, m) {
+  m.impl("ambiguous_multi_key", &ambiguous_xpu_fn);
+}
+
+// ---------------------------------------------------------------------------
+// Helper: simulate the priority-fallback lookup used by get_op_with_args
+// ---------------------------------------------------------------------------
+
+static decltype(torch::OperatorRegistry::instance()
+                    .find_operator("")
+                    ->implementations.end())
+pick_impl(torch::OperatorRegistration* op) {
+  using DK = c10::DispatchKey;
+  const std::vector<DK> preferred_keys = {
+      DK::CPU, DK::BackendSelect, DK::CatchAll};
+  auto chosen = op->implementations.end();
+  for (const auto& key : preferred_keys) {
+    chosen = op->implementations.find(key);
+    if (chosen != op->implementations.end()) break;
+  }
+  // Mirror the production rule: allow exactly-one-impl, reject ambiguous.
+  if (chosen == op->implementations.end() && op->implementations.size() == 1) {
+    chosen = op->implementations.begin();
+  }
+  return chosen;
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+// An operator registered only under BackendSelect must be queryable under
+// that key and must NOT appear under CPU.
+TEST(CompatTorchDispatchTest, BackendSelectOnlyRegistration) {
+  const auto qname = "compat_dispatch_test_lib::backend_select_only";
+  auto* op = torch::OperatorRegistry::instance().find_operator(qname);
+  ASSERT_NE(op, nullptr);
+
+  EXPECT_EQ(op->implementations.find(c10::DispatchKey::CPU),
+            op->implementations.end());
+
+  auto bs_it = op->implementations.find(c10::DispatchKey::BackendSelect);
+  ASSERT_NE(bs_it, op->implementations.end());
+
+  torch::FunctionArgs args;
+  args.add_arg(torch::IValue(int64_t(32)));
+  auto result = bs_it->second.call_with_args(args);
+  ASSERT_TRUE(result.get_value().is_int());
+  EXPECT_EQ(result.get_value().to_int(), 42);  // 32 + 10
+}
+
+// When CPU and BackendSelect are both registered, the priority lookup must
+// pick CPU (higher priority in get_op_with_args).
+TEST(CompatTorchDispatchTest, CpuPreferredOverBackendSelect) {
+  const auto qname = "compat_dispatch_test_lib::backend_select_and_cpu";
+  auto* op = torch::OperatorRegistry::instance().find_operator(qname);
+  ASSERT_NE(op, nullptr);
+
+  ASSERT_NE(op->implementations.find(c10::DispatchKey::CPU),
+            op->implementations.end());
+  ASSERT_NE(op->implementations.find(c10::DispatchKey::BackendSelect),
+            op->implementations.end());
+
+  auto chosen = pick_impl(op);
+  ASSERT_NE(chosen, op->implementations.end());
+  EXPECT_EQ(chosen->first, c10::DispatchKey::CPU);
+
+  torch::FunctionArgs args;
+  args.add_arg(torch::IValue(int64_t(41)));
+  auto result = chosen->second.call_with_args(args);
+  ASSERT_TRUE(result.get_value().is_int());
+  EXPECT_EQ(result.get_value().to_int(), 42);  // CPU impl: x + 1
+}
+
+// When CPU is absent, the priority lookup must fall through to BackendSelect.
+TEST(CompatTorchDispatchTest, BackendSelectPickedWhenCpuAbsent) {
+  const auto qname = "compat_dispatch_test_lib::backend_select_only";
+  auto* op = torch::OperatorRegistry::instance().find_operator(qname);
+  ASSERT_NE(op, nullptr);
+
+  auto chosen = pick_impl(op);
+  ASSERT_NE(chosen, op->implementations.end());
+  EXPECT_EQ(chosen->first, c10::DispatchKey::BackendSelect);
+
+  torch::FunctionArgs args;
+  args.add_arg(torch::IValue(int64_t(32)));
+  auto result = chosen->second.call_with_args(args);
+  ASSERT_TRUE(result.get_value().is_int());
+  EXPECT_EQ(result.get_value().to_int(), 42);  // BackendSelect impl: x + 10
+}
+
+// An operator registered only under one non-preferred key (e.g. CUDA) must
+// still be reachable when it's the sole implementation (deterministic).
+TEST(CompatTorchDispatchTest, UniqueNonPreferredKeyIsCallable) {
+  const auto qname = "compat_dispatch_test_lib::unique_non_preferred";
+  auto* op = torch::OperatorRegistry::instance().find_operator(qname);
+  ASSERT_NE(op, nullptr);
+  ASSERT_EQ(op->implementations.size(), 1UL);
+
+  auto chosen = pick_impl(op);
+  ASSERT_NE(chosen, op->implementations.end());
+
+  torch::FunctionArgs args;
+  args.add_arg(torch::IValue(int64_t(35)));
+  auto result = chosen->second.call_with_args(args);
+  ASSERT_TRUE(result.get_value().is_int());
+  EXPECT_EQ(result.get_value().to_int(), 42);  // unique impl: x + 7
+}
+
+// An operator with multiple non-preferred keys (CUDA + XPU) must produce
+// end() from pick_impl (the production code would raise an Ambiguous error).
+TEST(CompatTorchDispatchTest, AmbiguousMultiKeyProducesEnd) {
+  const auto qname = "compat_dispatch_test_lib::ambiguous_multi_key";
+  auto* op = torch::OperatorRegistry::instance().find_operator(qname);
+  ASSERT_NE(op, nullptr);
+  // Registered under CUDA and XPU – neither is in the preferred list.
+  ASSERT_GE(op->implementations.size(), 2UL);
+  EXPECT_EQ(op->implementations.find(c10::DispatchKey::CPU),
+            op->implementations.end());
+  EXPECT_EQ(op->implementations.find(c10::DispatchKey::BackendSelect),
+            op->implementations.end());
+
+  auto chosen = pick_impl(op);
+  // Must not resolve to any implementation – ambiguous.
+  EXPECT_EQ(chosen, op->implementations.end());
+}