Propagate modified reshape extents at concretization by jacobhinkle · Pull Request #630 · NVIDIA/Fuser

jacobhinkle · 2023-07-20T14:45:45Z

This PR propagates extent scalars after concretizing reshape. Previously the following fusion

  auto tv0 = makeSymbolicTensor(4);
  fusion->addInput(tv0);
  auto s0 = IrBuilder::create<Val>(DataType::Int);
  fusion->addInput(s0);

  auto sh = tensor_sizes(tv0);
  auto tv1 = reshape(tv0, {sh[0], div(sh[1], s0), s0, sh[2], sh[3]});
  // Reducing along axis 2 in tv1 is equivalent to a partial reduction across
  // axis 1 of tv0.
  auto vm = variance_mean(tv1, {2, 3, 4}, 0, true);
  fusion->addOutput(vm.mean);
  fusion->addOutput(vm.var);

would be concretized like this:

Inputs:
  T0_g[ iS0{i0}, iS1{i2}, iS2{i3}, iS3{i4} ], float
  i5, int64_t
Outputs:
  T7_g[ iS49{i0}, iS50{( i2 / i5 )}, bS37{1}, bS38{1}, bS39{1} ], float
  T6_g[ iS55{i0}, iS56{( i2 / i5 )}, bS32{1}, bS33{1}, bS34{1} ], float

%kernel_math {
T8_l[ iS40{i0}, iS45{4}rf, iS46{( ceilDiv(i2, 4) )}rf, iS42{i3}, iS43{i4} ] = view( T0_g[ iS0{i0}, iS1{i2}, iS2{i3}, iS3{i4} ] )
T2_l[ iS47{i0}, iS48{( i2 / i5 )}, rS15{i5}, rS16{i3}, rS17{i4} ](Avg),
T3_l[ iS51{i0}, iS52{( i2 / i5 )}, rS20{i5}, rS21{i3}, rS22{i4} ](Var),
T4_l[ iS57{i0}, iS58{( i2 / i5 )}, rS25{i5}, rS26{i3}, rS27{i4} ](Count)
 = Welford ( T8_l[ iS40{i0}, iS45{4}rf, iS46{( ceilDiv(i2, 4) )}rf, iS42{i3}, iS43{i4} ](Avg),
  allreduce = false )
T7_g[ iS49{i0}, iS50{( i2 / i5 )}, bS37{1}, bS38{1}, bS39{1} ]
   = broadcast( T2_l[ iS47{i0}, iS48{( i2 / i5 )}, rS15{i5}, rS16{i3}, rS17{i4} ] )
d17 = (double)(i5);
d19 = double(1) * d17;
d21 = (double)(i3);
d23 = d19 * d21;
d25 = (double)(i4);
d27 = d23 * d25;
d33 = reciprocal(d27);
T5_l[ iS53{i0}, iS54{( i2 / i5 )} ]
   = T3_l[ iS51{i0}, iS52{( i2 / i5 )}, rS20{i5}, rS21{i3}, rS22{i4} ]
   * d33;
T6_g[ iS55{i0}, iS56{( i2 / i5 )}, bS32{1}, bS33{1}, bS34{1} ]
   = broadcast( T5_l[ iS53{i0}, iS54{( i2 / i5 )} ] )
}

Now this concretizes as

Inputs:
  T0_g[ iS0{i0}, iS1{i2}, iS2{i3}, iS3{i4} ], float
  i5, int64_t
Outputs:
  T7_g[ iS52{i0}, iS53{4}, bS37{1}, bS38{1}, bS39{1} ], float
  T6_g[ iS62{i0}, iS63{4}, bS32{1}, bS33{1}, bS34{1} ], float

%kernel_math {
T8_l[ iS40{i0}, iS45{4}rf, iS46{( ceilDiv(i2, 4) )}rf, iS42{i3}, iS43{i4} ] = view( T0_g[ iS0{i0}, iS1{i2}, iS2{i3}, iS3{i4} ] )
T2_l[ iS49{i0}, iS50{4}, rS48{( ceilDiv(i2, 4) )}, rS16{i3}, rS17{i4} ](Avg),
T3_l[ iS56{i0}, iS57{4}, rS55{( ceilDiv(i2, 4) )}, rS21{i3}, rS22{i4} ](Var),
T4_l[ iS66{i0}, iS67{4}, rS65{( ceilDiv(i2, 4) )}, rS26{i3}, rS27{i4} ](Count)
 = Welford ( T8_l[ iS40{i0}, iS45{4}rf, iS46{( ceilDiv(i2, 4) )}rf, iS42{i3}, iS43{i4} ](Avg),
  allreduce = false )
T7_g[ iS52{i0}, iS53{4}, bS37{1}, bS38{1}, bS39{1} ]
   = broadcast( T2_l[ iS49{i0}, iS50{4}, rS48{( ceilDiv(i2, 4) )}, rS16{i3}, rS17{i4} ] )
d17 = (double)(i5);
d19 = double(1) * d17;
d21 = (double)(i3);
d23 = d19 * d21;
d25 = (double)(i4);
d27 = d23 * d25;
d33 = reciprocal(d27);
T5_l[ iS59{i0}, iS60{4} ]
   = T3_l[ iS56{i0}, iS57{4}, rS55{( ceilDiv(i2, 4) )}, rS21{i3}, rS22{i4} ]
   * d33;
T6_g[ iS62{i0}, iS63{4}, bS32{1}, bS33{1}, bS34{1} ]
   = broadcast( T5_l[ iS59{i0}, iS60{4} ] )
}

This helps ensure that no scalars get lost during segmentation, which could previously occur if the reshape output became a segmentation edge (see #629).

I also took this opportunity to remove TensorView::convertRfactorToRootDomain(). It was only used at segmentation and did a traversal for each segment input TV. Now, we call it once and that replaces all extents for all input TVs using a single traversal.

Fixes #629 and fixes #418.

This allows us to just do one traversal for the whole segment to replace vals instead of one traversal per input TV.

jacobhinkle · 2023-07-20T14:47:32Z

Note that an alternative approach would be to reset the extents of the concretized reshaped TV to the original expressions. I went with this approach since it more closely resembles the extents we would see using static reshapes.

jacobhinkle · 2023-07-20T14:49:34Z

csrc/tensor_view.cpp

-        TORCH_INTERNAL_ASSERT(
-            !replacement_extents.empty() &&
-            getMaybeRFactorDomain().size() == replacement_extents.size());


I omitted the check that replacement_extents is not empty since it did not seem necessary.

jacobhinkle · 2023-07-20T16:32:11Z

csrc/dynamic_transform.cpp

+      /*traverse_members*/ true,
+      /*traverse_attributes*/ true,


Since we propagate scalars, we should do this anywhere they might occur; namely in members and attributes. This was not necessary when we were only replacing IterTypes since IterDomains are traversed regardless of these settings.

Didn't we discuss duplicated replacements if members were also traversed? Would it be a concern?

naoyam · 2023-07-21T04:12:37Z

csrc/tensor_view.cpp

  return toString(indent_size);
 }

-void TensorView::convertRfactorToRootDomain() {


Is this change completely unrelated with the propagation of concretized vals?

Yes. At first I thought it was going to be necessary. I realized it was not necessary for this PR, but it is a little faster and reduces the TensorView interface so I figured I would keep it. I can split it into another PR if you prefer.

jacobhinkle · 2023-07-21T11:27:43Z

!build

naoyam

The PR looks good to me, but I have one question on the rfactor-root replacement

naoyam · 2023-07-21T15:40:51Z

csrc/dynamic_transform.cpp

+      /*traverse_members*/ true,
+      /*traverse_attributes*/ true,


Didn't we discuss duplicated replacements if members were also traversed? Would it be a concern?

naoyam · 2023-07-21T15:56:39Z

csrc/fusion_segmenter.cpp

+    for (const auto& id : rfactor) {
+      if (id->isRFactorProduct()) {
+        // Create new symbolic extents for rfactor iterDomains
+        auto domain_extent = (!tv_is_concrete)


I think this function is fine as it's mostly just copied from tensor_view.cpp. However, I don't know why we do use all symbolic extents if tv_is_concrete. It seems this means that if there's any IterDomain with a symbolic extent, all IterDomains would have symbolic extents.

This is not about this PR itself, but does it make sense?

Didn't we discuss duplicated replacements if members were also traversed? Would it be a concern?

This is an important point since it could lead to us concretizing one aspect of a Val but losing a previously concretized aspect (for example losing the concretized extent when we later concretize IterType). But I think we are OK if we are careful to use maybeMutated when creating replacements. In one case we do already do multiple mutations; we call OptOutMutator::mutate(id) at the beginning of our mutate(IterDomain* id) override.

It seems this means that if there's any IterDomain with a symbolic extent, all IterDomains would have symbolic extents.

Yeah that is strange. All IterDomains in the rfactor domain are checked to determine concreteness, but then the condition is only applied to rfactor products. The purpose of this function is to make the root a "standalone" domain so that we can bind input shapes to it, so it seems like all extents regardless of if they were already rfactor products should be either constant ints or pure symbolic.

We always need to create a new IterDomain for rfactor products in order to cut the connection to the root domain. It seems like a simpler condition would be

auto domain_extent = id->extent()->isConstScalar() ? id->extent() : IrBuilder::create<Val>(DataType::Int);

And in light of the above, we could move this outside of the id->isRFactorProduct() check, so that we would replace IDs that are not rfactor products if they have non-constant derived extents.

Please create a separate issue for this. Let's merge this PR as is.

Sounds good. Merging without that change. I will do that in another PR.

This reverts commit 5996e7a.

jacobhinkle added 3 commits July 20, 2023 10:20

Modify Issue 418 test to include validation

3d73387

Propagate reshape extents at concretization

a207bae

Move TensorView::convertRfactorToRootDomain to fusion_segmenter.cpp

a70be81

This allows us to just do one traversal for the whole segment to replace vals instead of one traversal per input TV.

jacobhinkle changed the title ~~Seg lost scalars~~ Propagate modified reshape extents at concretization Jul 20, 2023

jacobhinkle commented Jul 20, 2023

View reviewed changes

jacobhinkle requested review from kevinstephano and naoyam July 20, 2023 14:50

jacobhinkle marked this pull request as ready for review July 20, 2023 14:50

jacobhinkle marked this pull request as draft July 20, 2023 15:06

jacobhinkle added 2 commits July 20, 2023 12:29

Only propagate derived extents. Traverse members+attrs

98066d2

Update comment about dynamic -1 in ReshapeReductionShmoo

deeabf6

jacobhinkle commented Jul 20, 2023

View reviewed changes

Merge branch 'main' into seg_lost_scalars

0d3850b

jacobhinkle marked this pull request as ready for review July 20, 2023 19:32

naoyam reviewed Jul 21, 2023

View reviewed changes

Merge branch 'main' into seg_lost_scalars

92ba13b

naoyam approved these changes Jul 21, 2023

View reviewed changes

jacobhinkle and others added 4 commits July 21, 2023 13:23

Clean up and explain condition for replacing rfactor IDs

5996e7a

Revert "Clean up and explain condition for replacing rfactor IDs"

d05d1e7

This reverts commit 5996e7a.

Merge branch 'main' into seg_lost_scalars

fb9d2e3

Merge branch 'main' into seg_lost_scalars

536b59e

jacobhinkle merged commit d25e8c5 into main Jul 21, 2023

jacobhinkle deleted the seg_lost_scalars branch July 21, 2023 22:18

This was referenced Jul 24, 2023

Resolve broadcasts resulting from a PadOp #610

Merged

More lost scalars during segmentation for dynamic reshape #656

Closed

jacobhinkle mentioned this pull request Sep 26, 2025

Dead code #5210

Draft

Conversation

jacobhinkle commented Jul 20, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jacobhinkle commented Jul 20, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jacobhinkle commented Jul 21, 2023

Uh oh!

naoyam left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jacobhinkle commented Jul 20, 2023 •

edited

Loading