[Mlir-commits] [mlir] [mlir] [linalg] Fold reduce(broadcast(x)) max/min (PR #213190)

[Chuanqi Xu]

================
@@ -1,5 +1,88 @@
 // RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file | FileCheck %s
 
+// CHECK-LABEL: func.func @reduce_broadcast_maxsi
+func.func @reduce_broadcast_maxsi(%input: tensor<3xi32>,
+    %broadcast_init: tensor<3x4xi32>)
+    -> tensor<3xi32> {
+  %reduce_init = arith.constant dense<-2147483648> : tensor<3xi32>
+  %broadcasted = linalg.broadcast
+      ins(%input : tensor<3xi32>)
+      outs(%broadcast_init : tensor<3x4xi32>) dimensions = [1]
+  %result = linalg.reduce
+      ins(%broadcasted : tensor<3x4xi32>)
+      outs(%reduce_init : tensor<3xi32>) dimensions = [1]
+      (%in: i32, %out: i32) {
+      %max = arith.maxsi %in, %out : i32
+      linalg.yield %max : i32
+  }
+  return %result : tensor<3xi32>
+}
+// CHECK-NEXT: return %{{.*}} : tensor<3xi32>
+
+// -----
+
+// CHECK-LABEL: func.func @reduce_broadcast_minui
+func.func @reduce_broadcast_minui(%input: tensor<3xi8>,
+    %broadcast_init: tensor<2x3x4xi8>)
+    -> tensor<3xi8> {
+  %reduce_init = arith.constant dense<255> : tensor<3xi8>
+  %broadcasted = linalg.broadcast
+      ins(%input : tensor<3xi8>)
+      outs(%broadcast_init : tensor<2x3x4xi8>) dimensions = [0, 2]
+  %result = linalg.reduce
+      ins(%broadcasted : tensor<2x3x4xi8>)
+      outs(%reduce_init : tensor<3xi8>) dimensions = [0, 2]
+      (%in: i8, %out: i8) {
+      %min = arith.minui %in, %out : i8
+      linalg.yield %min : i8
+  }
+  return %result : tensor<3xi8>
+}
+// CHECK-NEXT: return %{{.*}} : tensor<3xi8>
+
+// -----
+
+// CHECK-LABEL: func.func @reduce_broadcast_wrong_dimension
+// CHECK: linalg.broadcast
+// CHECK: linalg.reduce
+func.func @reduce_broadcast_wrong_dimension(%input: tensor<3xi32>,
+    %broadcast_init: tensor<3x4xi32>, %reduce_init: tensor<4xi32>)
+    -> tensor<4xi32> {
+  %broadcasted = linalg.broadcast
+      ins(%input : tensor<3xi32>)
+      outs(%broadcast_init : tensor<3x4xi32>) dimensions = [1]
+  %result = linalg.reduce
+      ins(%broadcasted : tensor<3x4xi32>)
+      outs(%reduce_init : tensor<4xi32>) dimensions = [0]
+      (%in: i32, %out: i32) {
+      %max = arith.maxsi %in, %out : i32
+      linalg.yield %max : i32
+  }
+  return %result : tensor<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @reduce_broadcast_non_identity_init
+// CHECK: linalg.broadcast
+// CHECK: linalg.reduce
+func.func @reduce_broadcast_non_identity_init(%input: tensor<3xi32>,
+    %broadcast_init: tensor<3x4xi32>)
+    -> tensor<3xi32> {
+  %reduce_init = arith.constant dense<0> : tensor<3xi32>
+  %broadcasted = linalg.broadcast
+      ins(%input : tensor<3xi32>)
+      outs(%broadcast_init : tensor<3x4xi32>) dimensions = [1]
+  %result = linalg.reduce
+      ins(%broadcasted : tensor<3x4xi32>)
+      outs(%reduce_init : tensor<3xi32>) dimensions = [1]
+      (%in: i32, %out: i32) {
+      %max = arith.maxsi %in, %out : i32
+      linalg.yield %max : i32
+  }
+  return %result : tensor<3xi32>
+}
+
----------------
ChuanqiXu9 wrote:

Done

https://github.com/llvm/llvm-project/pull/213190