[MLIR][TOSA] Fix Conv3D bias dim check

[RoboTux]

verifyConvOpErrorIf() assumes output channel is the 4th dimension of the
output type but this is wrong for conv3d which now uses that verifier.
Use rank - 1 which works accross the operations using this verifier
(conv2d, conv3d and depthwise_conv3d).

[llvmbot]

@llvm/pr-subscribers-mlir-tosa

@llvm/pr-subscribers-mlir

Author: Thomas Preud'homme (RoboTux)

Changes

verifyConvOpErrorIf() assumes output channel is the 4th dimension of the
output type but this is wrong for conv3d which now uses that verifier.
Use rank - 1 which works accross the operations using this verifier
(conv2d, conv3d and depthwise_conv3d).

---

Full diff: https://github.com/llvm/llvm-project/pull/137296.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Tosa/IR/TosaOps.cpp (+2-1)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 751ae785bda6f..17873444b2d71 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -558,7 +558,8 @@ static LogicalResult verifyConvOpErrorIf(T op) {
     return success();
 
   const int64_t biasChannels = biasType.getDimSize(0);
-  const int64_t outputChannels = outputType.getDimSize(3);
+  const int64_t outputChannels =
+      outputType.getDimSize(outputType.getRank() - 1);
   if (biasChannels == ShapedType::kDynamic ||
       outputChannels == ShapedType::kDynamic)
     // Skip following checks if biasChannels or outputChannels is dynamic dim

[lhutton1]

Thanks for the fix @RoboTux!