Decouple loop and allocation during preseg for multi-GPU

[wujingyue]

This post tries to summarize why we need to decouple loop domains and allocation domains during preseg for multi-GPU. In addition to the philosophical reason that they represent different concepts, we have practical reasons for that.

Context: This comes from several offline discussions with @Priya2698. Although we are on the same page, I think it's worth writing down our thought process for posterity.

First principles

At the start of segmentation,

1. Loop and allocation may differ in the following ways:
   a. Loop can be more parallelized than allocation.
   b. Loop and allocation can be in different orders.
2. Allocation of all segmentation boundaries must be finalized.

Rationales

1(a) is needed for stream parallelization. For example,

TensorView* tv0 = makeContigTensor(2);
TensorView* tv1 = set(tv0);
TensorView* tv2 = set(tv1);

tv1->axis(0)->parallelize(ParallelType::Stream);
tv2->axis(1)->parallelize(ParallelType::Stream);

where tv1 is produced in a row-wise host loop and consumed in a column-wise host loop.
However, unlike loop domains, tv1 needs to be allocated fully, i.e., its allocation domain is the same as logical without parallelization.

1(b) is needed for convenience. The allocation domain of a tensor should respect its stride order by definition. However, the loop domain should prioritize placing Streams/DIDs in the front so we do things like inlineMost and reorder(..., num_device_dims). It's better for other, non-parallel IterDomains to follow the logical order to minimize disruption. Alternatively, we can let the host IR lowering, all get*Heuristicss and all schedule*s reorder locally, but this won't be as convenient.

2 is needed because it's too late to change allocation during or after segmentation. Segmentation can't change allocation because it's read-only. Schedulers can't change allocation of segmentation boundaries because it only sees one segment at a time and doesn't know the implications of changing the allocation of a boundary tensor.

Implications

Preseg passes need to decouple loop and allocation.

This doesn't mean loop and allocation have to be decoupled from the start. They just need to be decoupled at the end of the preseg stage.

[wujingyue]

cc @Priya2698 is this a good summary?

[Priya2698]

cc @Priya2698 is this a good summary?

LGTM. I will update with the implications and required changes when I experiment with this. Thank you

[Priya2698]

We can also consider still keep the split by stream in allocation even if it is not parallelized and hence, fully allocated. That is, loop = [Stream (s), i0/s, i1], allocation = [s, i0/s, i1]. This can be an intermediate step to decoupling loop and allocation where we allow them to be permutations of each other.

[wujingyue]

We can also consider still keep the split by stream in allocation even if it is not parallelized and hence, fully allocated. That is, loop = [Stream (s), i0/s, i1], allocation = [s, i0/s, i1]. This can be an intermediate step to decoupling loop and allocation where we allow them to be permutations of each other.

Yes. This doesn't violate any first principles and doesn't change the behavior of the program compared to [i0,i1]. It may lead to suboptimal kernel due to extra index computation.

[Priya2698]

Some comments from team discussion:

1. There is no explicit contract between loop and allocation at the beginning of the scheduling. There may be existing assumptions that we break when loop and allocation diverge.
2. Vectorization

• For schedulers that do not explicitly check for allocation domain maybe assuming allocation == logical and silently affect vectorization.
• Analysis may benefit from device dimensions being outermost in allocation domain. Alternatively, it can be made smarter to consider their extent as 1.

3. Design choice for retaining transforms in allocation domain even when fully replicated: Consider a fusion --

tv1 = set (tv0) // [i0, i1]
fusion->addInput(tv0)
fusion->addOutput(tv1)

tv1->outer_split(1, s)
tv1->parallelize(ParallelType::Stream)

The loop domain for tv1 = [i0, iStream{s}, i1/s]. The allocation domain can be represented as (no transforms)[i0, i1] or (retain transforms) [i0, s, i1/s].

There is no clear motivation to pick one design over the other, and it may be beneficial to start out by not restricting ourselves to one of them.

CC: @jjsjann123

[wujingyue]

cc @naoyam and @zasdfgbnm

Principle 1(a) is one of the examples of how loop and allocation can diverge. The loop domain is parallelized but the allocation domain isn't. This representation is the best I can think of without inventing another memory type.

However, if we were to follow the convention in single-GPU schedulers, we could create a memory type called "Stream". This way, loop=[iStream{i0}, iS{i1}], memory_type=Stream is allocated within the host for loop and has allocation size of i1. loop=[iStream{i0}, iS{i1}], memory_type=Global instead is allocated outside the host for loop and has allocation size of i0*i1.

That's why I said in the meeting that multi-GPU will use more of allocation domains and less of memory types. To me, for the multi-GPU use cases, allocation domains are more expressive than memory types. For example, it can also represent a stride order (cf. Principle 1(b)). So I'll likely stick with the current plan without having to worry about memory type, another concept. But I'm happy to hear your thoughts on this.

[naoyam]

I don't think I get what you meant, for example, "allocation domains are more expressive than memory types", I don't understand what that means. Allocation domains and memory types represent different things, so I'm not sure why you're favoring one over another.

Let's talk offline.