[Mlir-commits] [mlir] eb1b55c - [mlir][tosa] Add tosa.reduce_any and tosa.reduce_all linalg lowering
Rob Suderman
llvmlistbot at llvm.org
Fri Apr 2 14:33:28 PDT 2021
Author: Rob Suderman
Date: 2021-04-02T14:32:18-07:00
New Revision: eb1b55c652a78f17a913ffbc4369b75cea25f23f
URL: https://github.com/llvm/llvm-project/commit/eb1b55c652a78f17a913ffbc4369b75cea25f23f
DIFF: https://github.com/llvm/llvm-project/commit/eb1b55c652a78f17a913ffbc4369b75cea25f23f.diff
LOG: [mlir][tosa] Add tosa.reduce_any and tosa.reduce_all linalg lowering
Added lowerings for Tosa's reduce boolean operations. This includes a fix to
maintain the output rank of reduce operations.
Reviewed By: silvas
Differential Revision: https://reviews.llvm.org/D99228
Added:
Modified:
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index dbc7654818b3c..3fdbc6f89733b 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -514,6 +514,12 @@ static Attribute createInitialValueForReduceOp(Operation *op, Type elementTy,
return rewriter.getIntegerAttr(
elementTy, APInt::getSignedMinValue(elementTy.getIntOrFloatBitWidth()));
+ if (isa<tosa::ReduceAllOp>(op) && elementTy.isInteger(1))
+ return rewriter.getIntegerAttr(elementTy, APInt::getAllOnesValue(1));
+
+ if (isa<tosa::ReduceAnyOp>(op) && elementTy.isInteger(1))
+ return rewriter.getIntegerAttr(elementTy, APInt::getNullValue(1));
+
if (isa<tosa::ArgMaxOp>(op) && elementTy.isa<FloatType>())
return rewriter.getFloatAttr(
elementTy, APFloat::getLargest(
@@ -573,6 +579,12 @@ static Value createLinalgBodyCalculationForReduceOp(Operation *op,
return rewriter.create<mlir::SelectOp>(loc, predicate, args[0], args[1]);
}
+ if (isa<tosa::ReduceAllOp>(op) && elementTy.isInteger(1))
+ return rewriter.create<mlir::AndOp>(loc, args);
+
+ if (isa<tosa::ReduceAnyOp>(op) && elementTy.isInteger(1))
+ return rewriter.create<mlir::OrOp>(loc, args);
+
return {};
}
@@ -613,6 +625,8 @@ static LogicalResult reduceMatchAndRewriteHelper(Operation *op, uint64_t axis,
: getParallelIteratorTypeName());
if (axis != i)
dstExprs.push_back(mlir::getAffineDimExpr(i, rewriter.getContext()));
+ else
+ dstExprs.push_back(rewriter.getAffineConstantExpr(0));
}
bool didEncounterError = false;
@@ -1419,7 +1433,8 @@ void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
PointwiseConverter<tosa::CeilOp>, PointwiseConverter<tosa::FloorOp>,
PointwiseConverter<tosa::ClampOp>, PointwiseConverter<tosa::ReluNOp>,
PointwiseConverter<tosa::SigmoidOp>, IdentityNConverter<tosa::IdentityOp>,
- IdentityNConverter<tosa::IdentityNOp>, ReduceConverter<tosa::ReduceMinOp>,
+ IdentityNConverter<tosa::IdentityNOp>, ReduceConverter<tosa::ReduceAllOp>,
+ ReduceConverter<tosa::ReduceAnyOp>, ReduceConverter<tosa::ReduceMinOp>,
ReduceConverter<tosa::ReduceMaxOp>, ReduceConverter<tosa::ReduceSumOp>,
ReduceConverter<tosa::ReduceProdOp>, ArgMaxConverter, ConcatConverter,
PadConverter, ReshapeConverter, RescaleConverter, ReverseConverter,
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index 83ec9222cf3c0..1bc4d6a8d0132 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -442,98 +442,124 @@ func @test_transpose(%arg0: tensor<1x2x3xi32>) -> () {
// -----
// 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
}
More information about the Mlir-commits
mailing list