iree-org / iree-org/iree

Transform matmul ops to transpose_matmul before dispatch region creation

Open
#6,512 5 comments 0 reactions 0 assignees View on GitHub
codegen
Dominant language
C++
Stars
3.9k
Forks
1k
Avg merge
4d 16h
Merged PRs (30d)
47

Description

In order to get better memory access patterns we should convert A*B into two dispatch regions, (At = transpose(A)) followed by (transpose(At)*B) .
When we vectorize this allows us to do block load along M and N dimensions, do outer product and accumulate on each SIMD lane.
The steps to get there are:
* Add a new named op for to represent transpose(A)*B
* Add a pass before dispatch region creation and convert all matmul to this format.
* Add a second pattern to convert transpose(A*B) into transpose(B)*transpose(A) in order to try to cancel some of the transpose.

The assumption is that it should benefit both CPU and GPU as they both use outer product to lower the contract.
This should be parallel to the mmt4d effort and should only apply to cases where we don't use the mmt4d route. If we enable mmt4d on all target this may become irrelevant but this will be useful in the meantime and will allow us to do a more meaningful comparison.

Contributor guide

Open the contributing guide

Research direction

Start by locating the matmul handling and the dispatch region creation pass described in the issue. Read the parallel mmt4d effort and trace where non-mmt4d matmul cases are processed. Done means the named transpose_matmul operation and the requested conversion patterns are implemented before dispatch region creation, with transpose cancellation covered.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.