mlir/test/mlir-vulkan-runner/umul_extended.mlir

   1 // Make sure that unsigned extended multiplication produces expected results
   2 // with and without expansion to primitive mul/add ops for WebGPU.
   3
   4 // RUN: mlir-opt %s -test-vulkan-runner-pipeline \
   5 // RUN:   | mlir-vulkan-runner - \
   6 // RUN:     --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
   7 // RUN:     --entry-point-result=void | FileCheck %s
   8
   9 // RUN: mlir-opt %s -test-vulkan-runner-pipeline -spirv-webgpu-prepare \
  10 // RUN:   | mlir-vulkan-runner - \
  11 // RUN:     --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
  12 // RUN:     --entry-point-result=void | FileCheck %s
  13
  14 // CHECK: [0, 1, -2,  1, 1048560, -87620295, -131071, -49]
  15 // CHECK: [0, 0,  1, -2,       0,     65534, -131070,   6]
  16 module 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 x
  25       %lhs = memref.load %arg0[%0] : memref<8xi32>
  26       %rhs = memref.load %arg1[%0] : memref<8xi32>
  27       %low, %hi = arith.mului_extended %lhs, %rhs : i32
  28       memref.store %low, %arg2[%0] : memref<8xi32>
  29       memref.store %hi, %arg3[%0] : memref<8xi32>
  30       gpu.return
  31     }
  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 : i32
  40
  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 : index
  48     %idx_1 = arith.constant 1 : index
  49     %idx_8 = arith.constant 8 : index
  50
  51     // Initialize input buffers.
  52     %lhs_vals = arith.constant dense<[0, 1, -1,  -1,  65535,  65535, -65535,  7]> : vector<8xi32>
  53     %rhs_vals = arith.constant dense<[0, 1,  2,  -1,     16,  -1337, -65535, -7]> : 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_add
  58         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     return
  65   }
  66   func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)
  67   func.func private @printMemrefI32(%ptr : memref<*xi32>)
  68 }