aleph-oh · aleph-oh · Jul 4, 2024 · Jul 4, 2024 · Jul 4, 2024 · Jul 4, 2024
diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
@@ -1552,6 +1552,7 @@ pub enum ExprKind {
 pub enum ForLoopKind {
     For,
     ForAwait,
+    CilkFor,
 }
 
 /// Used to differentiate between `async {}` blocks and `gen {}` blocks.

diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
@@ -1606,10 +1606,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
         };
 
         // Some(<pat>) => <body>,
+        // For cilk_for we instead perform Some(<pat>) => cilk_spawn <body>.
         let some_arm = {
             let some_pat = self.pat_some(pat_span, pat);
             let body_block = self.with_loop_scope(e.id, |this| this.lower_block(body, false));
-            let body_expr = self.arena.alloc(self.expr_block(body_block));
+            let body_expr = if matches!(loop_kind, ForLoopKind::CilkFor) {
+                self.arena.alloc(self.expr_spawn_block(body_block))
+            } else {
+                self.arena.alloc(self.expr_block(body_block))
+            };
             self.arm(some_pat, body_expr)
         };
 
@@ -1621,7 +1626,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let match_expr = {
             let iter = self.expr_ident(head_span, iter, iter_pat_nid);
             let next_expr = match loop_kind {
-                ForLoopKind::For => {
+                ForLoopKind::For | ForLoopKind::CilkFor => {
                     // `Iterator::next(&mut iter)`
                     let ref_mut_iter = self.expr_mut_addr_of(head_span, iter);
                     self.expr_call_lang_item_fn(
@@ -1660,10 +1665,17 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let loop_block = self.block_all(for_span, arena_vec![self; match_stmt], None);
 
         // `[opt_ident]: loop { ... }`
+        // If we have a cilk_for we want to indicate this in the loop so the header can be
+        // annotated correctly later.
+        let loop_source = if matches!(loop_kind, ForLoopKind::CilkFor) {
+            hir::LoopSource::CilkFor
+        } else {
+            hir::LoopSource::ForLoop
+        };
         let kind = hir::ExprKind::Loop(
             loop_block,
             self.lower_label(opt_label),
-            hir::LoopSource::ForLoop,
+            loop_source,
             self.lower_span(for_span.with_hi(head.span.hi())),
         );
         let loop_expr =
@@ -1700,6 +1712,17 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 let iter = self.arena.alloc(self.expr_unsafe(iter));
                 iter
             }
+            ForLoopKind::CilkFor => {
+                // FIXME(jhilton): this should call something for converting to a random-access iterator
+                // or otherwise semantically check that the expression is a random-access iterator (or
+                // only work on ranges for now).
+                // `::std::iter::IntoIterator::into_iter(<head>)`
+                self.expr_call_lang_item_fn(
+                    head_span,
+                    hir::LangItem::IntoIterIntoIter,
+                    arena_vec![self; head],
+                )
+            }
         };
 
         let match_expr = self.arena.alloc(self.expr_match(
@@ -2059,6 +2082,11 @@ impl<'hir> LoweringContext<'_, 'hir> {
         self.expr(b.span, hir::ExprKind::Block(b, None))
     }
 
+    fn expr_spawn_block(&mut self, b: &'hir hir::Block<'hir>) -> hir::Expr<'hir> {
+        let block = self.arena.alloc(self.expr_block(b));
+        self.expr(b.span, hir::ExprKind::CilkSpawn(block))
+    }
+
     pub(super) fn expr_array_ref(
         &mut self,
         span: Span,

diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -186,10 +186,8 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         }
     }
 
-    fn br(&mut self, dest: &'ll BasicBlock) {
-        unsafe {
-            llvm::LLVMBuildBr(self.llbuilder, dest);
-        }
+    fn br(&mut self, dest: &'ll BasicBlock) -> &'ll Value {
+        unsafe { llvm::LLVMBuildBr(self.llbuilder, dest) }
     }
 
     fn cond_br(
@@ -681,6 +679,42 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         }
     }
 
+    fn tapir_loop_spawn_strategy_metadata(&mut self, branch: &'ll Value) {
+        // NOTE(jhilton): I think we can't use llvm.loop.mustprogress since ranges can be unbounded.
+        const TAPIR_LOOP_SPAWN_STRATEGY: &'static str = "tapir.loop.spawn.strategy";
+        const LOOP_DIVIDE_AND_CONQUER: i32 = 1;
+        unsafe {
+            let temp = llvm::TemporaryMetadataNode::new(self.cx.llcx);
+            // First we need to make the metadata for tapir.loop.spawn.strategy.
+            let metadata_name = llvm::LLVMMDStringInContext2(
+                self.cx.llcx,
+                TAPIR_LOOP_SPAWN_STRATEGY.as_ptr() as *const c_char,
+                TAPIR_LOOP_SPAWN_STRATEGY.as_bytes().len(),
+            );
+            // Now we need to make the value for the divide-and-conquer strategy.
+            let metadata_value = llvm::LLVMValueAsMetadata(self.const_i32(LOOP_DIVIDE_AND_CONQUER));
+            // Bundle them together into a single node.
+            let tapir_metadata = [metadata_name, metadata_value];
+            let tapir_metadata = llvm::LLVMMDNodeInContext2(
+                self.cx.llcx,
+                tapir_metadata.as_ptr(),
+                tapir_metadata.len(),
+            );
+
+            // Now group the loop with all its associated nodes and assign the node to the branch.
+            let v = [&temp, tapir_metadata];
+            let node = llvm::LLVMMDNodeInContext2(self.cx.llcx, v.as_ptr(), v.len());
+            // Make the node self-referential: this is what loops expect.
+            llvm::LLVMRustReplaceMDOperandWith(node, 0, node);
+
+            llvm::LLVMSetMetadata(
+                branch,
+                llvm::MD_loop as c_uint,
+                llvm::LLVMMetadataAsValue(self.cx.llcx, node),
+            )
+        }
+    }
+
     fn store(&mut self, val: &'ll Value, ptr: &'ll Value, align: Align) -> &'ll Value {
         self.store_with_flags(val, ptr, align, MemFlags::empty())
     }

diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -428,6 +428,7 @@ pub enum MetadataType {
     MD_mem_parallel_loop_access = 10,
     MD_nonnull = 11,
     MD_align = 17,
+    MD_loop = 18,
     MD_type = 19,
     MD_vcall_visibility = 28,
     MD_noundef = 29,
@@ -899,6 +900,17 @@ extern "C" {
     pub fn LLVMSetMetadata<'a>(Val: &'a Value, KindID: c_uint, Node: &'a Value);
     pub fn LLVMGlobalSetMetadata<'a>(Val: &'a Value, KindID: c_uint, Metadata: &'a Metadata);
     pub fn LLVMValueAsMetadata(Node: &Value) -> &Metadata;
+    pub fn LLVMRustMDGetTemporary(C: &Context) -> &mut Metadata;
+    pub fn LLVMRustMDDeleteTemporary(Metadata: &mut Metadata);
+    /// SAFETY: there should exist exactly one reference to the value passed as [Metadata],
+    /// unless a self-reference is being created, in which case [NewMetadata] may be equal to
+    /// [Metadata].
+    /// We can't use an exclusive borrow to model this constraint.
+    pub fn LLVMRustReplaceMDOperandWith<'a>(
+        Metadata: &'a Metadata,
+        Index: size_t,
+        NewMetadata: &'a Metadata,
+    );
 
     // Operations on constants of any type
     pub fn LLVMConstNull(Ty: &Type) -> &Value;

diff --git a/compiler/rustc_codegen_llvm/src/llvm/mod.rs b/compiler/rustc_codegen_llvm/src/llvm/mod.rs
@@ -13,6 +13,7 @@ use rustc_data_structures::small_c_str::SmallCStr;
 use rustc_llvm::RustString;
 use std::cell::RefCell;
 use std::ffi::{CStr, CString};
+use std::ops::Deref;
 use std::str::FromStr;
 use std::string::FromUtf8Error;
 
@@ -324,3 +325,33 @@ impl Drop for OperandBundleDef<'_> {
         }
     }
 }
+
+pub struct TemporaryMetadataNode<'a> {
+    raw: core::ptr::NonNull<Metadata>,
+    phantom: core::marker::PhantomData<&'a mut Metadata>,
+}
+
+impl<'a> TemporaryMetadataNode<'a> {
+    pub fn new(llcx: &'a Context) -> Self {
+        let temp = unsafe { ffi::LLVMRustMDGetTemporary(llcx) };
+        Self { raw: temp.into(), phantom: core::marker::PhantomData }
+    }
+}
+
+impl Deref for TemporaryMetadataNode<'_> {
+    type Target = Metadata;
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { self.raw.as_ref() }
+    }
+}
+
+impl Drop for TemporaryMetadataNode<'_> {
+    fn drop(&mut self) {
+        // SAFETY: raw originated from LLVMRustMDGetTemporary so it should be aligned.
+        // Should also be dereferencable, initialized, and non-aliasing since we don't
+        // expose raw outside the module..
+        let raw = unsafe { self.raw.as_mut() };
+        unsafe { ffi::LLVMRustMDDeleteTemporary(raw) }
+    }
+}
diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -33,6 +33,11 @@ enum MergingSucc {
     True,
 }
 
+enum AddParallelLoopMetadata {
+    False,
+    True,
+}
+
 /// Used by `FunctionCx::codegen_terminator` for emitting common patterns
 /// e.g., creating a basic block, calling a function, etc.
 struct TerminatorCodegenHelper<'tcx> {
@@ -124,6 +129,23 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
         bx: &mut Bx,
         target: mir::BasicBlock,
         mergeable_succ: bool,
+    ) -> MergingSucc {
+        self.funclet_br_maybe_add_metadata(
+            fx,
+            bx,
+            target,
+            mergeable_succ,
+            AddParallelLoopMetadata::False,
+        )
+    }
+
+    fn funclet_br_maybe_add_metadata<Bx: BuilderMethods<'a, 'tcx>>(
+        &self,
+        fx: &mut FunctionCx<'a, 'tcx, Bx>,
+        bx: &mut Bx,
+        target: mir::BasicBlock,
+        mergeable_succ: bool,
+        add_parallel_loop_metadata: AddParallelLoopMetadata,
     ) -> MergingSucc {
         let (needs_landing_pad, is_cleanupret) = self.llbb_characteristics(fx, target);
         if mergeable_succ && !needs_landing_pad && !is_cleanupret {
@@ -139,7 +161,10 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
                 // to a trampoline.
                 bx.cleanup_ret(self.funclet(fx).unwrap(), Some(lltarget));
             } else {
-                bx.br(lltarget);
+                let inst = bx.br(lltarget);
+                if matches!(add_parallel_loop_metadata, AddParallelLoopMetadata::True) {
+                    bx.tapir_loop_spawn_strategy_metadata(inst);
+                }
             }
             MergingSucc::False
         }
@@ -1224,7 +1249,24 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             }
 
             mir::TerminatorKind::Goto { target } => {
-                helper.funclet_br(self, bx, target, mergeable_succ())
+                let add_tapir_metadata = if self.parallel_back_edges.contains(bb) {
+                    AddParallelLoopMetadata::True
+                } else {
+                    AddParallelLoopMetadata::False
+                };
+                // NOTE(jhilton): we attach the metadata to the parallel loop back edge.
+                // This is because in the structure of a loop, the back edge is more
+                // fundamental than the header and because the way LLVM detects loop metadata
+                // is by canonicalizing then checking the back edge. By adding the metadata to
+                // the back edge, we don't rely on the canonicalization succeeding, so it's a
+                // little less fragile (although I expect the canonicalization is very robust).
+                helper.funclet_br_maybe_add_metadata(
+                    self,
+                    bx,
+                    target,
+                    mergeable_succ(),
+                    add_tapir_metadata,
+                )
             }
 
             mir::TerminatorKind::SwitchInt { ref discr, ref targets } => {

diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs
@@ -119,6 +119,10 @@ pub struct FunctionCx<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> {
     /// we add a statement to compute the sync region.
     sync_region: Option<Bx::Value>,
 
+    /// We need to know the basic blocks terminated by back-edges that go into
+    /// parallel loop headers. This is so we can annotate with the right metadata.
+    parallel_back_edges: BitSet<mir::BasicBlock>,
+
     /// A stack of values returned from `tapir_runtime_start` for use in their corresponding `tapir_runtime_stop`
     /// call. We might need a more complex data structure than this if the token should ever be reused, but it's
     /// my impression that that isn't the case.
@@ -221,6 +225,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         per_local_var_debug_info: None,
         caller_location: None,
         sync_region: None,
+        parallel_back_edges: BitSet::new_empty(mir.basic_blocks.len()),
         runtime_hint_stack: SmallVec::new(),
     };
 
@@ -282,9 +287,30 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         })
     };
 
+    let parallel_back_edges = || {
+        let mut seen = rustc_data_structures::fx::FxHashSet::default();
+        mir::traversal::reverse_postorder(mir).filter_map(move |(bb, bb_data)| {
+            seen.insert(bb);
+            if let mir::TerminatorKind::Goto { target } = bb_data.terminator().kind {
+                // If the target of the jump is a parallel loop header and we've already observed
+                // it, we know this must be a loop back-edge.
+                if mir.basic_blocks[target].is_parallel_loop_header && seen.contains(&target) {
+                    return Some(bb);
+                }
+            }
+            None
+        })
+    };
+
     if Bx::supports_tapir() && uses_cilk_control_flow() {
         // Add a sync region at the top of the function, so we can use it later.
         fx.sync_region = Some(start_bx.sync_region_start());
+
+        // Let's figure out the parallel back-edges. These are edges into parallel
+        // loop headers.
+        parallel_back_edges().for_each(|bb| {
+            fx.parallel_back_edges.insert(bb);
+        });
     }
 
     // The builders will be created separately for each basic block at `codegen_block`.

diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs
@@ -58,7 +58,7 @@ pub trait BuilderMethods<'a, 'tcx>:
 
     fn ret_void(&mut self);
     fn ret(&mut self, v: Self::Value);
-    fn br(&mut self, dest: Self::BasicBlock);
+    fn br(&mut self, dest: Self::BasicBlock) -> Self::Value;
     fn cond_br(
         &mut self,
         cond: Self::Value,
@@ -170,6 +170,7 @@ pub trait BuilderMethods<'a, 'tcx>:
 
     fn range_metadata(&mut self, load: Self::Value, range: WrappingRange);
     fn nonnull_metadata(&mut self, load: Self::Value);
+    fn tapir_loop_spawn_strategy_metadata(&mut self, branch: Self::Value);
 
     fn store(&mut self, val: Self::Value, ptr: Self::Value, align: Align) -> Self::Value;
     fn store_with_flags(

diff --git a/compiler/rustc_hir/src/hir.rs b/compiler/rustc_hir/src/hir.rs
@@ -2055,6 +2055,8 @@ pub enum LoopSource {
     While,
     /// A `for _ in _ { .. }` loop.
     ForLoop,
+    /// A `cilk_for _ in _ { .. }` loop.
+    CilkFor,
 }
 
 impl LoopSource {
@@ -2063,6 +2065,7 @@ impl LoopSource {
             LoopSource::Loop => "loop",
             LoopSource::While => "while",
             LoopSource::ForLoop => "for",
+            LoopSource::CilkFor => "cilk_for",
         }
     }
 }

diff --git a/compiler/rustc_hir_typeck/src/expr.rs b/compiler/rustc_hir_typeck/src/expr.rs
@@ -1296,14 +1296,16 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         expected: Expectation<'tcx>,
         expr: &'tcx hir::Expr<'tcx>,
     ) -> Ty<'tcx> {
+        // FIXME(jhilton): might be a good place to enforce the semantic rules for a cilk_for.
+        // Not sure how to implement `continue` here, disallowing is the conservative option.
         let coerce = match source {
             // you can only use break with a value from a normal `loop { }`
             hir::LoopSource::Loop => {
                 let coerce_to = expected.coercion_target_type(self, body.span);
                 Some(CoerceMany::new(coerce_to))
             }
 
-            hir::LoopSource::While | hir::LoopSource::ForLoop => None,
+            hir::LoopSource::While | hir::LoopSource::ForLoop | hir::LoopSource::CilkFor => None,
         };
 
         let ctxt = BreakableCtxt {