Add type constraints

Author: kuhar

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

@@ -907,15 +907,14 @@ def GPU_AllReduceOperation : I32EnumAttr<"AllReduceOperation",
   let genSpecializedAttr = 0;
   let cppNamespace = "::mlir::gpu";
 }
+
+def AnyIntegerOrFloat : AnyTypeOf<[AnySignlessInteger, AnyFloat], "Integer or Float">;
+
 def GPU_AllReduceOperationAttr : EnumAttr<GPU_Dialect, GPU_AllReduceOperation,
                                           "all_reduce_op">;
 
 def GPU_AllReduceOp : GPU_Op<"all_reduce",
-    [SameOperandsAndResultType, IsolatedFromAbove]>,
-    Arguments<(ins AnyType:$value,
-               OptionalAttr<GPU_AllReduceOperationAttr>:$op,
-               UnitAttr:$uniform)>,
-    Results<(outs AnyType)> {
+    [SameOperandsAndResultType, IsolatedFromAbove]> {
   let summary = "Reduce values among workgroup.";
   let description = [{
     The `all_reduce` op reduces the value of every work item across a local
@@ -943,6 +942,14 @@ def GPU_AllReduceOp : GPU_Op<"all_reduce",
     If `uniform` flag is set either none or all work items of a workgroup
     need to execute this op in convergence.
   }];
+
+  let arguments = (ins
+    AnyIntegerOrFloat:$value,
+    OptionalAttr<GPU_AllReduceOperationAttr>:$op,
+    UnitAttr:$uniform
+  );
+  let results = (outs AnyIntegerOrFloat:$result);
+
   let regions = (region AnyRegion:$body);
   let assemblyFormat = [{ custom<AllReduceOperation>($op) $value
                           (`uniform` $uniform^)? $body attr-dict
@@ -952,12 +959,7 @@ def GPU_AllReduceOp : GPU_Op<"all_reduce",
   let hasRegionVerifier = 1;
 }
 
-def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce",
-    [SameOperandsAndResultType]>,
-    Arguments<(ins AnyType:$value,
-               GPU_AllReduceOperationAttr:$op,
-               UnitAttr:$uniform)>,
-    Results<(outs AnyType)> {
+def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce", [SameOperandsAndResultType]> {
   let summary = "Reduce values among subgroup.";
   let description = [{
     The `subgroup_reduce` op reduces the value of every work item across a
@@ -977,6 +979,14 @@ def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce",
     *  Floating point types: `add`, `mul`, `minf`, `maxf`, `minimumf`,
        `maximumf`
   }];
+
+  let arguments = (ins
+    AnyIntegerOrFloat:$value,
+    GPU_AllReduceOperationAttr:$op,
+    UnitAttr:$uniform
+  );
+  let results = (outs AnyIntegerOrFloat:$result);
+
   let assemblyFormat = [{ custom<AllReduceOperation>($op) $value
                           (`uniform` $uniform^)? attr-dict
                           `:` functional-type(operands, results) }];

mlir/include/mlir/IR/CommonTypeConstraints.td

@@ -34,7 +34,7 @@ def IsFixedVectorTypePred : CPred<[{::llvm::isa<::mlir::VectorType>($_self) &&
                                   !::llvm::cast<VectorType>($_self).isScalable()}]>;
 
 // Whether a type is a scalable VectorType.
-def IsVectorTypeWithAnyDimScalablePred 
+def IsVectorTypeWithAnyDimScalablePred
         : CPred<[{::llvm::isa<::mlir::VectorType>($_self) &&
                   ::llvm::cast<VectorType>($_self).isScalable()}]>;
 

mlir/test/Dialect/GPU/invalid.mlir

@@ -210,6 +210,14 @@ module attributes {gpu.container_module} {
 
 // -----
 
+func.func @reduce_bad_type(%arg0 : vector<4xf32>) {
+  // expected-error@+1 {{'gpu.all_reduce' op operand #0 must be Integer or Float}}
+  %res = gpu.all_reduce add %arg0 {} : (vector<4xf32>) -> vector<4xf32>
+  return
+}
+
+// -----
+
 func.func @reduce_no_op_no_body(%arg0 : f32) {
   // expected-error@+1 {{expected either an op attribute or a non-empty body}}
   %res = "gpu.all_reduce"(%arg0) ({}) : (f32) -> (f32)
@@ -325,6 +333,14 @@ func.func @reduce_invalid_op_type_maximumf(%arg0 : i32) {
 
 // -----
 
+func.func @subgroup_reduce_bad_type(%arg0 : vector<2xf32>) {
+  // expected-error@+1 {{'gpu.subgroup_reduce' op operand #0 must be Integer or Float}}
+  %res = gpu.subgroup_reduce add %arg0 : (vector<2xf32>) -> vector<2xf32>
+  return
+}
+
+// -----
+
 func.func @subgroup_reduce_invalid_op_type_and(%arg0 : f32) {
   // expected-error@+1 {{`and` reduction operation is not compatible with type 'f32'}}
   %res = gpu.subgroup_reduce and %arg0 : (f32) -> (f32)