How to Represent PyTorch's var_mean and getitem Operators with MLIR's Builtin Dialect?

MLIR
1

Hello everyone!
I am currently trying to represent PyTorch’s var_mean and getitem operators using MLIR Dialect, but I am facing some difficulties and would appreciate your guidance.

Specifically, the var_mean operator calculates the variance and mean of a tensor along a specified dimension, and then returns a tuple containing these two results. It’s important to note that the variance and mean returned by var_mean are also tensors. The getitem operator, on the other hand, retrieves a value from a tuple, tensor, etc., based on an index value. In my actual project scenario, getitem is used to retrieve a value from the result returned by the var_mean operator. It either retrieves the variance or the mean.

My confusion lies in how to represent these two operators with MLIR Dialect. I noticed that there is TupleType in MLIR, but I couldn’t find how to retrieve a value from a TupleType.

Any guidance or suggestions would be greatly appreciated. Thank you in advance for your help! :slight_smile:

2

There are intentionally no operations on the tuple type. Dialects are expected to define their own. So just go ahead an implement getitem the way you see fit for your project.

Also note that MLIR operations may have multiple results, so there is no need for tuple for MLIR purposes only for modeling external abstractions.

3

Yeah multiple results is the way to go.

Also have you looked at Torch-MLIR’s implementation of the VarMean op? It also implements the op with two results.

4

Thank you for your reply! I will take a look at how to use the multiple results feature :slight_smile:

5

Yes, I‘ve noticed the VarMean op of Torch-MLIR, but it seems that it’s implemented based on their own Torch dialect. For some reasons, in my project, I can only use MLIR’s builtin dialect. But I think I should be able to use the lowered code of Torch dialect as a reference. :smiley: