Improve documentation, add launch param attrs

Author: ashermancinelli

mlir/python/mlir/dialects/gpu/__init__.py

@@ -6,7 +6,7 @@
 from .._gpu_ops_gen import _Dialect
 from .._gpu_enum_gen import *
 from ..._mlir_libs._mlirDialectsGPU import *
-from typing import Callable, Sequence, Union, Optional
+from typing import Callable, Sequence, Union, Optional, List
 
 try:
     from ...ir import (
@@ -20,6 +20,7 @@
         Type,
         DictAttr,
         Attribute,
+        DenseI32ArrayAttr,
     )
     from .._ods_common import (
         get_default_loc_context as _get_default_loc_context,
@@ -29,26 +30,44 @@
     raise RuntimeError("Error loading imports from extension module") from e
 
 
-FUNCTION_TYPE_ATTRIBUTE_NAME = "function_type"
 KERNEL_ATTRIBUTE_NAME = "gpu.kernel"
+KNOWN_BLOCK_SIZE_ATTRIBUTE_NAME = "gpu.known_block_size"
+KNOWN_GRID_SIZE_ATTRIBUTE_NAME = "gpu.known_grid_size"
+
+FUNCTION_TYPE_ATTRIBUTE_NAME = "function_type"
 SYM_NAME_ATTRIBUTE_NAME = "sym_name"
 ARGUMENT_ATTRIBUTE_NAME = "arg_attrs"
 RESULT_ATTRIBUTE_NAME = "res_attrs"
 
-
 @_ods_cext.register_operation(_Dialect, replace=True)
 class GPUFuncOp(GPUFuncOp):
+    __doc__ = GPUFuncOp.__doc__
+
     def __init__(
         self,
         function_type: Union[FunctionType, TypeAttr],
         sym_name: Optional[Union[str, StringAttr]] = None,
         kernel: Optional[bool] = None,
         body_builder: Optional[Callable[[GPUFuncOp], None]] = None,
+        known_block_size: Optional[Union[List[int], DenseI32ArrayAttr]] = None,
+        known_grid_size: Optional[Union[List[int], DenseI32ArrayAttr]] = None,
         *args,
         loc=None,
         ip=None,
         **kwargs,
     ):
+        """
+        Create a GPUFuncOp with the provided `function_type`, `sym_name`, `kernel`, `body_builder`, `known_block_size`, and `known_grid_size`.
+        - `function_type` is a FunctionType or a TypeAttr.
+        - `sym_name` is a string or a StringAttr representing the function name.
+        - `kernel` is a boolean representing whether the function is a kernel.
+        - `body_builder` is an optional callback. When provided, a new entry block
+          is created and the callback is invoked with the new op as argument within
+          an InsertionPoint context already set for the block. The callback is
+          expected to insert a terminator in the block.
+        - `known_block_size` is an optional list of integers or a DenseI32ArrayAttr representing the known block size.
+        - `known_grid_size` is an optional list of integers or a DenseI32ArrayAttr representing the known grid size.
+        """
         function_type = (
             TypeAttr.get(function_type)
             if not isinstance(function_type, TypeAttr)
@@ -62,6 +81,20 @@ def __init__(
         if body_builder is not None:
             with InsertionPoint(self.add_entry_block()):
                 body_builder(self)
+        if known_block_size is not None:
+            if isinstance(known_block_size, list):
+                self.attributes[KNOWN_BLOCK_SIZE_ATTRIBUTE_NAME] = (
+                    DenseI32ArrayAttr.get(known_block_size)
+                )
+            else:
+                self.attributes[KNOWN_BLOCK_SIZE_ATTRIBUTE_NAME] = known_block_size
+        if known_grid_size is not None:
+            if isinstance(known_grid_size, list):
+                self.attributes[KNOWN_GRID_SIZE_ATTRIBUTE_NAME] = (
+                    DenseI32ArrayAttr.get(known_grid_size)
+                )
+            else:
+                self.attributes[KNOWN_GRID_SIZE_ATTRIBUTE_NAME] = known_grid_size
 
     @property
     def type(self) -> FunctionType:
@@ -77,6 +110,18 @@ def name(self) -> StringAttr:
     def is_kernel(self) -> bool:
         return KERNEL_ATTRIBUTE_NAME in self.attributes
 
+    @property
+    def known_block_size(self) -> Optional[DenseI32ArrayAttr]:
+        if KNOWN_BLOCK_SIZE_ATTRIBUTE_NAME not in self.attributes:
+            return None
+        return self.attributes[KNOWN_BLOCK_SIZE_ATTRIBUTE_NAME]
+
+    @property
+    def known_grid_size(self) -> Optional[DenseI32ArrayAttr]:
+        if KNOWN_GRID_SIZE_ATTRIBUTE_NAME not in self.attributes:
+            return None
+        return self.attributes[KNOWN_GRID_SIZE_ATTRIBUTE_NAME]
+
     def add_entry_block(self) -> Block:
         function_type = self.type
         return self.body.blocks.append(

mlir/test/python/dialects/gpu/dialect.py

@@ -1,10 +1,10 @@
 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
+import mlir.ir as ir
 import mlir.dialects.gpu as gpu
 import mlir.dialects.gpu.passes
 from mlir.passmanager import *
-import mlir.ir as ir
 
 
 def run(f):
@@ -70,15 +70,16 @@ def testObjectAttr():
 # CHECK-LABEL: testGPUFuncOp
 @run
 def testGPUFuncOp():
+    assert gpu.GPUFuncOp.__doc__ is not None
     module = Module.create()
     with InsertionPoint(module.body):
         gpu_module_name = StringAttr.get("gpu_module")
         gpumodule = gpu.GPUModuleOp(gpu_module_name)
         block = gpumodule.bodyRegion.blocks.append()
 
         def builder(func: gpu.GPUFuncOp) -> None:
-            _ = gpu.GlobalIdOp(gpu.Dimension.x)
-            _ = gpu.ReturnOp([])
+            gpu.GlobalIdOp(gpu.Dimension.x)
+            gpu.ReturnOp([])
 
         with InsertionPoint(block):
             name = StringAttr.get("kernel0")
@@ -96,6 +97,8 @@ def builder(func: gpu.GPUFuncOp) -> None:
                 sym_name="kernel1",
                 kernel=True,
                 body_builder=builder,
+                known_block_size=[1, 2, 3],
+                known_grid_size=DenseI32ArrayAttr.get([4, 5, 6]),
             )
 
             assert func.name.value == "kernel1"
@@ -104,13 +107,17 @@ def builder(func: gpu.GPUFuncOp) -> None:
             assert func.arguments == []
             assert func.entry_block == func.body.blocks[0]
             assert func.is_kernel
+            assert func.known_block_size == DenseI32ArrayAttr.get([1, 2, 3])
+            assert func.known_grid_size == DenseI32ArrayAttr.get([4, 5, 6])
 
-            non_kernel_func = gpu.GPUFuncOp(
+            func = gpu.GPUFuncOp(
                 func_type,
                 sym_name="non_kernel_func",
                 body_builder=builder,
             )
-            assert not non_kernel_func.is_kernel
+            assert not func.is_kernel
+            assert func.known_block_size is None
+            assert func.known_grid_size is None
 
     print(module)
 
@@ -119,7 +126,9 @@ def builder(func: gpu.GPUFuncOp) -> None:
     # CHECK:   %[[VAL_0:.*]] = gpu.global_id  x
     # CHECK:   gpu.return
     # CHECK: }
-    # CHECK: gpu.func @kernel1() kernel {
+    # CHECK: gpu.func @kernel1() kernel attributes
+    # CHECK-SAME: gpu.known_block_size = array<i32: 1, 2, 3>
+    # CHECK-SAME: gpu.known_grid_size = array<i32: 4, 5, 6>
     # CHECK:   %[[VAL_0:.*]] = gpu.global_id  x
     # CHECK:   gpu.return
     # CHECK: }