69 lines · plain
1// Make sure that signed extended multiplication produces expected results2// with and without expansion to primitive mul/add ops for WebGPU.3 4// RUN: mlir-opt %s -test-vulkan-runner-pipeline \5// RUN: | mlir-runner - \6// RUN: --shared-libs=%mlir_vulkan_runtime,%mlir_runner_utils \7// RUN: --entry-point-result=void | FileCheck %s8 9// RUN: mlir-opt %s -test-vulkan-runner-pipeline=spirv-webgpu-prepare \10// RUN: | mlir-runner - \11// RUN: --shared-libs=%mlir_vulkan_runtime,%mlir_runner_utils \12// RUN: --entry-point-result=void | FileCheck %s13 14// CHECK: [0, 1, -2, 1, 1048560, -87620295, -131071, 560969770]15// CHECK: [0, 0, -1, 0, 0, -1, 0, -499807318]16module attributes {17 gpu.container_module,18 spirv.target_env = #spirv.target_env<19 #spirv.vce<v1.4, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>20} {21 gpu.module @kernels {22 gpu.func @kernel_add(%arg0 : memref<8xi32>, %arg1 : memref<8xi32>, %arg2 : memref<8xi32>, %arg3 : memref<8xi32>)23 kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {24 %0 = gpu.block_id x25 %lhs = memref.load %arg0[%0] : memref<8xi32>26 %rhs = memref.load %arg1[%0] : memref<8xi32>27 %low, %hi = arith.mulsi_extended %lhs, %rhs : i3228 memref.store %low, %arg2[%0] : memref<8xi32>29 memref.store %hi, %arg3[%0] : memref<8xi32>30 gpu.return31 }32 }33 34 func.func @main() {35 %buf0 = memref.alloc() : memref<8xi32>36 %buf1 = memref.alloc() : memref<8xi32>37 %buf2 = memref.alloc() : memref<8xi32>38 %buf3 = memref.alloc() : memref<8xi32>39 %i32_0 = arith.constant 0 : i3240 41 // Initialize output buffers.42 %buf4 = memref.cast %buf2 : memref<8xi32> to memref<?xi32>43 %buf5 = memref.cast %buf3 : memref<8xi32> to memref<?xi32>44 call @fillResource1DInt(%buf4, %i32_0) : (memref<?xi32>, i32) -> ()45 call @fillResource1DInt(%buf5, %i32_0) : (memref<?xi32>, i32) -> ()46 47 %idx_0 = arith.constant 0 : index48 %idx_1 = arith.constant 1 : index49 %idx_8 = arith.constant 8 : index50 51 // Initialize input buffers.52 %lhs_vals = arith.constant dense<[0, 1, -1, -1, 65535, 65535, -65535, 2088183954]> : vector<8xi32>53 %rhs_vals = arith.constant dense<[0, 1, 2, -1, 16, -1337, -65535, -1028001427]> : vector<8xi32>54 vector.store %lhs_vals, %buf0[%idx_0] : memref<8xi32>, vector<8xi32>55 vector.store %rhs_vals, %buf1[%idx_0] : memref<8xi32>, vector<8xi32>56 57 gpu.launch_func @kernels::@kernel_add58 blocks in (%idx_8, %idx_1, %idx_1) threads in (%idx_1, %idx_1, %idx_1)59 args(%buf0 : memref<8xi32>, %buf1 : memref<8xi32>, %buf2 : memref<8xi32>, %buf3 : memref<8xi32>)60 %buf_low = memref.cast %buf4 : memref<?xi32> to memref<*xi32>61 %buf_hi = memref.cast %buf5 : memref<?xi32> to memref<*xi32>62 call @printMemrefI32(%buf_low) : (memref<*xi32>) -> ()63 call @printMemrefI32(%buf_hi) : (memref<*xi32>) -> ()64 return65 }66 func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)67 func.func private @printMemrefI32(%ptr : memref<*xi32>)68}69