Error occurs when Lowering Torch Backend IR -> StableHLO Backend IR
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 1.9k
- Forks
- 736
- Avg merge
- 5d 22h
- Merged PRs (30d)
- 15
Description
When I try to generate the forward and backward graphs of a simple model in StableHLO IR format, the forward graph can be converted correctly, but the following error occurs when converting the backward graph:
```shell
module {
func.func @M__0_forward_1(%arg0: tensor<10x32xf32>, %arg1: tensor<10xf32>, %arg2: tensor<8x8xf32>) -> (tensor<2x10xf32>, tensor<2x32xf32>) {
%cst = arith.constant dense<1> : tensor<1xi64>
%0 = stablehlo.reshape %arg2 : (tensor<8x8xf32>) -> tensor<2x32xf32>
%1 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<10x32xf32>) -> tensor<32x10xf32>
%2 = stablehlo.dot_general %0, %1, contracting_dims = [1] x [0] : (tensor<2x32xf32>, tensor<32x10xf32>) -> tensor<2x10xf32>
%3 = stablehlo.convert %cst : (tensor<1xi64>) -> tensor<1xf32>
%4 = stablehlo.reshape %3 : (tensor<1xf32>) -> tensor
%5 = stablehlo.broadcast_in_dim %arg1, dims = [0] : (tensor<10xf32>) -> tensor<10xf32>
%6 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor) -> tensor<10xf32>
%7 = stablehlo.multiply %5, %6 : tensor<10xf32>
%8 = stablehlo.broadcast_in_dim %7, dims = [1] : (tensor<10xf32>) -> tensor<2x10xf32>
%9 = stablehlo.broadcast_in_dim %2, dims = [0, 1] : (tensor<2x10xf32>) -> tensor<2x10xf32>
%10 = stablehlo.add %8, %9 : tensor<2x10xf32>
return %10, %0 : tensor<2x10xf32>, tensor<2x32xf32>
}
}
error: failed to legalize operation 'torch.constant.none'
error: Module does not conform to the Stablehlo backend contract. See dialect conversion legality information above.
Traceback (most recent call last):
File "/home/tomjiang/test/test-mlp/issue.py", line 54, in
loss.backward()
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch/_tensor.py", line 522, in backward
torch.autograd.backward(
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch/autograd/init.py", line 346, in backward
_engine_run_backward(
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch/autograd/graph.py", line 806, in _engine_run_backward
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch/autograd/function.py", line 306, in apply
return user_fn(self, *args)
^^^^^^^^^^^^^^^^^^^^
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1976, in backward
out = call_compiled_backward()
^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1908, in call_compiled_backward
CompiledFunction.compiled_bw = aot_config.bw_compiler(
^^^^^^^^^^^^^^^^^^^^^^^
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch/_dynamo/backends/common.py", line 47, in _wrapped_bw_compiler
return disable(disable(bw_compiler)(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 602, in _fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py", line 103, in f
out_f = compiler(fx_g, inps)
^^^^^^^^^^^^^^^^^^^^
File "/home/tomjiang/test/test-mlp/issue.py", line 32, in fx_import_aot_autograd_backend
m = fx.stateless_fx_import(gm, model_name=get_aot_graph_name(), output_type=OutputType.STABLEHLO)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch_mlir/fx.py", line 116, in stateless_fx_import
return _module_lowering(
^^^^^^^^^^^^^^^^^
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch_mlir/fx.py", line 47, in _module_lowering
return lower_mlir_module(verbose, output_type, torch_mod)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch_mlir/compiler_utils.py", line 167, in lower_mlir_module
run_pipeline_with_repro_report(
File "/home/tomjiang/miniconda3/envs/torch-mlir/lib/python3.11/site-packages/torch_mlir/compiler_utils.py", line 78, in run_pipeline_with_repro_report
raise TorchMlirCompilerError(trimmed_message) from None
torch_mlir.compiler_utils.TorchMlirCompilerError: Lowering Torch Backend IR -> StableHLO Backend IR failed with the following diagnostics:
python exception: Failure while executing pass pipeline
For Torch-MLIR developers, the error can be reproduced with:
$ torch-mlir-opt -pass-pipeline='builtin.module(torch-backend-to-stablehlo-backend-pipeline)' /tmp/UnnammedModule.mlir
Add '-mlir-print-ir-after-all -mlir-disable-threading' to get the IR dump for debugging purpose.
```
Code is:
```python
from typing import List
import torch
import torch.nn as nn
from torch._dynamo.backends.common import aot_autograd
from torch._functorch.aot_autograd import (
make_boxed_compiler,
get_aot_graph_name,
set_model_name,
)
from torch_mlir import fx
from torch_mlir.compiler_utils import OutputType
import torch_mlir
class M(nn.Module):
def init(self):
super().init()
self.l1 = nn.Linear(32, 10)
def forward(self, x):
x = x.view(-1, 32)
return self.l1(x)
@make_boxed_compiler
def fx_import_aot_autograd_backend(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
# print(gm.print_readable(False), flush=True)
# print(gm.code)
m = fx.stateless_fx_import(gm, model_name=get_aot_graph_name(), output_type=OutputType.STABLEHLO)
print(m, flush=True)
return gm
if name == 'main':
model = M()
criterion = nn.CrossEntropyLoss()
lr = 0.001
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
fx_import_backend = aot_autograd(fw_compiler=fx_import_aot_autograd_backend)
set_model_name("M")
model_opt = torch.compile(model, backend=fx_import_backend)
out = model_opt(torch.randn(8, 8))
labels = torch.randn(2, 10)
loss = criterion(out, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
```
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the reproduction in /home/tomjiang/test/test-mlp/issue.py, then run the reported torch-mlir-opt torch-backend-to-stablehlo-backend-pipeline command with IR dumping enabled. Trace the backward graph through fx.stateless_fx_import and the lowering pipeline; done means the backward graph converts without the torch.constant.none legalization failure and satisfies the StableHLO backend contract.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100