tint: spirv reader: detect and replace loads and stores of atomic variables with atomicLoad/Store

Bug: tint:1441
Change-Id: Iee89cb87ca063d8a98ff8ad789ba14dee65c036a
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/95140
Reviewed-by: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Antonio Maiorano <amaiorano@google.com>

src/tint/transform/spirv_atomic.cc

@@ -135,6 +135,10 @@ struct SpirvAtomic::State {
             });
         }
 
+        // Replace assignments and decls from atomic variables with atomicLoads, and assignments to
+        // atomic variables with atomicStores.
+        ReplaceLoadsAndStores();
+
         ctx.Clone();
     }
 
@@ -200,6 +204,70 @@ struct SpirvAtomic::State {
                 return nullptr;
             });
     }
+
+    void ReplaceLoadsAndStores() {
+        // Returns true if 'e' is a reference to an atomic variable or struct member
+        auto is_ref_to_atomic_var = [&](const sem::Expression* e) {
+            if (tint::Is<sem::Reference>(e->Type()) && e->SourceVariable() &&
+                (atomic_variables.count(e->SourceVariable()) != 0)) {
+                // If it's a struct member, make sure it's one we marked as atomic
+                if (auto* ma = e->As<sem::StructMemberAccess>()) {
+                    auto it = forked_structs.find(ma->Member()->Struct()->Declaration());
+                    if (it != forked_structs.end()) {
+                        auto& forked = it->second;
+                        return forked.atomic_members.count(ma->Member()->Index()) != 0;
+                    }
+                }
+                return true;
+            }
+            return false;
+        };
+
+        // Look for loads and stores via assignments and decls of atomic variables we've collected
+        // so far, and replace them with atomicLoad and atomicStore.
+        for (auto* atomic_var : atomic_variables) {
+            for (auto* vu : atomic_var->Users()) {
+                Switch(
+                    vu->Stmt()->Declaration(),
+                    [&](const ast::AssignmentStatement* assign) {
+                        auto* sem_lhs = ctx.src->Sem().Get(assign->lhs);
+                        if (is_ref_to_atomic_var(sem_lhs)) {
+                            ctx.Replace(assign, [=] {
+                                auto* lhs = ctx.CloneWithoutTransform(assign->lhs);
+                                auto* rhs = ctx.CloneWithoutTransform(assign->rhs);
+                                auto* call = b.Call(sem::str(sem::BuiltinType::kAtomicStore),
+                                                    b.AddressOf(lhs), rhs);
+                                return b.CallStmt(call);
+                            });
+                            return;
+                        }
+
+                        auto sem_rhs = ctx.src->Sem().Get(assign->rhs);
+                        if (is_ref_to_atomic_var(sem_rhs)) {
+                            ctx.Replace(assign->rhs, [=] {
+                                auto* rhs = ctx.CloneWithoutTransform(assign->rhs);
+                                return b.Call(sem::str(sem::BuiltinType::kAtomicLoad),
+                                              b.AddressOf(rhs));
+                            });
+                            return;
+                        }
+                    },
+                    [&](const ast::VariableDeclStatement* decl) {
+                        auto* var = decl->variable;
+                        if (auto* sem_ctor = ctx.src->Sem().Get(var->constructor)) {
+                            if (is_ref_to_atomic_var(sem_ctor)) {
+                                ctx.Replace(var->constructor, [=] {
+                                    auto* rhs = ctx.CloneWithoutTransform(var->constructor);
+                                    return b.Call(sem::str(sem::BuiltinType::kAtomicLoad),
+                                                  b.AddressOf(rhs));
+                                });
+                                return;
+                            }
+                        }
+                    });
+            }
+        }
+    }
 };
 
 SpirvAtomic::SpirvAtomic() = default;

src/tint/transform/spirv_atomic_test.cc

@@ -1018,5 +1018,366 @@ fn f() {
     EXPECT_EQ(expect, str(got));
 }
 
+TEST_F(SpirvAtomicTest, ReplaceAssignsAndDecls_Scaler) {
+    auto* src = R"(
+var<workgroup> wg : u32;
+
+fn f() {
+  stub_atomicAdd_u32(wg, 1u);
+
+  wg = 0u;
+  let a = wg;
+  var b : u32;
+  b = wg;
+}
+)";
+
+    auto* expect = R"(
+var<workgroup> wg : atomic<u32>;
+
+fn f() {
+  atomicAdd(&(wg), 1u);
+  atomicStore(&(wg), 0u);
+  let a = atomicLoad(&(wg));
+  var b : u32;
+  b = atomicLoad(&(wg));
+}
+)";
+
+    auto got = Run(src);
+
+    EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(SpirvAtomicTest, ReplaceAssignsAndDecls_Struct) {
+    auto* src = R"(
+struct S {
+  a : u32,
+}
+
+var<workgroup> wg : S;
+
+fn f() {
+  stub_atomicAdd_u32(wg.a, 1u);
+
+  wg.a = 0u;
+  let a = wg.a;
+  var b : u32;
+  b = wg.a;
+}
+)";
+
+    auto* expect = R"(
+struct S_atomic {
+  a : atomic<u32>,
+}
+
+struct S {
+  a : u32,
+}
+
+var<workgroup> wg : S_atomic;
+
+fn f() {
+  atomicAdd(&(wg.a), 1u);
+  atomicStore(&(wg.a), 0u);
+  let a = atomicLoad(&(wg.a));
+  var b : u32;
+  b = atomicLoad(&(wg.a));
+}
+)";
+
+    auto got = Run(src);
+
+    EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(SpirvAtomicTest, ReplaceAssignsAndDecls_NestedStruct) {
+    auto* src = R"(
+struct S0 {
+  a : u32,
+}
+
+struct S1 {
+  s0 : S0
+}
+
+var<workgroup> wg : S1;
+
+fn f() {
+  stub_atomicAdd_u32(wg.s0.a, 1u);
+
+  wg.s0.a = 0u;
+  let a = wg.s0.a;
+  var b : u32;
+  b = wg.s0.a;
+}
+)";
+
+    auto* expect = R"(
+struct S0_atomic {
+  a : atomic<u32>,
+}
+
+struct S0 {
+  a : u32,
+}
+
+struct S1_atomic {
+  s0 : S0_atomic,
+}
+
+struct S1 {
+  s0 : S0,
+}
+
+var<workgroup> wg : S1_atomic;
+
+fn f() {
+  atomicAdd(&(wg.s0.a), 1u);
+  atomicStore(&(wg.s0.a), 0u);
+  let a = atomicLoad(&(wg.s0.a));
+  var b : u32;
+  b = atomicLoad(&(wg.s0.a));
+}
+)";
+
+    auto got = Run(src);
+
+    EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(SpirvAtomicTest, ReplaceAssignsAndDecls_StructMultipleAtomics) {
+    auto* src = R"(
+struct S {
+  a : u32,
+  b : u32,
+  c : u32,
+}
+
+var<workgroup> wg : S;
+
+fn f() {
+  stub_atomicAdd_u32(wg.a, 1u);
+  stub_atomicAdd_u32(wg.b, 1u);
+
+  wg.a = 0u;
+  let a = wg.a;
+  var b : u32;
+  b = wg.a;
+
+  wg.b = 0u;
+  let c = wg.b;
+  var d : u32;
+  d = wg.b;
+
+  wg.c = 0u;
+  let e = wg.c;
+  var f : u32;
+  f = wg.c;
+}
+)";
+
+    auto* expect = R"(
+struct S_atomic {
+  a : atomic<u32>,
+  b : atomic<u32>,
+  c : u32,
+}
+
+struct S {
+  a : u32,
+  b : u32,
+  c : u32,
+}
+
+var<workgroup> wg : S_atomic;
+
+fn f() {
+  atomicAdd(&(wg.a), 1u);
+  atomicAdd(&(wg.b), 1u);
+  atomicStore(&(wg.a), 0u);
+  let a = atomicLoad(&(wg.a));
+  var b : u32;
+  b = atomicLoad(&(wg.a));
+  atomicStore(&(wg.b), 0u);
+  let c = atomicLoad(&(wg.b));
+  var d : u32;
+  d = atomicLoad(&(wg.b));
+  wg.c = 0u;
+  let e = wg.c;
+  var f : u32;
+  f = wg.c;
+}
+)";
+
+    auto got = Run(src);
+
+    EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(SpirvAtomicTest, ReplaceAssignsAndDecls_ArrayOfScalar) {
+    auto* src = R"(
+var<workgroup> wg : array<u32, 4>;
+
+fn f() {
+  stub_atomicAdd_u32(wg[1], 1u);
+
+  wg[1] = 0u;
+  let a = wg[1];
+  var b : u32;
+  b = wg[1];
+}
+)";
+
+    auto* expect = R"(
+var<workgroup> wg : array<atomic<u32>, 4u>;
+
+fn f() {
+  atomicAdd(&(wg[1]), 1u);
+  atomicStore(&(wg[1]), 0u);
+  let a = atomicLoad(&(wg[1]));
+  var b : u32;
+  b = atomicLoad(&(wg[1]));
+}
+)";
+
+    auto got = Run(src);
+
+    EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(SpirvAtomicTest, ReplaceAssignsAndDecls_ArrayOfStruct) {
+    auto* src = R"(
+struct S {
+  a : u32,
+}
+
+var<workgroup> wg : array<S, 4>;
+
+fn f() {
+  stub_atomicAdd_u32(wg[1].a, 1u);
+
+  wg[1].a = 0u;
+  let a = wg[1].a;
+  var b : u32;
+  b = wg[1].a;
+}
+)";
+
+    auto* expect = R"(
+struct S_atomic {
+  a : atomic<u32>,
+}
+
+struct S {
+  a : u32,
+}
+
+var<workgroup> wg : array<S_atomic, 4u>;
+
+fn f() {
+  atomicAdd(&(wg[1].a), 1u);
+  atomicStore(&(wg[1].a), 0u);
+  let a = atomicLoad(&(wg[1].a));
+  var b : u32;
+  b = atomicLoad(&(wg[1].a));
+}
+)";
+
+    auto got = Run(src);
+
+    EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(SpirvAtomicTest, ReplaceAssignsAndDecls_StructOfArray) {
+    auto* src = R"(
+struct S {
+  a : array<u32>,
+}
+
+@group(0) @binding(1) var<storage, read_write> s : S;
+
+fn f() {
+  stub_atomicAdd_u32(s.a[4], 1u);
+
+  s.a[4] = 0u;
+  let a = s.a[4];
+  var b : u32;
+  b = s.a[4];
+}
+)";
+
+    auto* expect = R"(
+struct S_atomic {
+  a : array<atomic<u32>>,
+}
+
+struct S {
+  a : array<u32>,
+}
+
+@group(0) @binding(1) var<storage, read_write> s : S_atomic;
+
+fn f() {
+  atomicAdd(&(s.a[4]), 1u);
+  atomicStore(&(s.a[4]), 0u);
+  let a = atomicLoad(&(s.a[4]));
+  var b : u32;
+  b = atomicLoad(&(s.a[4]));
+}
+)";
+
+    auto got = Run(src);
+
+    EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(SpirvAtomicTest, ReplaceAssignsAndDecls_ViaPtrLet) {
+    auto* src = R"(
+struct S {
+  i : u32,
+}
+
+@group(0) @binding(1) var<storage, read_write> s : S;
+
+fn f() {
+  let p0 = &(s);
+  let p1 : ptr<storage, u32, read_write> = &((*(p0)).i);
+  stub_atomicAdd_u32(*p1, 1u);
+
+  *p1 = 0u;
+  let a = *p1;
+  var b : u32;
+  b = *p1;
+}
+)";
+
+    auto* expect = R"(
+struct S_atomic {
+  i : atomic<u32>,
+}
+
+struct S {
+  i : u32,
+}
+
+@group(0) @binding(1) var<storage, read_write> s : S_atomic;
+
+fn f() {
+  let p0 = &(s);
+  let p1 : ptr<storage, atomic<u32>, read_write> = &((*(p0)).i);
+  atomicAdd(&(*(p1)), 1u);
+  atomicStore(&(*(p1)), 0u);
+  let a = atomicLoad(&(*(p1)));
+  var b : u32;
+  b = atomicLoad(&(*(p1)));
+}
+)";
+
+    auto got = Run(src);
+
+    EXPECT_EQ(expect, str(got));
+}
 }  // namespace
 }  // namespace tint::transform