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.

[guillon]

It's fine for me, @liamsemeria you can merge when ready.

Note for future work:

• may use the sdsl pad directly
• need to study performance issues with one-shot-bufferize

If there are further devs to do w.r.t. these points, we will make new PRs.

[liamsemeria]

@guillon I dont have the option to merge on my end, but I updated I rebased it with main and its ready to merge.

[guillon]

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

@qaco The PR is ready to merge, can you check?

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 93.54839% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/xtc/cli/explore.py	44.44%	5 Missing ⚠️
src/xtc/backends/mlir/MlirNodeBackend.py	33.33%	4 Missing ⚠️
src/xtc/backends/mlir/MlirOps.py	98.64%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[guillon]

@liamsemeria why do we need an additional setting using_tensors_hint ?
Shouldn't it be the default?

[liamsemeria]

@liamsemeria why do we need an additional setting using_tensors_hint ? Shouldn't it be the default?

@guillon It's needed to enable tensor specific changes to the compiler passes which are moving vectorization lowering patterns to after bufferization and using the FoldUnitDims pattern.

If it was enabled by default the major changes would be that there would be a lingering transform.sequence (the post vectorization patterns) in the IR Dump After Transform output of each test, and there is an annotation to the function that gets added in order to apply the vector patterns as well. Both are visible in this test. For the function annotation I was planning on double checking to see if there's a way around doing it at some point.

A way the using_tensors_hint can be removed is by doing a quick traversal of the IR to detect if tensors or memref are used. That way the changes I mentioned above would only apply to the tensor tests and not the memref tests. It could be done in the python bindings by looking at the function operands to see what dialect they are.

I can explain it more tomorrow if you need any clarifications.

[qaco]

Great job!

[liamsemeria]

@guillon @qaco its ready to merge unless you guys have any other changes in mind. Thanks for taking a look.

[guillon]

@guillon @qaco its ready to merge unless you guys have any other changes in mind. Thanks for taking a look.

I need to understand first the using_tensors_hint changes and the discussion on the buffer/vectorize post passes, I do not catch yet and I do not think it's good to have an additional flag which seems not "user" oriented.

[qaco]

@guillon @qaco its ready to merge unless you guys have any other changes in mind. Thanks for taking a look.

I need to understand first the using_tensors_hint changes and the discussion on the buffer/vectorize post passes, I do not catch yet and I do not think it's good to have an additional flag which seems not "user" oriented.

Should we make the integration of tensor transparent ? Like : if nothing happens at tensor level, the tensors are immediately bufferized, and then the buffer-level transformations occur

[guillon]

@guillon @qaco its ready to merge unless you guys have any other changes in mind. Thanks for taking a look.

I need to understand first the using_tensors_hint changes and the discussion on the buffer/vectorize post passes, I do not catch yet and I do not think it's good to have an additional flag which seems not "user" oriented.

Should we make the integration of tensor transparent ? Like : if nothing happens at tensor level, the tensors are immediately bufferized, and then the buffer-level transformations occur

Good point. Do you think it's feasible to have a single transform pass @liamsemeria?