diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -148,6 +148,7 @@ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/GPUTransformOps.td SOURCES + dialects/_gpu_transform_ops_ext.py dialects/transform/gpu.py DIALECT_NAME transform EXTENSION_NAME gpu_transform) diff --git a/mlir/python/mlir/dialects/_gpu_transform_ops_ext.py b/mlir/python/mlir/dialects/_gpu_transform_ops_ext.py new file mode 100644 --- /dev/null +++ b/mlir/python/mlir/dialects/_gpu_transform_ops_ext.py @@ -0,0 +1,70 @@ +# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +try: + from ..ir import * + from ._ods_common import get_op_result_or_value as _get_op_result_or_value + from ..dialects import pdl, transform +except ImportError as e: + raise RuntimeError("Error loading imports from extension module") from e + +from typing import List, Optional, Sequence, Tuple, Union, overload + + +class MapForallToBlocks: + """Specialization for MapForallToBlocks class.""" + + @overload + def __init__( + self, + result_type: Type, + target: Union[Operation, OpView, Value], + *, + grid_dims: Optional[Sequence[int]] = None, + generate_gpu_launch: Optional[bool] = None, + loc=None, + ip=None + ): + ... + + @overload + def __init__( + self, + target: Union[Operation, OpView, Value], + *, + grid_dims: Optional[Sequence[int]] = None, + generate_gpu_launch: Optional[bool] = None, + loc=None, + ip=None + ): + ... + + def __init__( + self, + result_type_or_target: Union[Operation, OpView, Type, Value], + target_or_none: Optional[Union[Operation, OpView, Value]] = None, + *, + grid_dims: Optional[Sequence[int]] = None, + generate_gpu_launch: Optional[bool] = None, + loc=None, + ip=None + ): + if isinstance(result_type_or_target, Type): + result_type = result_type_or_target + target = target_or_none + else: + result_type = transform.AnyOpType.get() + target = result_type_or_target + + if grid_dims is not None and not isinstance(grid_dims, ArrayAttr): + grid_dims = DenseI64ArrayAttr.get(grid_dims) + + super().__init__( + result_type, + target, + grid_dims=grid_dims, + generate_gpu_launch=generate_gpu_launch, + loc=loc, + ip=ip, + ) diff --git a/mlir/test/python/dialects/transform_gpu_ext.py b/mlir/test/python/dialects/transform_gpu_ext.py new file mode 100644 --- /dev/null +++ b/mlir/test/python/dialects/transform_gpu_ext.py @@ -0,0 +1,60 @@ +# RUN: %PYTHON %s | FileCheck %s + +from mlir.ir import * +from mlir.dialects import transform +from mlir.dialects import pdl +from mlir.dialects.transform import gpu + + +def run(f): + with Context(), Location.unknown(): + module = Module.create() + with InsertionPoint(module.body): + print("\nTEST:", f.__name__) + f() + print(module) + return f + + +@run +def testMapForallToBlocksCompact(): + sequence = transform.SequenceOp( + transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get() + ) + with InsertionPoint(sequence.body): + gpu.MapForallToBlocks(sequence.bodyTarget) + transform.YieldOp() + # CHECK-LABEL: TEST: testMapForallToBlocksCompact + # CHECK: = transform.gpu.map_forall_to_blocks + # CHECK-NOT: grid_dims + # CHECK-SAME: (!transform.any_op) -> !transform.any_op + # CHECK-NOT: grid_dims + + +@run +def testMapForallToBlocksTyped(): + sequence = transform.SequenceOp( + transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get() + ) + with InsertionPoint(sequence.body): + gpu.MapForallToBlocks( + transform.OperationType.get("test.dummy"), sequence.bodyTarget + ) + transform.YieldOp() + # CHECK-LABEL: TEST: testMapForallToBlocksTyped + # CHECK: = transform.gpu.map_forall_to_blocks + # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy"> + + +@run +def testMapForallToBlocksGridDims(): + sequence = transform.SequenceOp( + transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get() + ) + with InsertionPoint(sequence.body): + gpu.MapForallToBlocks(sequence.bodyTarget, grid_dims=[4, 2]) + transform.YieldOp() + # CHECK-LABEL: TEST: testMapForallToBlocksGridDims + # CHECK: = transform.gpu.map_forall_to_blocks + # CHECK-SAME: grid_dims = [4, 2] + # CHECK-SAME: (!transform.any_op) -> !transform.any_op