[Mlir-commits] [mlir] [mlir] [dataflow] Refactoring the definition of program points in dat… (PR #105656)

[donald chen]

cxy-1993 wrote:

> > In my opinion, this doesn't fundamentally deviate from the essence of current data flow analysis. When we update flat symbol attrs or op result, we still rely on initializeGlobalSlotsOp/GlobalSlotGetOp to update operand states. It's just that the current updates only modify the states of some operands in initializeGlobalSlotsOp(visit op will update all). Therefore, I suggest modifying the existing visit op result at (https://github.com/llvm/torch-mlir/blob/98e08023bbf71e00ab81e980eac9f7c96f1f24b4/lib/Dialect/Torch/Transforms/InlineGlobalSlots.cpp#L193C31-L193C46) to visit GlobalSlotGetOp, and modifying the visit flatSymbolRefPoint at (https://github.com/llvm/torch-mlir/blob/98e08023bbf71e00ab81e980eac9f7c96f1f24b4/lib/Dialect/Torch/Transforms/InlineGlobalSlots.cpp#L210) to visit initializeGlobalSlotsOp. This would align with the current dataflow analysis's definition of a program point.
> > If you have any further questions, please don't hesitate to discuss them with me. Additionally, I'm more than happy to assist you in modifying this section of code if needed.
> 
> Thanks for your suggestion @cxy-1993! I am currently looking into how to adapt torch-mlir but I am unfortunately not familiar with the data-flow analysis and might have missed something. Therefore please allow to ask for further clarification.
> 
> While the argument to the `visit(ProgramPoint point)` method still is of a `ProgramPoint`, the latter was changed and now essentially is a `PointerUnion`. Before this change, `point` could be casted to a `Value` (within `visit()`) which is no longer possible. To get the value, one could instead cast `point` to an `Operation` and obtain the result value from this op. Furthermore, `getProgramPoint` needs to be replaced with `getLatticeAnchor`, thus
> 
> ```c++
> auto *flatSymbolRefPoint = getProgramPoint<FlatSymbolRefProgramPoint>(globalSlotGet.getSlotAttr());
> auto *valueState = getOrCreateFor<InlineGlobalSlotsAnalysisState>(flatSymbolRefPoint, globalSlotGet.getResult()):
> ```
> 
> becomes
> 
> ```c++
> auto *flatSymbolRefPoint = getLatticeAnchor<FlatSymbolRefProgramPoint>(globalSlotGet.getSlotAttr());
> auto *valueState = getOrCreateFor<InlineGlobalSlotsAnalysisState>(*flatSymbolRefPoint, globalSlotGet.getResult()):
> ```
> 
> What is failing here is `getOrCreateFor` as the first argument needs to be a `ProgramPoint` but now is a `FlatSymbolRefProgramPoint`, while assuming that the constructor of `InlineGlobalSlotsAnalysisState` needs to be changed from
> 
> ```c++
> InlineGlobalSlotsAnalysisState(ProgramPoint point) : AnalysisState(point) {
> ```
> 
> to
> 
> ```c++
> InlineGlobalSlotsAnalysisState(LatticeAnchor anchor) : AnalysisState(anchor) {
> ```
> 
> Any pointers would be appreciated!

I have a limited understanding of the specific functionality of this pass. Please point out any issues if you find any. For the sake of clarity, let's call the program point before this patch the 'ori program point'. My overall modification idea is as follows:

The InlineGlobalSlotsAnalysis::visit function might accept two types of ori program points and update the lattice state: one is the result of Torch::GlobalSlotGetOp, and the other is FlatSymbolRefProgramPoint. They will respectively use Torch::GlobalSlotGetOp and Torch::InitializeGlobalSlotsOp to update the states of other lattices.

Therefore, I suggest modifying the original code:

getOrCreateFor<InlineGlobalSlotsAnalysisState>(flatSymbolRefPoint, xxx) to getOrCreateFor<InlineGlobalSlotsAnalysisState>(initializeGlobalSlotsOp, xxx).
getOrCreateFor<InlineGlobalSlotsAnalysisState>(value, xx) to something like getOrCreateFor<InlineGlobalSlotsAnalysisState>(value.getDefiningOp<Torch::GlobalSlotGetOp>(), flatSymbolRefPoint).

At this point, the visit's program point will only have to deal initializeGlobalSlotsOp and GlobalSlotGetOp, and we can use the information from these two ops to update the lattice state related to these op.

https://github.com/llvm/llvm-project/pull/105656