tint: Support @diagnostic on loop and loop body

Bug: tint:1809
Change-Id: Ib3ccfd823f9cccb67bebbf04927d54f193a4e281
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/124321
Commit-Queue: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>

src/tint/ast/loop_statement.cc

@@ -14,6 +14,8 @@
 
 #include "src/tint/ast/loop_statement.h"
 
+#include <utility>
+
 #include "src/tint/program_builder.h"
 
 TINT_INSTANTIATE_TYPEINFO(tint::ast::LoopStatement);
@@ -24,11 +26,16 @@ LoopStatement::LoopStatement(ProgramID pid,
                              NodeID nid,
                              const Source& src,
                              const BlockStatement* b,
-                             const BlockStatement* cont)
+                             const BlockStatement* cont,
-    : Base(pid, nid, src), body(b), continuing(cont) {
+                             utils::VectorRef<const ast::Attribute*> attrs)
+    : Base(pid, nid, src), body(b), continuing(cont), attributes(std::move(attrs)) {
     TINT_ASSERT(AST, body);
     TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, body, program_id);
     TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, continuing, program_id);
+    for (auto* attr : attributes) {
+        TINT_ASSERT(AST, attr);
+        TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, attr, program_id);
+    }
 }
 
 LoopStatement::~LoopStatement() = default;
@@ -38,7 +45,8 @@ const LoopStatement* LoopStatement::Clone(CloneContext* ctx) const {
     auto src = ctx->Clone(source);
     auto* b = ctx->Clone(body);
     auto* cont = ctx->Clone(continuing);
-    return ctx->dst->create<LoopStatement>(src, b, cont);
+    auto attrs = ctx->Clone(attributes);
+    return ctx->dst->create<LoopStatement>(src, b, cont, std::move(attrs));
 }
 
 }  // namespace tint::ast

src/tint/ast/loop_statement.h

@@ -28,12 +28,13 @@ class LoopStatement final : public Castable<LoopStatement, Statement> {
     /// @param source the loop statement source
     /// @param body the body statements
     /// @param continuing the continuing statements
+    /// @param attributes the while statement attributes
     LoopStatement(ProgramID pid,
                   NodeID nid,
                   const Source& source,
                   const BlockStatement* body,
-                  const BlockStatement* continuing);
+                  const BlockStatement* continuing,
-
+                  utils::VectorRef<const ast::Attribute*> attributes);
     /// Destructor
     ~LoopStatement() override;
 
@@ -48,6 +49,9 @@ class LoopStatement final : public Castable<LoopStatement, Statement> {
 
     /// The continuing statements
     const BlockStatement* const continuing;
+
+    /// The attribute list
+    const utils::Vector<const Attribute*, 1> attributes;
 };
 
 }  // namespace tint::ast

src/tint/ast/loop_statement_test.cc

@@ -14,6 +14,7 @@
 
 #include "src/tint/ast/loop_statement.h"
 
+#include "gmock/gmock.h"
 #include "gtest/gtest-spi.h"
 #include "src/tint/ast/discard_statement.h"
 #include "src/tint/ast/if_statement.h"
@@ -30,7 +31,7 @@ TEST_F(LoopStatementTest, Creation) {
 
     auto* continuing = Block(create<DiscardStatement>());
 
-    auto* l = create<LoopStatement>(body, continuing);
+    auto* l = create<LoopStatement>(body, continuing, utils::Empty);
     ASSERT_EQ(l->body->statements.Length(), 1u);
     EXPECT_EQ(l->body->statements[0], b);
     ASSERT_EQ(l->continuing->statements.Length(), 1u);
@@ -42,21 +43,32 @@ TEST_F(LoopStatementTest, Creation_WithSource) {
 
     auto* continuing = Block(create<DiscardStatement>());
 
-    auto* l = create<LoopStatement>(Source{Source::Location{20, 2}}, body, continuing);
+    auto* l =
+        create<LoopStatement>(Source{Source::Location{20, 2}}, body, continuing, utils::Empty);
     auto src = l->source;
     EXPECT_EQ(src.range.begin.line, 20u);
     EXPECT_EQ(src.range.begin.column, 2u);
 }
 
+TEST_F(LoopStatementTest, Creation_WithAttributes) {
+    auto* attr1 = DiagnosticAttribute(builtin::DiagnosticSeverity::kOff, "foo");
+    auto* attr2 = DiagnosticAttribute(builtin::DiagnosticSeverity::kOff, "bar");
+
+    auto* body = Block(Return());
+    auto* l = create<LoopStatement>(body, nullptr, utils::Vector{attr1, attr2});
+
+    EXPECT_THAT(l->attributes, testing::ElementsAre(attr1, attr2));
+}
+
 TEST_F(LoopStatementTest, IsLoop) {
-    auto* l = create<LoopStatement>(Block(), Block());
+    auto* l = create<LoopStatement>(Block(), Block(), utils::Empty);
     EXPECT_TRUE(l->Is<LoopStatement>());
 }
 
 TEST_F(LoopStatementTest, HasContinuing_WithoutContinuing) {
     auto* body = Block(create<DiscardStatement>());
 
-    auto* l = create<LoopStatement>(body, nullptr);
+    auto* l = create<LoopStatement>(body, nullptr, utils::Empty);
     EXPECT_FALSE(l->continuing);
 }
 
@@ -65,7 +77,7 @@ TEST_F(LoopStatementTest, HasContinuing_WithContinuing) {
 
     auto* continuing = Block(create<DiscardStatement>());
 
-    auto* l = create<LoopStatement>(body, continuing);
+    auto* l = create<LoopStatement>(body, continuing, utils::Empty);
     EXPECT_TRUE(l->continuing);
 }
 
@@ -73,7 +85,7 @@ TEST_F(LoopStatementTest, Assert_Null_Body) {
     EXPECT_FATAL_FAILURE(
         {
             ProgramBuilder b;
-            b.create<LoopStatement>(nullptr, nullptr);
+            b.create<LoopStatement>(nullptr, nullptr, utils::Empty);
         },
         "internal compiler error");
 }
@@ -83,7 +95,7 @@ TEST_F(LoopStatementTest, Assert_DifferentProgramID_Body) {
         {
             ProgramBuilder b1;
             ProgramBuilder b2;
-            b1.create<LoopStatement>(b2.Block(), b1.Block());
+            b1.create<LoopStatement>(b2.Block(), b1.Block(), utils::Empty);
         },
         "internal compiler error");
 }
@@ -93,7 +105,7 @@ TEST_F(LoopStatementTest, Assert_DifferentProgramID_Continuing) {
         {
             ProgramBuilder b1;
             ProgramBuilder b2;
-            b1.create<LoopStatement>(b1.Block(), b2.Block());
+            b1.create<LoopStatement>(b1.Block(), b2.Block(), utils::Empty);
         },
         "internal compiler error");
 }

src/tint/program_builder.h

@@ -3264,20 +3264,26 @@ class ProgramBuilder {
     /// @param source the source information
     /// @param body the loop body
     /// @param continuing the optional continuing block
+    /// @param attributes optional attributes
     /// @returns the loop statement pointer
-    const ast::LoopStatement* Loop(const Source& source,
+    const ast::LoopStatement* Loop(
-                                   const ast::BlockStatement* body,
+        const Source& source,
-                                   const ast::BlockStatement* continuing = nullptr) {
+        const ast::BlockStatement* body,
-        return create<ast::LoopStatement>(source, body, continuing);
+        const ast::BlockStatement* continuing = nullptr,
+        utils::VectorRef<const ast::Attribute*> attributes = utils::Empty) {
+        return create<ast::LoopStatement>(source, body, continuing, std::move(attributes));
     }
 
     /// Creates a ast::LoopStatement with input body and optional continuing
     /// @param body the loop body
     /// @param continuing the optional continuing block
+    /// @param attributes optional attributes
     /// @returns the loop statement pointer
-    const ast::LoopStatement* Loop(const ast::BlockStatement* body,
+    const ast::LoopStatement* Loop(
-                                   const ast::BlockStatement* continuing = nullptr) {
+        const ast::BlockStatement* body,
-        return create<ast::LoopStatement>(body, continuing);
+        const ast::BlockStatement* continuing = nullptr,
+        utils::VectorRef<const ast::Attribute*> attributes = utils::Empty) {
+        return create<ast::LoopStatement>(body, continuing, std::move(attributes));
     }
 
     /// Creates a ast::ForLoopStatement with input body and optional initializer, condition,

src/tint/reader/spirv/function.cc

@@ -743,7 +743,7 @@ struct LoopStatementBuilder final : public Castable<LoopStatementBuilder, Statem
     /// @param builder the program builder
     /// @returns the built ast::LoopStatement
     ast::LoopStatement* Build(ProgramBuilder* builder) const override {
-        return builder->create<ast::LoopStatement>(Source{}, body, continuing);
+        return builder->create<ast::LoopStatement>(Source{}, body, continuing, utils::Empty);
     }
 
     /// Loop-statement block body

src/tint/reader/wgsl/parser_impl.cc

@@ -1281,7 +1281,7 @@ Maybe<const ast::Statement*> ParserImpl::statement() {
         return sw.value;
     }
 
-    auto loop = loop_statement();
+    auto loop = loop_statement(attrs.value);
     if (loop.errored) {
         return Failure::kErrored;
     }
@@ -1731,13 +1731,18 @@ Maybe<const ast::CaseSelector*> ParserImpl::case_selector() {
 }
 
 // loop_statement
-//   : LOOP BRACKET_LEFT statements continuing_statement? BRACKET_RIGHT
+//   : attribute* LOOP attribute* BRACKET_LEFT statements continuing_statement? BRACKET_RIGHT
-Maybe<const ast::LoopStatement*> ParserImpl::loop_statement() {
+Maybe<const ast::LoopStatement*> ParserImpl::loop_statement(AttributeList& attrs) {
     Source source;
     if (!match(Token::Type::kLoop, &source)) {
         return Failure::kNoMatch;
     }
 
+    auto body_attrs = attribute_list();
+    if (body_attrs.errored) {
+        return Failure::kErrored;
+    }
+
     Maybe<const ast::BlockStatement*> continuing(Failure::kErrored);
     auto body_start = peek().source();
     auto body = expect_brace_block("loop", [&]() -> Maybe<StatementList> {
@@ -1757,11 +1762,12 @@ Maybe<const ast::LoopStatement*> ParserImpl::loop_statement() {
     }
     auto body_end = last_source();
 
+    TINT_DEFER(attrs.Clear());
     return create<ast::LoopStatement>(
         source,
         create<ast::BlockStatement>(Source::Combine(body_start, body_end), body.value,
-                                    utils::Empty),
+                                    std::move(body_attrs.value)),
-        continuing.value);
+        continuing.value, std::move(attrs));
 }
 
 ForHeader::ForHeader(const ast::Statement* init,

src/tint/reader/wgsl/parser_impl.h

@@ -512,9 +512,10 @@ class ParserImpl {
     /// Parses a `func_call_statement` grammar element
     /// @returns the parsed function call or nullptr
     Maybe<const ast::CallStatement*> func_call_statement();
-    /// Parses a `loop_statement` grammar element
+    /// Parses a `loop_statement` grammar element, with the attribute list provided as `attrs`.
+    /// @param attrs the list of attributes for the statement
     /// @returns the parsed loop or nullptr
-    Maybe<const ast::LoopStatement*> loop_statement();
+    Maybe<const ast::LoopStatement*> loop_statement(AttributeList& attrs);
     /// Parses a `for_header` grammar element, erroring on parse failure.
     /// @returns the parsed for header or nullptr
     Expect<std::unique_ptr<ForHeader>> expect_for_header();

src/tint/reader/wgsl/parser_impl_loop_stmt_test.cc

@@ -20,7 +20,8 @@ namespace {
 
 TEST_F(ParserImplTest, LoopStmt_BodyNoContinuing) {
     auto p = parser("loop { discard; }");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_TRUE(e.matched);
     EXPECT_FALSE(e.errored);
     EXPECT_FALSE(p->has_error()) << p->error();
@@ -39,7 +40,8 @@ TEST_F(ParserImplTest, LoopStmt_BodyNoContinuing) {
 
 TEST_F(ParserImplTest, LoopStmt_BodyWithContinuing) {
     auto p = parser("loop { discard; continuing { discard; }}");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_TRUE(e.matched);
     EXPECT_FALSE(e.errored);
     EXPECT_FALSE(p->has_error()) << p->error();
@@ -64,7 +66,8 @@ TEST_F(ParserImplTest, LoopStmt_BodyWithContinuing) {
 
 TEST_F(ParserImplTest, LoopStmt_NoBodyNoContinuing) {
     auto p = parser("loop { }");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_TRUE(e.matched);
     EXPECT_FALSE(e.errored);
     EXPECT_FALSE(p->has_error()) << p->error();
@@ -75,7 +78,8 @@ TEST_F(ParserImplTest, LoopStmt_NoBodyNoContinuing) {
 
 TEST_F(ParserImplTest, LoopStmt_NoBodyWithContinuing) {
     auto p = parser("loop { continuing { discard; }}");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_TRUE(e.matched);
     EXPECT_FALSE(e.errored);
     EXPECT_FALSE(p->has_error()) << p->error();
@@ -85,9 +89,35 @@ TEST_F(ParserImplTest, LoopStmt_NoBodyWithContinuing) {
     EXPECT_TRUE(e->continuing->statements[0]->Is<ast::DiscardStatement>());
 }
 
+TEST_F(ParserImplTest, LoopStmt_StmtAttributes) {
+    auto p = parser("@diagnostic(off, derivative_uniformity) loop { }");
+    auto attrs = p->attribute_list();
+    auto l = p->loop_statement(attrs.value);
+    EXPECT_FALSE(p->has_error()) << p->error();
+    EXPECT_FALSE(l.errored);
+    ASSERT_TRUE(l.matched);
+
+    EXPECT_TRUE(attrs->IsEmpty());
+    ASSERT_EQ(l->attributes.Length(), 1u);
+    EXPECT_TRUE(l->attributes[0]->Is<ast::DiagnosticAttribute>());
+}
+
+TEST_F(ParserImplTest, LoopStmt_BodyAttributes) {
+    auto p = parser("loop @diagnostic(off, derivative_uniformity) { }");
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
+    EXPECT_TRUE(e.matched);
+    EXPECT_FALSE(e.errored);
+    EXPECT_FALSE(p->has_error()) << p->error();
+    ASSERT_NE(e.value, nullptr);
+    ASSERT_EQ(e->body->attributes.Length(), 1u);
+    EXPECT_TRUE(e->body->attributes[0]->Is<ast::DiagnosticAttribute>());
+}
+
 TEST_F(ParserImplTest, LoopStmt_MissingBracketLeft) {
     auto p = parser("loop discard; }");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_EQ(e.value, nullptr);
@@ -97,7 +127,8 @@ TEST_F(ParserImplTest, LoopStmt_MissingBracketLeft) {
 
 TEST_F(ParserImplTest, LoopStmt_MissingBracketRight) {
     auto p = parser("loop { discard; ");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_EQ(e.value, nullptr);
@@ -107,7 +138,8 @@ TEST_F(ParserImplTest, LoopStmt_MissingBracketRight) {
 
 TEST_F(ParserImplTest, LoopStmt_InvalidStatements) {
     auto p = parser("loop { discard }");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_EQ(e.value, nullptr);
@@ -117,7 +149,8 @@ TEST_F(ParserImplTest, LoopStmt_InvalidStatements) {
 
 TEST_F(ParserImplTest, LoopStmt_InvalidContinuing) {
     auto p = parser("loop { continuing { discard }}");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_EQ(e.value, nullptr);
@@ -127,7 +160,8 @@ TEST_F(ParserImplTest, LoopStmt_InvalidContinuing) {
 
 TEST_F(ParserImplTest, LoopStmt_Continuing_BreakIf) {
     auto p = parser("loop { continuing { break if 1 + 2 < 5; }}");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_TRUE(e.matched);
     EXPECT_FALSE(e.errored);
     EXPECT_FALSE(p->has_error()) << p->error();
@@ -139,7 +173,8 @@ TEST_F(ParserImplTest, LoopStmt_Continuing_BreakIf) {
 
 TEST_F(ParserImplTest, LoopStmt_Continuing_BreakIf_MissingExpr) {
     auto p = parser("loop { continuing { break if; }}");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_TRUE(p->has_error());
@@ -149,7 +184,8 @@ TEST_F(ParserImplTest, LoopStmt_Continuing_BreakIf_MissingExpr) {
 
 TEST_F(ParserImplTest, LoopStmt_Continuing_BreakIf_InvalidExpr) {
     auto p = parser("loop { continuing { break if switch; }}");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_TRUE(p->has_error());
@@ -159,7 +195,8 @@ TEST_F(ParserImplTest, LoopStmt_Continuing_BreakIf_InvalidExpr) {
 
 TEST_F(ParserImplTest, LoopStmt_NoContinuing_BreakIf) {
     auto p = parser("loop { break if true; }");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_TRUE(p->has_error());
@@ -169,7 +206,8 @@ TEST_F(ParserImplTest, LoopStmt_NoContinuing_BreakIf) {
 
 TEST_F(ParserImplTest, LoopStmt_Continuing_BreakIf_MissingSemicolon) {
     auto p = parser("loop { continuing { break if 1 + 2 < 5 }}");
-    auto e = p->loop_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->loop_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_TRUE(p->has_error());

src/tint/reader/wgsl/parser_impl_statement_test.cc

@@ -350,6 +350,18 @@ TEST_F(ParserImplTest, Statement_ConsumedAttributes_If) {
     EXPECT_EQ(s->attributes.Length(), 1u);
 }
 
+TEST_F(ParserImplTest, Statement_ConsumedAttributes_Loop) {
+    auto p = parser("@diagnostic(off, derivative_uniformity) loop {}");
+    auto e = p->statement();
+    ASSERT_FALSE(p->has_error()) << p->error();
+    EXPECT_TRUE(e.matched);
+    EXPECT_FALSE(e.errored);
+
+    auto* s = As<ast::LoopStatement>(e.value);
+    ASSERT_NE(s, nullptr);
+    EXPECT_EQ(s->attributes.Length(), 1u);
+}
+
 TEST_F(ParserImplTest, Statement_ConsumedAttributes_Switch) {
     auto p = parser("@diagnostic(off, derivative_uniformity) switch (0) { default{} }");
     auto e = p->statement();

src/tint/resolver/attribute_validation_test.cc

@@ -1144,6 +1144,39 @@ INSTANTIATE_TEST_SUITE_P(ResolverAttributeValidationTest,
                                          TestParams{AttributeKind::kWorkgroup, false},
                                          TestParams{AttributeKind::kBindingAndGroup, false}));
 
+using LoopStatementAttributeTest = TestWithParams;
+TEST_P(LoopStatementAttributeTest, IsValid) {
+    auto& params = GetParam();
+
+    WrapInFunction(
+        Loop(Block(Return()), Block(), createAttributes(Source{{12, 34}}, *this, params.kind)));
+
+    if (params.should_pass) {
+        EXPECT_TRUE(r()->Resolve()) << r()->error();
+    } else {
+        EXPECT_FALSE(r()->Resolve());
+        EXPECT_EQ(r()->error(), "12:34 error: attribute is not valid for loop statements");
+    }
+}
+INSTANTIATE_TEST_SUITE_P(ResolverAttributeValidationTest,
+                         LoopStatementAttributeTest,
+                         testing::Values(TestParams{AttributeKind::kAlign, false},
+                                         TestParams{AttributeKind::kBinding, false},
+                                         TestParams{AttributeKind::kBuiltin, false},
+                                         TestParams{AttributeKind::kDiagnostic, true},
+                                         TestParams{AttributeKind::kGroup, false},
+                                         TestParams{AttributeKind::kId, false},
+                                         TestParams{AttributeKind::kInterpolate, false},
+                                         TestParams{AttributeKind::kInvariant, false},
+                                         TestParams{AttributeKind::kLocation, false},
+                                         TestParams{AttributeKind::kMustUse, false},
+                                         TestParams{AttributeKind::kOffset, false},
+                                         TestParams{AttributeKind::kSize, false},
+                                         TestParams{AttributeKind::kStage, false},
+                                         TestParams{AttributeKind::kStride, false},
+                                         TestParams{AttributeKind::kWorkgroup, false},
+                                         TestParams{AttributeKind::kBindingAndGroup, false}));
+
 using WhileStatementAttributeTest = TestWithParams;
 TEST_P(WhileStatementAttributeTest, IsValid) {
     auto& params = GetParam();
@@ -1225,6 +1258,13 @@ TEST_P(BlockStatementTest, ForStatementBody) {
          });
     Check();
 }
+TEST_P(BlockStatementTest, LoopStatementBody) {
+    Func("foo", utils::Empty, ty.void_(),
+         utils::Vector{
+             Loop(Block(utils::Vector{Break()}, createAttributes({}, *this, GetParam().kind))),
+         });
+    Check();
+}
 TEST_P(BlockStatementTest, WhileStatementBody) {
     Func("foo", utils::Empty, ty.void_(),
          utils::Vector{

src/tint/resolver/resolver.cc

@@ -4281,6 +4281,9 @@ SEM* Resolver::StatementScope(const ast::Statement* ast, SEM* sem, F&& callback)
                 return handle_attributes(f, sem, "for statements");
             },
             [&](const ast::IfStatement* i) { return handle_attributes(i, sem, "if statements"); },
+            [&](const ast::LoopStatement* l) {
+                return handle_attributes(l, sem, "loop statements");
+            },
             [&](const ast::SwitchStatement* s) {
                 return handle_attributes(s, sem, "switch statements");
             },

src/tint/resolver/uniformity_test.cc

@@ -8547,6 +8547,114 @@ fn foo() {
     }
 }
 
+TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnLoopStatement_CallInBody) {
+    auto& param = GetParam();
+    utils::StringStream ss;
+    ss << R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+fn foo() {
+  )"
+       << "@diagnostic(" << param << ", derivative_uniformity)"
+       << R"(loop {
+    _ = dpdx(1.0);
+    continuing {
+      break if non_uniform == 0;
+    }
+  }
+}
+)";
+
+    RunTest(ss.str(), param != builtin::DiagnosticSeverity::kError);
+    if (param == builtin::DiagnosticSeverity::kOff) {
+        EXPECT_TRUE(error_.empty());
+    } else {
+        utils::StringStream err;
+        err << ToStr(param) << ": 'dpdx' must only be called";
+        EXPECT_THAT(error_, ::testing::HasSubstr(err.str()));
+    }
+}
+
+TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnLoopStatement_CallInContinuing) {
+    auto& param = GetParam();
+    utils::StringStream ss;
+    ss << R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+fn foo() {
+  )"
+       << "@diagnostic(" << param << ", derivative_uniformity)"
+       << R"(loop {
+    continuing {
+      _ = dpdx(1.0);
+      break if non_uniform == 0;
+    }
+  }
+}
+)";
+
+    RunTest(ss.str(), param != builtin::DiagnosticSeverity::kError);
+    if (param == builtin::DiagnosticSeverity::kOff) {
+        EXPECT_TRUE(error_.empty());
+    } else {
+        utils::StringStream err;
+        err << ToStr(param) << ": 'dpdx' must only be called";
+        EXPECT_THAT(error_, ::testing::HasSubstr(err.str()));
+    }
+}
+
+TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnLoopBody_CallInBody) {
+    auto& param = GetParam();
+    utils::StringStream ss;
+    ss << R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+fn foo() {
+  loop )"
+       << "@diagnostic(" << param << ", derivative_uniformity)"
+       << R"( {
+    _ = dpdx(1.0);
+    continuing {
+      break if non_uniform == 0;
+    }
+  }
+}
+)";
+
+    RunTest(ss.str(), param != builtin::DiagnosticSeverity::kError);
+    if (param == builtin::DiagnosticSeverity::kOff) {
+        EXPECT_TRUE(error_.empty());
+    } else {
+        utils::StringStream err;
+        err << ToStr(param) << ": 'dpdx' must only be called";
+        EXPECT_THAT(error_, ::testing::HasSubstr(err.str()));
+    }
+}
+
+TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnLoopBody_CallInContinuing) {
+    auto& param = GetParam();
+    utils::StringStream ss;
+    ss << R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+fn foo() {
+  loop )"
+       << "@diagnostic(" << param << ", derivative_uniformity)"
+       << R"( {
+    continuing {
+      _ = dpdx(1.0);
+      break if non_uniform == 0;
+    }
+  }
+}
+)";
+
+    RunTest(ss.str(), param != builtin::DiagnosticSeverity::kError);
+    if (param == builtin::DiagnosticSeverity::kOff) {
+        EXPECT_TRUE(error_.empty());
+    } else {
+        utils::StringStream err;
+        err << ToStr(param) << ": 'dpdx' must only be called";
+        EXPECT_THAT(error_, ::testing::HasSubstr(err.str()));
+    }
+}
+
 TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnSwitchStatement_CallInCondition) {
     auto& param = GetParam();
     utils::StringStream ss;

src/tint/sem/diagnostic_severity_test.cc

@@ -55,7 +55,11 @@ class DiagnosticSeverityTest : public TestHelper {
         //     for (var i = 0; false; i++) @diagnostic(warning, chromium_unreachable_code) {
         //       return;
         //     }
-        //   }
+        //
+        //     @diagnostic(warning, chromium_unreachable_code)
+        //     loop @diagnostic(off, chromium_unreachable_code) {
+        //       return;
+        //     }
         //
         //     @diagnostic(error, chromium_unreachable_code)
         //     while (false) @diagnostic(warning, chromium_unreachable_code) {
@@ -77,6 +81,8 @@ class DiagnosticSeverityTest : public TestHelper {
         auto case_severity = builtin::DiagnosticSeverity::kWarning;
         auto for_severity = builtin::DiagnosticSeverity::kError;
         auto for_body_severity = builtin::DiagnosticSeverity::kWarning;
+        auto loop_severity = builtin::DiagnosticSeverity::kWarning;
+        auto loop_body_severity = builtin::DiagnosticSeverity::kOff;
         auto while_severity = builtin::DiagnosticSeverity::kError;
         auto while_body_severity = builtin::DiagnosticSeverity::kWarning;
         auto attr = [&](auto severity) {
@@ -90,6 +96,7 @@ class DiagnosticSeverityTest : public TestHelper {
         auto* return_foo_case = Return();
         auto* return_foo_default = Return();
         auto* return_foo_for = Return();
+        auto* return_foo_loop = Return();
         auto* return_foo_while = Return();
         auto* else_stmt = Block(utils::Vector{return_foo_else}, attr(else_body_severity));
         auto* elseif = If(Expr(false), Block(return_foo_elseif), Else(else_stmt));
@@ -102,10 +109,12 @@ class DiagnosticSeverityTest : public TestHelper {
         auto* fl =
             For(Decl(Var("i", ty.i32())), false, Increment("i"),
                 Block(utils::Vector{return_foo_for}, attr(for_body_severity)), attr(for_severity));
+        auto* l = Loop(Block(utils::Vector{return_foo_loop}, attr(loop_body_severity)), Block(),
+                       attr(loop_severity));
         auto* wl = While(false, Block(utils::Vector{return_foo_while}, attr(while_body_severity)),
                          attr(while_severity));
         auto* block_1 =
-            Block(utils::Vector{if_foo, return_foo_block, swtch, fl, wl}, attr(block_severity));
+            Block(utils::Vector{if_foo, return_foo_block, swtch, fl, l, wl}, attr(block_severity));
         auto* func_attr = DiagnosticAttribute(func_severity, "chromium_unreachable_code");
         auto* foo = Func("foo", {}, ty.void_(), utils::Vector{block_1}, utils::Vector{func_attr});
 
@@ -139,6 +148,9 @@ class DiagnosticSeverityTest : public TestHelper {
         EXPECT_EQ(p.Sem().DiagnosticSeverity(fl->continuing, rule), for_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(fl->body, rule), for_body_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_for, rule), for_body_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(l, rule), loop_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(l->body, rule), loop_body_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_loop, rule), loop_body_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(wl, rule), while_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(wl->condition, rule), while_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(wl->body, rule), while_body_severity);

src/tint/writer/wgsl/generator_impl.cc

@@ -1025,7 +1025,21 @@ bool GeneratorImpl::EmitDiscard(const ast::DiscardStatement*) {
 }
 
 bool GeneratorImpl::EmitLoop(const ast::LoopStatement* stmt) {
-    line() << "loop {";
+    {
+        auto out = line();
+
+        if (!stmt->attributes.IsEmpty()) {
+            if (!EmitAttributes(out, stmt->attributes)) {
+                return false;
+            }
+            out << " ";
+        }
+
+        out << "loop ";
+        if (!EmitBlockHeader(out, stmt->body)) {
+            return false;
+        }
+    }
     increment_indent();
 
     if (!EmitStatements(stmt->body->statements)) {

test/tint/diagnostic_filtering/loop_attribute.wgsl

@@ -0,0 +1,10 @@
+@fragment
+fn main(@location(0) x : f32) {
+  @diagnostic(warning, derivative_uniformity)
+  loop {
+    _ = dpdx(1.0);
+    continuing {
+      break if x > 0.0;
+    }
+  }
+}

test/tint/diagnostic_filtering/loop_attribute.wgsl.expected.dxc.hlsl

@@ -0,0 +1,28 @@
+diagnostic_filtering/loop_attribute.wgsl:5:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+struct tint_symbol_1 {
+  float x : TEXCOORD0;
+};
+
+void main_inner(float x) {
+  while (true) {
+    {
+      if ((x > 0.0f)) { break; }
+    }
+  }
+}
+
+void main(tint_symbol_1 tint_symbol) {
+  main_inner(tint_symbol.x);
+  return;
+}

test/tint/diagnostic_filtering/loop_attribute.wgsl.expected.fxc.hlsl

@@ -0,0 +1,28 @@
+diagnostic_filtering/loop_attribute.wgsl:5:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+struct tint_symbol_1 {
+  float x : TEXCOORD0;
+};
+
+void main_inner(float x) {
+  while (true) {
+    {
+      if ((x > 0.0f)) { break; }
+    }
+  }
+}
+
+void main(tint_symbol_1 tint_symbol) {
+  main_inner(tint_symbol.x);
+  return;
+}

test/tint/diagnostic_filtering/loop_attribute.wgsl.expected.glsl

@@ -0,0 +1,28 @@
+diagnostic_filtering/loop_attribute.wgsl:5:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+#version 310 es
+precision highp float;
+
+layout(location = 0) in float x_1;
+void tint_symbol(float x) {
+  while (true) {
+    {
+      if ((x > 0.0f)) { break; }
+    }
+  }
+}
+
+void main() {
+  tint_symbol(x_1);
+  return;
+}

test/tint/diagnostic_filtering/loop_attribute.wgsl.expected.msl

@@ -0,0 +1,32 @@
+diagnostic_filtering/loop_attribute.wgsl:5:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+#include <metal_stdlib>
+
+using namespace metal;
+struct tint_symbol_2 {
+  float x [[user(locn0)]];
+};
+
+void tint_symbol_inner(float x) {
+  while (true) {
+    {
+      if ((x > 0.0f)) { break; }
+    }
+  }
+}
+
+fragment void tint_symbol(tint_symbol_2 tint_symbol_1 [[stage_in]]) {
+  tint_symbol_inner(tint_symbol_1.x);
+  return;
+}
+

test/tint/diagnostic_filtering/loop_attribute.wgsl.expected.spvasm

@@ -0,0 +1,55 @@
+diagnostic_filtering/loop_attribute.wgsl:5:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 21
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %main "main" %x_1
+               OpExecutionMode %main OriginUpperLeft
+               OpName %x_1 "x_1"
+               OpName %main_inner "main_inner"
+               OpName %x "x"
+               OpName %main "main"
+               OpDecorate %x_1 Location 0
+      %float = OpTypeFloat 32
+%_ptr_Input_float = OpTypePointer Input %float
+        %x_1 = OpVariable %_ptr_Input_float Input
+       %void = OpTypeVoid
+          %4 = OpTypeFunction %void %float
+         %13 = OpConstantNull %float
+       %bool = OpTypeBool
+         %16 = OpTypeFunction %void
+ %main_inner = OpFunction %void None %4
+          %x = OpFunctionParameter %float
+          %8 = OpLabel
+               OpBranch %9
+          %9 = OpLabel
+               OpLoopMerge %10 %11 None
+               OpBranch %12
+         %12 = OpLabel
+               OpBranch %11
+         %11 = OpLabel
+         %14 = OpFOrdGreaterThan %bool %x %13
+               OpBranchConditional %14 %10 %9
+         %10 = OpLabel
+               OpReturn
+               OpFunctionEnd
+       %main = OpFunction %void None %16
+         %18 = OpLabel
+         %20 = OpLoad %float %x_1
+         %19 = OpFunctionCall %void %main_inner %20
+               OpReturn
+               OpFunctionEnd

test/tint/diagnostic_filtering/loop_attribute.wgsl.expected.wgsl

@@ -0,0 +1,22 @@
+diagnostic_filtering/loop_attribute.wgsl:5:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_attribute.wgsl:7:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+@fragment
+fn main(@location(0) x : f32) {
+  @diagnostic(warning, derivative_uniformity) loop {
+    _ = dpdx(1.0);
+
+    continuing {
+      break if (x > 0.0);
+    }
+  }
+}

test/tint/diagnostic_filtering/loop_body_attribute.wgsl

@@ -0,0 +1,9 @@
+@fragment
+fn main(@location(0) x : f32) {
+  loop @diagnostic(warning, derivative_uniformity) {
+    _ = dpdx(1.0);
+    continuing {
+      break if x > 0.0;
+    }
+  }
+}

test/tint/diagnostic_filtering/loop_body_attribute.wgsl.expected.dxc.hlsl

@@ -0,0 +1,28 @@
+diagnostic_filtering/loop_body_attribute.wgsl:4:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+struct tint_symbol_1 {
+  float x : TEXCOORD0;
+};
+
+void main_inner(float x) {
+  while (true) {
+    {
+      if ((x > 0.0f)) { break; }
+    }
+  }
+}
+
+void main(tint_symbol_1 tint_symbol) {
+  main_inner(tint_symbol.x);
+  return;
+}

test/tint/diagnostic_filtering/loop_body_attribute.wgsl.expected.fxc.hlsl

@@ -0,0 +1,28 @@
+diagnostic_filtering/loop_body_attribute.wgsl:4:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+struct tint_symbol_1 {
+  float x : TEXCOORD0;
+};
+
+void main_inner(float x) {
+  while (true) {
+    {
+      if ((x > 0.0f)) { break; }
+    }
+  }
+}
+
+void main(tint_symbol_1 tint_symbol) {
+  main_inner(tint_symbol.x);
+  return;
+}

test/tint/diagnostic_filtering/loop_body_attribute.wgsl.expected.glsl

@@ -0,0 +1,28 @@
+diagnostic_filtering/loop_body_attribute.wgsl:4:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+#version 310 es
+precision highp float;
+
+layout(location = 0) in float x_1;
+void tint_symbol(float x) {
+  while (true) {
+    {
+      if ((x > 0.0f)) { break; }
+    }
+  }
+}
+
+void main() {
+  tint_symbol(x_1);
+  return;
+}

test/tint/diagnostic_filtering/loop_body_attribute.wgsl.expected.msl

@@ -0,0 +1,32 @@
+diagnostic_filtering/loop_body_attribute.wgsl:4:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+#include <metal_stdlib>
+
+using namespace metal;
+struct tint_symbol_2 {
+  float x [[user(locn0)]];
+};
+
+void tint_symbol_inner(float x) {
+  while (true) {
+    {
+      if ((x > 0.0f)) { break; }
+    }
+  }
+}
+
+fragment void tint_symbol(tint_symbol_2 tint_symbol_1 [[stage_in]]) {
+  tint_symbol_inner(tint_symbol_1.x);
+  return;
+}
+

test/tint/diagnostic_filtering/loop_body_attribute.wgsl.expected.spvasm

@@ -0,0 +1,55 @@
+diagnostic_filtering/loop_body_attribute.wgsl:4:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 21
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %main "main" %x_1
+               OpExecutionMode %main OriginUpperLeft
+               OpName %x_1 "x_1"
+               OpName %main_inner "main_inner"
+               OpName %x "x"
+               OpName %main "main"
+               OpDecorate %x_1 Location 0
+      %float = OpTypeFloat 32
+%_ptr_Input_float = OpTypePointer Input %float
+        %x_1 = OpVariable %_ptr_Input_float Input
+       %void = OpTypeVoid
+          %4 = OpTypeFunction %void %float
+         %13 = OpConstantNull %float
+       %bool = OpTypeBool
+         %16 = OpTypeFunction %void
+ %main_inner = OpFunction %void None %4
+          %x = OpFunctionParameter %float
+          %8 = OpLabel
+               OpBranch %9
+          %9 = OpLabel
+               OpLoopMerge %10 %11 None
+               OpBranch %12
+         %12 = OpLabel
+               OpBranch %11
+         %11 = OpLabel
+         %14 = OpFOrdGreaterThan %bool %x %13
+               OpBranchConditional %14 %10 %9
+         %10 = OpLabel
+               OpReturn
+               OpFunctionEnd
+       %main = OpFunction %void None %16
+         %18 = OpLabel
+         %20 = OpLoad %float %x_1
+         %19 = OpFunctionCall %void %main_inner %20
+               OpReturn
+               OpFunctionEnd

test/tint/diagnostic_filtering/loop_body_attribute.wgsl.expected.wgsl

@@ -0,0 +1,22 @@
+diagnostic_filtering/loop_body_attribute.wgsl:4:9 warning: 'dpdx' must only be called from uniform control flow
+    _ = dpdx(1.0);
+        ^^^^^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:7 note: control flow depends on possibly non-uniform value
+      break if x > 0.0;
+      ^^^^^
+
+diagnostic_filtering/loop_body_attribute.wgsl:6:16 note: user-defined input 'x' of 'main' may be non-uniform
+      break if x > 0.0;
+               ^
+
+@fragment
+fn main(@location(0) x : f32) {
+  loop @diagnostic(warning, derivative_uniformity) {
+    _ = dpdx(1.0);
+
+    continuing {
+      break if (x > 0.0);
+    }
+  }
+}