mirror of https://github.com/llvm/torch-mlir
a83e106f92
This is a first step towards reworking the scalarize-shapes pass which has been integral to our ONNX frontend path detangling shape computations. ## Purpose: 1. Restrict the scope of the pass to only apply to op sequences which are used to compute shapes. 2. Make the pass more efficient by applying patterns in an appropriate order for scalarization propagation. 3. Report failed scalarization patterns for easier debugging (Not yet implemented). I can't seem to find a good path for this right now to capture the right diagnostics. I'd like to defer this addition to a later patch so we can add some high-value patterns to this pass in the meantime. With these changes, some reworking of the conversions themselves will be necessary. 1. The removal of the SqueezeDim fold pattern was an appropriate fix to avoid folding a pattern that may be needed to propagate further. The reversal of pattern application order uncovered this bug. The addition of rank 0 item logic was added to replace the functionality needed from the squeeze dim pattern. 2. Rework getListFromTensor to modify a `SmallVector<OpFoldResult>` to allow processing value tensor literals without immediately materializing the ints. This should factor out a significant portion of code that was used in specific cases to handle constants. ## RFC 1: Currently, we are going to add all prim list of int ops to the worklist. Can anyone identify problems with uniformly anchoring on prim lists of ints? E.g. Does there exist a Torch Op satisfying all of the following conditions: 1. Accepts a list of constant ints, LIST, as an input 2. The role of LIST is **not** shape related. All the examples I can think of are indeed shape related: padding ints passed to a pad op, kernel size ints passed to a conv op, size ints passed to a view op, etc. 4. The LIST is not gotten entirely from scalars already. If there does not exist a torch op satisfying all three of those conditions, I think it will be safe to "anchor" on prim lists of ints. ### Conclusion for RFC 1: I just scanned through the `GeneratedTorchOps.td` and `TorchOps.td` for all references of `AnyTorchListOfTorchIntType` and verified this will not be problematic to apply in any of those cases. ## RFC 2: What should I use to report failed scalarization? Like my dumb idea was just to walk back through the func op after applying the passes and check if anything in the worklist is still a tensor. If so, emit/log a warning. It certainly works, since you can just look at the warnings and start debugging from the last printed warning upwards, but there has to be a better way to handle this without walking back through the func.func op. ### Conclusion for RFC 2: I tried a few things without much success. The fundamental problem is that identifying the cause of a failed scalarization could be myriad: 1. We could be missing a pattern for an op entirely: E.g., a pattern we need is scalarizing rank0 arithmetic ops (e.g. AtenMulTensorOp -> AtenMulIntOp). 2. We could fail a scalarization pattern because it should fold instead. This is specifically the case for rank0 where.self ops. These ops MUST fold, or we need to have custom lowering logic for the rank 0 case. 3. Walking through the func op a second time and emiting a warning for ops that have tensor result types seems to give locations that are inconsistent or hard to track in the converted IR. Doing this on IR that doesn't apply any patterns seems to give decent information, but it's still dramatically insufficient considering how complex these patterns can get, and still takes manually reading IR to try and figure out what is really blocking the simplification. I'd like to skip out on fleshing out the error reporting for now and come back to it after iterating a few time on the patterns. |
||
|---|---|---|
| .. | ||
| CAPI | ||
| Conversion | ||
| Dialect | ||
| RefBackend | ||
| CMakeLists.txt | ||
| InitAll.cpp | ||