add write_stride to avoid redundant write when output domain is merge…

[liqiangxl]

Fixes csarofeen/pytorch#2125
Before this fix was 340~370 GB/s on A100-80G, after this fix increased to 1.19e+03 GB/s on the same GPU.

[csarofeen]

I'm hesitant to take this change in. We need to think a bit more carefully about this optimization.

[csarofeen]

Could you try to make this comment a bit more clear. I don't understand what the key-val pairs mean by this comment.

[csarofeen]

I think I understand what you're doing, but it seems dangerous to me. I think this is fine in your example in the test, but I don't think this would be fine in more complex situations of broadcast concretization with arbitrary splits/merges. This seems hard coded for patterns of merges then splits. Which could be fine, but we'd need a better check to ensure that's the pattern being experienced. It's a common pattern for sure, but not the only one.

[liqiangxl]

You are right. This PR only fix patterns where a domain is isParallelTypeThread(pt) && merged_from_broadcast. It won't touch other broadcast concretization with arbitrary splits/merges, actually, thread_predict can't handle these general cases. To achieve this, we need to apply the write_stride to the index of the tensor, which I looked and it seems to be more complex than this thread_predict approach. I can revist that more general approach.

[csarofeen]

If you can explain to me how it only gets applied in this more isolated instance (likely I just forgot the code paths), then I'm happy to proceed and merge this PR. I just don't understand how the codebase avoids applying this rule in other situations.

[liqiangxl]

Only when a domain is isParallelTypeThread(pt) && merged_from_broadcast., it will be added to the map. e.g. iblockIdx.x103{( i0 * ( 1 * 1 ) )}, which is parallelized by blockIdx.x and merged from broadcast domains. It won't add other domains to the map, e.g. iS112{( i0 * ( 1 * 1 ) )} becase it is not parallelized by thread/blocks, iblockIdx.x113{( ceilDiv(( i0 * ( 1 * 1 ) ), blockDim.y) )} becase it is not a merge from broadcast domains.
see the following code:

        auto domain = merge->out();
        auto pt = domain->getParallelType();
        bool from_broadcast =
            merge->outer()->isBroadcast() || merge->inner()->isBroadcast();
        if (isParallelTypeThread(pt) && from_broadcast) {
          merged_broadcast_domains.emplace_back(std::make_pair(out_tv, domain));
          tensor_with_domain_merged_from_broadcast = true;
          break;
        }

[csarofeen]

Sorry, forgot to push this review.

[csarofeen]

You should check that it's a concretized domain. Likely we should even check that it's concretized and inlined.

[csarofeen]

It seems the check here is starting at a leaf node, check its definition, if that definition is a merge and one side of it is a concretized broadcast, then we want to modify its write predicate. This check seems to effectively do that, but that structure is not clearly written.

[csarofeen]

i.e. I don't understand why we need to process all expressions between rfactor and leaf domains, instead of just each definition of each leaf.

[csarofeen]

This also doesn't seem to work if we had something like:
i2 = merge(i0, b1)
i3 = merge(i2, b2)
i4 = merge(i3, i4)

[csarofeen]

These constraints are generally fine for now, but I want the code to be a bit more precise in exactly what it's WAR-ing, and that it is a perf WAR. Basically what you're trying to do is predicate the writes of a producer to only write on the first unique time a producer value is generated.

This is straight forward to do in such a simple case, but is challenging to implement more generally. For example what if for some reason the tensor of concern is an intermediate tensor and not a terminating output. Then what if there's non-trivial communication going on with its consumer. We might not just be able to write once, as we might have complex producer-consumer synchronization expectations.

I want to make sure we limit the scope of this to as narrow as possible to make sure it's not accidentally triggered in a more complex situation.

[csarofeen]

I also wonder if this is worth doing on anything shared or local memory, or if it should be limited to global writes.

[liqiangxl]

I revised the algorithm, it can capture the case you mensioned. Current alg:
(1) starting at an output tensor, if one of its leaf node's definition is a merge and parallelized by thread/block, move to (2)
(2) check all the root domains, if it is a broadcast, map it to its concretized domain (or concretized domain in its exact map)
(3) Loop over all leaf nodes whose definition is a merge and parallelized by thread/block
---(3.1) find all root domains that are merged to this leaf domain.
---(3.2) Set stride when use extent of merged root domains to index leaf domain.

  // e.g. Root: [I1,B2,B3] -> Leaf: [I1*B2*B3], the merged_root_domains =
  // {B3,B2,I1}. root_stride = {1, len(B3), len(B3) * len(B2)}. For broadcast
  // root domain, use the extent of its concretized domain. 

---(3.3) set write stride

    // write_stride_mod will generate condition: index % pair(1) < pair(2)
    std::unordered_map<ParallelType, std::pair<Val*, Val*>> write_stride_mod;
    // write_stride_less will generate condition: index < write_stride_less
    std::unordered_map<ParallelType, Val*> write_stride_less;

7 addtional test cases are added:

  // Test case where [B1,I2,I3] is merged to [B1I2I3] and paralled by blockIdx.x
  // The write pattern should be: write only the first len(I2)*len(I3) blocks,
  // where len(I) is the extent of the concretized domain.
  // write_stride_less= ( ( 1 * T0.size[1] ) * T0.size[0] )
  // condition = blockIdx.x < ( ( 1 * T0.size[1] ) * T0.size[0] )
  runTest({true, false, false, false});

  // Test case where [I1,B2,I3] is merged to [I1B2I3] and paralled by blockIdx.x
  // The write pattern should be: len(I3) blocks write, in every len(B2)*len(I3)
  // blocks, write_stride_mod= {len(B2)*len(I3), len(I3)} condition = blockIdx.x
  // % (len(B2)*len(I3)) < len(I3)
  runTest({false, true, false, false});

  // Test case where [I1,I2,B3] is merged to [I1I2B3] and paralled by blockIdx.x
  // The write pattern should be: write every len(B3) blocks,
  // where len(B) is the extent of the domain.
  // write_stride_mod= ( 1 * T1.size[2] )
  runTest({false, false, true, false});

  // Test case where [I1,B2,B3] is merged to [I1B2B3] and paralled by blockIdx.x
  // The write pattern should be: write every len(B2)*len(B3) blocks,
  // where len(B) is the extent of the concretized domain.
  // write_stride_mod= ( ( 1 * T1.size[2] ) * T1.size[1] )
  runTest({false, true, true, false});

  // Test case where [B1,I2,B3] is merged to [B1I2B3] and paralled by blockIdx.x
  // The write pattern should be: write every len(B3) blocks of the first
  // len(I2) * len(B3) blocks write_stride_less= ( ( 1 * T1.size[2] ) *
  // T0.size[0] ) write_stride_mod= ( 1 * T1.size[2] )
  runTest({true, false, true, false});

  // Test case where [B1,B2,I3] is merged to [B1B2I3] and paralled by blockIdx.x
  // The write pattern should be: write only the first len(I3) blocks
  // write_stride_less= ( 1 * T0.size[0] )
  runTest({true, true, false, false});

  // Test case where [B1,B2,B3] is merged to [B1B2B3] and paralled by blockIdx.x
  // The write pattern should be: write only the first block
  // write_stride_less= 1 (write when blockIDx.x < 1)
  runTest({true, true, true, false});

[csarofeen]

What's T1 and T0 here?

[liqiangxl]

T0 is a tv with first 3 dims = [Broadcast_00, Iter_01, Broadcast_02], T1 is another tv with first 3 dims = [Iter_10, Iter_11, Iter_12], Broadcast_00 is concretized to Iter_10 and Broadcast_02 is concretized to Iter_12 (T1.size[2]). I should revise this comment without using T0 and T1.

[csarofeen]

Pushing some questions. I think there's only one thing that I'd consider blocker which is making sure the merges are ordered. I believe the rest of the logic is safe, but that's the one potential issue I could find.

I will note that this could be made a much more general pass but is difficult to do so today. What we seem to be looking for here is having an "inlined" (inlined here could also mean shared parallelization dimensions) resolved broadcast merged with a non-broadcast dimension. When we have these merges, what we really want to do is predicate away the indexing across the broadcast dimension so we don't write multiple times. Inlined here is really an interesting concept, where it could also include parallelizations that exact map from producer-consumers. i.e. for this pass it doesn't matter if the block parallelized dimension is inlined or if it simply maps from producer-consumer with this resolved broadcast. I think there's some interesting parallelization x memory location aspect to the applicability of a more general implementation of this optimization.

Today it would be cumbersome to implement a more generic version of this pass. However, it could become easier to implement with #32

[csarofeen]

I assume by len(B3) you mean the size of the dimension the broadcast is resolved to, not 1?

[liqiangxl]

yes!

[csarofeen]

Does the predicate generated depend on the order of merges?

[csarofeen]

I think this comment is indicating that we're thinking of cases where we have a merge with the outer being broadcast and inner not being broadcast. The point being if we have this broadcast axis getting peeled off (a broadcast axis that's resolved inlined) then we only need to write for the first pass of the inner side of that merge.

I think this only works if the iter domain doesn't have "zero merged in", so as long as this is currently limited to global memory instances then it should be valid but wouldn't necessarily be valid, for example, if it was done on shared memory.

[liqiangxl]

The generated predication does not depend on the order of the merge, it only depends on the order of the root domains, e.g. [broadcast, iter] or [iter, broadcast]. Current fix is only for global memory. what's the meaning of "zero merged in"?

[csarofeen]

zero merged in is currently only applicable to smem and local, but for multi-gpu it will get extended. It's indicative of the removal of allocation associated to locality across parallelization. i.e. if something is local memory it will remove the portion of the allocation bound to block and grid dims.

[csarofeen]

Does this have to be a recursive merge? i.e. all expressions between root domains and this leaf domain must all be merge?

[liqiangxl]

Yes. This is the current assumption. If there is a split in the path, predicate away the indexing across the broadcast dimension becomes difficult.

[csarofeen]

But this isn't actually recursive right? Doesn't this need to grab the expressions of ld to domain, and make sure they're all merges?

[csarofeen]

Or are you just saying it's the final expression you want to improve the write predication on?

[liqiangxl]

You are right. I need to make sure all are merges. so the following return will be called if a non-merge expr is detected.

        // current analysis of predication is only valid if all the exprs
        // between this lead domain and root domains are merge
        return std::vector<IterDomain*>();

[csarofeen]

I meant sequential merges, so a merge operation going into another. Something along the lines of for(auto expr : StmtSort::getExprs(fusion, ld)){...

[csarofeen]

Sorry, it should actually be similar to: auto all_exp = DependencyCheck::getAllExprsBetween( {rfactor_domain.begin(), rfactor_domain.end()}, {ld}); as it should be bounded by the allocated domain.

[liqiangxl]

(1) changed getRootDomain to getMaybeRFactorDomain();
(2) the code can process a merge operation going into another, e.g. [I,B,B] --> [IB, B] --> [IB*B]

[csarofeen]

Is this needed because your merged_root_domain is reversed order from the root ordering?

[liqiangxl]

yes.

[csarofeen]

Does it need to be written in reverse order?

[liqiangxl]

reverse order is convinent for the calculation of strides. Of course, we can also store it in unreversed order and visit backward when calculate stride.

[csarofeen]

I kind of like the latter approach, as it keeps what's stored in a 1:1 "zippable" order. Generally changing order of iterating on a data structure makes a lot of sense to me, flipping the order of a data structure relative to how it's generally been stored is a bit more "mentally complex" in my opinion and should be done if storing as such improves cache locallity/performance (i.e. when we change layout of tensors to improve locality of tiles as they're accessed).

[csarofeen]

This will now only work on a single root domain -> merge -> parallelization, no?

If you have I0, I1, I2 and merge(merge(I0, I1), I2) the outer of the second merge is not in the root domain because it's an intermediate.

I think you need to start with the root domain, but as you're processing the above case basically update the domain you're tracking so:
Root [I0, I1, I2]
intermediate [I0I1, I2]
Final [I0I1*I2]
Then on the intermediate you also know the distance between the outer and inner is 1.

[liqiangxl]

The if-statement is trying to stop applying this thread/block predication if there is a merge where the merged root domains are not neighbors, e.g. [I1, B2, B3] with merges (I1*B3) will be rejected because I1 is a root domain, B3 is a root domain, and their distance is not 1.
It won't influence cases where one of the merged domain is not a root domain. e.g. in the case you mensioned, I0, I1, I2 and merge(merge(I0, I1), I2), the if-statement test is false. so the case will not be rejcted. and the predication will be added for this case.

Your suggested approach is more general and I implemented it in the reivsed version.

[csarofeen]

Could you clean up your uses of lambda functions with [&] except where necessary. I also think your lambda functions are getting a bit long, and it might be nice to take some out and put them in anonymous namespaces as a function.

[csarofeen]

isn't the only thing you're passing in root_domain? [&] seems excessive here.

[liqiangxl]

changed to functions.

[csarofeen]

I think you should assert that all the c_ids returned by this are in the same group in the IdMappingMode::EXACT graph.

[liqiangxl]

added check.

[naoyam]

@csarofeen Not sure why this assertion should hold. A broadcast ID can be concretized to different non-broadcast IDs and they don't need to be exactly mapped. They should be only permissively mapped.

[naoyam]

Also, if we really want to check this, shouldn't we look at all IDs in mapped_rd_exact_set instead of just one of them? Should we have the break at line 588?

[liqiangxl]

For a mapped_rd in mapped_rd_exact_set, we take its allConcretizedDomains and store to all_cids.
The check is to check all_cids are fromSameExactGroup.
The break at line 588 is to break from the loop over other mapped_rd in mapped_rd_exact_set as we already find a concretized domain, which can be used to represent the size of the broadcast root domain.

[liqiangxl]

If this happens, the fusion will be segmented by our heuristics, right? I added a test case FusionAvoidRedundantWriteDifferentConcretizedDomains_CUDA and tested pointwise and reduciton and the fusion is segmented in both cases. Just in case, I removed the break and added additional check:

            // make sure this broadcast root domain is concretized to the same domain
            // otherwise, treat as not concretized becase we don't know which one to use.
            auto iter = concretized_broadcast_root_domains_.find(rd);
            if(iter != concretized_broadcast_root_domains_.end() && iter->second != *all_cids.begin()) {
              concretized_broadcast_root_domains_.erase(iter);
              // break to move to the next broadcast root domain
              break;
            }else{
              concretized_broadcast_root_domains_[rd] = *all_cids.begin();
            }

[naoyam]

This is a part of fusion lowering, which is designed to be independent from how the schedulers and segmenter work, so we should not assume how they are segmented. In fact, we could just define a fusion like the example I posted above and pass it to GpuLower, and it should still work or at least error out if it cannot be handled without segmentation.

[naoyam]

Also, I still don't understand why concretization matters. As shown here, #11 (comment), the issue happens even with non-concretized broadcast domains.

[liqiangxl]

This is a part of fusion lowering, which is designed to be independent from how the schedulers and segmenter work, so we should not assume how they are segmented. In fact, we could just define a fusion like the example I posted above and pass it to GpuLower, and it should still work or at least error out if it cannot be handled without segmentation.

Added to test case FusionAvoidRedundantWriteDifferentConcretizedDomains_CUDA, where the non-segmented fusion is directly passed to GpuLower.

[liqiangxl]

Also, I still don't understand why concretization matters. As shown here, #11 (comment), the issue happens even with non-concretized broadcast domains.

Marked as a case can't be handled by this PR.

[csarofeen]

I'm skeptical that changing the position of merges would work as it would change the indexing.
T1[I0, I1]->merge(0, 1)
is not indexed the same as
T1[I0, I1]->merge(1, 0)

[csarofeen]

Does line 633 indicate that the predicate is invariant to the order of the root domains going into the merge?

[liqiangxl]

order of the root domains going into the merge doesn't influence the indexing. Here, we are indexing between multi-dimension neighboring root domains (indices) and one leaf domain (linear_index). e.g. [I0, I1] --> [I0*I1], linear_index = indices[0]*stride[0] + indices[1]*stride[1]; stride=[len(I1), 1]; The range of indices[i] is from 0 to len(i)-1

[csarofeen]

should this and the instance below be std::distance(root_domain.begin(), inner_iter)?

[liqiangxl]

corrected.

[csarofeen]

@naoyam does a concretized domain mean it's inline concretized? I'm not sure if this PR is assuming that concretized domains it's processing are shared inlined domains of consumers.

[naoyam]

No, concretization has nothing to do with inlining.

[csarofeen]

I kind of like the latter approach, as it keeps what's stored in a 1:1 "zippable" order. Generally changing order of iterating on a data structure makes a lot of sense to me, flipping the order of a data structure relative to how it's generally been stored is a bit more "mentally complex" in my opinion and should be done if storing as such improves cache locallity/performance (i.e. when we change layout of tensors to improve locality of tiles as they're accessed).

[csarofeen]

I guess this is what surprises me about my comment on line 598

[liqiangxl]

I changed the algorithm. The previous is difficult to understand and extend.

[csarofeen]

Please don't use three letter acronyms, previous_id or current_id would be fine, cur_r_id also would be okay, but crd is a bit too short of a variable name for my liking. Code should read like normal text, and acronyms should be used sparingly unless they're extensive/standard in the codebase i.e. id is used frequently as IterDomain so that's generally fine as it's so common one should learn the convention to read the codebase.

[liqiangxl]

make sense.

[csarofeen]

This is confusing to me, can you expand the comments in this section a bit?

[liqiangxl]

I changed the algorithm.

// e.g. [I2,I4] -> [B1,I2,B3,I4,B5] -> [B1*I2*B3*I4*B5]_blockIdx.x where Bx are concretized with extents larger than 1.
// In this case, the leaf domain {B1*I2*B3*I4*B5} is written redundantly because len(B1)*len(I2)*len(B3)*len(I4)*len(B5) blocks are writing to len(I2)*len(I4) locations.
// This class will set a map from ParallelType to the index that needs to write.
// The method is: (1) Thinking the leaf domain as a 5-D array with extended_stride:
// extended_stride = [len(I2)*len(B3)*len(I4)*len(B5), len(B3)*len(I4)*len(B5), len(I4)*len(B5), len(B5), 1]
// (2) calculate the index of each dimension: linear_index = blockIdx.x; index[i] = (linear_index / extended_stride[i]); linear_index %= extended_stride[i];
// (3) only the non-broadcasted dimensions need to write, dim-1(I2) and dim-3(I4). So the write index is:
// write_index = index[1] * extended_stride[1] + index[3] * extended_stride[3] 
// Finally, the write condition is: if (blockIdx.x == write_index)

[liqiangxl]

I forgot to push my local change on last Monday, updated now.

[liqiangxl]

!build

[naoyam]

Instead of generating the index that corresponds to the leaf ID, it seems it's more intuitive to just use the index of the broadcast root ID and create a predicate as index_of_the_broadcast_root_id == 0

[liqiangxl]

Sounds like if we have multiple merged broadcast root IDs, then the condition will be index_of_the_broadcast_root_id0 == 0 && index_of_the_broadcast_root_id1 == 0 && .... Is this prefered?

[naoyam]

The logic of generating the predicate would be more straightforward. The generated predicate may end up having multiple conditions, and it appears to me easier to reason about.

[liqiangxl]

sure. I'll change to that.

[naoyam]

LGTM. Please confirm again the perf regression is fixed as the PR changed a lot.

[naoyam]

Just as a safety measure, assert that:

concretized_broadcast_root_domains_.emplace(rd, concrete_root_id).second == true

[liqiangxl]

!build

Naoya did the review.