PReLU op

[LucQueen]

I added prelu op, and i run the code below, it is passed,

import torch
import torch_mlir
m = torch.nn.PReLU()
x = torch.randn([3,4,3])
res = m(x)
print(res.shape)
print("success")
module = torch_mlir.compile(m, x, output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, use_tracing=True)

but i get error, when i run the other code

import torch

import torch_mlir

class PReLU(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x, w):
        return torch.nn.functional.prelu(x, w)

m = PReLU()
x = torch.randn([3,4,3])
w = torch.randn([4])
res = m(x, w)
module = torch_mlir.compile(m, [x, w], output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, use_tracing=True)

module = torch_mlir.compile(m, [x, w], output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, use_tracing=True)
  File "/mnt/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir/torch_mlir/__init__.py", line 194, in compile
    run_pipeline_with_repro_report(
  File "/mnt/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir/torch_mlir/compiler_utils.py", line 49, in run_pipeline_with_repro_report
    raise Exception(f"""
Exception:
Lowering Torch Backend IR -> Linalg-on-Tensors Backend IR failed with the following diagnostics:
error: 'linalg.generic' op inferred input/output operand #1 has shape's dimension #0 to be 3, but found 4
note: see current operation: %13 = "linalg.generic"(%0, %1, %12) ({
^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
  %16 = "arith.constant"() {value = 0.000000e+00 : f32} : () -> f32
  %17 = "arith.cmpf"(%arg2, %16) {predicate = 9 : i64} : (f32, f32) -> i1
  %18 = "arith.select"(%17, %arg2, %16) : (i1, f32, f32) -> f32
  %19 = "arith.select"(%17, %16, %arg2) : (i1, f32, f32) -> f32
  %20 = "arith.mulf"(%16, %19) : (f32, f32) -> f32
  %21 = "arith.addf"(%18, %20) : (f32, f32) -> f32
  "linalg.yield"(%21) : (f32) -> ()
}) {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"], operand_segment_sizes = dense<[2, 1]> : vector<2xi32>} : (tensor<3x4x3xf32>, tensor<4xf32>, tensor<3x4x3xf32>) -> tensor<3x4x3xf32>


Error can be reproduced with:
$ torch-mlir-opt -pass-pipeline='torch-backend-to-linalg-on-tensors-backend-pipeline' /tmp/PReLU.mlir
Add '-mlir-print-ir-after-all -mlir-disable-threading' to get the IR dump for debugging purpose.

i am very confused，can someone help me？

[ramiro050]

Hi @LucQueen,

I just wanted to let you know that we have torch-MLIR office hours every Thursday where I can show you my approach for debugging issues like this one. If the time doesn't work you for you, I'm happy to help through github issues. I just wanted to make sure you were aware of that option! 😄