wehub-resource-sync
201 lines
7.9 KiB
C++
201 lines
7.9 KiB
C++
// 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());
|
||
}
|