alibaba / alibaba/BladeDISC

Introduce MLIR transform dialect to BladeDISC

Open
#787 2 comments 1 reaction 0 assignees View on GitHub
Transform Dialect
Dominant language
C++
Stars
933
Forks
169
PR merge metrics
No merged PRs in 30d

Description

We'll start to explore using MLIR [transform dialect ](https://mlir.llvm.org/docs/Dialects/Transform) to do codegen for (fused) compute-intensive pattern. The initial target is to support gemm codegen on ARM platform to address the dynamic shape problem of Arm Compute Library.

The initial plan is:

- [x] Step 1, enhance the fusion decision pass. We’ll add a new fusion kind `kTransform` for the transform-based fusion pattern.
- [x] Step 2, lower the lmhlo fusion op to linalg on tensor.
- [x] Step 3, transform the linalg computation to loops using transform dialect.
- [x] Step 4, refined the transformed loop to make it suitable for BladeDISC runtime.
- [x] Step 5, add a new pass to the disc pass pipeline to drive the above process.
- [x] Step 6, weight pre-packing support
- [x] add `disc_linalg.multi_level_pack` op, used for doing packing.
- [x] add `transform.disc.cache_read` transform op, relying on `disc_linalg.multi_level_pack` op.
- [x] add folding support for `disc_linalg.multi_level_pack`.
- [x] lower `disc_linalg.multi_level_pack` to loop if it can not be folded.
- [x] fuse const weight op into the `kTransform` fusion pattern, lower it to linalg and then schedule it.
- [x] Step 7, assign a default schedule for each `kTransform` pattern.
- [x] Step 8, schedule selection logic injection
- [x] Step 9, initial model level testing: bert (albert).
- [x] Step 10, support nt, tn, tt format GEMM.
- [ ] Step 11, support batch matmul
- [x] Step 12, support GEMM epilogue fusion.
- [ ] Step 13, performance optimization

Contributor guide

No contributing guide indexed for this repository

Research direction

No source files or tests are named. Start by tracing the fusion decision pass and the disc pass pipeline, then review the lmhlo-to-linalg and MLIR transform-dialect stages described in the plan. The remaining scope is batch matmul support and performance optimization, but the issue does not define completion criteria for either.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
compilers, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
20/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.