diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp --- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp +++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp @@ -380,6 +380,12 @@ return rewriter.getIntegerAttr( elementTy, APInt::getSignedMinValue(elementTy.getIntOrFloatBitWidth())); + if (isa(op) && elementTy.isInteger(1)) + return rewriter.getIntegerAttr(elementTy, APInt::getAllOnesValue(1)); + + if (isa(op) && elementTy.isInteger(1)) + return rewriter.getIntegerAttr(elementTy, APInt::getNullValue(1)); + return {}; } @@ -430,6 +436,12 @@ return rewriter.create(loc, predicate, args[0], args[1]); } + if (isa(op) && elementTy.isInteger(1)) + return rewriter.create(loc, args); + + if (isa(op) && elementTy.isInteger(1)) + return rewriter.create(loc, args); + return {}; } @@ -470,6 +482,8 @@ : getParallelIteratorTypeName()); if (axis != i) dstExprs.push_back(mlir::getAffineDimExpr(i, rewriter.getContext())); + else + dstExprs.push_back(rewriter.getAffineConstantExpr(0)); } bool didEncounterError = false; @@ -877,7 +891,8 @@ PointwiseConverter, PointwiseConverter, PointwiseConverter, PointwiseConverter, IdentityNConverter, - IdentityNConverter, ReduceConverter, + IdentityNConverter, ReduceConverter, + ReduceConverter, ReduceConverter, ReduceConverter, ReduceConverter, ReduceConverter, ConcatOpConversion, ReshapeOpConverter, TransposeConverter, RescaleOpConverter>(context); diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir --- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir +++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir @@ -339,98 +339,124 @@ // ----- // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> -// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> -// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)> +// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (0, d1)> +// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, 0)> // CHECK-LABEL: @reduce_float // CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xf32> func @reduce_float(%arg0: tensor<5x4xf32>) -> () { - // CHECK: [[INIT:%.+]] = linalg.init_tensor [4] + // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 4] // CHECK: [[CST0:%.+]] = constant 0.0 // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST0]]) - // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<4xf32>) + // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<1x4xf32>) // CHECK: ^bb0(%arg1: f32, %arg2: f32) // CHECK: [[RES:%.+]] = addf %arg1, %arg2 : f32 // CHECK: linalg.yield [[RES]] : f32 - %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<4xf32> + %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32> - // CHECK: [[INIT:%.+]] = linalg.init_tensor [5] + // CHECK: [[INIT:%.+]] = linalg.init_tensor [5, 1] // CHECK: [[CST0:%.+]] = constant 0.0 // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST0]]) - // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<5xf32>) + // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<5x1xf32>) // CHECK: ^bb0(%arg1: f32, %arg2: f32) // CHECK: [[RES:%.+]] = addf %arg1, %arg2 : f32 // CHECK: linalg.yield [[RES]] : f32 - %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xf32>) -> tensor<5xf32> + %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xf32>) -> tensor<5x1xf32> // CHECK: constant 1.0 // CHECK: linalg.fill // CHECK: linalg.generic // CHECK: mulf - %2 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<4xf32> + %2 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32> // CHECK: constant 3.40282347E+38 : f32 // CHECK: linalg.fill // CHECK: linalg.generic // CHECK: cmpf olt // CHECK: select - %3 = "tosa.reduce_min"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<4xf32> + %3 = "tosa.reduce_min"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32> // CHECK: constant -3.40282347E+38 : f32 // CHECK: linalg.fill // CHECK: linalg.generic // CHECK: cmpf ogt // CHECK: select - %4 = "tosa.reduce_max"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<4xf32> + %4 = "tosa.reduce_max"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32> return } // ----- // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> -// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> -// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)> +// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (0, d1)> +// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, 0)> // CHECK-LABEL: @reduce_int // CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xi32> func @reduce_int(%arg0: tensor<5x4xi32>) -> () { - // CHECK: [[INIT:%.+]] = linalg.init_tensor [4] + // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 4] // CHECK: [[CST0:%.+]] = constant 0 // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST0]]) - // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<4xi32>) + // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<1x4xi32>) // CHECK: ^bb0(%arg1: i32, %arg2: i32) // CHECK: [[RES:%.+]] = addi %arg1, %arg2 : i32 // CHECK: linalg.yield [[RES]] : i32 - %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<4xi32> + %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32> - // CHECK: [[INIT:%.+]] = linalg.init_tensor [5] + // CHECK: [[INIT:%.+]] = linalg.init_tensor [5, 1] // CHECK: [[CST0:%.+]] = constant 0 // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST0]]) - // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<5xi32>) + // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<5x1xi32>) // CHECK: ^bb0(%arg1: i32, %arg2: i32) // CHECK: [[RES:%.+]] = addi %arg1, %arg2 : i32 // CHECK: linalg.yield [[RES]] : i32 - %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5xi32> + %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5x1xi32> // CHECK: constant 1 // CHECK: linalg.fill // CHECK: linalg.generic // CHECK: muli - %2 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<4xi32> + %2 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32> // CHECK: constant 2147483647 : i32 // CHECK: linalg.fill // CHECK: linalg.generic // CHECK: cmpi slt // CHECK: select - %3 = "tosa.reduce_min"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<4xi32> + %3 = "tosa.reduce_min"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32> // CHECK: constant -2147483648 : i32 // CHECK: linalg.fill // CHECK: linalg.generic // CHECK: cmpi sgt // CHECK: select - %4 = "tosa.reduce_max"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<4xi32> + %4 = "tosa.reduce_max"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32> + return +} + +// ----- + +// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> +// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (0, d1)> + +// CHECK-LABEL: @reduce_bool +// CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xi1> +func @reduce_bool(%arg0: tensor<5x4xi1>) -> () { + // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 4] + // CHECK: [[CST0:%.+]] = constant true + // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST0]]) + // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi1>) outs([[FILL]] : tensor<1x4xi1>) + // CHECK: ^bb0(%arg1: i1, %arg2: i1) + // CHECK: [[RES:%.+]] = and %arg1, %arg2 : i1 + // CHECK: linalg.yield [[RES]] : i1 + %0 = "tosa.reduce_all"(%arg0) {axis = 0 : i64} : (tensor<5x4xi1>) -> tensor<1x4xi1> + + // CHECK: constant false + // CHECK: linalg.fill + // CHECK: linalg.generic + // CHECK: or + %1 = "tosa.reduce_any"(%arg0) {axis = 0 : i64} : (tensor<5x4xi1>) -> tensor<1x4xi1> + return }