[mlir][spirv] Add GroupNonUniform min and max operations.

Add GroupNonUniform atihtmetic operations: FMax, FMin, SMax, SMin,
UMax, UMin.

Differential Revision: https://reviews.llvm.org/D73563

diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td
index 44141bf..79e5158 100644 (file)
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td
@@ -3161,6 +3161,12 @@ def SPV_OC_OpGroupNonUniformIAdd       : I32EnumAttrCase<"OpGroupNonUniformIAdd"
 def SPV_OC_OpGroupNonUniformFAdd       : I32EnumAttrCase<"OpGroupNonUniformFAdd", 350>;
 def SPV_OC_OpGroupNonUniformIMul       : I32EnumAttrCase<"OpGroupNonUniformIMul", 351>;
 def SPV_OC_OpGroupNonUniformFMul       : I32EnumAttrCase<"OpGroupNonUniformFMul", 352>;
+def SPV_OC_OpGroupNonUniformSMin       : I32EnumAttrCase<"OpGroupNonUniformSMin", 353>;
+def SPV_OC_OpGroupNonUniformUMin       : I32EnumAttrCase<"OpGroupNonUniformUMin", 354>;
+def SPV_OC_OpGroupNonUniformFMin       : I32EnumAttrCase<"OpGroupNonUniformFMin", 355>;
+def SPV_OC_OpGroupNonUniformSMax       : I32EnumAttrCase<"OpGroupNonUniformSMax", 356>;
+def SPV_OC_OpGroupNonUniformUMax       : I32EnumAttrCase<"OpGroupNonUniformUMax", 357>;
+def SPV_OC_OpGroupNonUniformFMax       : I32EnumAttrCase<"OpGroupNonUniformFMax", 358>;
 def SPV_OC_OpSubgroupBallotKHR         : I32EnumAttrCase<"OpSubgroupBallotKHR", 4421>;
 
 def SPV_OpcodeAttr :
@@ -3210,6 +3216,9 @@ def SPV_OpcodeAttr :
       SPV_OC_OpGroupNonUniformElect, SPV_OC_OpGroupNonUniformBallot,
       SPV_OC_OpGroupNonUniformIAdd, SPV_OC_OpGroupNonUniformFAdd,
       SPV_OC_OpGroupNonUniformIMul, SPV_OC_OpGroupNonUniformFMul,
+      SPV_OC_OpGroupNonUniformSMin, SPV_OC_OpGroupNonUniformUMin,
+      SPV_OC_OpGroupNonUniformFMin, SPV_OC_OpGroupNonUniformSMax,
+      SPV_OC_OpGroupNonUniformUMax, SPV_OC_OpGroupNonUniformFMax,
       SPV_OC_OpSubgroupBallotKHR
     ]>;
 

diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td
index 96da3cc..b1478d0 100644 (file)
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td
@@ -200,6 +200,114 @@ def SPV_GroupNonUniformFAddOp :
 
 // -----
 
+def SPV_GroupNonUniformFMaxOp :
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformFMax", SPV_Float, []> {
+  let summary = [{
+    A floating point maximum group operation of all Value operands
+    contributed by active invocations in by group.
+  }];
+
+  let description = [{
+    Result Type  must be a scalar or vector of floating-point type.
+
+    Execution must be Workgroup or Subgroup Scope.
+
+    The identity I for Operation is -INF. If Operation is ClusteredReduce,
+    ClusterSize must be specified.
+
+     The type of Value must be the same as Result Type.  The method used to
+    perform the group operation on the contributed Value(s) from active
+    invocations is implementation defined. From the set of Value(s) provided
+    by active invocations within a subgroup, if for any two Values one of
+    them is a NaN, the other is chosen. If all Value(s) that are used by the
+    current invocation are NaN, then the result is an undefined value.
+
+    ClusterSize is the size of cluster to use. ClusterSize must be a scalar
+    of integer type, whose Signedness operand is 0. ClusterSize must come
+    from a constant instruction. ClusterSize must be at least 1, and must be
+    a power of 2. If ClusterSize is greater than the declared SubGroupSize,
+    executing this instruction results in undefined behavior.
+
+    ### Custom assembly form
+
+    ```
+    scope ::= `"Workgroup"` | `"Subgroup"`
+    operation ::= `"Reduce"` | `"InclusiveScan"` | `"ExclusiveScan"` | ...
+    float-scalar-vector-type ::= float-type |
+                                 `vector<` integer-literal `x` float-type `>`
+    non-uniform-fmax-op ::= ssa-id `=` `spv.GroupNonUniformFMax` scope operation
+                            ssa-use ( `cluster_size` `(` ssa_use `)` )?
+                            `:` float-scalar-vector-type
+    ```
+
+    For example:
+
+    ```
+    %four = spv.constant 4 : i32
+    %scalar = ... : f32
+    %vector = ... : vector<4xf32>
+    %0 = spv.GroupNonUniformFMax "Workgroup" "Reduce" %scalar : f32
+    %1 = spv.GroupNonUniformFMax "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xf32>
+    ```
+  }];
+}
+
+// -----
+
+def SPV_GroupNonUniformFMinOp :
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformFMin", SPV_Float, []> {
+  let summary = [{
+    A floating point minimum group operation of all Value operands
+    contributed by active invocations in the group.
+  }];
+
+  let description = [{
+    Result Type  must be a scalar or vector of floating-point type.
+
+    Execution must be Workgroup or Subgroup Scope.
+
+    The identity I for Operation is +INF. If Operation is ClusteredReduce,
+    ClusterSize must be specified.
+
+     The type of Value must be the same as Result Type.  The method used to
+    perform the group operation on the contributed Value(s) from active
+    invocations is implementation defined. From the set of Value(s) provided
+    by active invocations within a subgroup, if for any two Values one of
+    them is a NaN, the other is chosen. If all Value(s) that are used by the
+    current invocation are NaN, then the result is an undefined value.
+
+    ClusterSize is the size of cluster to use. ClusterSize must be a scalar
+    of integer type, whose Signedness operand is 0. ClusterSize must come
+    from a constant instruction. ClusterSize must be at least 1, and must be
+    a power of 2. If ClusterSize is greater than the declared SubGroupSize,
+    executing this instruction results in undefined behavior.
+
+    ### Custom assembly form
+
+    ```
+    scope ::= `"Workgroup"` | `"Subgroup"`
+    operation ::= `"Reduce"` | `"InclusiveScan"` | `"ExclusiveScan"` | ...
+    float-scalar-vector-type ::= float-type |
+                                 `vector<` integer-literal `x` float-type `>`
+    non-uniform-fmin-op ::= ssa-id `=` `spv.GroupNonUniformFMin` scope operation
+                            ssa-use ( `cluster_size` `(` ssa_use `)` )?
+                            `:` float-scalar-vector-type
+    ```
+
+    For example:
+
+    ```
+    %four = spv.constant 4 : i32
+    %scalar = ... : f32
+    %vector = ... : vector<4xf32>
+    %0 = spv.GroupNonUniformFMin "Workgroup" "Reduce" %scalar : f32
+    %1 = spv.GroupNonUniformFMin "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xf32>
+    ```
+  }];
+}
+
+// -----
+
 def SPV_GroupNonUniformFMulOp :
     SPV_GroupNonUniformArithmeticOp<"GroupNonUniformFMul", SPV_Float, []> {
   let summary = [{
@@ -349,4 +457,202 @@ def SPV_GroupNonUniformIMulOp :
 
 // -----
 
+def SPV_GroupNonUniformSMaxOp :
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformSMax", SPV_Integer, []> {
+  let summary = [{
+    A signed integer maximum group operation of all Value operands
+    contributed by active invocations in the group.
+  }];
+
+  let description = [{
+    Result Type  must be a scalar or vector of integer type.
+
+    Execution must be Workgroup or Subgroup Scope.
+
+    The identity I for Operation is INT_MIN. If Operation is
+    ClusteredReduce, ClusterSize must be specified.
+
+     The type of Value must be the same as Result Type.
+
+    ClusterSize is the size of cluster to use. ClusterSize must be a scalar
+    of integer type, whose Signedness operand is 0. ClusterSize must come
+    from a constant instruction. ClusterSize must be at least 1, and must be
+    a power of 2. If ClusterSize is greater than the declared SubGroupSize,
+    executing this instruction results in undefined behavior.
+
+    ### Custom assembly form
+
+    ```
+    scope ::= `"Workgroup"` | `"Subgroup"`
+    operation ::= `"Reduce"` | `"InclusiveScan"` | `"ExclusiveScan"` | ...
+    integer-scalar-vector-type ::= integer-type |
+                                 `vector<` integer-literal `x` integer-type `>`
+    non-uniform-smax-op ::= ssa-id `=` `spv.GroupNonUniformSMax` scope operation
+                            ssa-use ( `cluster_size` `(` ssa_use `)` )?
+                            `:` integer-scalar-vector-type
+    ```
+
+    For example:
+
+    ```
+    %four = spv.constant 4 : i32
+    %scalar = ... : i32
+    %vector = ... : vector<4xi32>
+    %0 = spv.GroupNonUniformSMax "Workgroup" "Reduce" %scalar : i32
+    %1 = spv.GroupNonUniformSMax "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xi32>
+    ```
+  }];
+}
+
+// -----
+
+def SPV_GroupNonUniformSMinOp :
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformSMin", SPV_Integer, []> {
+  let summary = [{
+    A signed integer minimum group operation of all Value operands
+    contributed by active invocations in the group.
+  }];
+
+  let description = [{
+    Result Type  must be a scalar or vector of integer type.
+
+    Execution must be Workgroup or Subgroup Scope.
+
+    The identity I for Operation is INT_MAX. If Operation is
+    ClusteredReduce, ClusterSize must be specified.
+
+     The type of Value must be the same as Result Type.
+
+    ClusterSize is the size of cluster to use. ClusterSize must be a scalar
+    of integer type, whose Signedness operand is 0. ClusterSize must come
+    from a constant instruction. ClusterSize must be at least 1, and must be
+    a power of 2. If ClusterSize is greater than the declared SubGroupSize,
+    executing this instruction results in undefined behavior.
+
+    ### Custom assembly form
+
+    ```
+    scope ::= `"Workgroup"` | `"Subgroup"`
+    operation ::= `"Reduce"` | `"InclusiveScan"` | `"ExclusiveScan"` | ...
+    integer-scalar-vector-type ::= integer-type |
+                                 `vector<` integer-literal `x` integer-type `>`
+    non-uniform-smin-op ::= ssa-id `=` `spv.GroupNonUniformSMin` scope operation
+                            ssa-use ( `cluster_size` `(` ssa_use `)` )?
+                            `:` integer-scalar-vector-type
+    ```
+
+    For example:
+
+    ```
+    %four = spv.constant 4 : i32
+    %scalar = ... : i32
+    %vector = ... : vector<4xi32>
+    %0 = spv.GroupNonUniformSMin "Workgroup" "Reduce" %scalar : i32
+    %1 = spv.GroupNonUniformSMin "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xi32>
+    ```
+  }];
+}
+
+// -----
+
+def SPV_GroupNonUniformUMaxOp :
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformUMax", SPV_Integer, []> {
+  let summary = [{
+    An unsigned integer maximum group operation of all Value operands
+    contributed by active invocations in the group.
+  }];
+
+  let description = [{
+    Result Type  must be a scalar or vector of integer type, whose
+    Signedness operand is 0.
+
+    Execution must be Workgroup or Subgroup Scope.
+
+    The identity I for Operation is 0. If Operation is ClusteredReduce,
+    ClusterSize must be specified.
+
+     The type of Value must be the same as Result Type.
+
+    ClusterSize is the size of cluster to use. ClusterSize must be a scalar
+    of integer type, whose Signedness operand is 0. ClusterSize must come
+    from a constant instruction. ClusterSize must be at least 1, and must be
+    a power of 2. If ClusterSize is greater than the declared SubGroupSize,
+    executing this instruction results in undefined behavior.
+
+    ### Custom assembly form
+
+    ```
+    scope ::= `"Workgroup"` | `"Subgroup"`
+    operation ::= `"Reduce"` | `"InclusiveScan"` | `"ExclusiveScan"` | ...
+    integer-scalar-vector-type ::= integer-type |
+                                 `vector<` integer-literal `x` integer-type `>`
+    non-uniform-umax-op ::= ssa-id `=` `spv.GroupNonUniformUMax` scope operation
+                            ssa-use ( `cluster_size` `(` ssa_use `)` )?
+                            `:` integer-scalar-vector-type
+    ```
+
+    For example:
+
+    ```
+    %four = spv.constant 4 : i32
+    %scalar = ... : i32
+    %vector = ... : vector<4xi32>
+    %0 = spv.GroupNonUniformUMax "Workgroup" "Reduce" %scalar : i32
+    %1 = spv.GroupNonUniformUMax "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xi32>
+    ```
+  }];
+}
+
+// -----
+
+def SPV_GroupNonUniformUMinOp :
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformUMin", SPV_Integer, []> {
+  let summary = [{
+    An unsigned integer minimum group operation of all Value operands
+    contributed by active invocations in the group.
+  }];
+
+  let description = [{
+    Result Type  must be a scalar or vector of integer type, whose
+    Signedness operand is 0.
+
+    Execution must be Workgroup or Subgroup Scope.
+
+    The identity I for Operation is UINT_MAX. If Operation is
+    ClusteredReduce, ClusterSize must be specified.
+
+     The type of Value must be the same as Result Type.
+
+    ClusterSize is the size of cluster to use. ClusterSize must be a scalar
+    of integer type, whose Signedness operand is 0. ClusterSize must come
+    from a constant instruction. ClusterSize must be at least 1, and must be
+    a power of 2. If ClusterSize is greater than the declared SubGroupSize,
+    executing this instruction results in undefined behavior.
+
+    ### Custom assembly form
+
+    ```
+    scope ::= `"Workgroup"` | `"Subgroup"`
+    operation ::= `"Reduce"` | `"InclusiveScan"` | `"ExclusiveScan"` | ...
+    integer-scalar-vector-type ::= integer-type |
+                                 `vector<` integer-literal `x` integer-type `>`
+    non-uniform-umin-op ::= ssa-id `=` `spv.GroupNonUniformUMin` scope operation
+                            ssa-use ( `cluster_size` `(` ssa_use `)` )?
+                            `:` integer-scalar-vector-type
+    ```
+
+    For example:
+
+    ```
+    %four = spv.constant 4 : i32
+    %scalar = ... : i32
+    %vector = ... : vector<4xi32>
+    %0 = spv.GroupNonUniformUMin "Workgroup" "Reduce" %scalar : i32
+    %1 = spv.GroupNonUniformUMin "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xi32>
+    ```
+  }];
+}
+
+// -----
+
 #endif // SPIRV_NON_UNIFORM_OPS

diff --git a/mlir/test/Dialect/SPIRV/Serialization/non-uniform-ops.mlir b/mlir/test/Dialect/SPIRV/Serialization/non-uniform-ops.mlir
index 9037f55..225fbf7 100644 (file)
--- a/mlir/test/Dialect/SPIRV/Serialization/non-uniform-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/Serialization/non-uniform-ops.mlir
@@ -22,6 +22,20 @@ spv.module "Logical" "GLSL450" {
     spv.ReturnValue %0: f32
   }
 
+  // CHECK-LABEL: @group_non_uniform_fmax_reduce
+  func @group_non_uniform_fmax_reduce(%val: f32) -> f32 {
+    // CHECK: %{{.+}} = spv.GroupNonUniformFMax "Workgroup" "Reduce" %{{.+}} : f32
+    %0 = spv.GroupNonUniformFMax "Workgroup" "Reduce" %val : f32
+    spv.ReturnValue %0: f32
+  }
+
+  // CHECK-LABEL: @group_non_uniform_fmin_reduce
+  func @group_non_uniform_fmin_reduce(%val: f32) -> f32 {
+    // CHECK: %{{.+}} = spv.GroupNonUniformFMin "Workgroup" "Reduce" %{{.+}} : f32
+    %0 = spv.GroupNonUniformFMin "Workgroup" "Reduce" %val : f32
+    spv.ReturnValue %0: f32
+  }
+
   // CHECK-LABEL: @group_non_uniform_fmul_reduce
   func @group_non_uniform_fmul_reduce(%val: f32) -> f32 {
     // CHECK: %{{.+}} = spv.GroupNonUniformFMul "Workgroup" "Reduce" %{{.+}} : f32
@@ -51,4 +65,31 @@ spv.module "Logical" "GLSL450" {
     spv.ReturnValue %0: i32
   }
 
+  // CHECK-LABEL: @group_non_uniform_smax_reduce
+  func @group_non_uniform_smax_reduce(%val: i32) -> i32 {
+    // CHECK: %{{.+}} = spv.GroupNonUniformSMax "Workgroup" "Reduce" %{{.+}} : i32
+    %0 = spv.GroupNonUniformSMax "Workgroup" "Reduce" %val : i32
+    spv.ReturnValue %0: i32
+  }
+
+  // CHECK-LABEL: @group_non_uniform_smin_reduce
+  func @group_non_uniform_smin_reduce(%val: i32) -> i32 {
+    // CHECK: %{{.+}} = spv.GroupNonUniformSMin "Workgroup" "Reduce" %{{.+}} : i32
+    %0 = spv.GroupNonUniformSMin "Workgroup" "Reduce" %val : i32
+    spv.ReturnValue %0: i32
+  }
+
+  // CHECK-LABEL: @group_non_uniform_umax_reduce
+  func @group_non_uniform_umax_reduce(%val: i32) -> i32 {
+    // CHECK: %{{.+}} = spv.GroupNonUniformUMax "Workgroup" "Reduce" %{{.+}} : i32
+    %0 = spv.GroupNonUniformUMax "Workgroup" "Reduce" %val : i32
+    spv.ReturnValue %0: i32
+  }
+
+  // CHECK-LABEL: @group_non_uniform_umin_reduce
+  func @group_non_uniform_umin_reduce(%val: i32) -> i32 {
+    // CHECK: %{{.+}} = spv.GroupNonUniformUMin "Workgroup" "Reduce" %{{.+}} : i32
+    %0 = spv.GroupNonUniformUMin "Workgroup" "Reduce" %val : i32
+    spv.ReturnValue %0: i32
+  }
 }

diff --git a/mlir/test/Dialect/SPIRV/non-uniform-ops.mlir b/mlir/test/Dialect/SPIRV/non-uniform-ops.mlir
index e059233..15337e3 100644 (file)
--- a/mlir/test/Dialect/SPIRV/non-uniform-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/non-uniform-ops.mlir
@@ -82,6 +82,32 @@ func @group_non_uniform_fmul_clustered_reduce(%val: vector<2xf32>) -> vector<2xf
 // -----
 
 //===----------------------------------------------------------------------===//
+// spv.GroupNonUniformFMax
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @group_non_uniform_fmax_reduce
+func @group_non_uniform_fmax_reduce(%val: f32) -> f32 {
+  // CHECK: %{{.+}} = spv.GroupNonUniformFMax "Workgroup" "Reduce" %{{.+}} : f32
+  %0 = spv.GroupNonUniformFMax "Workgroup" "Reduce" %val : f32
+  return %0: f32
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GroupNonUniformFMin
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @group_non_uniform_fmin_reduce
+func @group_non_uniform_fmin_reduce(%val: f32) -> f32 {
+  // CHECK: %{{.+}} = spv.GroupNonUniformFMin "Workgroup" "Reduce" %{{.+}} : f32
+  %0 = spv.GroupNonUniformFMin "Workgroup" "Reduce" %val : f32
+  return %0: f32
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // spv.GroupNonUniformIAdd
 //===----------------------------------------------------------------------===//
 
@@ -153,3 +179,55 @@ func @group_non_uniform_imul_clustered_reduce(%val: vector<2xi32>) -> vector<2xi
   %0 = spv.GroupNonUniformIMul "Workgroup" "ClusteredReduce" %val cluster_size(%four) : vector<2xi32>
   return %0: vector<2xi32>
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GroupNonUniformSMax
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @group_non_uniform_smax_reduce
+func @group_non_uniform_smax_reduce(%val: i32) -> i32 {
+  // CHECK: %{{.+}} = spv.GroupNonUniformSMax "Workgroup" "Reduce" %{{.+}} : i32
+  %0 = spv.GroupNonUniformSMax "Workgroup" "Reduce" %val : i32
+  return %0: i32
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GroupNonUniformSMin
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @group_non_uniform_smin_reduce
+func @group_non_uniform_smin_reduce(%val: i32) -> i32 {
+  // CHECK: %{{.+}} = spv.GroupNonUniformSMin "Workgroup" "Reduce" %{{.+}} : i32
+  %0 = spv.GroupNonUniformSMin "Workgroup" "Reduce" %val : i32
+  return %0: i32
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GroupNonUniformUMax
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @group_non_uniform_umax_reduce
+func @group_non_uniform_umax_reduce(%val: i32) -> i32 {
+  // CHECK: %{{.+}} = spv.GroupNonUniformUMax "Workgroup" "Reduce" %{{.+}} : i32
+  %0 = spv.GroupNonUniformUMax "Workgroup" "Reduce" %val : i32
+  return %0: i32
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GroupNonUniformUMin
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @group_non_uniform_umin_reduce
+func @group_non_uniform_umin_reduce(%val: i32) -> i32 {
+  // CHECK: %{{.+}} = spv.GroupNonUniformUMin "Workgroup" "Reduce" %{{.+}} : i32
+  %0 = spv.GroupNonUniformUMin "Workgroup" "Reduce" %val : i32
+  return %0: i32
+}