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

[llvmlistbot at llvm.org]

Author: Luke Hutton
Date: 2026-01-27T18:26:55Z
New Revision: 3f5a5d45d18a514f086f3e07c9676ca5fb95bbe9

URL: https://github.com/llvm/llvm-project/commit/3f5a5d45d18a514f086f3e07c9676ca5fb95bbe9
DIFF: https://github.com/llvm/llvm-project/commit/3f5a5d45d18a514f086f3e07c9676ca5fb95bbe9.diff

LOG: [mlir][tosa] Add constant folding support for `tosa.dim` (#176975)

This enhances shape inference.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
    mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
    mlir/test/Dialect/Tosa/constant_folding.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
index 1783a5ef7c961..57fd1d2d20aa8 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
@@ -139,6 +139,8 @@ def Tosa_DimOp : Tosa_ShapeOp<"dim", [Pure]> {
   let results = (outs Tosa_Shape:$output);
 
   let hasVerifier = 1;
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 5f41c8c3f300f..07a0b6742d48a 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -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 auto resultAttrTy =
+      RankedTensorType::get(/*rank=*/1, builder.getIndexType());
+  return DenseElementsAttr::get(resultAttrTy, dimSize);
+}

diff  --git a/mlir/test/Dialect/Tosa/constant_folding.mlir b/mlir/test/Dialect/Tosa/constant_folding.mlir
index 1007af6c8bd82..4cb9a46e5d049 100644
--- a/mlir/test/Dialect/Tosa/constant_folding.mlir
+++ b/mlir/test/Dialect/Tosa/constant_folding.mlir
@@ -683,3 +683,30 @@ func.func @test_no_fold_add_shape_negative_overflow() -> !tosa.shape<6> {
   %c = tosa.add_shape %a, %b : (!tosa.shape<6>, !tosa.shape<6>) -> !tosa.shape<6>
   return %c : !tosa.shape<6>
 }
+
+// -----
+
+// CHECK-LABEL: @test_fold_dim
+// CHECK: tosa.const_shape  {values = dense<6> : tensor<1xindex>} : () -> !tosa.shape<1>
+func.func @test_fold_dim(%arg0: tensor<6xi32>) -> !tosa.shape<1> {
+  %dim = tosa.dim %arg0 {axis = 0 : i32} : (tensor<6xi32>) -> !tosa.shape<1>
+  return %dim : !tosa.shape<1>
+}
+
+// -----
+
+// CHECK-LABEL: @test_no_fold_dim_unranked_input
+// CHECK: tosa.dim
+func.func @test_no_fold_dim_unranked_input(%arg0: tensor<*xi32>) -> !tosa.shape<1> {
+  %dim = tosa.dim %arg0 {axis = 0 : i32} : (tensor<*xi32>) -> !tosa.shape<1>
+  return %dim : !tosa.shape<1>
+}
+
+// -----
+
+// CHECK-LABEL: @test_no_fold_dim_dynamic
+// CHECK: tosa.dim
+func.func @test_no_fold_dim_dynamic(%arg0: tensor<4x?xi32>) -> !tosa.shape<1> {
+  %dim = tosa.dim %arg0 {axis = 1 : i32} : (tensor<4x?xi32>) -> !tosa.shape<1>
+  return %dim : !tosa.shape<1>
+}