[spirv-reader][ir] Convert load/store usages of atomics

With the atomic types changed, we need to update any load/store of the
original SPIR-V value to the atomic equivalents.

Bug: 404501988
Change-Id: Idc48fa529bbf593a5a5ffda0693fe0ba355712e3
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/236954
Commit-Queue: dan sinclair <dsinclair@chromium.org>
Reviewed-by: James Price <jrprice@google.com>

src/tint/lang/spirv/reader/lower/atomics.cc

diff --git a/src/tint/lang/spirv/reader/lower/atomics.cc b/src/tint/lang/spirv/reader/lower/atomics.cc
index 23fe305..8b25e2b 100644
--- a/src/tint/lang/spirv/reader/lower/atomics.cc
+++ b/src/tint/lang/spirv/reader/lower/atomics.cc
@@ -58,8 +58,14 @@
     /// The `ir::Value`s to be converted to atomics
     Vector<core::ir::Value*, 8> values_to_convert_{};
 
+    /// The `ir::Values`s which have had their types changed, they then need to have their
+    /// loads/stores updated to match. This maps to the root FunctionParam or Var for each atomic.
+    Hashset<core::ir::Value*, 8> values_to_fix_usages_{};
+
     /// The `ir::Value`s which have been converted
     Hashset<core::ir::Value*, 8> converted_{};
+    /// The `ir::Value`s which have had their usages updated
+    Hashset<core::ir::Value*, 8> usages_updated_{};
 
     struct ForkedStruct {
         const core::type::Struct* src_struct = nullptr;
@@ -264,14 +270,19 @@
                 auto* atomic_ty = AtomicTypeFor(val, orig_ty);
                 res->SetType(atomic_ty);
 
-                tint::Switch(            //
-                    res->Instruction(),  //
+                tint::Switch(
+                    res->Instruction(),
                     [&](core::ir::Access* a) {
                         CheckForStructForking(a);
                         values_to_convert_.Push(a->Object());
-                    },                                                               //
-                    [&](core::ir::Let* l) { values_to_convert_.Push(l->Value()); },  //
-                    [&](core::ir::Var*) {},                                          //
+                    },
+                    [&](core::ir::Let* l) { values_to_convert_.Push(l->Value()); },
+                    [&](core::ir::Var* v) {
+                        auto* var_res = v->Result();
+                        if (usages_updated_.Add(var_res)) {
+                            ConvertUsages(var_res);
+                        }
+                    },
                     TINT_ICE_ON_NO_MATCH);
             },
             [&](core::ir::FunctionParam* param) {
@@ -279,6 +290,10 @@
                 auto* atomic_ty = AtomicTypeFor(val, orig_ty);
                 param->SetType(atomic_ty);
 
+                if (usages_updated_.Add(param)) {
+                    ConvertUsages(param);
+                }
+
                 for (auto& usage : param->Function()->UsagesUnsorted()) {
                     if (usage->instruction->Is<core::ir::Return>()) {
                         continue;
@@ -292,6 +307,28 @@
             TINT_ICE_ON_NO_MATCH);
     }
 
+    void ConvertUsages(core::ir::Value* val) {
+        val->ForEachUseUnsorted([&](const core::ir::Usage& usage) {
+            auto* inst = usage.instruction;
+
+            tint::Switch(  //
+                inst,
+                [&](core::ir::Load* ld) {
+                    b.InsertBefore(ld, [&] {
+                        b.CallWithResult(ld->DetachResult(), core::BuiltinFn::kAtomicLoad,
+                                         ld->From());
+                    });
+                    ld->Destroy();
+                },
+                [&](core::ir::Store* st) {
+                    b.InsertBefore(st, [&] {
+                        b.Call(ty.void_(), core::BuiltinFn::kAtomicStore, st->To(), st->From());
+                    });
+                    st->Destroy();
+                });
+        });
+    }
+
     void CheckForStructForking(core::ir::Access* access) {
         auto* cur_ty = access->Object()->Type()->UnwrapPtr();
         for (auto* idx : access->Indices()) {

src/tint/lang/spirv/reader/lower/atomics_test.cc

diff --git a/src/tint/lang/spirv/reader/lower/atomics_test.cc b/src/tint/lang/spirv/reader/lower/atomics_test.cc
index aef1ce2..3a163d4 100644
--- a/src/tint/lang/spirv/reader/lower/atomics_test.cc
+++ b/src/tint/lang/spirv/reader/lower/atomics_test.cc
@@ -1973,7 +1973,7 @@
     ASSERT_EQ(expect, str());
 }
 
-TEST_F(SpirvReader_AtomicsTest, DISABLED_ReplaceAssignsAndDecls_Scalar) {
+TEST_F(SpirvReader_AtomicsTest, ReplaceAssignsAndDecls_Scalar) {
     auto* f = b.ComputeFunction("main");
 
     core::ir::Var* wg = nullptr;
@@ -2013,7 +2013,22 @@
     Run(Atomics);
 
     auto* expect = R"(
-UNIMPLEMENTED
+$B1: {  # root
+  %wg:ptr<workgroup, atomic<u32>, read_write> = var undef
+}
+
+%main = @compute @workgroup_size(1u, 1u, 1u) func():void {
+  $B2: {
+    %b:ptr<function, u32, read_write> = var undef
+    %4:u32 = atomicAdd %wg, 0u
+    %5:void = atomicStore %wg, 0u
+    %6:u32 = atomicLoad %wg
+    %7:u32 = let %6
+    %8:u32 = atomicLoad %wg
+    store %b, %8
+    ret
+  }
+}
 )";
     ASSERT_EQ(expect, str());
 }