[Mlir-commits] [mlir] [mlir][SPIR-V] Lower AND/OR/XOR vector reductions (PR #192293)

[Igor Wodiany]

================
@@ -856,6 +856,105 @@ func.func @reduction_minui(%v : vector<3xi32>, %s: i32) -> i32 {
 
 // -----
 
+// CHECK-LABEL: func @reduction_and
+//  CHECK-SAME: (%[[V:.+]]: vector<4xi32>)
+//       CHECK:   %[[S0:.+]] = spirv.CompositeExtract %[[V]][0 : i32] : vector<4xi32>
+//       CHECK:   %[[S1:.+]] = spirv.CompositeExtract %[[V]][1 : i32] : vector<4xi32>
+//       CHECK:   %[[S2:.+]] = spirv.CompositeExtract %[[V]][2 : i32] : vector<4xi32>
+//       CHECK:   %[[S3:.+]] = spirv.CompositeExtract %[[V]][3 : i32] : vector<4xi32>
+//       CHECK:   %[[AND0:.+]] = spirv.BitwiseAnd %[[S0]], %[[S1]]
+//       CHECK:   %[[AND1:.+]] = spirv.BitwiseAnd %[[AND0]], %[[S2]]
+//       CHECK:   %[[AND2:.+]] = spirv.BitwiseAnd %[[AND1]], %[[S3]]
+//       CHECK:   return %[[AND2]]
+func.func @reduction_and(%v : vector<4xi32>) -> i32 {
+  %reduce = vector.reduction <and>, %v : vector<4xi32> into i32
+  return %reduce : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @reduction_and_acc
+//  CHECK-SAME: (%[[V:.+]]: vector<3xi32>, %[[S:.+]]: i32)
+//       CHECK:   %[[S0:.+]] = spirv.CompositeExtract %[[V]][0 : i32] : vector<3xi32>
+//       CHECK:   %[[S1:.+]] = spirv.CompositeExtract %[[V]][1 : i32] : vector<3xi32>
+//       CHECK:   %[[S2:.+]] = spirv.CompositeExtract %[[V]][2 : i32] : vector<3xi32>
+//       CHECK:   %[[AND0:.+]] = spirv.BitwiseAnd %[[S0]], %[[S1]]
+//       CHECK:   %[[AND1:.+]] = spirv.BitwiseAnd %[[AND0]], %[[S2]]
+//       CHECK:   %[[AND2:.+]] = spirv.BitwiseAnd %[[AND1]], %[[S]]
+//       CHECK:   return %[[AND2]]
+func.func @reduction_and_acc(%v : vector<3xi32>, %s: i32) -> i32 {
+  %reduce = vector.reduction <and>, %v, %s : vector<3xi32> into i32
+  return %reduce : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @reduction_or
+//  CHECK-SAME: (%[[V:.+]]: vector<4xi32>)
+//       CHECK:   %[[S0:.+]] = spirv.CompositeExtract %[[V]][0 : i32] : vector<4xi32>
+//       CHECK:   %[[S1:.+]] = spirv.CompositeExtract %[[V]][1 : i32] : vector<4xi32>
+//       CHECK:   %[[S2:.+]] = spirv.CompositeExtract %[[V]][2 : i32] : vector<4xi32>
+//       CHECK:   %[[S3:.+]] = spirv.CompositeExtract %[[V]][3 : i32] : vector<4xi32>
+//       CHECK:   %[[OR0:.+]] = spirv.BitwiseOr %[[S0]], %[[S1]]
+//       CHECK:   %[[OR1:.+]] = spirv.BitwiseOr %[[OR0]], %[[S2]]
+//       CHECK:   %[[OR2:.+]] = spirv.BitwiseOr %[[OR1]], %[[S3]]
+//       CHECK:   return %[[OR2]]
+func.func @reduction_or(%v : vector<4xi32>) -> i32 {
+  %reduce = vector.reduction <or>, %v : vector<4xi32> into i32
+  return %reduce : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @reduction_or_acc
+//  CHECK-SAME: (%[[V:.+]]: vector<3xi32>, %[[S:.+]]: i32)
+//       CHECK:   %[[S0:.+]] = spirv.CompositeExtract %[[V]][0 : i32] : vector<3xi32>
+//       CHECK:   %[[S1:.+]] = spirv.CompositeExtract %[[V]][1 : i32] : vector<3xi32>
+//       CHECK:   %[[S2:.+]] = spirv.CompositeExtract %[[V]][2 : i32] : vector<3xi32>
+//       CHECK:   %[[OR0:.+]] = spirv.BitwiseOr %[[S0]], %[[S1]]
+//       CHECK:   %[[OR1:.+]] = spirv.BitwiseOr %[[OR0]], %[[S2]]
+//       CHECK:   %[[OR2:.+]] = spirv.BitwiseOr %[[OR1]], %[[S]]
+//       CHECK:   return %[[OR2]]
+func.func @reduction_or_acc(%v : vector<3xi32>, %s: i32) -> i32 {
+  %reduce = vector.reduction <or>, %v, %s : vector<3xi32> into i32
+  return %reduce : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @reduction_xor
+//  CHECK-SAME: (%[[V:.+]]: vector<4xi32>)
+//       CHECK:   %[[S0:.+]] = spirv.CompositeExtract %[[V]][0 : i32] : vector<4xi32>
+//       CHECK:   %[[S1:.+]] = spirv.CompositeExtract %[[V]][1 : i32] : vector<4xi32>
+//       CHECK:   %[[S2:.+]] = spirv.CompositeExtract %[[V]][2 : i32] : vector<4xi32>
+//       CHECK:   %[[S3:.+]] = spirv.CompositeExtract %[[V]][3 : i32] : vector<4xi32>
+//       CHECK:   %[[XOR0:.+]] = spirv.BitwiseXor %[[S0]], %[[S1]]
+//       CHECK:   %[[XOR1:.+]] = spirv.BitwiseXor %[[XOR0]], %[[S2]]
+//       CHECK:   %[[XOR2:.+]] = spirv.BitwiseXor %[[XOR1]], %[[S3]]
----------------
IgWod wrote:

Same.

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