tint: Support @diagnostic on for loops

Bug: tint:1809
Change-Id: I0c6ce482a6d91ddfbcd3e69938ff09de79823e05
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/124101
Reviewed-by: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: James Price <jrprice@google.com>
diff --git a/src/tint/ast/for_loop_statement.cc b/src/tint/ast/for_loop_statement.cc
index aba956d..b2a5470 100644
--- a/src/tint/ast/for_loop_statement.cc
+++ b/src/tint/ast/for_loop_statement.cc
@@ -14,6 +14,8 @@
 
 #include "src/tint/ast/for_loop_statement.h"
 
+#include <utility>
+
 #include "src/tint/program_builder.h"
 
 TINT_INSTANTIATE_TYPEINFO(tint::ast::ForLoopStatement);
@@ -26,14 +28,24 @@
                                    const Statement* init,
                                    const Expression* cond,
                                    const Statement* cont,
-                                   const BlockStatement* b)
-    : Base(pid, nid, src), initializer(init), condition(cond), continuing(cont), body(b) {
+                                   const BlockStatement* b,
+                                   utils::VectorRef<const ast::Attribute*> attrs)
+    : Base(pid, nid, src),
+      initializer(init),
+      condition(cond),
+      continuing(cont),
+      body(b),
+      attributes(std::move(attrs)) {
     TINT_ASSERT(AST, body);
 
     TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, initializer, program_id);
     TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, condition, program_id);
     TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, continuing, program_id);
     TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, body, program_id);
+    for (auto* attr : attributes) {
+        TINT_ASSERT(AST, attr);
+        TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, attr, program_id);
+    }
 }
 
 ForLoopStatement::ForLoopStatement(ForLoopStatement&&) = default;
@@ -48,7 +60,8 @@
     auto* cond = ctx->Clone(condition);
     auto* cont = ctx->Clone(continuing);
     auto* b = ctx->Clone(body);
-    return ctx->dst->create<ForLoopStatement>(src, init, cond, cont, b);
+    auto attrs = ctx->Clone(attributes);
+    return ctx->dst->create<ForLoopStatement>(src, init, cond, cont, b, std::move(attrs));
 }
 
 }  // namespace tint::ast
diff --git a/src/tint/ast/for_loop_statement.h b/src/tint/ast/for_loop_statement.h
index 59e0587..6063dda 100644
--- a/src/tint/ast/for_loop_statement.h
+++ b/src/tint/ast/for_loop_statement.h
@@ -32,13 +32,15 @@
     /// @param condition the optional loop condition expression
     /// @param continuing the optional continuing statement
     /// @param body the loop body
+    /// @param attributes the while statement attributes
     ForLoopStatement(ProgramID pid,
                      NodeID nid,
                      const Source& source,
                      const Statement* initializer,
                      const Expression* condition,
                      const Statement* continuing,
-                     const BlockStatement* body);
+                     const BlockStatement* body,
+                     utils::VectorRef<const ast::Attribute*> attributes);
     /// Move constructor
     ForLoopStatement(ForLoopStatement&&);
     ~ForLoopStatement() override;
@@ -60,6 +62,9 @@
 
     /// The loop body block
     const BlockStatement* const body;
+
+    /// The attribute list
+    const utils::Vector<const Attribute*, 1> attributes;
 };
 
 }  // namespace tint::ast
diff --git a/src/tint/ast/for_loop_statement_test.cc b/src/tint/ast/for_loop_statement_test.cc
index 5e2a000..ad9050e 100644
--- a/src/tint/ast/for_loop_statement_test.cc
+++ b/src/tint/ast/for_loop_statement_test.cc
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "gmock/gmock.h"
 #include "gtest/gtest-spi.h"
 #include "src/tint/ast/binary_expression.h"
 #include "src/tint/ast/test_helper.h"
@@ -50,6 +51,15 @@
     EXPECT_EQ(l->body, body);
 }
 
+TEST_F(ForLoopStatementTest, Creation_WithAttributes) {
+    auto* attr1 = DiagnosticAttribute(builtin::DiagnosticSeverity::kOff, "foo");
+    auto* attr2 = DiagnosticAttribute(builtin::DiagnosticSeverity::kOff, "bar");
+    auto* body = Block(Return());
+    auto* l = For(nullptr, nullptr, nullptr, body, utils::Vector{attr1, attr2});
+
+    EXPECT_THAT(l->attributes, testing::ElementsAre(attr1, attr2));
+}
+
 TEST_F(ForLoopStatementTest, Assert_Null_Body) {
     EXPECT_FATAL_FAILURE(
         {
diff --git a/src/tint/program_builder.h b/src/tint/program_builder.h
index dc50acf..d5576b8 100644
--- a/src/tint/program_builder.h
+++ b/src/tint/program_builder.h
@@ -3278,37 +3278,44 @@
         return create<ast::LoopStatement>(body, continuing);
     }
 
-    /// Creates a ast::ForLoopStatement with input body and optional initializer,
-    /// condition and continuing.
+    /// Creates a ast::ForLoopStatement with input body and optional initializer, condition,
+    /// continuing, and attributes.
     /// @param source the source information
     /// @param init the optional loop initializer
     /// @param cond the optional loop condition
     /// @param cont the optional loop continuing
     /// @param body the loop body
+    /// @param attributes optional attributes
     /// @returns the for loop statement pointer
     template <typename COND>
-    const ast::ForLoopStatement* For(const Source& source,
-                                     const ast::Statement* init,
-                                     COND&& cond,
-                                     const ast::Statement* cont,
-                                     const ast::BlockStatement* body) {
+    const ast::ForLoopStatement* For(
+        const Source& source,
+        const ast::Statement* init,
+        COND&& cond,
+        const ast::Statement* cont,
+        const ast::BlockStatement* body,
+        utils::VectorRef<const ast::Attribute*> attributes = utils::Empty) {
         return create<ast::ForLoopStatement>(source, init, Expr(std::forward<COND>(cond)), cont,
-                                             body);
+                                             body, std::move(attributes));
     }
 
-    /// Creates a ast::ForLoopStatement with input body and optional initializer,
-    /// condition and continuing.
+    /// Creates a ast::ForLoopStatement with input body and optional initializer, condition,
+    /// continuing, and attributes.
     /// @param init the optional loop initializer
     /// @param cond the optional loop condition
     /// @param cont the optional loop continuing
     /// @param body the loop body
+    /// @param attributes optional attributes
     /// @returns the for loop statement pointer
     template <typename COND>
-    const ast::ForLoopStatement* For(const ast::Statement* init,
-                                     COND&& cond,
-                                     const ast::Statement* cont,
-                                     const ast::BlockStatement* body) {
-        return create<ast::ForLoopStatement>(init, Expr(std::forward<COND>(cond)), cont, body);
+    const ast::ForLoopStatement* For(
+        const ast::Statement* init,
+        COND&& cond,
+        const ast::Statement* cont,
+        const ast::BlockStatement* body,
+        utils::VectorRef<const ast::Attribute*> attributes = utils::Empty) {
+        return create<ast::ForLoopStatement>(init, Expr(std::forward<COND>(cond)), cont, body,
+                                             std::move(attributes));
     }
 
     /// Creates a ast::WhileStatement with input body, condition, and optional attributes.
diff --git a/src/tint/reader/wgsl/parser_impl.cc b/src/tint/reader/wgsl/parser_impl.cc
index 7936d65..c1a379e 100644
--- a/src/tint/reader/wgsl/parser_impl.cc
+++ b/src/tint/reader/wgsl/parser_impl.cc
@@ -1288,7 +1288,7 @@
         return loop.value;
     }
 
-    auto stmt_for = for_statement();
+    auto stmt_for = for_statement(attrs.value);
     if (stmt_for.errored) {
         return Failure::kErrored;
     }
@@ -1881,7 +1881,7 @@
 
 // for_statement
 //   : FOR PAREN_LEFT for_header PAREN_RIGHT compound_statement
-Maybe<const ast::ForLoopStatement*> ParserImpl::for_statement() {
+Maybe<const ast::ForLoopStatement*> ParserImpl::for_statement(AttributeList& attrs) {
     Source source;
     if (!match(Token::Type::kFor, &source)) {
         return Failure::kNoMatch;
@@ -1897,8 +1897,9 @@
         return Failure::kErrored;
     }
 
+    TINT_DEFER(attrs.Clear());
     return create<ast::ForLoopStatement>(source, header->initializer, header->condition,
-                                         header->continuing, body.value);
+                                         header->continuing, body.value, std::move(attrs));
 }
 
 // while_statement
diff --git a/src/tint/reader/wgsl/parser_impl.h b/src/tint/reader/wgsl/parser_impl.h
index 57550e0..9693ea4 100644
--- a/src/tint/reader/wgsl/parser_impl.h
+++ b/src/tint/reader/wgsl/parser_impl.h
@@ -517,9 +517,10 @@
     /// 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();
-    /// Parses a `for_statement` grammar element
+    /// Parses a `for_statement` grammar element, with the attribute list provided as `attrs`.
+    /// @param attrs the list of attributes for the statement
     /// @returns the parsed for loop or nullptr
-    Maybe<const ast::ForLoopStatement*> for_statement();
+    Maybe<const ast::ForLoopStatement*> for_statement(AttributeList& attrs);
     /// Parses a `while_statement` grammar element, with the attribute list provided as `attrs`.
     /// @param attrs the list of attributes for the statement
     /// @returns the parsed while loop or nullptr
diff --git a/src/tint/reader/wgsl/parser_impl_for_stmt_test.cc b/src/tint/reader/wgsl/parser_impl_for_stmt_test.cc
index 5af8f22..528be9f 100644
--- a/src/tint/reader/wgsl/parser_impl_for_stmt_test.cc
+++ b/src/tint/reader/wgsl/parser_impl_for_stmt_test.cc
@@ -24,7 +24,8 @@
 // Test an empty for loop.
 TEST_F(ForStmtTest, Empty) {
     auto p = parser("for (;;) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -37,7 +38,8 @@
 // Test a for loop with non-empty body.
 TEST_F(ForStmtTest, Body) {
     auto p = parser("for (;;) { discard; }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -51,7 +53,8 @@
 // Test a for loop declaring a variable in the initializer statement.
 TEST_F(ForStmtTest, InitializerStatementDecl) {
     auto p = parser("for (var i: i32 ;;) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -68,7 +71,8 @@
 // statement.
 TEST_F(ForStmtTest, InitializerStatementDeclEqual) {
     auto p = parser("for (var i: i32 = 0 ;;) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -84,7 +88,8 @@
 // Test a for loop declaring a const variable in the initializer statement.
 TEST_F(ForStmtTest, InitializerStatementConstDecl) {
     auto p = parser("for (let i: i32 = 0 ;;) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -100,7 +105,8 @@
 // Test a for loop assigning a variable in the initializer statement.
 TEST_F(ForStmtTest, InitializerStatementAssignment) {
     auto p = parser("for (i = 0 ;;) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -113,7 +119,8 @@
 // Test a for loop incrementing a variable in the initializer statement.
 TEST_F(ForStmtTest, InitializerStatementIncrement) {
     auto p = parser("for (i++;;) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -126,7 +133,8 @@
 // Test a for loop calling a function in the initializer statement.
 TEST_F(ForStmtTest, InitializerStatementFuncCall) {
     auto p = parser("for (a(b,c) ;;) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -139,7 +147,8 @@
 // Test a for loop with a break condition
 TEST_F(ForStmtTest, BreakCondition) {
     auto p = parser("for (; 0 == 1;) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -152,7 +161,8 @@
 // Test a for loop assigning a variable in the continuing statement.
 TEST_F(ForStmtTest, ContinuingAssignment) {
     auto p = parser("for (;; x = 2) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -165,7 +175,8 @@
 // Test a for loop with an increment statement as the continuing statement.
 TEST_F(ForStmtTest, ContinuingIncrement) {
     auto p = parser("for (;; x++) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -178,7 +189,8 @@
 // Test a for loop calling a function in the continuing statement.
 TEST_F(ForStmtTest, ContinuingFuncCall) {
     auto p = parser("for (;; a(b,c)) { }");
-    auto fl = p->for_statement();
+    ParserImpl::AttributeList attrs;
+    auto fl = p->for_statement(attrs);
     EXPECT_FALSE(p->has_error()) << p->error();
     EXPECT_FALSE(fl.errored);
     ASSERT_TRUE(fl.matched);
@@ -188,11 +200,26 @@
     EXPECT_TRUE(fl->body->Empty());
 }
 
+// Test a for loop with attributes.
+TEST_F(ForStmtTest, WithAttributes) {
+    auto p = parser("@diagnostic(off, derivative_uniformity) for (;;) { }");
+    auto attrs = p->attribute_list();
+    auto fl = p->for_statement(attrs.value);
+    EXPECT_FALSE(p->has_error()) << p->error();
+    EXPECT_FALSE(fl.errored);
+    ASSERT_TRUE(fl.matched);
+
+    EXPECT_TRUE(attrs->IsEmpty());
+    ASSERT_EQ(fl->attributes.Length(), 1u);
+    EXPECT_TRUE(fl->attributes[0]->Is<ast::DiagnosticAttribute>());
+}
+
 class ForStmtErrorTest : public ParserImplTest {
   public:
     void TestForWithError(std::string for_str, std::string error_str) {
         auto p_for = parser(for_str);
-        auto e_for = p_for->for_statement();
+        ParserImpl::AttributeList attrs;
+        auto e_for = p_for->for_statement(attrs);
 
         EXPECT_FALSE(e_for.matched);
         EXPECT_TRUE(e_for.errored);
diff --git a/src/tint/reader/wgsl/parser_impl_statement_test.cc b/src/tint/reader/wgsl/parser_impl_statement_test.cc
index d10e549..8c920a7b 100644
--- a/src/tint/reader/wgsl/parser_impl_statement_test.cc
+++ b/src/tint/reader/wgsl/parser_impl_statement_test.cc
@@ -326,6 +326,18 @@
     EXPECT_EQ(s->attributes.Length(), 1u);
 }
 
+TEST_F(ParserImplTest, Statement_ConsumedAttributes_For) {
+    auto p = parser("@diagnostic(off, derivative_uniformity) for (;false;) {}");
+    auto e = p->statement();
+    ASSERT_FALSE(p->has_error()) << p->error();
+    EXPECT_TRUE(e.matched);
+    EXPECT_FALSE(e.errored);
+
+    auto* s = As<ast::ForLoopStatement>(e.value);
+    ASSERT_NE(s, nullptr);
+    EXPECT_EQ(s->attributes.Length(), 1u);
+}
+
 TEST_F(ParserImplTest, Statement_ConsumedAttributes_If) {
     auto p = parser("@diagnostic(off, derivative_uniformity) if true {}");
     auto e = p->statement();
diff --git a/src/tint/resolver/attribute_validation_test.cc b/src/tint/resolver/attribute_validation_test.cc
index 0b80300..23f82a2 100644
--- a/src/tint/resolver/attribute_validation_test.cc
+++ b/src/tint/resolver/attribute_validation_test.cc
@@ -1111,6 +1111,39 @@
                                          TestParams{AttributeKind::kWorkgroup, false},
                                          TestParams{AttributeKind::kBindingAndGroup, false}));
 
+using ForStatementAttributeTest = TestWithParams;
+TEST_P(ForStatementAttributeTest, IsValid) {
+    auto& params = GetParam();
+
+    WrapInFunction(For(nullptr, Expr(false), nullptr, 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 for statements");
+    }
+}
+INSTANTIATE_TEST_SUITE_P(ResolverAttributeValidationTest,
+                         ForStatementAttributeTest,
+                         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();
diff --git a/src/tint/resolver/resolver.cc b/src/tint/resolver/resolver.cc
index 2425f51..48fa4b1 100644
--- a/src/tint/resolver/resolver.cc
+++ b/src/tint/resolver/resolver.cc
@@ -4274,6 +4274,9 @@
             [&](const ast::BlockStatement* block) {
                 return handle_attributes(block, sem, "block statements");
             },
+            [&](const ast::ForLoopStatement* f) {
+                return handle_attributes(f, sem, "for statements");
+            },
             [&](const ast::IfStatement* i) { return handle_attributes(i, sem, "if statements"); },
             [&](const ast::SwitchStatement* s) {
                 return handle_attributes(s, sem, "switch statements");
diff --git a/src/tint/resolver/uniformity_test.cc b/src/tint/resolver/uniformity_test.cc
index 274e5ad..e166a06 100644
--- a/src/tint/resolver/uniformity_test.cc
+++ b/src/tint/resolver/uniformity_test.cc
@@ -8376,6 +8376,101 @@
     }
 }
 
+TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnForStatement_CallInInitializer) {
+    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"(for (var b = (non_uniform == 42 && dpdx(1.0) > 0.0); false;) {
+  }
+}
+)";
+
+    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, AttributeOnForStatement_CallInCondition) {
+    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"(for (; non_uniform == 42 && dpdx(1.0) > 0.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, AttributeOnForStatement_CallInIncrement) {
+    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"(for (var b = false; false; b = (non_uniform == 42 && dpdx(1.0) > 0.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, AttributeOnForStatement_CallInBody) {
+    auto& param = GetParam();
+    utils::StringStream ss;
+    ss << R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+@group(0) @binding(1) var t : texture_2d<f32>;
+@group(0) @binding(2) var s : sampler;
+fn foo() {
+  )"
+       << "@diagnostic(" << param << ", derivative_uniformity)"
+       << R"(for (; non_uniform == 42;) {
+    let color = textureSample(t, s, vec2(0, 0));
+  }
+}
+)";
+
+    RunTest(ss.str(), param != builtin::DiagnosticSeverity::kError);
+    if (param == builtin::DiagnosticSeverity::kOff) {
+        EXPECT_TRUE(error_.empty());
+    } else {
+        utils::StringStream err;
+        err << ToStr(param) << ": 'textureSample' must only be called";
+        EXPECT_THAT(error_, ::testing::HasSubstr(err.str()));
+    }
+}
+
 TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnIfStatement_CallInCondition) {
     auto& param = GetParam();
     utils::StringStream ss;
diff --git a/src/tint/sem/diagnostic_severity_test.cc b/src/tint/sem/diagnostic_severity_test.cc
index b487663..a22b43e 100644
--- a/src/tint/sem/diagnostic_severity_test.cc
+++ b/src/tint/sem/diagnostic_severity_test.cc
@@ -52,6 +52,12 @@
         //     }
         //
         //     @diagnostic(error, chromium_unreachable_code)
+        //     for (var i = 0; false; i++) @diagnostic(warning, chromium_unreachable_code) {
+        //       return;
+        //     }
+        //   }
+        //
+        //     @diagnostic(error, chromium_unreachable_code)
         //     while (false) @diagnostic(warning, chromium_unreachable_code) {
         //       return;
         //     }
@@ -69,6 +75,8 @@
         auto else_body_severity = builtin::DiagnosticSeverity::kInfo;
         auto switch_severity = builtin::DiagnosticSeverity::kError;
         auto case_severity = builtin::DiagnosticSeverity::kWarning;
+        auto for_severity = builtin::DiagnosticSeverity::kError;
+        auto for_body_severity = builtin::DiagnosticSeverity::kWarning;
         auto while_severity = builtin::DiagnosticSeverity::kError;
         auto while_body_severity = builtin::DiagnosticSeverity::kWarning;
         auto attr = [&](auto severity) {
@@ -81,6 +89,7 @@
         auto* return_foo_block = Return();
         auto* return_foo_case = Return();
         auto* return_foo_default = Return();
+        auto* return_foo_for = 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));
@@ -90,10 +99,13 @@
             Case(CaseSelector(0_a), Block(utils::Vector{return_foo_case}, attr(case_severity)));
         auto* swtch = Switch(42_a, utils::Vector{case_stmt, DefaultCase(Block(return_foo_default))},
                              attr(switch_severity));
+        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* 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, wl}, attr(block_severity));
+            Block(utils::Vector{if_foo, return_foo_block, swtch, fl, 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});
 
@@ -121,6 +133,12 @@
         EXPECT_EQ(p.Sem().DiagnosticSeverity(case_stmt->body, rule), case_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_case, rule), case_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_default, rule), switch_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(fl, rule), while_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(fl->initializer, rule), for_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(fl->condition, rule), for_severity);
+        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(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);
diff --git a/src/tint/writer/wgsl/generator_impl.cc b/src/tint/writer/wgsl/generator_impl.cc
index 7446d35..0421616 100644
--- a/src/tint/writer/wgsl/generator_impl.cc
+++ b/src/tint/writer/wgsl/generator_impl.cc
@@ -1059,6 +1059,14 @@
 
     {
         auto out = line();
+
+        if (!stmt->attributes.IsEmpty()) {
+            if (!EmitAttributes(out, stmt->attributes)) {
+                return false;
+            }
+            out << " ";
+        }
+
         out << "for";
         {
             ScopedParen sp(out);
diff --git a/test/tint/diagnostic_filtering/for_loop_attribute.wgsl b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl
new file mode 100644
index 0000000..e46c2c9
--- /dev/null
+++ b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl
@@ -0,0 +1,7 @@
+@fragment
+fn main(@location(0) x : f32) {
+  var v = vec4<f32>(0);
+  @diagnostic(warning, derivative_uniformity)
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+  }
+}
diff --git a/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.dxc.hlsl b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.dxc.hlsl
new file mode 100644
index 0000000..560765d
--- /dev/null
+++ b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.dxc.hlsl
@@ -0,0 +1,33 @@
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 warning: 'dpdx' must only be called from uniform control flow
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:3 note: control flow depends on possibly non-uniform value
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+  ^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 note: return value of 'dpdx' may be non-uniform
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+struct tint_symbol_1 {
+  float x : TEXCOORD0;
+};
+
+void main_inner(float x) {
+  float4 v = (0.0f).xxxx;
+  {
+    while (true) {
+      bool tint_tmp = (x > v.x);
+      if (tint_tmp) {
+        tint_tmp = (ddx(1.0f) > 0.0f);
+      }
+      if (!((tint_tmp))) { break; }
+    }
+  }
+}
+
+void main(tint_symbol_1 tint_symbol) {
+  main_inner(tint_symbol.x);
+  return;
+}
diff --git a/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.fxc.hlsl b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.fxc.hlsl
new file mode 100644
index 0000000..560765d
--- /dev/null
+++ b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.fxc.hlsl
@@ -0,0 +1,33 @@
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 warning: 'dpdx' must only be called from uniform control flow
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:3 note: control flow depends on possibly non-uniform value
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+  ^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 note: return value of 'dpdx' may be non-uniform
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+struct tint_symbol_1 {
+  float x : TEXCOORD0;
+};
+
+void main_inner(float x) {
+  float4 v = (0.0f).xxxx;
+  {
+    while (true) {
+      bool tint_tmp = (x > v.x);
+      if (tint_tmp) {
+        tint_tmp = (ddx(1.0f) > 0.0f);
+      }
+      if (!((tint_tmp))) { break; }
+    }
+  }
+}
+
+void main(tint_symbol_1 tint_symbol) {
+  main_inner(tint_symbol.x);
+  return;
+}
diff --git a/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.glsl b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.glsl
new file mode 100644
index 0000000..d2565f0
--- /dev/null
+++ b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.glsl
@@ -0,0 +1,33 @@
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 warning: 'dpdx' must only be called from uniform control flow
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:3 note: control flow depends on possibly non-uniform value
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+  ^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 note: return value of 'dpdx' may be non-uniform
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+#version 310 es
+precision highp float;
+
+layout(location = 0) in float x_1;
+void tint_symbol(float x) {
+  vec4 v = vec4(0.0f);
+  {
+    while (true) {
+      bool tint_tmp = (x > v.x);
+      if (tint_tmp) {
+        tint_tmp = (dFdx(1.0f) > 0.0f);
+      }
+      if (!((tint_tmp))) { break; }
+    }
+  }
+}
+
+void main() {
+  tint_symbol(x_1);
+  return;
+}
diff --git a/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.msl b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.msl
new file mode 100644
index 0000000..bf82a9a
--- /dev/null
+++ b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.msl
@@ -0,0 +1,30 @@
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 warning: 'dpdx' must only be called from uniform control flow
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:3 note: control flow depends on possibly non-uniform value
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+  ^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 note: return value of 'dpdx' may be non-uniform
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+#include <metal_stdlib>
+
+using namespace metal;
+struct tint_symbol_2 {
+  float x [[user(locn0)]];
+};
+
+void tint_symbol_inner(float x) {
+  float4 v = float4(0.0f);
+  for(; ((x > v[0]) && (dfdx(1.0f) > 0.0f)); ) {
+  }
+}
+
+fragment void tint_symbol(tint_symbol_2 tint_symbol_1 [[stage_in]]) {
+  tint_symbol_inner(tint_symbol_1.x);
+  return;
+}
+
diff --git a/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.spvasm b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.spvasm
new file mode 100644
index 0000000..302b2a4
--- /dev/null
+++ b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.spvasm
@@ -0,0 +1,81 @@
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 warning: 'dpdx' must only be called from uniform control flow
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:3 note: control flow depends on possibly non-uniform value
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+  ^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 note: return value of 'dpdx' may be non-uniform
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 39
+; 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 %v "v"
+               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
+    %v4float = OpTypeVector %float 4
+         %10 = OpConstantNull %v4float
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+       %uint = OpTypeInt 32 0
+     %uint_0 = OpConstant %uint 0
+%_ptr_Function_float = OpTypePointer Function %float
+       %bool = OpTypeBool
+    %float_1 = OpConstant %float 1
+         %29 = OpConstantNull %float
+         %34 = OpTypeFunction %void
+ %main_inner = OpFunction %void None %4
+          %x = OpFunctionParameter %float
+          %8 = OpLabel
+          %v = OpVariable %_ptr_Function_v4float Function %10
+               OpStore %v %10
+               OpBranch %13
+         %13 = OpLabel
+               OpLoopMerge %14 %15 None
+               OpBranch %16
+         %16 = OpLabel
+         %21 = OpAccessChain %_ptr_Function_float %v %uint_0
+         %22 = OpLoad %float %21
+         %23 = OpFOrdGreaterThan %bool %x %22
+               OpSelectionMerge %25 None
+               OpBranchConditional %23 %26 %25
+         %26 = OpLabel
+         %27 = OpDPdx %float %float_1
+         %30 = OpFOrdGreaterThan %bool %27 %29
+               OpBranch %25
+         %25 = OpLabel
+         %31 = OpPhi %bool %23 %16 %30 %26
+         %17 = OpLogicalNot %bool %31
+               OpSelectionMerge %32 None
+               OpBranchConditional %17 %33 %32
+         %33 = OpLabel
+               OpBranch %14
+         %32 = OpLabel
+               OpBranch %15
+         %15 = OpLabel
+               OpBranch %13
+         %14 = OpLabel
+               OpReturn
+               OpFunctionEnd
+       %main = OpFunction %void None %34
+         %36 = OpLabel
+         %38 = OpLoad %float %x_1
+         %37 = OpFunctionCall %void %main_inner %38
+               OpReturn
+               OpFunctionEnd
diff --git a/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.wgsl b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.wgsl
new file mode 100644
index 0000000..cec593b
--- /dev/null
+++ b/test/tint/diagnostic_filtering/for_loop_attribute.wgsl.expected.wgsl
@@ -0,0 +1,18 @@
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 warning: 'dpdx' must only be called from uniform control flow
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:3 note: control flow depends on possibly non-uniform value
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+  ^^^
+
+diagnostic_filtering/for_loop_attribute.wgsl:5:21 note: return value of 'dpdx' may be non-uniform
+  for (; x > v.x && dpdx(1.0) > 0.0; ) {
+                    ^^^^^^^^^
+
+@fragment
+fn main(@location(0) x : f32) {
+  var v = vec4<f32>(0);
+  @diagnostic(warning, derivative_uniformity) for(; ((x > v.x) && (dpdx(1.0) > 0.0)); ) {
+  }
+}