xref: /llvm-project/mlir/test/Conversion/MemRefToSPIRV/atomic.mlir (revision c63febb1025564b078a5c8e52e6df638e8a1d808) 
1// RUN: mlir-opt -split-input-file -convert-memref-to-spirv %s -o - | FileCheck %s
2
3module attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader], []>, #spirv.resource_limits<>>} {
4
5//      CHECK: func.func @atomic_addi_storage_buffer
6// CHECK-SAME: (%[[VAL:.+]]: i32,
7func.func @atomic_addi_storage_buffer(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>, %i0: index, %i1: index, %i2: index) -> i32 {
8  // CHECK: %[[AC:.+]] = spirv.AccessChain
9  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicIAdd <Device> <AcquireRelease> %[[AC]], %[[VAL]] : !spirv.ptr<i32, StorageBuffer>
10  // CHECK: return %[[ATOMIC]]
11  %0 = memref.atomic_rmw "addi" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>) -> i32
12  return %0: i32
13}
14
15//      CHECK: func.func @atomic_maxs_workgroup
16// CHECK-SAME: (%[[VAL:.+]]: i32,
17func.func @atomic_maxs_workgroup(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<Workgroup>>, %i0: index, %i1: index, %i2: index) -> i32 {
18  // CHECK: %[[AC:.+]] = spirv.AccessChain
19  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicSMax <Workgroup> <AcquireRelease> %[[AC]], %[[VAL]] : !spirv.ptr<i32, Workgroup>
20  // CHECK: return %[[ATOMIC]]
21  %0 = memref.atomic_rmw "maxs" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<Workgroup>>) -> i32
22  return %0: i32
23}
24
25//      CHECK: func.func @atomic_maxu_storage_buffer
26// CHECK-SAME: (%[[VAL:.+]]: i32,
27func.func @atomic_maxu_storage_buffer(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>, %i0: index, %i1: index, %i2: index) -> i32 {
28  // CHECK: %[[AC:.+]] = spirv.AccessChain
29  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicUMax <Device> <AcquireRelease> %[[AC]], %[[VAL]] : !spirv.ptr<i32, StorageBuffer>
30  // CHECK: return %[[ATOMIC]]
31  %0 = memref.atomic_rmw "maxu" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>) -> i32
32  return %0: i32
33}
34
35//      CHECK: func.func @atomic_mins_workgroup
36// CHECK-SAME: (%[[VAL:.+]]: i32,
37func.func @atomic_mins_workgroup(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<Workgroup>>, %i0: index, %i1: index, %i2: index) -> i32 {
38  // CHECK: %[[AC:.+]] = spirv.AccessChain
39  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicSMin <Workgroup> <AcquireRelease> %[[AC]], %[[VAL]] : !spirv.ptr<i32, Workgroup>
40  // CHECK: return %[[ATOMIC]]
41  %0 = memref.atomic_rmw "mins" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<Workgroup>>) -> i32
42  return %0: i32
43}
44
45//      CHECK: func.func @atomic_minu_storage_buffer
46// CHECK-SAME: (%[[VAL:.+]]: i32,
47func.func @atomic_minu_storage_buffer(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>, %i0: index, %i1: index, %i2: index) -> i32 {
48  // CHECK: %[[AC:.+]] = spirv.AccessChain
49  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicUMin <Device> <AcquireRelease> %[[AC]], %[[VAL]] : !spirv.ptr<i32, StorageBuffer>
50  // CHECK: return %[[ATOMIC]]
51  %0 = memref.atomic_rmw "minu" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>) -> i32
52  return %0: i32
53}
54
55//      CHECK: func.func @atomic_ori_workgroup
56// CHECK-SAME: (%[[VAL:.+]]: i32,
57func.func @atomic_ori_workgroup(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<Workgroup>>, %i0: index, %i1: index, %i2: index) -> i32 {
58  // CHECK: %[[AC:.+]] = spirv.AccessChain
59  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicOr <Workgroup> <AcquireRelease> %[[AC]], %[[VAL]] : !spirv.ptr<i32, Workgroup>
60  // CHECK: return %[[ATOMIC]]
61  %0 = memref.atomic_rmw "ori" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<Workgroup>>) -> i32
62  return %0: i32
63}
64
65//      CHECK: func.func @atomic_andi_storage_buffer
66// CHECK-SAME: (%[[VAL:.+]]: i32,
67func.func @atomic_andi_storage_buffer(%value: i32, %memref: memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>, %i0: index, %i1: index, %i2: index) -> i32 {
68  // CHECK: %[[AC:.+]] = spirv.AccessChain
69  // CHECK: %[[ATOMIC:.+]] = spirv.AtomicAnd <Device> <AcquireRelease> %[[AC]], %[[VAL]] : !spirv.ptr<i32, StorageBuffer>
70  // CHECK: return %[[ATOMIC]]
71  %0 = memref.atomic_rmw "andi" %value, %memref[%i0, %i1, %i2] : (i32, memref<2x3x4xi32, #spirv.storage_class<StorageBuffer>>) -> i32
72  return %0: i32
73}
74
75}
76
77