76 lines · plain
1// RUN: mlir-opt %s \2// RUN: | mlir-opt -pass-pipeline='builtin.module(cse,func.func(gpu-async-region),xevm-attach-target,gpu.module(convert-gpu-to-llvm-spv{use-64bit-index=true},convert-xevm-to-llvm,cse))' \3// RUN: | mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-vector-to-llvm -convert-arith-to-llvm \4// RUN: | mlir-opt -gpu-to-llvm -reconcile-unrealized-casts -cse -gpu-module-to-binary \5// RUN: | mlir-runner \6// RUN: --shared-libs=%mlir_levelzero_runtime \7// RUN: --shared-libs=%mlir_runner_utils \8// RUN: --shared-libs=%mlir_c_runner_utils \9// RUN: --entry-point-result=void \10// RUN: | FileCheck %s11 12module @gemm attributes {gpu.container_module} {13 14 gpu.module @kernel {15 gpu.func @store_constant(%ptr: !llvm.ptr<1>) kernel {16 %const_val = arith.constant 42.0 : f3217 %thread_x = gpu.lane_id18 %thread_x_i64 = arith.index_cast %thread_x : index to i6419 %ptr_next_1 = llvm.getelementptr %ptr[%thread_x_i64] : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, i3220 llvm.store %const_val, %ptr_next_1 : f32, !llvm.ptr<1>21 gpu.return22 }23 }24 func.func @test(%src : memref<8x16xf32>) -> memref<8x16xf32> attributes {llvm.emit_c_interface} {25 %c1 = arith.constant 1 : index26 %c16 = arith.constant 16 : index27 %memref_0 = gpu.alloc() : memref<8x16xf32>28 gpu.memcpy %memref_0, %src : memref<8x16xf32>, memref<8x16xf32>29 %0 = memref.extract_aligned_pointer_as_index %memref_0 : memref<8x16xf32> -> index30 %1 = arith.index_cast %0 : index to i6431 %2 = llvm.inttoptr %1 : i64 to !llvm.ptr32 %src_casted = llvm.addrspacecast %2 : !llvm.ptr to !llvm.ptr<1>33 gpu.launch_func @kernel::@store_constant blocks in (%c1, %c1, %c1) threads in (%c16, %c1, %c1)34 args(%src_casted : !llvm.ptr<1>)35 %dst = memref.alloc() : memref<8x16xf32>36 gpu.memcpy %dst, %memref_0 : memref<8x16xf32>, memref<8x16xf32>37 gpu.dealloc %memref_0 : memref<8x16xf32>38 39 return %dst : memref<8x16xf32>40 }41 42 func.func @main() attributes {llvm.emit_c_interface} {43 %A = memref.alloc() : memref<8x16xf32>44 %c0 = arith.constant 0 : index45 %c1 = arith.constant 1 : index46 %c8 = arith.constant 8 : index47 %c16 = arith.constant 16 : index48 %c11_f32 = arith.constant 11.11 : f3249 scf.for %i = %c0 to %c8 step %c1 {50 scf.for %j = %c0 to %c16 step %c1 {51 memref.store %c11_f32, %A[%i, %j] : memref<8x16xf32>52 }53 }54 %B = call @test(%A) : (memref<8x16xf32>) -> memref<8x16xf32>55 %B_cast = memref.cast %B : memref<8x16xf32> to memref<*xf32>56 %A_cast = memref.cast %A : memref<8x16xf32> to memref<*xf32>57 call @printMemrefF32(%A_cast) : (memref<*xf32>) -> ()58 call @printMemrefF32(%B_cast) : (memref<*xf32>) -> ()59 60 // CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}61 // CHECK-NEXT: [11.11{{.*}}]62 // CHECK-COUNT-96: 11.1163 // CHECK-NEXT: [11.11{{.*}}]64 65 // CHECK-NEXT: Unranked Memref base@ = 0x{{[0-9a-f]+}}66 // CHECK-NEXT: [42, 42, 42, 42, 42, 42, 42, 42, 42, 42, 42, 42, 42, 42, 42, 42]67 // CHECK-COUNT-96: 11.1168 // CHECK-NEXT: [11.11{{.*}}]69 70 memref.dealloc %A : memref<8x16xf32>71 memref.dealloc %B : memref<8x16xf32>72 return73 }74 func.func private @printMemrefF32(%ptr : memref<*xf32>)75}76