80 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects import transform5from mlir.dialects.transform import gpu6 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 f23 24 25@run26def testMapForallToBlocksCompact(target):27 gpu.MapForallToBlocks(target)28 # CHECK-LABEL: TEST: testMapForallToBlocksCompact29 # CHECK: = transform.gpu.map_forall_to_blocks30 # CHECK-NOT: grid_dims31 # CHECK-SAME: (!transform.any_op) -> !transform.any_op32 # CHECK-NOT: grid_dims33 34 35@run36def testMapForallToBlocksTyped(target):37 gpu.MapForallToBlocks(transform.OperationType.get("test.dummy"), target)38 # CHECK-LABEL: TEST: testMapForallToBlocksTyped39 # CHECK: = transform.gpu.map_forall_to_blocks40 # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">41 42 43@run44def testMapForallToBlocksGridDims(target):45 gpu.MapForallToBlocks(target, grid_dims=[4, 2, 1])46 # CHECK-LABEL: TEST: testMapForallToBlocksGridDims47 # CHECK: = transform.gpu.map_forall_to_blocks48 # CHECK-SAME: grid_dims = [4, 2, 1]49 # CHECK-SAME: (!transform.any_op) -> !transform.any_op50 51 52@run53def testMapNestedForallToThreadsCompact(target):54 gpu.MapNestedForallToThreads(target)55 # CHECK-LABEL: TEST: testMapNestedForallToThreadsCompact56 # CHECK: transform.gpu.map_nested_forall_to_threads57 # CHECK-SAME: block_dims = []58 # CHECK-SAME: (!transform.any_op) -> !transform.any_op59 60 61@run62def testMapNestedForallToThreadsTyped(target):63 gpu.MapNestedForallToThreads(transform.OperationType.get("test.dummy"), target)64 # CHECK-LABEL: TEST: testMapNestedForallToThreadsTyped65 # CHECK: transform.gpu.map_nested_forall_to_threads66 # CHECK-SAME: block_dims = []67 # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">68 69 70@run71def testMapNestedForallToThreadsAttributes(target):72 gpu.MapNestedForallToThreads(73 target, block_dims=[4, 2], warp_size=64, sync_after_distribute=False74 )75 # CHECK-LABEL: TEST: testMapNestedForallToThreadsAttributes76 # CHECK: transform.gpu.map_nested_forall_to_threads77 # CHECK-SAME: block_dims = [4, 2]78 # CHECK-SAME: sync_after_distribute = false79 # CHECK-SAME: warp_size = 6480