Mlir Backend: Tensor Dialect Support

[liamsemeria]

Motivation

Support for ops in the tensor dialect allows for tracking of producer-consumer relationships and broadcasting, which allow for operator fusion and element-wise operations respectively.

Description

The mlir backend now has an option use_tensor_dialect that causes ops to be generated in the tensor dialect. The tensor dialect gets lowered into memref by a new bufferization pass after the transform pass is applied (can be printed with print_bufferization_ir=True).

Discussion

How the Tensor Dialect Affects the IR:

• matmul and conv2d:
  The bufferization results in the exact same lowered mlir as the memref dialect ops (at least for unscheduled).

• relu:
  Collapsing the shape to 1 dim requires the tensor to be expanded (unlike the memref), resulting in an extra memory allocation after bufferization. So the relu for the tensor dialect is non-collapsing, which is also required for consumer fusion to work properly in the future.

• pad and unpad:
  The tensor implementation uses a linalg.generic which is needed for fusion. It has dynamic dims which requires mlir: updated extra-tools version #70 an update to the extra tools for the c backend to work properly.

[qaco]

It is really cool ! Congrats !

When you say that you wait for an xDSL release, is it because you know that the feature will be added ?

[liamsemeria]

It is really cool ! Congrats !

When you say that you wait for an xDSL release, is it because you know that the feature will be added ?

Thanks! yeah I contributed tensor.pad but its probably easiest to just wait for the next release to add it to xtc.

[qaco]

I will need some time to review this! Pls ping me when it's ready!

[liamsemeria]

I will need some time to review this! Pls ping me when it's ready!

@qaco Sorry for the delay during my break I thought I had an idea to make the pad better but turns out it wouldn't work as I expected. The PR is ready to review. I'm currently working on making the compile time not insanely terrible but the majority of the tensor dialect PR would be unchanged.