tosa.reduce_sum: Support any rank

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

@@ -1224,12 +1224,12 @@ def Tosa_ReduceAllOp : Tosa_Op<"reduce_all", [
   }];
 
   let arguments = (ins
-    Tosa_Tensor1Dto4D:$input,
+    Tosa_Tensor:$input,
     I64Attr:$axis
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
@@ -1253,12 +1253,12 @@ def Tosa_ReduceAnyOp : Tosa_Op<"reduce_any", [
   }];
 
   let arguments = (ins
-    Tosa_Tensor1Dto4D:$input,
+    Tosa_Tensor:$input,
     I64Attr:$axis
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
@@ -1282,12 +1282,12 @@ def Tosa_ReduceMaxOp : Tosa_Op<"reduce_max", [
   }];
 
   let arguments = (ins
-    Tosa_Tensor1Dto4D:$input,
+    Tosa_Tensor:$input,
     I64Attr:$axis
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
@@ -1311,12 +1311,12 @@ def Tosa_ReduceMinOp : Tosa_Op<"reduce_min", [
   }];
 
   let arguments = (ins
-    Tosa_Tensor1Dto4D:$input,
+    Tosa_Tensor:$input,
     I64Attr:$axis
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
@@ -1340,12 +1340,12 @@ def Tosa_ReduceProdOp : Tosa_Op<"reduce_prod", [
   }];
 
   let arguments = (ins
-    Tosa_Tensor1Dto4D:$input,
+    Tosa_Tensor:$input,
     I64Attr:$axis
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
@@ -1369,12 +1369,12 @@ def Tosa_ReduceSumOp : Tosa_Op<"reduce_sum", [
   }];
 
   let arguments = (ins
-    Tosa_Tensor1Dto4D:$input,
+    Tosa_Tensor:$input,
     I64Attr:$axis
   );
 
   let results = (outs
-    Tosa_Tensor1Dto4D:$output
+    Tosa_Tensor:$output
   );
 
   let hasFolder = 1;

mlir/test/Dialect/Tosa/reduce_rank5.mlir

@@ -0,0 +1,44 @@
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+// AMD extention to allow any rank in tosa reduction operations
+
+// -----
+// CHECK-LABEL: reduce_all
+func.func @test_reduce_all(%arg0: tensor<13x21x3x5x9xi1>) -> tensor<1x21x3x5x9xi1> {
+  %0 = "tosa.reduce_all"(%arg0) {axis = 0 : i64} : (tensor<13x21x3x5x9xi1>) -> tensor<1x21x3x5x9xi1>
+  return %0 : tensor<1x21x3x5x9xi1>
+}
+
+// -----
+// CHECK-LABEL: reduce_any
+func.func @test_reduce_any(%arg0: tensor<13x21x3x5x9xi1>) -> tensor<1x21x3x5x9xi1> {
+  %0 = "tosa.reduce_any"(%arg0) {axis = 0 : i64} : (tensor<13x21x3x5x9xi1>) -> tensor<1x21x3x5x9xi1>
+  return %0 : tensor<1x21x3x5x9xi1>
+}
+
+// -----
+// CHECK-LABEL: reduce_max
+func.func @test_reduce_max(%arg0: tensor<13x21x3x5x9xf32>) -> tensor<1x21x3x5x9xf32> {
+  %0 = "tosa.reduce_max"(%arg0) {axis = 0 : i64} : (tensor<13x21x3x5x9xf32>) -> tensor<1x21x3x5x9xf32>
+  return %0 : tensor<1x21x3x5x9xf32>
+}
+
+// -----
+// CHECK-LABEL: reduce_min
+func.func @test_reduce_min(%arg0: tensor<13x21x3x5x9xf32>) -> tensor<1x21x3x5x9xf32> {
+  %0 = "tosa.reduce_min"(%arg0) {axis = 0 : i64} : (tensor<13x21x3x5x9xf32>) -> tensor<1x21x3x5x9xf32>
+  return %0 : tensor<1x21x3x5x9xf32>
+}
+
+// -----
+// CHECK-LABEL: reduce_product
+func.func @test_reduce_product(%arg0: tensor<13x21x3x5x9xf32>) -> tensor<1x21x3x5x9xf32> {
+  %0 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<13x21x3x5x9xf32>) -> tensor<1x21x3x5x9xf32>
+  return %0 : tensor<1x21x3x5x9xf32>
+}
+
+// -----
+// CHECK-LABEL: reduce_sum
+func.func @test_reduce_sum(%arg0: tensor<13x21x3x5x9xf32>) -> tensor<1x21x3x5x9xf32> {
+  %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<13x21x3x5x9xf32>) -> tensor<1x21x3x5x9xf32>
+  return %0 : tensor<1x21x3x5x9xf32>
+}