15f4f9220SIngo Müller# RUN: %PYTHON %s | FileCheck %s 25f4f9220SIngo Müller 35f4f9220SIngo Müllerfrom mlir.ir import * 45f4f9220SIngo Müllerfrom mlir.dialects import transform 55f4f9220SIngo Müllerfrom mlir.dialects.transform import gpu 65f4f9220SIngo Müller 75f4f9220SIngo Müller 85f4f9220SIngo Müllerdef run(f): 95f4f9220SIngo Müller with Context(), Location.unknown(): 105f4f9220SIngo Müller module = Module.create() 115f4f9220SIngo Müller with InsertionPoint(module.body): 12b17acc08SAlex Zinenko sequence = transform.SequenceOp( 1392233062Smax transform.FailurePropagationMode.Propagate, 14b17acc08SAlex Zinenko [], 15b17acc08SAlex Zinenko transform.AnyOpType.get(), 16b17acc08SAlex Zinenko ) 17b17acc08SAlex Zinenko with InsertionPoint(sequence.body): 18b17acc08SAlex Zinenko f(sequence.bodyTarget) 19b17acc08SAlex Zinenko transform.YieldOp() 205f4f9220SIngo Müller print("\nTEST:", f.__name__) 215f4f9220SIngo Müller print(module) 225f4f9220SIngo Müller return f 235f4f9220SIngo Müller 245f4f9220SIngo Müller 255f4f9220SIngo Müller@run 26b17acc08SAlex Zinenkodef testMapForallToBlocksCompact(target): 27b17acc08SAlex Zinenko gpu.MapForallToBlocks(target) 285f4f9220SIngo Müller # CHECK-LABEL: TEST: testMapForallToBlocksCompact 295f4f9220SIngo Müller # CHECK: = transform.gpu.map_forall_to_blocks 305f4f9220SIngo Müller # CHECK-NOT: grid_dims 315f4f9220SIngo Müller # CHECK-SAME: (!transform.any_op) -> !transform.any_op 325f4f9220SIngo Müller # CHECK-NOT: grid_dims 335f4f9220SIngo Müller 345f4f9220SIngo Müller 355f4f9220SIngo Müller@run 36b17acc08SAlex Zinenkodef testMapForallToBlocksTyped(target): 37b17acc08SAlex Zinenko gpu.MapForallToBlocks(transform.OperationType.get("test.dummy"), target) 385f4f9220SIngo Müller # CHECK-LABEL: TEST: testMapForallToBlocksTyped 395f4f9220SIngo Müller # CHECK: = transform.gpu.map_forall_to_blocks 405f4f9220SIngo Müller # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy"> 415f4f9220SIngo Müller 425f4f9220SIngo Müller 435f4f9220SIngo Müller@run 44b17acc08SAlex Zinenkodef testMapForallToBlocksGridDims(target): 45*92f088d3SNicolas Vasilache gpu.MapForallToBlocks(target, grid_dims=[4, 2, 1]) 465f4f9220SIngo Müller # CHECK-LABEL: TEST: testMapForallToBlocksGridDims 475f4f9220SIngo Müller # CHECK: = transform.gpu.map_forall_to_blocks 48*92f088d3SNicolas Vasilache # CHECK-SAME: grid_dims = [4, 2, 1] 495f4f9220SIngo Müller # CHECK-SAME: (!transform.any_op) -> !transform.any_op 50b17acc08SAlex Zinenko 51b17acc08SAlex Zinenko 52b17acc08SAlex Zinenko@run 53b17acc08SAlex Zinenkodef testMapNestedForallToThreadsCompact(target): 54b17acc08SAlex Zinenko gpu.MapNestedForallToThreads(target) 55b17acc08SAlex Zinenko # CHECK-LABEL: TEST: testMapNestedForallToThreadsCompact 56b17acc08SAlex Zinenko # CHECK: transform.gpu.map_nested_forall_to_threads 57b17acc08SAlex Zinenko # CHECK-SAME: block_dims = [] 58b17acc08SAlex Zinenko # CHECK-SAME: (!transform.any_op) -> !transform.any_op 59b17acc08SAlex Zinenko 60b17acc08SAlex Zinenko 61b17acc08SAlex Zinenko@run 62b17acc08SAlex Zinenkodef testMapNestedForallToThreadsTyped(target): 63b17acc08SAlex Zinenko gpu.MapNestedForallToThreads(transform.OperationType.get("test.dummy"), target) 64b17acc08SAlex Zinenko # CHECK-LABEL: TEST: testMapNestedForallToThreadsTyped 65b17acc08SAlex Zinenko # CHECK: transform.gpu.map_nested_forall_to_threads 66b17acc08SAlex Zinenko # CHECK-SAME: block_dims = [] 67b17acc08SAlex Zinenko # CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy"> 68b17acc08SAlex Zinenko 69b17acc08SAlex Zinenko 70b17acc08SAlex Zinenko@run 71b17acc08SAlex Zinenkodef testMapNestedForallToThreadsAttributes(target): 72b17acc08SAlex Zinenko gpu.MapNestedForallToThreads( 73b17acc08SAlex Zinenko target, block_dims=[4, 2], warp_size=64, sync_after_distribute=False 74b17acc08SAlex Zinenko ) 75b17acc08SAlex Zinenko # CHECK-LABEL: TEST: testMapNestedForallToThreadsAttributes 76b17acc08SAlex Zinenko # CHECK: transform.gpu.map_nested_forall_to_threads 77b17acc08SAlex Zinenko # CHECK-SAME: block_dims = [4, 2] 78b17acc08SAlex Zinenko # CHECK-SAME: sync_after_distribute = false 79b17acc08SAlex Zinenko # CHECK-SAME: warp_size = 64 80