[CodePartition] Support arbitrary data channel

[htyu]

Previously data channels are created for data flows originating from a load operation to its consumers only. This change adds fundamental support for channeling from any operation to its consumers. This will be useful for the cases where data partition isn't available, such as when BLOCK_M is smaller than 128.

To support arbitrary data channel where SMEM isn't naturally, a local_store is created on the producer side to dump the register data into SMEM and a local_load use created on the consumer side to load data back into registers.

The patch also extends the previous data channeling logics to handle one-producer multi-consumer mode, where only one buffer is needed.

Some refactoring is also done to the existing code base. For example, async copy creating code is separated from async communication emission. The data structure of Channel is refactored mainly to get the case where producer operation can be deallocated/reallocated, e.g, a forOp.

Although the current auto partition heuristics still only considers DotLike ops in addition to previous LoadOp and ExperimentalDescriptorLoadOp, the logics there can be extended to work on more anchor ops.

[manman-ren]

I wonder why we need to change from (srcOp, dstOp, srcOperand) to (op, operandIdx). Will it be better if we keep dstOp instead of op?

[htyu]

This is because the source operations can be changing. For example it could be a loop, or an if. The consumers are relatively stable. op now means dstOp, and the class now provides setter functions to access srcOp,dstOp, srcOperand, dstOperand.

[manman-ren]

Is there a case where we need to get transitive users?

[htyu]

Yes, like below which is from fusing one kernel into one gemm kernel:

for (...)
  c = dot(a,b,c)

op1(c) 
op2(c)

There will be two channels created. Both has the loop as the producer.

[manman-ren]

Is this the only place we update producerChannels? the rest of the function seems to update consumerChannels which existed prior to the change (i.e to minimize number of sync points).

[htyu]

Yes. The function computes producerChannels first followed by the computation of consumerChannels.

[manman-ren]

Oh lowering for non-tma loads are moved out of this function.

[htyu]

Yeah as a part of the barrier insertion, as they are coupled.

[manman-ren]

Will this have any impact on the default mode where producer is a loadOp? I guess it can, if the load has multiple consumers. Will we create multiple async_copy?

[htyu]

Yes it affects load with multiple consumers where we haven't seen so far. If the load has multiple consumers, we will create multiple async_copy without the current change. With only one buffer is created for channels that can be compressed (i.e, one master dominating channel).