Update workgroup count op syntax (iree-org#21656)

The current syntax for workgroup count ops has a trailing operand list,
which is problematic when the following op has results, e.g.:
```mlir
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice
%foo = arith.constant 1 : index
```
gets parsed as:
```mlir
%x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice %foo
= arith.constant 1 : index
```
and errors out.

So far this hasn't been an issue, as the workgroup count ops are
typically followed by a `return`, which has no results, but with the
introduction of `workgroup_count_split_reduction_modifier` we've begun
hitting this case more often. This updates the syntax to have additional
literals which remove the ambiguity.

Syntax change:
Before:
```mlir
iree_tensor_ext.dispatch.workgroup_count_from_slice
iree_tensor_ext.dispatch.workgroup_count_split_reduction_modifier(%x, %y, %z),
iree_tensor_ext.dispatch.workgroup_count_from_dag_root
```
After:
```mlir
iree_tensor_ext.dispatch.workgroup_count_from_slice()
iree_tensor_ext.dispatch.workgroup_count_split_reduction_modifier workgroups(%x, %y, %z) workload()
iree_tensor_ext.dispatch.workgroup_count_from_dag_root()
```

Signed-off-by: keshavvinayak01 <[email protected]>

Author: rkayaith

compiler/plugins/target/CUDA/test/smoketest.mlir

@@ -11,7 +11,7 @@ module attributes {
 
 stream.executable public @add_dispatch_executable {
   stream.executable.export @add_dispatch workgroups(%arg0 : index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   }
   builtin.module  {
@@ -36,7 +36,7 @@ stream.executable public @add_dispatch_executable {
 
 stream.executable public @mul_dispatch_executable {
   stream.executable.export @mul_dispatch workgroups(%arg0 : index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   }
   builtin.module  {

compiler/plugins/target/LLVMCPU/test/smoketest_embedded.mlir

@@ -13,7 +13,7 @@ module attributes {
 
 stream.executable public @add_dispatch_0 {
   stream.executable.export @add_dispatch_0 workgroups(%arg0 : index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   }
   builtin.module  {

compiler/plugins/target/LLVMCPU/test/smoketest_system.mlir

@@ -15,7 +15,7 @@ module attributes {
 
 stream.executable public @add_dispatch_0 {
   stream.executable.export @add_dispatch_0 workgroups(%arg0 : index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   }
   builtin.module  {

compiler/plugins/target/MetalSPIRV/test/smoketest.mlir

@@ -15,7 +15,7 @@ module attributes {
 
 stream.executable public @reduce_dispatch {
   stream.executable.export @reduce_dispatch workgroups(%arg0 : index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   }
   builtin.module {

compiler/plugins/target/ROCM/builtins/tuning/test/spec_gfx942.mlir

@@ -31,7 +31,7 @@ hal.executable public @main {
     iree_codegen.default_tuning_spec = #rocm.builtin.tuning_module<"iree_default_tuning_spec_gfx942.mlir">
   }>) {
     hal.executable.export public @matmul_transpose_b ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device) -> (index, index, index) {
-      %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice
+      %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice()
       hal.return %x, %y, %z : index, index, index
     }
     builtin.module {
@@ -82,7 +82,7 @@ hal.executable public @main {
     iree_codegen.default_tuning_spec = #rocm.builtin.tuning_module<"iree_default_tuning_spec_gfx942.mlir">
   }>) {
     hal.executable.export public @attention ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device) -> (index, index, index) {
-      %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice
+      %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice()
       hal.return %x, %y, %z : index, index, index
     }
     builtin.module {
@@ -138,7 +138,7 @@ hal.executable public @main {
     iree_codegen.default_tuning_spec = #rocm.builtin.tuning_module<"iree_default_tuning_spec_gfx942.mlir">
   }>) {
     hal.executable.export public @attention ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device) -> (index, index, index) {
-      %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice
+      %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice()
       hal.return %x, %y, %z : index, index, index
     }
     builtin.module {

compiler/plugins/target/ROCM/test/gpu_encoding_attrs.mlir

@@ -21,7 +21,7 @@
 
 stream.executable public @main {
   stream.executable.export @main workgroups(%arg0: index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   }
   builtin.module {

compiler/plugins/target/ROCM/test/lowering_strategy_from_tuning_spec.mlir

@@ -28,7 +28,7 @@
 hal.executable public @main {
   hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb">) {
     hal.executable.export public @matmul_transpose_b ordinal(0) layout(#pipeline_layout) count(%arg0: !hal.device) -> (index, index, index) {
-      %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice
+      %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice()
       hal.return %x, %y, %z : index, index, index
     }
     builtin.module {

compiler/plugins/target/ROCM/test/smoketest.mlir

@@ -10,7 +10,7 @@ module attributes {
 
 stream.executable public @add_dispatch_executable {
   stream.executable.export @add_dispatch workgroups(%arg0 : index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   }
   builtin.module  {
@@ -35,7 +35,7 @@ stream.executable public @add_dispatch_executable {
 
 stream.executable public @mul_dispatch_executable {
   stream.executable.export @mul_dispatch workgroups(%arg0 : index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   }
   builtin.module  {
@@ -80,7 +80,7 @@ module attributes {
 
 stream.executable public @executable {
   stream.executable.export @export workgroups(%arg0 : index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   } loc(#loc)
   builtin.module  {

compiler/plugins/target/ROCM/test/target_device_features.mlir

@@ -83,7 +83,7 @@
 
 stream.executable public @reduce_dispatch {
   stream.executable.export @reduce_dispatch workgroups(%arg0: index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   }
   builtin.module {

compiler/plugins/target/VMVX/test/smoketest.mlir

@@ -10,7 +10,7 @@ module attributes {
 
 stream.executable public @add_dispatch_0 {
   stream.executable.export @add_dispatch_0 workgroups(%arg0 : index) -> (index, index, index) {
-    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root %arg0
+    %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_dag_root(%arg0)
     stream.return %x, %y, %z : index, index, index
   }
   builtin.module  {