[Mlir-commits] [mlir] [mlir][tosa] Add constant folding support for `tosa.dim` (PR #176975)

[Sayan Saha]

================
@@ -1764,3 +1764,18 @@ OpFoldResult tosa::AddShapeOp::fold(FoldAdaptor adaptor) {
   return binaryFolder<FoldAddAdaptor>(
       input1Attr, input2Attr, input1Attr.getType(), /*foldDenseValues=*/true);
 }
+
+OpFoldResult tosa::DimOp::fold(FoldAdaptor adaptor) {
+  const auto inputTy = llvm::dyn_cast<ShapedType>(getInput1().getType());
+  if (!inputTy || !inputTy.hasRank())
+    return {};
+  const int32_t axis = getAxis();
+  const int64_t dimSize = inputTy.getDimSize(axis);
+  if (ShapedType::isDynamic(dimSize))
+    return {};
+
+  OpBuilder builder(getContext());
+  const int64_t rank = cast<tosa::shapeType>(getResult().getType()).getRank();
----------------
sahas3 wrote:

Won't rank always be 1 irrespective of the input tensor shape since `tosa.dim` only accepts a scalar `axis` and returns the shape for that dimension?

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