[SYCL-MLIR] Pretty print SYCL grid operations (#8122)

- 1-D Grid operations: `attr-dict : type($result)`;
- 3-D Grid operations: `$dimension attr-dict : type($result)`;

Signed-off-by: Victor Perez <[email protected]>

Author: victor-eds

mlir-sycl/include/mlir/Dialect/SYCL/IR/SYCLOps.td

@@ -427,52 +427,57 @@ def VectorSplatArg : MemRefOf<[VectorEltTy]>;
 // SYCL GRID OPERATIONS
 ////////////////////////////////////////////////////////////////////////////////
 
-def SYCLNumWorkItemsOp : SYCL_Op<"num_work_items", [SYCLIndexSpaceGetRange]> {
+class SYCL1DGridOp<string mnemonic, list<Trait> traits = []>
+    : SYCL_Op<mnemonic, traits> {
+  let arguments = (ins);
+  let results = (outs I32:$result);
+  let assemblyFormat = "attr-dict `:` type($result)";
+}
+
+class SYCL3DGridOp<string mnemonic, list<Trait> traits = []>
+    : SYCL_Op<mnemonic, traits> {
+  let arguments = (ins Optional<I32>:$dimension);
+  let assemblyFormat = "$dimension attr-dict `:` type($result)";
+}
+
+class SYCL3DIDGridOp<string mnemonic>
+    : SYCL3DGridOp<mnemonic, [SYCLIndexSpaceGetID]> {
+  let results = (outs SYCLIndexSpaceGetIDResult:$result);
+}
+
+class SYCL3DRangeGridOp<string mnemonic>
+    : SYCL3DGridOp<mnemonic, [SYCLIndexSpaceGetRange]> {
+  let results = (outs SYCLIndexSpaceGetRangeResult:$result);
+}
+
+def SYCLNumWorkItemsOp : SYCL3DRangeGridOp<"num_work_items"> {
   let summary = "Retrieve the number of work-items.";
   let description = [{
     This operation returns the number of work-items in the index space. If the
     optional argument is passed, the number of work-items in the given dimension
     is returned.
   }];
-
-  let arguments = (ins Optional<I32>:$dimension);
-  let results = (outs SYCLIndexSpaceGetRangeResult:$result);
-  let assemblyFormat = [{
-    `(` operands `)` attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLGlobalIDOp : SYCL_Op<"global_id", [SYCLIndexSpaceGetID]> {
+def SYCLGlobalIDOp : SYCL3DIDGridOp<"global_id"> {
   let summary = "Retrieve the global ID of the item.";
   let description = [{
     This operation returns the global ID of the item in the index space. If the
     optional argument is passed, the global ID of the item in the given
     dimension is returned.
   }];
-
-  let arguments = (ins Optional<I32>:$dimension);
-  let results = (outs SYCLIndexSpaceGetIDResult:$result);
-  let assemblyFormat = [{
-    `(` operands `)` attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLLocalIDOp : SYCL_Op<"local_id", [SYCLIndexSpaceGetID]> {
+def SYCLLocalIDOp : SYCL3DIDGridOp<"local_id"> {
   let summary = "Retrieve the local ID of the item.";
   let description = [{
     This operation returns the local ID of the item within the work-group. If
     the optional argument is passed, the local ID of the item in the given
     dimension is returned.
   }];
-
-  let arguments = (ins Optional<I32>:$dimension);
-  let results = (outs SYCLIndexSpaceGetIDResult:$result);
-  let assemblyFormat = [{
-    `(` operands `)` attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLGlobalOffsetOp : SYCL_Op<"global_offset", [SYCLIndexSpaceGetID]>,
+def SYCLGlobalOffsetOp : SYCL3DIDGridOp<"global_offset">,
                        Deprecated<"deprecated in SYCL 2020"> {
   let summary = "Retrieve the global offset of the item.";
   let description = [{
@@ -482,121 +487,67 @@ def SYCLGlobalOffsetOp : SYCL_Op<"global_offset", [SYCLIndexSpaceGetID]>,
 
     Note that this is deprecated in SYCL 2020.
   }];
-
-  let arguments = (ins Optional<I32>:$dimension);
-  let results = (outs SYCLIndexSpaceGetIDResult:$result);
-  let assemblyFormat = [{
-    `(` operands `)` attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLNumWorkGroupsOp : SYCL_Op<"num_work_groups", [SYCLIndexSpaceGetRange]> {
+def SYCLNumWorkGroupsOp : SYCL3DRangeGridOp<"num_work_groups"> {
   let summary = "Retrieve the number of work-groups.";
   let description = [{
     This operation returns the number of work-groups in the index space. If the
     optional argument is passed, the number of work-groups in the given
     dimension is returned.
   }];
-
-  let arguments = (ins Optional<I32>:$dimension);
-  let results = (outs SYCLIndexSpaceGetRangeResult:$result);
-  let assemblyFormat = [{
-    `(` operands `)` attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLWorkGroupSizeOp : SYCL_Op<"work_group_size", [SYCLIndexSpaceGetRange]> {
+def SYCLWorkGroupSizeOp : SYCL3DRangeGridOp<"work_group_size"> {
   let summary = "Retrieve the number of work-items in a work-group.";
   let description = [{
     This operation returns the number of work-items per work-group. If the
     optional argument is passed, the number of work-items per work-group in the
     given dimension is returned.
   }];
-
-  let arguments = (ins Optional<I32>:$dimension);
-  let results = (outs SYCLIndexSpaceGetRangeResult:$result);
-  let assemblyFormat = [{
-    `(` operands `)` attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLWorkGroupIDOp : SYCL_Op<"work_group_id", [SYCLIndexSpaceGetID]> {
+def SYCLWorkGroupIDOp : SYCL3DIDGridOp<"work_group_id"> {
   let summary = "Retrieve the ID of the work-group.";
   let description = [{
     This operation returns the ID of the work-group. If the optional argument is
     passed, the ID of the work-group in the given dimension is returned.
   }];
-
-  let arguments = (ins Optional<I32>:$dimension);
-  let results = (outs SYCLIndexSpaceGetIDResult:$result);
-  let assemblyFormat = [{
-    `(` operands `)` attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLNumSubGroupsOp : SYCL_Op<"num_sub_groups"> {
+def SYCLNumSubGroupsOp : SYCL1DGridOp<"num_sub_groups"> {
   let summary = "Retrieve the number of sub-groups.";
   let description = [{
     This operation returns the number of sub-groups.
   }];
-
-  let arguments = (ins);
-  let results = (outs I32:$result);
-  let assemblyFormat = [{
-    attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLSubGroupMaxSizeOp : SYCL_Op<"sub_group_max_size"> {
+def SYCLSubGroupMaxSizeOp : SYCL1DGridOp<"sub_group_max_size"> {
   let summary = "Retrieve the maximum size of a sub-group.";
   let description = [{
     This operation returns the maximum size of a sub-group.
   }];
-
-  let arguments = (ins);
-  let results = (outs I32:$result);
-  let assemblyFormat = [{
-    attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLSubGroupSizeOp : SYCL_Op<"sub_group_size"> {
+def SYCLSubGroupSizeOp : SYCL1DGridOp<"sub_group_size"> {
   let summary = "Retrieve the sub-group size.";
   let description = [{
     This operation returns the sub-group size.
   }];
-
-  let arguments = (ins);
-  let results = (outs I32:$result);
-  let assemblyFormat = [{
-    attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLSubGroupIDOp : SYCL_Op<"sub_group_id"> {
+def SYCLSubGroupIDOp : SYCL1DGridOp<"sub_group_id"> {
   let summary = "Retrieve the ID of the sub-group.";
   let description = [{
     This operation returns the ID of the sub-group.
   }];
-
-  let arguments = (ins);
-  let results = (outs I32:$result);
-  let assemblyFormat = [{
-    attr-dict `:` functional-type(operands, results)
-  }];
 }
 
-def SYCLSubGroupLocalIDOp : SYCL_Op<"sub_group_local_id"> {
+def SYCLSubGroupLocalIDOp : SYCL1DGridOp<"sub_group_local_id"> {
   let summary = "Retrieve the local ID of the sub-group.";
   let description = [{
     This operation returns the local ID of the sub-group.
   }];
-
-  let arguments = (ins);
-  let results = (outs I32:$result);
-  let assemblyFormat = [{
-    attr-dict `:` functional-type(operands, results)
-  }];
 }
 
 ////////////////////////////////////////////////////////////////////////////////

mlir-sycl/test/Conversion/SYCLToGPU/grid-ops.mlir

@@ -19,7 +19,7 @@
 // CHECK-NEXT:      return %[[VAL_2]] : !sycl_range_1_
 // CHECK-NEXT:    }
 func.func @test_num_work_items() -> !sycl_range_1_ {
-  %0 = sycl.num_work_items() : () -> !sycl_range_1_
+  %0 = sycl.num_work_items : !sycl_range_1_
   return %0 : !sycl_range_1_
 }
 
@@ -40,7 +40,7 @@ func.func @test_num_work_items() -> !sycl_range_1_ {
 // CHECK-NEXT:      return %[[VAL_9]] : index
 // CHECK-NEXT:    }
 func.func @test_num_work_items_dim(%i: i32) -> index {
-  %0 = sycl.num_work_items(%i) : (i32) -> index
+  %0 = sycl.num_work_items %i : index
   return %0 : index
 }
 
@@ -61,7 +61,7 @@ func.func @test_num_work_items_dim(%i: i32) -> index {
 // CHECK-NEXT:      return %[[VAL_2]] : !sycl_id_2_
 // CHECK-NEXT:    }
 func.func @test_global_id() -> !sycl_id_2_ {
-  %0 = sycl.global_id() : () -> !sycl_id_2_
+  %0 = sycl.global_id : !sycl_id_2_
   return %0 : !sycl_id_2_
 }
 
@@ -82,7 +82,7 @@ func.func @test_global_id() -> !sycl_id_2_ {
 // CHECK-NEXT:      return %[[VAL_9]] : index
 // CHECK-NEXT:    }
 func.func @test_global_id_dim(%i: i32) -> index {
-  %0 = sycl.global_id(%i) : (i32) -> index
+  %0 = sycl.global_id %i : index
   return %0 : index
 }
 
@@ -108,7 +108,7 @@ func.func @test_global_id_dim(%i: i32) -> index {
 // CHECK-NEXT:      return %[[VAL_2]] : !sycl_id_3_
 // CHECK-NEXT:    }
 func.func @test_local_id() -> !sycl_id_3_ {
-  %0 = sycl.local_id() : () -> !sycl_id_3_
+  %0 = sycl.local_id : !sycl_id_3_
   return %0 : !sycl_id_3_
 }
 
@@ -129,7 +129,7 @@ func.func @test_local_id() -> !sycl_id_3_ {
 // CHECK-NEXT:      return %[[VAL_9]] : index
 // CHECK-NEXT:    }
 func.func @test_local_id_dim(%i: i32) -> index {
-  %0 = sycl.local_id(%i) : (i32) -> index
+  %0 = sycl.local_id %i : index
   return %0 : index
 }
 
@@ -155,7 +155,7 @@ func.func @test_local_id_dim(%i: i32) -> index {
 // CHECK-NEXT:      return %[[VAL_2]] : !sycl_range_3_
 // CHECK-NEXT:    }
 func.func @test_work_group_size() -> !sycl_range_3_ {
-  %0 = sycl.work_group_size() : () -> !sycl_range_3_
+  %0 = sycl.work_group_size : !sycl_range_3_
   return %0 : !sycl_range_3_
 }
 
@@ -176,7 +176,7 @@ func.func @test_work_group_size() -> !sycl_range_3_ {
 // CHECK-NEXT:      return %[[VAL_9]] : index
 // CHECK-NEXT:    }
 func.func @test_work_group_size_dim(%i: i32) -> index {
-  %0 = sycl.work_group_size(%i) : (i32) -> index
+  %0 = sycl.work_group_size %i : index
   return %0 : index
 }
 
@@ -192,7 +192,7 @@ func.func @test_work_group_size_dim(%i: i32) -> index {
 // CHECK-NEXT:      return %[[VAL_2]] : !sycl_id_1_
 // CHECK-NEXT:    }
 func.func @test_work_group_id() -> !sycl_id_1_ {
-  %0 = sycl.work_group_id() : () -> !sycl_id_1_
+  %0 = sycl.work_group_id : !sycl_id_1_
   return %0 : !sycl_id_1_
 }
 
@@ -213,7 +213,7 @@ func.func @test_work_group_id() -> !sycl_id_1_ {
 // CHECK-NEXT:      return %[[VAL_9]] : index
 // CHECK-NEXT:    }
 func.func @test_work_group_id_dim(%i: i32) -> index {
-  %0 = sycl.work_group_id(%i) : (i32) -> index
+  %0 = sycl.work_group_id %i : index
   return %0 : index
 }
 
@@ -223,7 +223,7 @@ func.func @test_work_group_id_dim(%i: i32) -> index {
 // CHECK-NEXT:      return %[[VAL_1]] : i32
 // CHECK-NEXT:    }
 func.func @test_num_sub_groups() -> i32 {
-  %0 = sycl.num_sub_groups : () -> i32
+  %0 = sycl.num_sub_groups : i32
   return %0 : i32
 }
 
@@ -233,7 +233,7 @@ func.func @test_num_sub_groups() -> i32 {
 // CHECK-NEXT:      return %[[VAL_1]] : i32
 // CHECK-NEXT:    }
 func.func @test_sub_group_size() -> i32 {
-  %0 = sycl.sub_group_size : () -> i32
+  %0 = sycl.sub_group_size : i32
   return %0 : i32
 }
 
@@ -243,6 +243,6 @@ func.func @test_sub_group_size() -> i32 {
 // CHECK-NEXT:      return %[[VAL_1]] : i32
 // CHECK-NEXT:    }
 func.func @test_sub_group_id() -> i32 {
-  %0 = sycl.sub_group_id : () -> i32
+  %0 = sycl.sub_group_id : i32
   return %0 : i32
 }

mlir-sycl/test/Conversion/SYCLToSPIRV/grid-ops.mlir

@@ -26,7 +26,7 @@ module attributes {gpu.container_module} {
     // CHECK-NEXT           }
     gpu.func @test_global_offset() kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<>} {
-      %0 = sycl.global_offset() : () -> !sycl_id_1_
+      %0 = sycl.global_offset : !sycl_id_1_
       gpu.return
     }
 
@@ -53,7 +53,7 @@ module attributes {gpu.container_module} {
     // CHECK-NEXT           }
     gpu.func @test_global_offset_dim(%i: i32) kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<>} {
-      %0 = sycl.global_offset(%i) : (i32) -> index
+      %0 = sycl.global_offset %i : index
       gpu.return
     }
 
@@ -77,7 +77,7 @@ module attributes {gpu.container_module} {
     // CHECK-NEXT           }
     gpu.func @test_num_work_groups() kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<>} {
-      %0 = sycl.num_work_groups() : () -> !sycl_range_2_
+      %0 = sycl.num_work_groups : !sycl_range_2_
       gpu.return
     }
 
@@ -104,7 +104,7 @@ module attributes {gpu.container_module} {
     // CHECK-NEXT           }
     gpu.func @test_num_work_groups_dim(%i: i32) kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<>} {
-      %0 = sycl.num_work_groups(%i) : (i32) -> index
+      %0 = sycl.num_work_groups %i : index
       gpu.return
     }
 
@@ -115,7 +115,7 @@ module attributes {gpu.container_module} {
     // CHECK-NEXT           }
     gpu.func @test_sub_group_max_size() kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<>} {
-      %0 = sycl.sub_group_max_size : () -> i32
+      %0 = sycl.sub_group_max_size : i32
       gpu.return
     }
 
@@ -126,7 +126,7 @@ module attributes {gpu.container_module} {
     // CHECK-NEXT           }
     gpu.func @test_sub_group_local_id() kernel
       attributes {spirv.entry_point_abi = #spirv.entry_point_abi<>} {
-      %0 = sycl.sub_group_local_id : () -> i32
+      %0 = sycl.sub_group_local_id : i32
       gpu.return
     }
   }