diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td index df067397558ee7..5823f5447522c5 100644 --- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td +++ b/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; diff --git a/mlir/test/Dialect/Tosa/reduce_rank5.mlir b/mlir/test/Dialect/Tosa/reduce_rank5.mlir new file mode 100644 index 00000000000000..8d0172942ecaab --- /dev/null +++ b/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> +}