71 lines · plain
1// Make sure that addition with carry produces expected results2// with and without expansion to primitive add/cmp 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, 42, 0, 42]15// CHECK: [1, 0, 1, 1]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<4xi32>, %arg1 : memref<4xi32>, %arg2 : memref<4xi32>, %arg3 : memref<4xi32>)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<4xi32>26 %rhs = memref.load %arg1[%0] : memref<4xi32>27 %sum, %carry = arith.addui_extended %lhs, %rhs : i32, i128 29 %carry_i32 = arith.extui %carry : i1 to i3230 31 memref.store %sum, %arg2[%0] : memref<4xi32> memref.store %carry_i32, %arg3[%0] : memref<4xi32>32 gpu.return33 }34 }35 36 func.func @main() {37 %buf0 = memref.alloc() : memref<4xi32>38 %buf1 = memref.alloc() : memref<4xi32>39 %buf2 = memref.alloc() : memref<4xi32>40 %buf3 = memref.alloc() : memref<4xi32>41 %i32_0 = arith.constant 0 : i3242 43 // Initialize output buffers.44 %buf4 = memref.cast %buf2 : memref<4xi32> to memref<?xi32>45 %buf5 = memref.cast %buf3 : memref<4xi32> to memref<?xi32>46 call @fillResource1DInt(%buf4, %i32_0) : (memref<?xi32>, i32) -> ()47 call @fillResource1DInt(%buf5, %i32_0) : (memref<?xi32>, i32) -> ()48 49 %idx_0 = arith.constant 0 : index50 %idx_1 = arith.constant 1 : index51 %idx_4 = arith.constant 4 : index52 53 // Initialize input buffers.54 %lhs_vals = arith.constant dense<[-1, 24, 4294967295, 43]> : vector<4xi32>55 %rhs_vals = arith.constant dense<[1, 18, 1, 4294967295]> : vector<4xi32>56 vector.store %lhs_vals, %buf0[%idx_0] : memref<4xi32>, vector<4xi32>57 vector.store %rhs_vals, %buf1[%idx_0] : memref<4xi32>, vector<4xi32>58 59 gpu.launch_func @kernels::@kernel_add60 blocks in (%idx_4, %idx_1, %idx_1) threads in (%idx_1, %idx_1, %idx_1)61 args(%buf0 : memref<4xi32>, %buf1 : memref<4xi32>, %buf2 : memref<4xi32>, %buf3 : memref<4xi32>)62 %buf_sum = memref.cast %buf4 : memref<?xi32> to memref<*xi32>63 %buf_carry = memref.cast %buf5 : memref<?xi32> to memref<*xi32>64 call @printMemrefI32(%buf_sum) : (memref<*xi32>) -> ()65 call @printMemrefI32(%buf_carry) : (memref<*xi32>) -> ()66 return67 }68 func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)69 func.func private @printMemrefI32(%ptr : memref<*xi32>)70}71