Skip to content

fix: preserve implementation lookups under optimization #929

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
voltjia wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix: preserve implementation lookups under optimization #929

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 27 additions & 9 deletions scripts/generate_wrappers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1534,18 +1534,36 @@ def _append_unique(declaration, definition):
declarations.append(declaration)
definitions.append(definition)

_append_unique(
f"extern template std::vector<std::size_t> "
f"Operator<{op_type}>::active_implementation_indices(Device::Type);",
f"template std::vector<std::size_t> "
f"Operator<{op_type}>::active_implementation_indices(Device::Type);",
active_implementation_declaration = (
"template <>\n"
"std::vector<std::size_t>\n"
f"Operator<{op_type}>::active_implementation_indices(Device::Type);"
)
active_implementation_definition = (
"template <>\n"
"std::vector<std::size_t>\n"
f"Operator<{op_type}>::active_implementation_indices("
"Device::Type dev_type) {\n"
f" return detail::ActiveImplementationIndices<{op_type}>(dev_type);\n"
"}"
)
_append_unique(active_implementation_declaration, active_implementation_definition)

default_implementation_declaration = (
"template <>\n"
"std::size_t\n"
f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);"
)
default_implementation_definition = (
"template <>\n"
"std::size_t\n"
f"Operator<{op_type}>::DefaultImplementationIndex("
"Device::Type dev_type) {\n"
f" return detail::DefaultImplementationIndex<{op_type}>(dev_type);\n"
"}"
)
_append_unique(
f"extern template std::size_t "
f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);",
f"template std::size_t "
f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);",
default_implementation_declaration, default_implementation_definition
)

for call in operator.calls:
Expand Down
96 changes: 59 additions & 37 deletions src/operator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,15 @@ struct CacheKeyBuilder {
template <typename Key, Device::Type kDev>
struct ActiveImplementations;

namespace detail {

template <typename Key>
std::vector<std::size_t> ActiveImplementationIndices(Device::Type dev_type);
template <typename Key>
std::size_t DefaultImplementationIndex(Device::Type dev_type);

} // namespace detail

class OperatorBase {
public:
virtual ~OperatorBase() = default;
Expand Down Expand Up @@ -297,20 +306,7 @@ class Operator : public OperatorBase {

static std::vector<std::size_t> active_implementation_indices(
Device::Type dev_type) {
if (!detail::ListContains(dev_type, ActiveDevices<Key>{})) {
return {};
}

std::vector<std::size_t> result;
DispatchFunc<ActiveDevices<Key>>(
dev_type,
[&](auto device_tag) {
constexpr Device::Type kDev = decltype(device_tag)::value;
result = detail::ListToVector(
typename ActiveImplementations<Key, kDev>::type{});
},
"Operator::active_implementation_indices");
return result;
return detail::ActiveImplementationIndices<Key>(dev_type);
}

template <typename... Args>
Expand All @@ -334,30 +330,8 @@ class Operator : public OperatorBase {
static constexpr std::size_t implementation_index_{implementation_index};

private:
template <auto first, auto... rest>
static constexpr std::size_t FirstActiveImplementationIndex(
List<first, rest...>) {
return static_cast<std::size_t>(first);
}

static std::size_t FirstActiveImplementationIndex(List<>) {
assert(false && "operator has no active implementation for this device");
std::abort();
}

static std::size_t DefaultImplementationIndex(Device::Type dev_type) {
std::size_t default_index{0};

DispatchFunc<ActiveDevices<Key>>(
dev_type,
[&](auto device_tag) {
constexpr Device::Type kDev = decltype(device_tag)::value;
default_index = FirstActiveImplementationIndex(
typename ActiveImplementations<Key, kDev>::type{});
},
"Operator::DefaultImplementationIndex");

return default_index;
return detail::DefaultImplementationIndex<Key>(dev_type);
}

static Config DefaultConfig(Device::Type dev_type) {
Expand Down Expand Up @@ -458,6 +432,54 @@ struct ActiveImplementations {
Key, kDev, std::make_index_sequence<kMaxImplementations>>::type;
};

namespace detail {

template <auto first, auto... rest>
constexpr std::size_t FirstActiveImplementationIndex(List<first, rest...>) {
return static_cast<std::size_t>(first);
}

inline std::size_t FirstActiveImplementationIndex(List<>) {
assert(false && "operator has no active implementation for this device");
std::abort();
}

template <typename Key>
std::size_t DefaultImplementationIndex(Device::Type dev_type) {
std::size_t default_index{0};

DispatchFunc<ActiveDevices<Key>>(
dev_type,
[&](auto device_tag) {
constexpr Device::Type kDev = decltype(device_tag)::value;
default_index = FirstActiveImplementationIndex(
typename ActiveImplementations<Key, kDev>::type{});
},
"Operator::DefaultImplementationIndex");

return default_index;
}

template <typename Key>
std::vector<std::size_t> ActiveImplementationIndices(Device::Type dev_type) {
if (!ListContains(dev_type, ActiveDevices<Key>{})) {
return {};
}

std::vector<std::size_t> result;
DispatchFunc<ActiveDevices<Key>>(
dev_type,
[&](auto device_tag) {
constexpr Device::Type kDev = decltype(device_tag)::value;
result =
ListToVector(typename ActiveImplementations<Key, kDev>::type{});
},
"Operator::active_implementation_indices");
return result;
}

} // namespace detail

} // namespace infini::ops

#endif
1 change: 1 addition & 0 deletions tests/test_cpp_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ def test_cpp_operator_call_instantiation_smoke(tmp_path):
[
_compiler("CXX", "c++"),
"-std=c++17",
"-O3",
"-Werror",
f"-I{include_dir}",
str(source),
Expand Down
36 changes: 28 additions & 8 deletions tests/test_generate_wrappers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,12 +97,22 @@ def test_operator_call_instantiations_externalize_default_implementation_lookup(
operator
)

signature = (
"std::size_t "
declaration = (
"template <>\n"
"std::size_t\n"
"Operator<::infini::ops::Abs>::DefaultImplementationIndex(Device::Type);"
)
assert f"extern template {signature}" in declarations
assert f"template {signature}" in definitions
definition = (
"template <>\n"
"std::size_t\n"
"Operator<::infini::ops::Abs>::DefaultImplementationIndex("
"Device::Type dev_type) {\n"
" return detail::DefaultImplementationIndex<::infini::ops::Abs>("
"dev_type);\n"
"}"
)
assert declaration in declarations
assert definition in definitions


def test_operator_call_instantiations_externalize_active_implementation_query():
Expand All @@ -113,12 +123,22 @@ def test_operator_call_instantiations_externalize_active_implementation_query():
operator
)

signature = (
"std::vector<std::size_t> "
declaration = (
"template <>\n"
"std::vector<std::size_t>\n"
"Operator<::infini::ops::Add>::active_implementation_indices(Device::Type);"
)
assert f"extern template {signature}" in declarations
assert f"template {signature}" in definitions
definition = (
"template <>\n"
"std::vector<std::size_t>\n"
"Operator<::infini::ops::Add>::active_implementation_indices("
"Device::Type dev_type) {\n"
" return detail::ActiveImplementationIndices<::infini::ops::Add>("
"dev_type);\n"
"}"
)
assert declaration in declarations
assert definition in definitions


def test_operator_call_instantiations_keep_scalar_and_optional_tensor_overloads_distinct(
Expand Down
Loading