1// RUN: mlir-spirv-cpu-runner %s -e main --entry-point-result=void --shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext,%spirv_wrapper_library_dir/libmlir_test_spirv_cpu_runner_c_wrappers%shlibext 2 3// CHECK: [[[7.7, 0, 0], [7.7, 0, 0], [7.7, 0, 0]], [[0, 7.7, 0], [0, 7.7, 0], [0, 7.7, 0]], [[0, 0, 7.7], [0, 0, 7.7], [0, 0, 7.7]]] 4module attributes { 5 gpu.container_module, 6 spv.target_env = #spv.target_env< 7 #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class, SPV_KHR_8bit_storage]>, 8 {max_compute_workgroup_invocations = 128 : i32, 9 max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}> 10} { 11 gpu.module @kernels { 12 gpu.func @sum(%arg0 : memref<3xf32>, %arg1 : memref<3x3xf32>, %arg2 : memref<3x3x3xf32>) 13 kernel attributes { spv.entry_point_abi = {local_size = dense<[1, 1, 1]>: vector<3xi32>}} { 14 %i0 = constant 0 : index 15 %i1 = constant 1 : index 16 %i2 = constant 2 : index 17 18 %x = load %arg0[%i0] : memref<3xf32> 19 %y = load %arg1[%i0, %i0] : memref<3x3xf32> 20 %sum = addf %x, %y : f32 21 22 store %sum, %arg2[%i0, %i0, %i0] : memref<3x3x3xf32> 23 store %sum, %arg2[%i0, %i1, %i0] : memref<3x3x3xf32> 24 store %sum, %arg2[%i0, %i2, %i0] : memref<3x3x3xf32> 25 store %sum, %arg2[%i1, %i0, %i1] : memref<3x3x3xf32> 26 store %sum, %arg2[%i1, %i1, %i1] : memref<3x3x3xf32> 27 store %sum, %arg2[%i1, %i2, %i1] : memref<3x3x3xf32> 28 store %sum, %arg2[%i2, %i0, %i2] : memref<3x3x3xf32> 29 store %sum, %arg2[%i2, %i1, %i2] : memref<3x3x3xf32> 30 store %sum, %arg2[%i2, %i2, %i2] : memref<3x3x3xf32> 31 gpu.return 32 } 33 } 34 35 func @main() { 36 %input1 = alloc() : memref<3xf32> 37 %input2 = alloc() : memref<3x3xf32> 38 %output = alloc() : memref<3x3x3xf32> 39 %0 = constant 0.0 : f32 40 %3 = constant 3.4 : f32 41 %4 = constant 4.3 : f32 42 %input1_casted = memref_cast %input1 : memref<3xf32> to memref<?xf32> 43 %input2_casted = memref_cast %input2 : memref<3x3xf32> to memref<?x?xf32> 44 %output_casted = memref_cast %output : memref<3x3x3xf32> to memref<?x?x?xf32> 45 call @fillF32Buffer1D(%input1_casted, %3) : (memref<?xf32>, f32) -> () 46 call @fillF32Buffer2D(%input2_casted, %4) : (memref<?x?xf32>, f32) -> () 47 call @fillF32Buffer3D(%output_casted, %0) : (memref<?x?x?xf32>, f32) -> () 48 49 %one = constant 1 : index 50 gpu.launch_func @kernels::@sum 51 blocks in (%one, %one, %one) threads in (%one, %one, %one) 52 args(%input1 : memref<3xf32>, %input2 : memref<3x3xf32>, %output : memref<3x3x3xf32>) 53 %result = memref_cast %output : memref<3x3x3xf32> to memref<*xf32> 54 call @print_memref_f32(%result) : (memref<*xf32>) -> () 55 return 56 } 57 func private @fillF32Buffer1D(%arg0 : memref<?xf32>, %arg1 : f32) 58 func private @fillF32Buffer2D(%arg0 : memref<?x?xf32>, %arg1 : f32) 59 func private @fillF32Buffer3D(%arg0 : memref<?x?x?xf32>, %arg1 : f32) 60 func private @print_memref_f32(%arg0 : memref<*xf32>) 61} 62