Skip to content

Rewrite ParallelDimensionMap with expr simplifier#86

Merged
zasdfgbnm merged 6 commits intomainfrom
parallel-dim
Mar 29, 2023
Merged

Rewrite ParallelDimensionMap with expr simplifier#86
zasdfgbnm merged 6 commits intomainfrom
parallel-dim

Conversation

@zasdfgbnm
Copy link
Collaborator

@zasdfgbnm zasdfgbnm commented Mar 28, 2023

I am seeing trivial thread predicate in one matmul schedule #87, and that trivial thread predicate causes a 10% slowdown. So I start my work on removing that predicate. My first step is a rewrite of ParallelDimensionMap. The new code is much shorter in lines of code, and should have better result (for example, in FusionParallelDimensionMap3_CUDA, blockDim.x becomes 20).

This PR uses the following equation to calculate the extent of a parallel dim:

parallel_dim = simplifyExpr(max(extent1, extent2, ..., extentN));

Future work:

  • Simplify trivial predicate in Implement with split merge mmigdal-nv/Fuser#1, what we need to do is:
    • When simplifying all expressions, we should assume parallel_index < parallel_dim in additional to
      axioms.push_back(SimplifyingIrBuilder::ltExpr(pidx, pdim));
    • Omit the predicate not only when isExact(ptype), but also when the loop extent equals parallel_dim.
  • Simplify the following to use the extent computed in this PR:
    for (auto& entry : simplified_parallel_iter_extents) {

@zasdfgbnm zasdfgbnm marked this pull request as ready for review March 28, 2023 22:08
@zasdfgbnm zasdfgbnm requested a review from naoyam March 28, 2023 22:08
}

const auto tidx_pt = ParallelType::TIDx;
auto warp_size = at::cuda::warp_size();
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

at::cuda::warp_size(); was added to support ROCm, which we no longer care anymore.

@zasdfgbnm zasdfgbnm changed the title Rewrite ParallelDimensionMap with expr simplifier Rewrite ParallelDimensionMap with expr simplifier Mar 28, 2023
Copy link
Collaborator

@naoyam naoyam left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice cleanup!

@zasdfgbnm zasdfgbnm requested a review from naoyam March 29, 2023 17:28
@zasdfgbnm
Copy link
Collaborator Author

!build

Copy link
Collaborator

@naoyam naoyam left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants