Skip to content

fix: Add generate wrapper change #796

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
baominghelly wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix: Add generate wrapper change #796

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
132 changes: 116 additions & 16 deletions scripts/generate_wrappers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,6 +242,12 @@ def __call__(self, op_name):
elif node.kind == CursorKind.CXX_METHOD and node.spelling == "operator()":
calls.append(node)

parsed_operator = _parse_operator_header(op_name)
constructors = _prefer_header_type_spellings(
constructors, parsed_operator.constructors
)
calls = _prefer_header_type_spellings(calls, parsed_operator.calls)

return _Operator(op_name, constructors, calls)

@staticmethod
Expand Down Expand Up @@ -275,6 +281,27 @@ def _parse_operator_header(op_name):
return _Operator(op_name, constructors, calls)


def _prefer_header_type_spellings(functions, header_functions):
header_by_names = {}

for function in header_functions:
names = tuple(arg.spelling for arg in function.get_arguments())
header_by_names.setdefault(names, []).append(function)

aligned = []

for function in functions:
names = tuple(arg.spelling for arg in function.get_arguments())
matches = header_by_names.get(names, ())

if len(matches) == 1:
aligned.append(matches[0])
else:
aligned.append(function)

return aligned


def _strip_cpp_comments(source):
source = re.sub(r"/\*.*?\*/", "", source, flags=re.DOTALL)
return re.sub(r"//.*", "", source)
Expand Down Expand Up @@ -443,6 +470,22 @@ def _find_vector_int64_params(op_name):
return set(re.findall(r"std::vector<int64_t>\s+(\w+)", source))


def _find_optional_vector_int64_params(op_name):
"""Return params declared as `std::optional<std::vector<int64_t>>`.

libclang can report these as plain `int`, just like non-optional vector
params, so scan the base header to preserve the generated API signature.
"""
source = _find_base_header(op_name).read_text()

return set(
re.findall(
r"std::optional\s*<\s*std::vector\s*<\s*int64_t\s*>\s*>\s+(\w+)",
source,
)
)


def _find_tensor_params(op_name):
source = _find_base_header(op_name).read_text()

Expand Down Expand Up @@ -474,6 +517,7 @@ def _generate_pybind11(operator):
optional_non_tensor_params = _find_optional_non_tensor_params(operator.name)
vector_tensor_params = _find_vector_tensor_params(operator.name)
vector_int64_params = _find_vector_int64_params(operator.name)
optional_vector_int64_params = _find_optional_vector_int64_params(operator.name)
params_with_defaults = _find_params_with_defaults(operator.name)

def _is_optional_tensor(arg):
Expand All @@ -490,6 +534,9 @@ def _is_optional_tensor(arg):
def _is_optional(arg):
return "std::optional" in arg.type.spelling

def _is_optional_vector_int64(arg):
return arg.spelling in optional_vector_int64_params

def _is_vector_tensor(arg):
if arg.spelling in vector_tensor_params:
return True
Expand All @@ -510,6 +557,10 @@ def _generate_params(node):
parts.append(f"std::optional<py::object> {arg.spelling}")
elif _is_vector_tensor(arg):
parts.append(f"std::vector<py::object> {arg.spelling}")
elif _is_optional_vector_int64(arg):
parts.append(
f"const std::optional<std::vector<int64_t>> {arg.spelling}"
)
elif _is_vector_int64(arg):
parts.append(f"const std::vector<int64_t> {arg.spelling}")
else:
Expand Down Expand Up @@ -908,6 +959,7 @@ def _generate_generated_dispatch_entries(operator):
tensor_params = _find_tensor_params(operator.name)
vector_tensor_params = _find_vector_tensor_params(operator.name)
vector_int64_params = _find_vector_int64_params(operator.name)
optional_vector_int64_params = _find_optional_vector_int64_params(operator.name)

def _is_optional_tensor(arg):
spelling = arg.type.spelling
Expand All @@ -920,6 +972,9 @@ def _is_optional_tensor(arg):

return arg.spelling in optional_tensor_params

def _is_optional_vector_int64(arg):
return arg.spelling in optional_vector_int64_params

def _is_vector_tensor(arg):
if arg.spelling in vector_tensor_params:
return True
Expand All @@ -933,10 +988,13 @@ def _is_tensor(arg):
if arg.spelling in optional_non_tensor_params:
return False

if arg.spelling in tensor_params:
if "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling:
return True

return "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling
if _is_known_non_tensor_type(arg.type.spelling):
return False

return arg.spelling in tensor_params

def _generate_params(node):
parts = []
Expand All @@ -949,6 +1007,8 @@ def _generate_params(node):
parts.append(f"std::optional<Tensor> {arg.spelling}")
elif _is_vector_tensor(arg):
parts.append(f"std::vector<Tensor> {arg.spelling}")
elif _is_optional_vector_int64(arg):
parts.append(f"std::optional<std::vector<int64_t>> {arg.spelling}")
elif _is_vector_int64(arg):
parts.append(f"std::vector<int64_t> {arg.spelling}")
elif _is_tensor(arg):
Expand Down Expand Up @@ -1108,12 +1168,32 @@ def _strip_top_level_const(type_spelling):
return type_spelling


def _is_known_non_tensor_type(type_spelling):
type_spelling = _strip_top_level_const(type_spelling)

if "Tensor" in type_spelling or "TensorView" in type_spelling:
return False

if type_spelling in {
"bool",
"double",
"float",
"int64_t",
"std::size_t",
"std::string",
}:
return True

return type_spelling.startswith("std::optional<")


def _generate_operator_call_instantiation_entries(operator):
optional_tensor_params = _find_optional_tensor_params(operator.name)
optional_non_tensor_params = _find_optional_non_tensor_params(operator.name)
tensor_params = _find_tensor_params(operator.name)
vector_tensor_params = _find_vector_tensor_params(operator.name)
vector_int64_params = _find_vector_int64_params(operator.name)
optional_vector_int64_params = _find_optional_vector_int64_params(operator.name)

def _is_optional_tensor(arg):
spelling = arg.type.spelling
Expand All @@ -1129,6 +1209,9 @@ def _is_optional_tensor(arg):

return False

def _is_optional_vector_int64(arg):
return arg.spelling in optional_vector_int64_params

def _is_vector_tensor(arg):
if arg.spelling in vector_tensor_params:
return True
Expand All @@ -1144,10 +1227,13 @@ def _is_tensor(arg):
if arg.spelling in optional_non_tensor_params:
return False

if arg.spelling in tensor_params:
if "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling:
return True

return "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling
if _is_known_non_tensor_type(arg.type.spelling):
return False

return arg.spelling in tensor_params

def _normalized_type(arg):
if _is_optional_tensor(arg):
Expand All @@ -1156,6 +1242,9 @@ def _normalized_type(arg):
if _is_vector_tensor(arg):
return "std::vector<Tensor>"

if _is_optional_vector_int64(arg):
return "std::optional<std::vector<int64_t>>"

if _is_vector_int64(arg):
return "std::vector<int64_t>"

Expand Down Expand Up @@ -1188,6 +1277,16 @@ def _append_optional_params(prefix, params):
declarations = []
definitions = []

seen_entries = set()

def _append_unique(declaration, definition):
if definition in seen_entries:
return

seen_entries.add(definition)
declarations.append(declaration)
definitions.append(definition)

for call in operator.calls:
template_arguments = _generate_template_arguments(call)
params = _generate_parameters(call)
Expand All @@ -1202,14 +1301,13 @@ def _append_optional_params(prefix, params):
)

make_template_arguments = ", ".join(
f"const {_strip_top_level_const(arg.type.spelling)}&" for arg in rest_args
f"const {_normalized_type(arg)}&" for arg in rest_args
)
make_params = ", ".join(
f"const {_strip_top_level_const(arg.type.spelling)}& {arg.spelling}"
for arg in rest_args
f"const {_normalized_type(arg)}& {arg.spelling}" for arg in rest_args
)
make_function_params = _append_optional_params(
f"const Config& config, const {_strip_top_level_const(first_arg.type.spelling)} {first_arg.spelling}",
f"const Config& config, const {_normalized_type(first_arg)} {first_arg.spelling}",
make_params,
)
make_instantiation = f"Operator<{op_type}>::Make<{make_template_arguments}>({make_function_params})"
Expand All @@ -1218,16 +1316,18 @@ def _append_optional_params(prefix, params):
f"Operator<{op_type}>::operator()<{template_arguments}>({params}) const"
)

declarations.append(f"extern template void {instantiation};")
declarations.append(
f"extern template std::unique_ptr<Operator<{op_type}>> {make_instantiation};"
_append_unique(
f"extern template void {instantiation};",
f"template void {instantiation};",
)
declarations.append(f"extern template void {operator_instantiation};")
definitions.append(f"template void {instantiation};")
definitions.append(
f"template std::unique_ptr<Operator<{op_type}>> {make_instantiation};"
_append_unique(
f"extern template std::unique_ptr<Operator<{op_type}>> {make_instantiation};",
f"template std::unique_ptr<Operator<{op_type}>> {make_instantiation};",
)
_append_unique(
f"extern template void {operator_instantiation};",
f"template void {operator_instantiation};",
)
definitions.append(f"template void {operator_instantiation};")

return declarations, definitions

Expand Down
65 changes: 65 additions & 0 deletions tests/test_generate_wrappers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,71 @@ class Clamp {
) in text


def test_extractor_prefers_header_types_for_reused_parameter_names(
monkeypatch, tmp_path
):
module = _load_generator_module()
base_header = tmp_path / "histogram.h"
base_header.write_text(
"""
class Histogram {
public:
virtual void operator()(const Tensor input, const Tensor bins,
const std::optional<Tensor> weight,
const bool density, Tensor hist,
Tensor bin_edges) const = 0;
virtual void operator()(const Tensor input, const int64_t bins,
const std::optional<std::vector<double>> range,
const std::optional<Tensor> weight,
const bool density, Tensor hist,
Tensor bin_edges) const = 0;
};
"""
)
monkeypatch.setattr(module, "_find_base_header", lambda op_name: base_header)

clang_calls = [
module._ParsedFunction(
[
module._ParsedArgument("const int", "input"),
module._ParsedArgument("const int", "bins"),
module._ParsedArgument("const int", "weight"),
module._ParsedArgument("const bool", "density"),
module._ParsedArgument("int", "hist"),
module._ParsedArgument("int", "bin_edges"),
]
),
module._ParsedFunction(
[
module._ParsedArgument("const int", "input"),
module._ParsedArgument("const int64_t", "bins"),
module._ParsedArgument("const int", "range"),
module._ParsedArgument("const int", "weight"),
module._ParsedArgument("const bool", "density"),
module._ParsedArgument("int", "hist"),
module._ParsedArgument("int", "bin_edges"),
]
),
]

operator = module._Operator(
"histogram",
constructors=[],
calls=module._prefer_header_type_spellings(
clang_calls, module._parse_operator_header("histogram").calls
),
)

declarations, _ = module._generate_operator_call_instantiation_entries(operator)
text = "\n".join(declarations)

assert (
"Call<Tensor, int64_t, std::optional<std::vector<double>>, "
"std::optional<Tensor>, bool, Tensor, Tensor>"
) in text
assert "Call<Tensor, Tensor, int, std::optional<Tensor>" not in text


def test_pybind_default_implementation_uses_first_active_index(monkeypatch, tmp_path):
module = _load_generator_module()
base_header = tmp_path / "mul.h"
Expand Down
Loading