75 lines · plain
1// RUN: mlir-opt %s -test-spirv-cpu-runner-pipeline \2// RUN: | mlir-runner - -e main --entry-point-result=void --shared-libs=%mlir_runner_utils,%mlir_spirv_cpu_runtime --link-nested-modules \3// RUN: | FileCheck %s4 5// CHECK: data =6// CHECK{LITERAL}: [[[7.7, 0, 0],7// CHECK{LITERAL}: [7.7, 0, 0],8// CHECK{LITERAL}: [7.7, 0, 0]],9// CHECK{LITERAL}: [[0, 7.7, 0],10// CHECK{LITERAL}: [0, 7.7, 0],11// CHECK{LITERAL}: [0, 7.7, 0]],12// CHECK{LITERAL}: [[0, 0, 7.7],13// CHECK{LITERAL}: [0, 0, 7.7],14// CHECK{LITERAL}: [0, 0, 7.7]]]15 16module attributes {17 gpu.container_module,18 spirv.target_env = #spirv.target_env<19 #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class, SPV_KHR_8bit_storage]>,20 #spirv.resource_limits<21 max_compute_workgroup_invocations = 128,22 max_compute_workgroup_size = [128, 128, 64]>>23} {24 gpu.module @kernels {25 gpu.func @sum(%arg0 : memref<3xf32>, %arg1 : memref<3x3xf32>, %arg2 : memref<3x3x3xf32>)26 kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {27 %i0 = arith.constant 0 : index28 %i1 = arith.constant 1 : index29 %i2 = arith.constant 2 : index30 31 %x = memref.load %arg0[%i0] : memref<3xf32>32 %y = memref.load %arg1[%i0, %i0] : memref<3x3xf32>33 %sum = arith.addf %x, %y : f3234 35 memref.store %sum, %arg2[%i0, %i0, %i0] : memref<3x3x3xf32>36 memref.store %sum, %arg2[%i0, %i1, %i0] : memref<3x3x3xf32>37 memref.store %sum, %arg2[%i0, %i2, %i0] : memref<3x3x3xf32>38 memref.store %sum, %arg2[%i1, %i0, %i1] : memref<3x3x3xf32>39 memref.store %sum, %arg2[%i1, %i1, %i1] : memref<3x3x3xf32>40 memref.store %sum, %arg2[%i1, %i2, %i1] : memref<3x3x3xf32>41 memref.store %sum, %arg2[%i2, %i0, %i2] : memref<3x3x3xf32>42 memref.store %sum, %arg2[%i2, %i1, %i2] : memref<3x3x3xf32>43 memref.store %sum, %arg2[%i2, %i2, %i2] : memref<3x3x3xf32>44 gpu.return45 }46 }47 48 func.func @main() {49 %input1 = memref.alloc() : memref<3xf32>50 %input2 = memref.alloc() : memref<3x3xf32>51 %output = memref.alloc() : memref<3x3x3xf32>52 %0 = arith.constant 0.0 : f3253 %3 = arith.constant 3.4 : f3254 %4 = arith.constant 4.3 : f3255 %input1_casted = memref.cast %input1 : memref<3xf32> to memref<?xf32>56 %input2_casted = memref.cast %input2 : memref<3x3xf32> to memref<?x?xf32>57 %output_casted = memref.cast %output : memref<3x3x3xf32> to memref<?x?x?xf32>58 call @fillF32Buffer1D(%input1_casted, %3) : (memref<?xf32>, f32) -> ()59 call @fillF32Buffer2D(%input2_casted, %4) : (memref<?x?xf32>, f32) -> ()60 call @fillF32Buffer3D(%output_casted, %0) : (memref<?x?x?xf32>, f32) -> ()61 62 %one = arith.constant 1 : index63 gpu.launch_func @kernels::@sum64 blocks in (%one, %one, %one) threads in (%one, %one, %one)65 args(%input1 : memref<3xf32>, %input2 : memref<3x3xf32>, %output : memref<3x3x3xf32>)66 %result = memref.cast %output : memref<3x3x3xf32> to memref<*xf32>67 call @printMemrefF32(%result) : (memref<*xf32>) -> ()68 return69 }70 func.func private @fillF32Buffer1D(%arg0 : memref<?xf32>, %arg1 : f32)71 func.func private @fillF32Buffer2D(%arg0 : memref<?x?xf32>, %arg1 : f32)72 func.func private @fillF32Buffer3D(%arg0 : memref<?x?x?xf32>, %arg1 : f32)73 func.func private @printMemrefF32(%arg0 : memref<*xf32>)74}75