46 lines · python
1# RUN: %PYTHON %s | FileCheck %s2# This is just a smoke test that the dialect is functional.3 4from mlir.ir import *5from mlir.dialects import amdgpu, func6 7 8def constructAndPrintInModule(f):9 print("\nTEST:", f.__name__)10 with Context(), Location.unknown():11 module = Module.create()12 with InsertionPoint(module.body):13 f()14 print(module)15 return f16 17 18# CHECK-LABEL: testSmoke19@constructAndPrintInModule20def testSmoke():21 # CHECK: amdgpu.lds_barrier22 amdgpu.LDSBarrierOp()23 24 25# CHECK-LABEL: testFatRawBufferCastOpParams26@constructAndPrintInModule27def testFatRawBufferCastOpParams():28 memref_type = MemRefType.get(29 [ShapedType.get_dynamic_size(), ShapedType.get_dynamic_size()],30 F32Type.get(),31 )32 f = func.FuncOp("test_raw_buffer_cast_params", ([memref_type], []))33 with InsertionPoint(f.add_entry_block()):34 block_args = f.arguments35 amdgpu.FatRawBufferCastOp(block_args[0])36 amdgpu.FatRawBufferCastOp(block_args[0], resetOffset=True)37 amdgpu.FatRawBufferCastOp(block_args[0], boundsCheck=False)38 amdgpu.FatRawBufferCastOp(block_args[0], boundsCheck=False, resetOffset=True)39 func.ReturnOp([])40 41 # CHECK: func.func @test_raw_buffer_cast_params(%[[ARG0:.+]]: memref<?x?xf32>) {42 # CHECK: amdgpu.fat_raw_buffer_cast %[[ARG0]] : memref<?x?xf32> to memref<?x?xf32, #amdgpu.address_space<fat_raw_buffer>>43 # CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] resetOffset : memref<?x?xf32> to memref<?x?xf32, #amdgpu.address_space<fat_raw_buffer>>44 # CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] boundsCheck(false) : memref<?x?xf32> to memref<?x?xf32, #amdgpu.address_space<fat_raw_buffer>>45 # CHECK-NEXT: amdgpu.fat_raw_buffer_cast %[[ARG0]] boundsCheck(false) resetOffset : memref<?x?xf32> to memref<?x?xf32, #amdgpu.address_space<fat_raw_buffer>>46