[TOSA] Add upsample_nearest2d, split_dim, outer, GELU tanh mode and misc #3886
Conversation
- Add Torch to TOSA lowering for the following ops: + torch.aten.upsample_nearest2d + torch.aten.upsample_nearest2d.vec + torch.aten.outer + torch.prims.split_dim - Add Tanh approximation mode for GELU lowering - Add different types support for compare ops - Add different input and output types support for linalg vector norm lowering - Update xfail with new e2e results - Add new LIT tests to basic.mlir Signed-off-by: Justin Ngo <[email protected]> Change-Id: I7b1d44d94319cf94fcc9d234cc07708ef9ce321e
| // "tanh" approximate | ||
| // GELU(x) = 0.5 * x * (1 + Tanh(sqrt(2/pi) * (x + 0.044715 * x^3)) | ||
| auto selfShape = selfType.getShape(); | ||
| auto numElem = std::accumulate(selfShape.begin(), selfShape.end(), 1, |
There was a problem hiding this comment.
This assumes the shape is all static right? Can you add a failure check for dynamic shape?
Alternatively, can't we rely on broadcasting semantics of TosaOps to correctly expand the shape even for dynamic dims if these constants are defined with size 1?
There was a problem hiding this comment.
I've added a static shape check for tanh approximate. I've tried using size 1 constant but for some reasons it crashed fx_importer_target (specifically the ElementwiseGeluApproximateTanhModule_basic test), so I will resort to static shape requirement for now since TOSA currently only supports static shape.
| self, rewriter.getDenseI64ArrayAttr(reshapedSelfShape)); | ||
|
|
||
| // Calculate PyTorch-styled gather indices | ||
| SmallVector<int32_t> targetIndicesVec; |
There was a problem hiding this comment.
Should this be int64_t as well since all the other types are int64_t?
There was a problem hiding this comment.
I'm currently setting it int32_t since other places in the code that use TOSA GatherOp or ScatterOp all have i32 indices conversion/definition (even in your AtenIndexTensorHackedTwin commit about a month ago, you also used i32 indices conversion). I will double-check this matter later and if it needs to be changed then I will change all occurrences in a later PR.
| static_cast<double>(outputWidth) / static_cast<double>(selfWidth); | ||
| } | ||
|
|
||
| if (isOutputSizeNone) { |
There was a problem hiding this comment.
#nit: can be merged with the previous if block since isOutputSizeNone implies !isScaleFactorsNone
There was a problem hiding this comment.
I've rewritten the if-else flow here to make it clearer and more concise. It should address this nit too.
Change-Id: I7b1d44d94319cf94fcc9d234cc07708ef9ce321e