llvm / torch-mlir

The Torch-MLIR project aims to provide first class support from the PyTorch ecosystem to the MLIR ecosystem.
Other
1.3k stars 480 forks source link

Unimplemented operator 'aten.upsample_nearest2d.vec' #2679

Open luo-yicong opened 9 months ago

luo-yicong commented 9 months ago

I encountered the following problem when converting vae decoder, part of stable diffusion

python exception: Failure while executing pass pipeline:
error: "__module.up_blocks.0/__module.up_blocks.0.upsamplers.0/aten::upsample_nearest2d"("/usr/local/lib/python3.11/dist-packages/torch/nn/functional.py":4001:0): unsupported by backend contract: Unimplemented operator 'aten.upsample_nearest2d.vec'
note: "__module.up_blocks.0/__module.up_blocks.0.upsamplers.0/aten::upsample_nearest2d"("/usr/local/lib/python3.11/dist-packages/torch/nn/functional.py":4001:0): see current operation: %498 = "torch.operator"(%496, %2, %497) <{name = "aten.upsample_nearest2d.vec"}> : (!torch.tensor<[2,512,64,64],f32>, !torch.none, !torch.list<float>) -> !torch.tensor<[2,512,128,128],f32>

The script is as follows:

import torch
import torch_mlir
from diffusers import StableDiffusionPipeline
model_id = "/data/models/stable-diffusion/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)

vae_decoder = pipe.vae.decoder
vae_decoder.eval()
module = torch_mlir.compile(vae_decoder, torch.rand(2,4,64,64), output_type=torch_mlir.OutputType.LINALG_ON_TENSORS,use_tracing=True)
with open("./pytorch_vae_decoder.mlir", "w", encoding="utf-8") as outf:
    outf.write(str(module))
ramiro050 commented 8 months ago

The error means that there is no conversion to linalg for aten.upsample_nearest2d.vec. A conversion pattern needs to be added to fix it.