[SWP] Attempt to move all scheduling logic to a scheduling pass

[manman-ren]

The main purpose is to move all scheduling related logic including scheduleDeps/scheduleDistOne/scheduleRemaining to a new loopScheduling pass. The new pass will generate (stage, cluster) attributes for each operation inside a loop.

For operations that are created during createAsyncOps, we manually add the (stage, cluster) based on the attributes prior to lowering. In most cases, the lowered operations should have (stage, cluster) of the original loadOp, or should have (stage, cluster) of the first use of the loadOp.

This patch also gets rid of prefetchCluster, instead it prefetches one stage before the actual use
setStageCluster(forOp, wait, stageForFirstUse - 1, clusterForFirstUse + 1);
setStageCluster(forOp, wait, stageForFirstUse - 1, clusterForFirstUse + 1);
comparing to
schedule.insert(wait, numStages - 2, prefetchCluster);
schedule.insert(viewLoad, numStages - 2, prefetchCluster);

At end of createAsyncOps, we make sure all operations inside the loop have (stage, cluster) attributes.

There is no need to maintain CoarseSchedule in SWP, instead we will just use (stage, cluster) attributes.

[manman-ren]

I am keeping the logic for loadOpToIndLvelAndUse inside assignMemoryLayout. It is not clear to me how to split the logic between loopScheduling and SWP. This function sets fields such as

  ttg::SharedEncodingAttr sharedEncoding = nullptr;
  ttg::BlockedEncodingAttr blockedEncoding = nullptr;
  bool loadIsMMAV3 = false;
  bool usedByDot = false;

The decision of whether or not to pipeline a loadOp depends on earlier decisions based on encoding etc. Maybe we can move this whole function to loopScheduling? But we will need to duplicate some of the logic to set the fields inside SWP.
CC @pawelszczerbuk @ThomasRaoux

[manman-ren]

This diff is not ready for detailed review. But I would appreciate looking at changes to MatmulLoopPipeline.cpp to see if the overall direction looks fine.

[manman-ren]

assignMemoryLayouts still uses loadOpToIndLevelAndUse in order to match functionality prior to the refactoring. But we can simplify it to find the first Dot use.

[manman-ren]

We have set (stage, cluster) for extractIdx to the first Use of all loadOps (see helper getExtractIdxStageCluster). If wait/viewLoad are put in the prefetch stage, cluster, we need to reset the attributes for extractIdx here.
It is not clear to me if we need multiple copies of extractIdx here.

[manman-ren]

@pawelszczerbuk @ThomasRaoux Part of assignMemoryLayouts ends up in LoopScheduling. It calls getSharedEncIfAllUsersAreDotEnc at line 136.

[manman-ren]

assignMemoryLayouts still uses loadOpToIndLevelAndUse in order to match functionality prior to the refactoring. But we can simplify it to find the first Dot use.

[manman-ren]

Instead of calling scheduleDependencies on CoarseSchedule, we can have a different helper scheduleDependenciesWithAttributes(forOp, numStages) to apply (stage, cluster) attributes on the dependent ops.

[pawelszczerbuk]

Should we have these strings defined in some shared file?

[ThomasRaoux]

What is the meaning of loop.load? Ideally we should only annotate with stage and schedule information otherwise we are adding ad hoc conventions between the different passes.

[manman-ren]

Yeah I haven't paid much attention to this yet. We should define them in some header.

The plan is to only annotate with (stage, cluster), but the last problem of separating out assignMemoryLayouts caused some complications. We are deciding which loads to pipeline in part1 of assignMemoryLayouts, which is in LoopScheduling, and the second part of assignMemoryLayouts in SWP needs to know which loads are going to be pipelined.