xref: /llvm-project/mlir/test/python/dialects/transform_gpu_ext.py (revision 92f088d335e4a55f10e74c541ac52ec2a4d4ceeb) 
1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4from mlir.dialects import transform
5from mlir.dialects.transform import gpu
6
7
8def run(f):
9    with Context(), Location.unknown():
10        module = Module.create()
11        with InsertionPoint(module.body):
12            sequence = transform.SequenceOp(
13                transform.FailurePropagationMode.Propagate,
14                [],
15                transform.AnyOpType.get(),
16            )
17            with InsertionPoint(sequence.body):
18                f(sequence.bodyTarget)
19                transform.YieldOp()
20        print("\nTEST:", f.__name__)
21        print(module)
22    return f
23
24
25@run
26def testMapForallToBlocksCompact(target):
27    gpu.MapForallToBlocks(target)
28    # CHECK-LABEL: TEST: testMapForallToBlocksCompact
29    # CHECK: = transform.gpu.map_forall_to_blocks
30    # CHECK-NOT: grid_dims
31    # CHECK-SAME: (!transform.any_op) -> !transform.any_op
32    # CHECK-NOT: grid_dims
33
34
35@run
36def testMapForallToBlocksTyped(target):
37    gpu.MapForallToBlocks(transform.OperationType.get("test.dummy"), target)
38    # CHECK-LABEL: TEST: testMapForallToBlocksTyped
39    # CHECK: = transform.gpu.map_forall_to_blocks
40    # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">
41
42
43@run
44def testMapForallToBlocksGridDims(target):
45    gpu.MapForallToBlocks(target, grid_dims=[4, 2, 1])
46    # CHECK-LABEL: TEST: testMapForallToBlocksGridDims
47    # CHECK: = transform.gpu.map_forall_to_blocks
48    # CHECK-SAME: grid_dims = [4, 2, 1]
49    # CHECK-SAME: (!transform.any_op) -> !transform.any_op
50
51
52@run
53def testMapNestedForallToThreadsCompact(target):
54    gpu.MapNestedForallToThreads(target)
55    # CHECK-LABEL: TEST: testMapNestedForallToThreadsCompact
56    # CHECK: transform.gpu.map_nested_forall_to_threads
57    # CHECK-SAME: block_dims = []
58    # CHECK-SAME: (!transform.any_op) -> !transform.any_op
59
60
61@run
62def testMapNestedForallToThreadsTyped(target):
63    gpu.MapNestedForallToThreads(transform.OperationType.get("test.dummy"), target)
64    # CHECK-LABEL: TEST: testMapNestedForallToThreadsTyped
65    # CHECK: transform.gpu.map_nested_forall_to_threads
66    # CHECK-SAME: block_dims = []
67    # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">
68
69
70@run
71def testMapNestedForallToThreadsAttributes(target):
72    gpu.MapNestedForallToThreads(
73        target, block_dims=[4, 2], warp_size=64, sync_after_distribute=False
74    )
75    # CHECK-LABEL: TEST: testMapNestedForallToThreadsAttributes
76    # CHECK: transform.gpu.map_nested_forall_to_threads
77    # CHECK-SAME: block_dims = [4, 2]
78    # CHECK-SAME: sync_after_distribute = false
79    # CHECK-SAME: warp_size = 64
80