tint: fix emitting duplicate structs for atomicCompareExchangeWeak

Bug: tint:1574
Change-Id: Id4ae2d2de9ac4678260f4ecfb3a0f779d170f9a4
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/92280
Reviewed-by: James Price <jrprice@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Antonio Maiorano <amaiorano@google.com>

src/tint/writer/glsl/generator_impl.cc

@@ -920,7 +920,7 @@ bool GeneratorImpl::EmitWorkgroupAtomicCall(std::ostream& out,
         case sem::BuiltinType::kAtomicCompareExchangeWeak: {
             // Emit the builtin return type unique to this overload. This does not
             // exist in the AST, so it will not be generated in Generate().
-            if (!EmitStructType(&helpers_, builtin->ReturnType()->As<sem::Struct>())) {
+            if (!EmitStructTypeOnce(&helpers_, builtin->ReturnType()->As<sem::Struct>())) {
                 return false;
             }
 
@@ -2822,6 +2822,14 @@ bool GeneratorImpl::EmitStructType(TextBuffer* b, const sem::Struct* str) {
     return true;
 }
 
+bool GeneratorImpl::EmitStructTypeOnce(TextBuffer* buffer, const sem::Struct* str) {
+    auto it = emitted_structs_.emplace(str);
+    if (!it.second) {
+        return true;
+    }
+    return EmitStructType(buffer, str);
+}
+
 bool GeneratorImpl::EmitStructMembers(TextBuffer* b, const sem::Struct* str, bool emit_offsets) {
     ScopedIndent si(b);
     for (auto* mem : str->Members()) {

src/tint/writer/glsl/generator_impl.h

@@ -411,6 +411,12 @@ class GeneratorImpl : public TextGenerator {
     /// @param ty the struct to generate
     /// @returns true if the struct is emitted
     bool EmitStructType(TextBuffer* buffer, const sem::Struct* ty);
+    /// Handles generating a structure declaration only the first time called. Subsequent calls are
+    /// a no-op and return true.
+    /// @param buffer the text buffer that the type declaration will be written to
+    /// @param ty the struct to generate
+    /// @returns true if the struct is emitted
+    bool EmitStructTypeOnce(TextBuffer* buffer, const sem::Struct* ty);
     /// Handles generating the members of a structure
     /// @param buffer the text buffer that the struct members will be written to
     /// @param ty the struct to generate
@@ -503,6 +509,7 @@ class GeneratorImpl : public TextGenerator {
     std::unordered_map<const sem::Vector*, std::string> dynamic_vector_write_;
     std::unordered_map<const sem::Vector*, std::string> int_dot_funcs_;
     std::unordered_map<const sem::Type*, std::string> float_modulo_funcs_;
+    std::unordered_set<const sem::Struct*> emitted_structs_;
     bool requires_oes_sample_variables_ = false;
     bool requires_default_precision_qualifier_ = false;
     Version version_;

src/tint/writer/hlsl/generator_impl.cc

@@ -1767,7 +1767,7 @@ bool GeneratorImpl::EmitWorkgroupAtomicCall(std::ostream& out,
         case sem::BuiltinType::kAtomicCompareExchangeWeak: {
             // Emit the builtin return type unique to this overload. This does not
             // exist in the AST, so it will not be generated in Generate().
-            if (!EmitStructType(&helpers_, builtin->ReturnType()->As<sem::Struct>())) {
+            if (!EmitStructTypeOnce(&helpers_, builtin->ReturnType()->As<sem::Struct>())) {
                 return false;
             }
 
@@ -3921,6 +3921,14 @@ bool GeneratorImpl::EmitStructType(TextBuffer* b, const sem::Struct* str) {
     return true;
 }
 
+bool GeneratorImpl::EmitStructTypeOnce(TextBuffer* buffer, const sem::Struct* str) {
+    auto it = emitted_structs_.emplace(str);
+    if (!it.second) {
+        return true;
+    }
+    return EmitStructType(buffer, str);
+}
+
 bool GeneratorImpl::EmitUnaryOp(std::ostream& out, const ast::UnaryOpExpression* expr) {
     switch (expr->op) {
         case ast::UnaryOp::kIndirection:

src/tint/writer/hlsl/generator_impl.h

@@ -411,6 +411,12 @@ class GeneratorImpl : public TextGenerator {
     /// @param ty the struct to generate
     /// @returns true if the struct is emitted
     bool EmitStructType(TextBuffer* buffer, const sem::Struct* ty);
+    /// Handles generating a structure declaration only the first time called. Subsequent calls are
+    /// a no-op and return true.
+    /// @param buffer the text buffer that the type declaration will be written to
+    /// @param ty the struct to generate
+    /// @returns true if the struct is emitted
+    bool EmitStructTypeOnce(TextBuffer* buffer, const sem::Struct* ty);
     /// Handles a unary op expression
     /// @param out the output of the expression stream
     /// @param expr the expression to emit
@@ -530,6 +536,7 @@ class GeneratorImpl : public TextGenerator {
     std::unordered_map<const sem::Matrix*, std::string> dynamic_matrix_vector_write_;
     std::unordered_map<const sem::Matrix*, std::string> dynamic_matrix_scalar_write_;
     std::unordered_map<const sem::Type*, std::string> value_or_one_if_zero_;
+    std::unordered_set<const sem::Struct*> emitted_structs_;
 };
 
 }  // namespace tint::writer::hlsl

src/tint/writer/msl/generator_impl.cc

@@ -826,46 +826,66 @@ bool GeneratorImpl::EmitAtomicCall(std::ostream& out,
             return call("atomic_exchange_explicit", true);
 
         case sem::BuiltinType::kAtomicCompareExchangeWeak: {
-            // Emit the builtin return type unique to this overload. This does not
-            // exist in the AST, so it will not be generated in Generate().
-            if (!EmitStructType(&helpers_, builtin->ReturnType()->As<sem::Struct>())) {
-                return false;
-            }
-
             auto* ptr_ty = TypeOf(expr->args[0])->UnwrapRef()->As<sem::Pointer>();
             auto sc = ptr_ty->StorageClass();
+            auto* str = builtin->ReturnType()->As<sem::Struct>();
 
-            auto func = utils::GetOrCreate(atomicCompareExchangeWeak_, sc, [&]() -> std::string {
-                auto name = UniqueIdentifier("atomicCompareExchangeWeak");
-                auto& buf = helpers_;
-
-                line(&buf) << "template <typename A, typename T>";
-                {
-                    auto f = line(&buf);
-                    auto str_name = StructName(builtin->ReturnType()->As<sem::Struct>());
-                    f << str_name << " " << name << "(";
-                    if (!EmitStorageClass(f, sc)) {
+            auto func = utils::GetOrCreate(
+                atomicCompareExchangeWeak_, ACEWKeyType{{sc, str}}, [&]() -> std::string {
+                    // Emit the builtin return type unique to this overload. This does not
+                    // exist in the AST, so it will not be generated in Generate().
+                    if (!EmitStructTypeOnce(&helpers_, builtin->ReturnType()->As<sem::Struct>())) {
                         return "";
                     }
-                    f << " A* atomic, T compare, T value) {";
-                }
 
-                buf.IncrementIndent();
-                TINT_DEFER({
-                    buf.DecrementIndent();
-                    line(&buf) << "}";
-                    line(&buf);
+                    auto name = UniqueIdentifier("atomicCompareExchangeWeak");
+                    auto& buf = helpers_;
+                    auto* atomic_ty = builtin->Parameters()[0]->Type();
+                    auto* arg_ty = builtin->Parameters()[1]->Type();
+
+                    {
+                        auto f = line(&buf);
+                        auto str_name = StructName(builtin->ReturnType()->As<sem::Struct>());
+                        f << str_name << " " << name << "(";
+                        if (!EmitTypeAndName(f, atomic_ty, "atomic")) {
+                            return "";
+                        }
+                        f << ", ";
+                        if (!EmitTypeAndName(f, arg_ty, "compare")) {
+                            return "";
+                        }
+                        f << ", ";
+                        if (!EmitTypeAndName(f, arg_ty, "value")) {
+                            return "";
+                        }
+                        f << ") {";
+                    }
+
+                    buf.IncrementIndent();
+                    TINT_DEFER({
+                        buf.DecrementIndent();
+                        line(&buf) << "}";
+                        line(&buf);
+                    });
+
+                    {
+                        auto f = line(&buf);
+                        if (!EmitTypeAndName(f, arg_ty, "old_value")) {
+                            return "";
+                        }
+                        f << " = compare;";
+                    }
+                    line(&buf) << "bool exchanged = "
+                                  "atomic_compare_exchange_weak_explicit(atomic, "
+                                  "&old_value, value, memory_order_relaxed, "
+                                  "memory_order_relaxed);";
+                    line(&buf) << "return {old_value, exchanged};";
+                    return name;
                 });
 
-                line(&buf) << "T old_value = compare;";
-                line(&buf) << "bool exchanged = "
-                              "atomic_compare_exchange_weak_explicit(atomic, "
-                              "&old_value, value, memory_order_relaxed, "
-                              "memory_order_relaxed);";
-                line(&buf) << "return {old_value, exchanged};";
-                return name;
-            });
-
+            if (func.empty()) {
+                return false;
+            }
             return call(func, false);
         }
 
@@ -2765,6 +2785,14 @@ bool GeneratorImpl::EmitStructType(TextBuffer* b, const sem::Struct* str) {
     return true;
 }
 
+bool GeneratorImpl::EmitStructTypeOnce(TextBuffer* buffer, const sem::Struct* str) {
+    auto it = emitted_structs_.emplace(str);
+    if (!it.second) {
+        return true;
+    }
+    return EmitStructType(buffer, str);
+}
+
 bool GeneratorImpl::EmitUnaryOp(std::ostream& out, const ast::UnaryOpExpression* expr) {
     // Handle `-e` when `e` is signed, so that we ensure that if `e` is the
     // largest negative value, it returns `e`.

src/tint/writer/msl/generator_impl.h

@@ -16,6 +16,7 @@
 #define SRC_TINT_WRITER_MSL_GENERATOR_IMPL_H_
 
 #include <string>
+#include <tuple>
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
@@ -332,6 +333,12 @@ class GeneratorImpl : public TextGenerator {
     /// @param str the struct to generate
     /// @returns true if the struct is emitted
     bool EmitStructType(TextBuffer* buffer, const sem::Struct* str);
+    /// Handles generating a structure declaration only the first time called. Subsequent calls are
+    /// a no-op and return true.
+    /// @param buffer the text buffer that the type declaration will be written to
+    /// @param ty the struct to generate
+    /// @returns true if the struct is emitted
+    bool EmitStructTypeOnce(TextBuffer* buffer, const sem::Struct* ty);
     /// Handles a unary op expression
     /// @param out the output of the expression stream
     /// @param expr the expression to emit
@@ -400,13 +407,13 @@ class GeneratorImpl : public TextGenerator {
     /// type.
     SizeAndAlign MslPackedTypeSizeAndAlign(const sem::Type* ty);
 
-    using StorageClassToString = std::unordered_map<ast::StorageClass, std::string>;
-
     std::function<bool()> emit_continuing_;
 
     /// Name of atomicCompareExchangeWeak() helper for the given pointer storage
-    /// class.
-    StorageClassToString atomicCompareExchangeWeak_;
+    /// class and struct return type
+    using ACEWKeyType =
+        utils::UnorderedKeyWrapper<std::tuple<ast::StorageClass, const sem::Struct*>>;
+    std::unordered_map<ACEWKeyType, std::string> atomicCompareExchangeWeak_;
 
     /// Unique name of the 'TINT_INVARIANT' preprocessor define. Non-empty only if
     /// an invariant attribute has been generated.
@@ -423,6 +430,7 @@ class GeneratorImpl : public TextGenerator {
     std::unordered_map<const sem::Builtin*, std::string> builtins_;
     std::unordered_map<const sem::Type*, std::string> unary_minus_funcs_;
     std::unordered_map<uint32_t, std::string> int_dot_funcs_;
+    std::unordered_set<const sem::Struct*> emitted_structs_;
 };
 
 }  // namespace tint::writer::msl