[TOSA] Add transposed conv support #4360
Conversation
Lower aten.conv_transpose2d into tosa.transpose_conv2d. Refresh FX importer TOSA xfails to drop the transpose-conv cases that now pass, and document the weight layout mapping. Change-Id: I23be2230a0948784402dca574597db1d979d5aee
|
|
||
| // Weight layout reference: | ||
| // Conv : PyTorch OIHW -> TOSA OHWI | ||
| // Depthwise : PyTorch OIHW* -> TOSA HWIM (*out = in * multiplier) |
There was a problem hiding this comment.
Thanks for adding this reference, though it's not clear what you are trying to imply with (*out = in * multiplier) here.
|
|
||
| SmallVector<int32_t> iohwToOhwi({1, 2, 3, 0}); | ||
| SmallVector<int64_t> ohwiWeightShape( | ||
| {weightShape[1], weightShape[2], weightShape[3], weightShape[0]}); |
There was a problem hiding this comment.
#nit: weightShape[iohwToOhwi[0]], weightShape[iohwToOhwi[1]], ... and so on will help readability I think
| // Result type is NHWC (we'll transpose back). | ||
| auto outNCHW = makeShapeTorchCompatible(outputTy.getShape()); | ||
| SmallVector<int64_t> outNHWC( | ||
| {outNCHW[0], outNCHW[2], outNCHW[3], outNCHW[1]}); |
There was a problem hiding this comment.
similar comment to above to use the nchwToNhwcDims instead of hardcoding 0,2,3,1
| @@ -1,8 +1,23 @@ | |||
| // RUN: torch-mlir-opt <%s -convert-torch-to-tosa -split-input-file -verify-diagnostics | |||
| // RUN: torch-mlir-opt <%s -convert-torch-to-tosa -split-input-file -verify-diagnostics | FileCheck %s | |||
There was a problem hiding this comment.
Is verify-diagnostics needed anymore unless you plan to add tests to verify the failure cases?
| // Quantized rescale. | ||
| Value rescaledResult = transposedOutput; | ||
| if (isa<quant::QuantizedType>(inputElemTy)) { | ||
| rescaledResult = tosa::buildRescaleOpConvOutput( |
There was a problem hiding this comment.
ConvTranspose2DQInt8_basic e2e test likely triggers this code path -- any idea why that is still failing? If that's not the correct test, are there any e2e tests that triggers this code?
2 participants
Lower aten.conv_transpose2d into tosa.transpose_conv2d. Refresh FX importer TOSA xfails to drop the transpose-conv cases that now pass, and document the weight layout mapping.