60 lines · plain
1// RUN: mlir-opt %s -pass-pipeline='builtin.module(spirv-attach-target{ver=v1.0 caps=Addresses,Int64,Kernel},convert-gpu-to-spirv{use-64bit-index=true},gpu.module(spirv.module(spirv-lower-abi-attrs,spirv-update-vce)),func.func(llvm-request-c-wrappers),convert-scf-to-cf,convert-to-llvm,gpu-to-llvm{use-bare-pointers-for-kernels=true},gpu-module-to-binary,expand-strided-metadata,lower-affine,reconcile-unrealized-casts)' \2// RUN: | mlir-runner \3// RUN: --shared-libs=%mlir_levelzero_runtime \4// RUN: --shared-libs=%mlir_runner_utils \5// RUN: --entry-point-result=void \6// RUN: | FileCheck %s7 8module @add attributes {gpu.container_module} {9 memref.global "private" constant @__constant_2x2x2xf32_0 : memref<2x2x2xf32> = dense<[[[1.1, 2.2], [3.3, 4.4]], [[5.5, 6.6], [7.7, 8.8 ]]]>10 memref.global "private" constant @__constant_2x2x2xf32 : memref<2x2x2xf32> = dense<[[[1.2, 2.3], [4.5, 5.8]], [[7.2, 8.3], [10.5, 11.8]]]>11 func.func @main() {12 %0 = memref.get_global @__constant_2x2x2xf32 : memref<2x2x2xf32>13 %1 = memref.get_global @__constant_2x2x2xf32_0 : memref<2x2x2xf32>14 %2 = call @test(%0, %1) : (memref<2x2x2xf32>, memref<2x2x2xf32>) -> memref<2x2x2xf32>15 %cast = memref.cast %2 : memref<2x2x2xf32> to memref<*xf32>16 call @printMemrefF32(%cast) : (memref<*xf32>) -> ()17 return18 }19 func.func private @printMemrefF32(memref<*xf32>)20 func.func @test(%arg0: memref<2x2x2xf32>, %arg1: memref<2x2x2xf32>) -> memref<2x2x2xf32> {21 %c2 = arith.constant 2 : index22 %c1 = arith.constant 1 : index23 %mem = gpu.alloc host_shared () : memref<2x2x2xf32>24 memref.copy %arg1, %mem : memref<2x2x2xf32> to memref<2x2x2xf32>25 %memref_0 = gpu.alloc host_shared () : memref<2x2x2xf32>26 memref.copy %arg0, %memref_0 : memref<2x2x2xf32> to memref<2x2x2xf32>27 %memref_2 = gpu.alloc host_shared () : memref<2x2x2xf32>28 %2 = gpu.wait async29 %3 = gpu.launch_func async [%2] @test_kernel::@test_kernel blocks in (%c2, %c2, %c2) threads in (%c1, %c1, %c1)30 args(%memref_0 : memref<2x2x2xf32>, %mem : memref<2x2x2xf32>, %memref_2 : memref<2x2x2xf32>)31 gpu.wait [%3]32 %alloc = memref.alloc() : memref<2x2x2xf32>33 memref.copy %memref_2, %alloc : memref<2x2x2xf32> to memref<2x2x2xf32>34 %4 = gpu.wait async35 %5 = gpu.dealloc async [%4] %memref_2 : memref<2x2x2xf32>36 %6 = gpu.dealloc async [%5] %memref_0 : memref<2x2x2xf32>37 %7 = gpu.dealloc async [%6] %mem : memref<2x2x2xf32>38 gpu.wait [%7]39 return %alloc : memref<2x2x2xf32>40 }41 gpu.module @test_kernel42 attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Addresses, Int64, Kernel], []>, api=OpenCL, #spirv.resource_limits<>>} {43 gpu.func @test_kernel(%arg0: memref<2x2x2xf32>, %arg1: memref<2x2x2xf32>, %arg2: memref<2x2x2xf32>) kernel44 attributes {gpu.known_block_size = array<i32: 1, 1, 1>, gpu.known_grid_size = array<i32: 2, 2, 2>, spirv.entry_point_abi = #spirv.entry_point_abi<>} {45 %0 = gpu.block_id x46 %1 = gpu.block_id y47 %2 = gpu.block_id z48 %3 = memref.load %arg0[%0, %1, %2] : memref<2x2x2xf32>49 %4 = memref.load %arg1[%0, %1, %2] : memref<2x2x2xf32>50 %5 = arith.addf %3, %4 : f3251 memref.store %5, %arg2[%0, %1, %2] : memref<2x2x2xf32>52 gpu.return53 }54 }55 // CHECK: [2.3, 4.5]56 // CHECK: [7.8, 10.2]57 // CHECK: [12.7, 14.9]58 // CHECK: [18.2, 20.6]59}60