New lowering pass: loop rotation #2500

zasdfgbnm · 2023-02-21T22:41:35Z

This PR adds a new lowering pass: loop rotation.

Loop rotation is a loop transformation that transforms the code like below

for i in range(n):
  statement1(i)
  statement2(i)
  statement3(i)
  statement4(i)

into something like:

if 0 < n:
  for i = 0:
    statement1(i)
    statement2(i)
for i in range(n):
  statement3(i)
  statement4(i)
  if i + 1 < n:
    statement1(i+1)
    statement2(i+1)

Right now, because existing predicates should already cover all illegal access, I am not materializing the 0 < n and i + 1 < n predicates, so I am actually generating:

for i = 0:
  statement1(i)
  statement2(i)
for i in range(n):
  statement3(i)
  statement4(i)
  if True:
    statement1(i+1)
    statement2(i+1)

Loop rotation is conceptually very simple regarding index and predicate lowering, and their changes are mostly mechanical. For the trivial loop for i = 0, our index and predicate lowering should already handle it correctly. For the exprs inside the if i + 1 < n block, the only change needed is to replace all the loop->index() with loop->index() + 1, which requires me to detect the if i + 1 < n block and mark its body as "rotated" during index and predicate lowering.

In order to use loop rotation, users can do fusion.rotateLoop(tv, dim, {tv1, tv2}), then during lowering, the loop tv->axis(dim) will be rotated, and the expressions allocating and computing tv1 and tv2 will be the rotated exprs.

This reverts commit 5eb594f.

…on-submit

third_party/nvfuser/csrc/lower_loop_rotation.h

third_party/nvfuser/test/test_gpu_loop_rotation.cpp

naoyam

Still working on, but some initial comments

third_party/nvfuser/csrc/fusion.h

third_party/nvfuser/csrc/index_compute.h

third_party/nvfuser/csrc/kernel_ir.cpp

third_party/nvfuser/csrc/fusion.h

third_party/nvfuser/csrc/type.h

naoyam · 2023-02-27T18:38:29Z

I get this build error (clang 13):

/home/nmaruyama/pytorch/debug2/third_party/nvfuser/csrc/ir_cloner.h:67:10: error: lambda capture 'this' is not used [-Werror,-Wunused-lambda-capture]
        [this](auto&... x) { return std::make_tuple<Ts...>(clone(x)...); },
         ^
/home/nmaruyama/pytorch/debug2/third_party/nvfuser/csrc/ir_cloner.h:49:22: note: in instantiation of function template specialization 'nvfuser::IrCloner::clone<nvfuser::TensorView *, long, std::unordered_set<nvfuser::Statement *>>' requested here
      copy.push_back(clone(p));
                     ^
/home/nmaruyama/pytorch/debug2/third_party/nvfuser/csrc/fusion.cpp:83:40: note: in instantiation of function template specialization 'nvfuser::IrCloner::clone<std::tuple<nvfuser::TensorView *, long, std::unordered_set<nvfuser::Statement *>>>' requested here
  to->loop_rotation_param_ = ir_cloner.clone(from->loop_rotation_param_);
                                       ^
1 error generated.

third_party/nvfuser/csrc/predicate_compute.cpp

third_party/nvfuser/csrc/lower_loop_rotation.h

third_party/nvfuser/csrc/index_compute.cpp

third_party/nvfuser/csrc/lower_index_compute.cpp

third_party/nvfuser/csrc/lower_loop_rotation.cpp

…on-submit

zasdfgbnm · 2023-02-28T00:17:26Z

I get this build error (clang 13):

/home/nmaruyama/pytorch/debug2/third_party/nvfuser/csrc/ir_cloner.h:67:10: error: lambda capture 'this' is not used [-Werror,-Wunused-lambda-capture]
        [this](auto&... x) { return std::make_tuple<Ts...>(clone(x)...); },
         ^
/home/nmaruyama/pytorch/debug2/third_party/nvfuser/csrc/ir_cloner.h:49:22: note: in instantiation of function template specialization 'nvfuser::IrCloner::clone<nvfuser::TensorView *, long, std::unordered_set<nvfuser::Statement *>>' requested here
      copy.push_back(clone(p));
                     ^
/home/nmaruyama/pytorch/debug2/third_party/nvfuser/csrc/fusion.cpp:83:40: note: in instantiation of function template specialization 'nvfuser::IrCloner::clone<std::tuple<nvfuser::TensorView *, long, std::unordered_set<nvfuser::Statement *>>>' requested here
  to->loop_rotation_param_ = ir_cloner.clone(from->loop_rotation_param_);
                                       ^
1 error generated.

I think this is a bug of clang, but I changed the code to this->clone instead of clone. Hope this will make clang happy.

naoyam · 2023-02-28T01:14:18Z

I get this build error (clang 13):

/home/nmaruyama/pytorch/debug2/third_party/nvfuser/csrc/ir_cloner.h:67:10: error: lambda capture 'this' is not used [-Werror,-Wunused-lambda-capture]
        [this](auto&... x) { return std::make_tuple<Ts...>(clone(x)...); },
         ^
/home/nmaruyama/pytorch/debug2/third_party/nvfuser/csrc/ir_cloner.h:49:22: note: in instantiation of function template specialization 'nvfuser::IrCloner::clone<nvfuser::TensorView *, long, std::unordered_set<nvfuser::Statement *>>' requested here
      copy.push_back(clone(p));
                     ^
/home/nmaruyama/pytorch/debug2/third_party/nvfuser/csrc/fusion.cpp:83:40: note: in instantiation of function template specialization 'nvfuser::IrCloner::clone<std::tuple<nvfuser::TensorView *, long, std::unordered_set<nvfuser::Statement *>>>' requested here
  to->loop_rotation_param_ = ir_cloner.clone(from->loop_rotation_param_);
                                       ^
1 error generated.

I think this is a bug of clang, but I changed the code to this->clone instead of clone. Hope this will make clang happy.

Yeah, looks like so. I don't see any error with the change.

naoyam · 2023-02-28T01:44:11Z

third_party/nvfuser/test/test_gpu_loop_rotation.cpp

+  inlineMost();
+  fusion.rotateLoop(tv4, -1, {tv1, tv2});


This looks inconsistent to me. Both are a fusion-wide operation, and one implicitly uses the fusion from a FusionGuard, but another is a method of Fusion. I'd say just rotateLoop(tv4, -1, {tv1, tv2}); would align well with the existing APIs.

naoyam · 2023-02-28T01:50:28Z

Just curious, can loop rotation be nested? Not that I think it's important.

third_party/nvfuser/csrc/lower_loop_rotation.cpp

zasdfgbnm · 2023-02-28T03:07:46Z

Just curious, can loop rotation be nested? Not that I think it's important.

I don't think so. To support nested loop rotation, we need to carefully handle the registerReplace and registerInsertBefore to make sure it does not change the abandoned main loop.

naoyam · 2023-02-28T03:17:14Z

Just curious, can loop rotation be nested? Not that I think it's important.

I don't think so. To support nested loop rotation, we need to carefully handle the registerReplace and registerInsertBefore to make sure it does not change the abandoned main loop.

OK. Maybe add a comment about it?

zasdfgbnm · 2023-02-28T07:03:42Z

Just curious, can loop rotation be nested? Not that I think it's important.

I don't think so. To support nested loop rotation, we need to carefully handle the registerReplace and registerInsertBefore to make sure it does not change the abandoned main loop.

OK. Maybe add a comment about it?

Actually, it is pretty easy to support it with very little extra work: I can just rotate one loop each pass.

I updated this PR to support this in 67f7df3

naoyam

Approving now. Thanks for the updates. This is still not something I like, #2500 (comment), but it's a question on code style, and it shouldn't block the PR.

zasdfgbnm added 18 commits February 21, 2023 14:36

New lowering pass: loop rotation

24dad4f

save

37ed0ed

cleanup

337d780

cleanup

b3b99cf

better comment

1789922

more cleanup

335e443

cleanup

5eb594f

Revert "cleanup"

c98b7d2

This reverts commit 5eb594f.

fix

d3f2e60

adjust tests

f517acd

better predicate printting

3a39b8e

fix canOmitStopPredicate

3cd238e

Merge branch 'devel' of github.com:csarofeen/pytorch into loop-rotati…

5445186

…on-submit

save

a806b13

test

a1311b2

save

d71980f

comment tests

d2b3f8e

save

1ee9a4b

zasdfgbnm commented Feb 22, 2023

View reviewed changes

third_party/nvfuser/csrc/lower_loop_rotation.h Show resolved Hide resolved

zasdfgbnm commented Feb 22, 2023

View reviewed changes

third_party/nvfuser/test/test_gpu_loop_rotation.cpp Show resolved Hide resolved

zasdfgbnm marked this pull request as ready for review February 22, 2023 22:32

zasdfgbnm requested a review from naoyam February 22, 2023 22:37

naoyam reviewed Feb 27, 2023

View reviewed changes

zasdfgbnm added 5 commits February 27, 2023 14:49

fix

abefa1a

update

655f5a7

comment

9c17d49

cleanup

836cbbd

Merge branch 'devel' of github.com:csarofeen/pytorch into loop-rotati…

b077be4

…on-submit

fix is_same

a86d11e

naoyam reviewed Feb 28, 2023

View reviewed changes

third_party/nvfuser/csrc/lower_loop_rotation.cpp Show resolved Hide resolved

zasdfgbnm added 3 commits February 27, 2023 19:02

check instead of warn

8ce278d

fix

f7257ed

save

d9c1efd

Support multiple nested

67f7df3

zasdfgbnm added 3 commits February 28, 2023 17:28

canOmitStopPredicate

80d7763

fix

896b867

IrBuilder

20564e3

This was referenced Mar 1, 2023

Rewrite reducePredicateRegisterUsage #2533

Merged

Add optional bool preserve_error for expr simplifier #2534

Merged

naoyam approved these changes Mar 1, 2023

View reviewed changes

zasdfgbnm added 2 commits March 1, 2023 10:32

scheduler_utils::rotateLoop

e083f84

fix include

9a85085

zasdfgbnm merged commit 591b8ae into devel Mar 1, 2023

zasdfgbnm deleted the loop-rotation-submit branch March 1, 2023 18:58

New lowering pass: loop rotation #2500

New lowering pass: loop rotation #2500

Uh oh!

Conversation

zasdfgbnm commented Feb 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

naoyam left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

naoyam commented Feb 27, 2023

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zasdfgbnm commented Feb 28, 2023

Uh oh!

naoyam commented Feb 28, 2023

Uh oh!

naoyam Feb 28, 2023

Choose a reason for hiding this comment

Uh oh!

naoyam commented Feb 28, 2023

Uh oh!

Uh oh!

zasdfgbnm commented Feb 28, 2023

Uh oh!

naoyam commented Feb 28, 2023

Uh oh!

zasdfgbnm commented Feb 28, 2023

Uh oh!

naoyam left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zasdfgbnm commented Feb 21, 2023 •

edited

Loading

naoyam left a comment •

edited

Loading