Stablehlo compiler #338
Draft
Stablehlo compiler #338
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add COMPILE_STABLEHLO_OP_BY_OP CompileDepth. Allow compilation/ execution starting from stablehlo.
|
||||||||||||||
❌ 4 Tests Failed:
View the top 3 failed test(s) by shortest run time
To view more test analytics, go to the Test Analytics Dashboard |
env/activate
Outdated
| fi | ||
| pip install $TT_TORCH_HOME/dist/torchvision*.whl | ||
| pip install --pre torch-mlir torchvision |
There was a problem hiding this comment.
what's torchvision for here?
tt_torch/tools/utils.py
Outdated
| @@ -275,7 +276,7 @@ def post_init(self): | |||
| else: | |||
| torch._dynamo.config.inline_inbuilt_nn_modules = True | |||
|
|
|||
| def save_unique_ops(self): | |||
| def save_unique_ops(self, mode=None): | |||
There was a problem hiding this comment.
perhaps default to mode="torch" as opposed to None.
There was a problem hiding this comment.
I could do that. Just want to confirm, by default we want the files to be named {self.results_path}{pytest_test}_torch_unique_ops.json or {self.results_path}{pytest_test}_unique_ops.json then? Right now, the way it is:
default --> {self.results_path}{pytest_test}_unique_ops.json
torch --> {self.results_path}{pytest_test}_torch_unique_ops.json
stablehlo --> {self.results_path}{pytest_test}_stablehlo_unique_ops.json
| executor.add_gm(gm, graph_constants) | ||
| return executor | ||
|
|
||
|
|
There was a problem hiding this comment.
This section is duplicated. Can you move it into a helper function?
There was a problem hiding this comment.
I think my latest commit should address this. By complicated, I assumed you meant the isinstance checks. I agree. I now check and assign parsed_module when initializing the Executor object.
| self.gm = gm | ||
| self.graph_constants = tuple(graph_constants) | ||
|
|
||
| def gm_op_by_op(self, *inputs): |
There was a problem hiding this comment.
So this is actually the same function from torch_backend, but stripped. This is actually going through a torch graph op by op, and only called during COMPILE_STABLEHLO_OP_BY_OP, since it's implied the input is a torch graph. Doing so, we create 2 json dumps - one for stablehlo and one for torch, during one run.
| # No conversion required. | ||
| new_inputs = new_inputs + ((input),) | ||
| inputs = new_inputs | ||
| if self.compiler_config.compile_depth == CompileDepth.EXECUTE: |
There was a problem hiding this comment.
isn't this handled by base backed?
There was a problem hiding this comment.
This is base_backend, which returns executor. This is in parallel with torch executor.
def _base_backend(gm_or_shlo, example_inputs, compiler_config):
# Called during EXECUTE
# input is a torch graph
if isinstance(gm_or_shlo, torch.fx.GraphModule):
shlo, executor, gm, graph_constants = torch_to_shlo(
gm_or_shlo, example_inputs, compiler_config
)
# input is a stablehlo string module
elif isinstance(gm_or_shlo, str):
shlo = parse_module_from_str(gm_or_shlo)
executor = StablehloExecutor(
parsed_module=shlo, compiler_config=compiler_config
)
else:
assert False, "Compiler input not valid"
binary = shlo_to_flatbuffer(shlo, compiler_config)
executor.set_binary(binary)
return executor
|
||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Ticket
#203
Problem description
Currently, we only support compilation or execution of torch graphs. We do not support compilation of stablehlo graphs because this would fail at the first unsupported stablehlo op. The solution is to create a module for each op. This can be done through stablehlo support.
What's changed
_base_backend-->_torch_backend_base_backendis for execution, for other compile depths using_torch_backendor_shlo_backendtt_torch/dynamo/backend.pyintott_torch/dynamo/backend.py,tt_torch/dynamo/shlo_backend.py,tt_torch/dynamo/torch_backend.pytt_torch/dynamo/backend.pycontents are mostly moved tott_torch/dynamo/torch_backend.py,tt_torch/dynamo/executor.pytt_torch/dynamo/executor.pyprovides base executor class to be used by both stablehlo and torch backendsenv/activateedited to add support for dependenciesTesting strategy: