Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fixed Grid Reduction Perf Issue with caching and FP16 usage. #243

Merged
merged 11 commits into from
Jul 29, 2020
Merged

Fixed Grid Reduction Perf Issue with caching and FP16 usage. #243

Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
8596f98
Apply reordering of Reduction Schedule in the magic scheduler on unro…
kevinstephano Jul 22, 2020
c33928b
Fix clang format.
kevinstephano Jul 22, 2020
b9c20a2
Added better computeAt() statements for rFactored Tensors of reductio…
kevinstephano Jul 23, 2020
a4da708
Adding a new struct LaunchParam that has boolean indicating whether a…
kevinstephano Jul 24, 2020
28a6bd0
Add equal and hash functions for caching of ReductionParams.
kevinstephano Jul 25, 2020
bfade59
Fix clang formatting.
kevinstephano Jul 25, 2020
cd9f818
Forgot to only check Launch Param equivalence for Reduction Parameter…
kevinstephano Jul 27, 2020
02e62eb
Fixed some debug prints
kevinstephano Jul 27, 2020
4cf0765
Made changes in style based on Leonard's suggestions.
kevinstephano Jul 28, 2020
4807070
Made clang format fixes.
kevinstephano Jul 28, 2020
86d56dc
Fixed Grid Reduction performance problem. I had placed the grid reduc…
kevinstephano Jul 29, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 14 additions & 15 deletions torch/csrc/jit/codegen/cuda/scheduler.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -488,20 +488,20 @@ c10::optional<ReductionParams> scheduleReduction(
} else {
if (rparams.cross_grid) {
// Reduction Splits
// [outputs, |rF-Leftover, rf-Unroll, X-Block, X-Grid, X-Warp|]
// [outputs, |rF-Leftover, rf-Unroll, X-Grid, X-Block, X-Warp|]
// Idx: 0 | 1(-5) 2(-4) 3(-3) 4(-2) 5(-1) |
// -------------------------------------------------
// Reduction Dimensions
red_tv->split(1, rparams.bdimx.value);
red_tv->split(1, rparams.gdimy.value);
red_tv->split(1, rparams.bdimy.value);
red_tv->split(1, rparams.gdimy.value);
red_tv->split(1, kLoopUnrollSplit);

// Reordering the Unroll dimension eases applying computeAt()
// for preceeding operations and the rFactored Tensor.
// |------ Reordered --------|
// V V
// [outputs, |rF-Leftover, X-Warp, X-Block, X-Grid, rf-Unroll|]
// [outputs, |rF-Leftover, X-Warp, X-Grid, X-Block, rf-Unroll|]
// Idx: 0 | 1(-5) 2(-4) 3(-3) 4(-2) 5(-1) |
// -------------------------------------------------
// Reduction Dimensions
Expand All @@ -512,31 +512,31 @@ c10::optional<ReductionParams> scheduleReduction(

// WARNING: computeAt will coalesce the rFactored dimensions
// rFactored Reduction Tensor after computeAt():
// [Outputs, |X-Warp, X-Block, X-Grid, rF-Leftover, rF-Unroll|]
// [Outputs, |X-Warp, X-Grid, X-Block, rF-Leftover, rF-Unroll|]
// Idx: 0 | 1(-5) 2(-4) 3(-3) 4(-2) 5(-1) |
// -------------------------------------------------
// Reduction Dimensions
red_tv_rf->computeAt(red_tv, -1);

// After the Reduction Tensor has rFactoring applied
// Reduction Output Tensor:
// [Outputs, X-Warp, X-Block, X-Grid]
// [Outputs, X-Warp, X-Grid, X-Block]
// Idx: 0 1(-3) 2(-2) 3(-1)

red_tv_rf->axis(-1)->parallelize(ParallelType::Unroll);

red_tv->axis(0)->parallelize(ParallelType::BIDx);
red_tv->axis(-3)->parallelize(ParallelType::TIDx);
red_tv->axis(-2)->parallelize(ParallelType::TIDy);
red_tv->axis(-1)->parallelize(ParallelType::BIDy);
red_tv->axis(-2)->parallelize(ParallelType::BIDy);
red_tv->axis(-1)->parallelize(ParallelType::TIDy);

// Bind Inputs to Reduction
// The computeAt is not to the inner most dimension of the rFactored
// tensor in order to force the creation of separate loop nests to cause
// Inputs to be separately read in their own loop.
// computeAt(-2)------|
// V
// [Outputs, X-Warp, X-Block, X-Grid, rF-Leftover,| rF-Unroll]
// [Outputs, X-Warp, X-Grid, X-Block, rF-Leftover,| rF-Unroll]
// Idx: 0 1(-5) 2(-4) 3(-3) 4(-2) 5(-1)
Val* input = fusion->origin(red_tv_rf)->as<ReductionOp>()->in();
if (!fusion->hasInput(input)) {
Expand Down Expand Up @@ -603,20 +603,19 @@ c10::optional<ReductionParams> scheduleReduction(
if (rparams.cross_block) {
if (rparams.cross_grid) {
// Reduction Splits
// [outputs, |rF-Leftover, rf-Unroll, X-Block, X-Grid|]
// [outputs, |rF-Leftover, rf-Unroll, X-Grid, X-Block|]
// Idx: 0 | 1(-4) 2(-3) 3(-2) 4(-1) |
// -----------------------------------------
// Reduction Dimensions
red_tv->split(1, rparams.bdimy.value);
red_tv->split(1, rparams.gdimy.value);
red_tv->split(1, rparams.bdimy.value);
red_tv->split(1, kLoopUnrollSplit);

// Reordering the Unroll dimension eases applying computeAt()
// for preceeding operations and the rFactored Tensor.
// |--- Reordered ----|
// V V
// [outputs, |rF-Leftover, X-Grid, X-Block, rF-Unroll|]
// [outputs, |rF-Leftover, X-Block, X-Grid, rF-Unroll|]
// Idx: 0 | 1(-4) 2(-3) 3(-2) 4(-1) |
// -----------------------------------------
// Reduction Dimensions
Expand All @@ -633,23 +632,23 @@ c10::optional<ReductionParams> scheduleReduction(

// WARNING: computeAt will coalesce the rFactored dimensions
// rFactored Reduction Tensor after computeAt():
// [<output dims>, |X-Grid, X-Block, rF-Leftover, rF-Unroll|]
// [<output dims>, |X-Block, X-Grid, rF-Leftover, rF-Unroll|]
// Idx: 0 -- 1 | 2(-4) 3(-3) 4(-2) 5(-1) |
// -----------------------------------------
// Reduction Dimensions
red_tv_rf->computeAt(red_tv, -1);

// After the Reduction Tensor has rFactoring applied
// Reduction Output Tensor:
// [Out-Leftover, Out-PerBlock, X-Grid, X-Block]
// [Out-Leftover, Out-PerBlock, X-Block, X-Grid]
// Idx: 0 1 2(-2) 3(-1)

red_tv_rf->axis(-1)->parallelize(ParallelType::Unroll);

red_tv->axis(0)->parallelize(ParallelType::BIDx);
red_tv->axis(1)->parallelize(ParallelType::TIDx);
red_tv->axis(-2)->parallelize(ParallelType::BIDy);
red_tv->axis(-1)->parallelize(ParallelType::TIDy);
red_tv->axis(-2)->parallelize(ParallelType::TIDy);
red_tv->axis(-1)->parallelize(ParallelType::BIDy);

// Bind Inputs to Reduction
// The computeAt is not to the inner most dimension of the rFactored
Expand Down