Skip to content
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.

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

Adapt to API break where resolve_literal takes an InputInfo arg. #126

Merged
merged 1 commit into from
Sep 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
5 changes: 4 additions & 1 deletion shark_turbine/aot/builtins/jittable.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
GraphNodeImporter,
FxImporter,
FxImporterHooks,
InputInfo,
)

from ...support.ir_imports import (
Expand Down Expand Up @@ -77,7 +78,9 @@ def __init__(self, module_builder: ModuleBuilder):
# symbols we have done this to here.
self.cloned_global_symbols: set[str] = set()

def resolve_literal(self, gni: GraphNodeImporter, literal: Any) -> Optional[Value]:
def resolve_literal(
self, gni: GraphNodeImporter, literal: Any, info: Optional[InputInfo] = None
) -> Optional[Value]:
module_builder = self.module_builder
cloned_global_symbols = self.cloned_global_symbols

Expand Down
4 changes: 3 additions & 1 deletion shark_turbine/aot/support/procedural/exported_program.py
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,9 @@ def store_produced_value(
).result
util_d.GlobalStoreOp(converted_value, materialized_global.symbol_name)

def resolve_literal(self, gni: GraphNodeImporter, literal: Any) -> Optional[Value]:
def resolve_literal(
self, gni: GraphNodeImporter, literal: Any, info: Optional[InputInfo] = None
) -> Optional[Value]:
# We support resolution of tracked reference types. Currently this
# only includes Tensors. All others we let the importer do what it
# is going to do.
Expand Down
Loading