use std::collections::HashMap;
use std::fmt;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JThreadState {
Unstarted,
Running,
StopRequested,
Finished,
Joined,
}
impl fmt::Display for JThreadState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Unstarted => write!(f, "unstarted"),
Self::Running => write!(f, "running"),
Self::StopRequested => write!(f, "stop_requested"),
Self::Finished => write!(f, "finished"),
Self::Joined => write!(f, "joined"),
}
}
}
#[derive(Debug, Clone)]
pub struct JThreadDescriptor {
pub callable: String,
pub args: Vec<String>,
pub has_stop_token: bool,
pub thread_id: Option<u64>,
pub state: JThreadState,
}
impl JThreadDescriptor {
pub fn new(callable: &str) -> Self {
Self {
callable: callable.to_string(),
args: Vec::new(),
has_stop_token: false,
thread_id: None,
state: JThreadState::Unstarted,
}
}
pub fn emit_create(&self) -> String {
let st = if self.has_stop_token {
", ptr %stop_token"
} else {
""
};
format!(
"call i32 @pthread_create(ptr %thread, ptr null, ptr @{}, ptr %args{})",
self.callable, st
)
}
pub fn emit_join(&self) -> String {
"call i32 @pthread_join(i64 %tid, ptr null)".to_string()
}
pub fn emit_request_stop(&self) -> String {
"call void @stop_source_request_stop(ptr %stop_source)".to_string()
}
pub fn emit_stop_requested_check(&self) -> String {
"call i1 @stop_token_stop_requested(ptr %stop_token)".to_string()
}
}
pub struct JThreadBuilder {
callable: String,
args: Vec<String>,
with_stop_token: bool,
}
impl JThreadBuilder {
pub fn new(callable: &str) -> Self {
Self {
callable: callable.to_string(),
args: Vec::new(),
with_stop_token: false,
}
}
pub fn with_arg(mut self, arg: &str) -> Self {
self.args.push(arg.to_string());
self
}
pub fn with_stop_token(mut self) -> Self {
self.with_stop_token = true;
self
}
pub fn build(self) -> JThreadDescriptor {
JThreadDescriptor {
callable: self.callable,
args: self.args,
has_stop_token: self.with_stop_token,
thread_id: None,
state: JThreadState::Unstarted,
}
}
}
#[derive(Debug, Clone)]
pub struct StopSource {
pub stop_requested: bool,
pub callbacks: Vec<StopCallback>,
pub id: u64,
}
impl StopSource {
pub fn new(id: u64) -> Self {
Self {
stop_requested: false,
callbacks: Vec::new(),
id,
}
}
pub fn request_stop(&mut self) {
self.stop_requested = true;
}
pub fn get_token(&self) -> StopToken {
StopToken {
source_id: self.id,
stop_requested: self.stop_requested,
}
}
pub fn stop_possible(&self) -> bool {
true
}
}
#[derive(Debug, Clone, Copy)]
pub struct StopToken {
pub source_id: u64,
pub stop_requested: bool,
}
impl StopToken {
pub fn stop_requested(&self) -> bool {
self.stop_requested
}
pub fn stop_possible(&self) -> bool {
true
}
}
#[derive(Debug, Clone)]
pub struct StopCallback {
pub function: String,
pub execute_on_request: bool,
}
impl StopCallback {
pub fn new(function: &str) -> Self {
Self {
function: function.to_string(),
execute_on_request: true,
}
}
pub fn register_on(&self, _token: &StopToken) {}
}
#[derive(Debug)]
pub struct StopState {
stop_requested: std::sync::atomic::AtomicBool,
}
impl StopState {
pub fn new() -> Self {
Self {
stop_requested: std::sync::atomic::AtomicBool::new(false),
}
}
pub fn request_stop(&self) -> bool {
self.stop_requested
.compare_exchange(
false,
true,
std::sync::atomic::Ordering::AcqRel,
std::sync::atomic::Ordering::Acquire,
)
.is_ok()
}
pub fn stop_requested(&self) -> bool {
self.stop_requested
.load(std::sync::atomic::Ordering::Acquire)
}
}
impl Default for StopState {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct Latch {
pub expected: usize,
pub current: usize,
pub max: usize,
}
impl Latch {
pub fn new(expected: usize) -> Self {
Self {
expected,
current: expected,
max: expected,
}
}
pub fn count_down(&mut self, n: usize) {
if n >= self.current {
self.current = 0;
} else {
self.current -= n;
}
}
pub fn try_wait(&self) -> bool {
self.current == 0
}
pub fn wait(&self) -> bool {
self.current == 0
}
pub fn arrive_and_wait(&mut self) {
self.count_down(1);
}
pub fn max_val(&self) -> usize {
self.max
}
pub fn emit_create(&self) -> String {
format!(
"call ptr @latch_create(i64 {})\n store i64 {}, ptr %latch_counter",
self.expected, self.expected
)
}
pub fn emit_count_down(&self, n: usize) -> String {
format!("call void @latch_count_down(ptr %latch, i64 {})", n)
}
pub fn emit_wait(&self) -> String {
"call void @latch_wait(ptr %latch)".to_string()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BarrierArrivalToken {
pub phase: u64,
}
impl BarrierArrivalToken {
pub fn new(phase: u64) -> Self {
Self { phase }
}
}
#[derive(Debug, Clone)]
pub struct Barrier {
pub expected: usize,
pub phase: u64,
pub arrival_count: usize,
pub completion_function: Option<String>,
}
impl Barrier {
pub fn new(expected: usize) -> Self {
Self {
expected,
phase: 0,
arrival_count: 0,
completion_function: None,
}
}
pub fn with_completion(expected: usize, completion: &str) -> Self {
Self {
expected,
phase: 0,
arrival_count: 0,
completion_function: Some(completion.to_string()),
}
}
pub fn arrive(&mut self) -> BarrierArrivalToken {
self.arrival_count += 1;
let current_phase = self.phase;
if self.arrival_count >= self.expected {
self.phase += 1;
self.arrival_count = 0;
}
BarrierArrivalToken::new(current_phase)
}
pub fn wait(&self, _token: BarrierArrivalToken) {}
pub fn arrive_and_wait(&mut self) {
let t = self.arrive();
self.wait(t);
}
pub fn arrive_and_drop(&mut self) {
self.arrive();
}
pub fn emit_create(&self) -> String {
let comp = self
.completion_function
.as_ref()
.map(|f| format!(", ptr @{}", f))
.unwrap_or_default();
format!("call ptr @barrier_create(i64 {}{})", self.expected, comp)
}
pub fn emit_arrive(&self) -> String {
"call i64 @barrier_arrive(ptr %barrier)".to_string()
}
pub fn emit_wait(&self) -> String {
"call void @barrier_wait(ptr %barrier, i64 %phase)".to_string()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SemaphoreKind {
Counting(usize),
Binary,
}
impl fmt::Display for SemaphoreKind {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Counting(n) => write!(f, "counting_semaphore<{}>", n),
Self::Binary => write!(f, "binary_semaphore"),
}
}
}
#[derive(Debug, Clone)]
pub struct Semaphore {
pub kind: SemaphoreKind,
pub counter: usize,
pub max: usize,
}
impl Semaphore {
pub fn new_counting(max: usize, initial: usize) -> Self {
Self {
kind: SemaphoreKind::Counting(max),
counter: initial.min(max),
max,
}
}
pub fn new_binary(initial: usize) -> Self {
let val = if initial > 0 { 1 } else { 0 };
Self {
kind: SemaphoreKind::Binary,
counter: val,
max: 1,
}
}
pub fn acquire(&mut self) {
if self.counter > 0 {
self.counter -= 1;
}
}
pub fn try_acquire(&mut self) -> bool {
if self.counter > 0 {
self.counter -= 1;
true
} else {
false
}
}
pub fn release(&mut self, update: usize) {
self.counter = (self.counter.saturating_add(update)).min(self.max);
}
pub fn max_val(&self) -> usize {
self.max
}
pub fn emit_acquire(&self) -> String {
"call void @semaphore_acquire(ptr %sem)".to_string()
}
pub fn emit_try_acquire(&self) -> String {
"call i1 @semaphore_try_acquire(ptr %sem)".to_string()
}
pub fn emit_release(&self, update: usize) -> String {
format!("call void @semaphore_release(ptr %sem, i64 {})", update)
}
}
#[derive(Debug, Clone)]
pub struct SyncBuf {
pub wrapped: String,
pub emit_atomically: bool,
buffer: Vec<String>,
}
impl SyncBuf {
pub fn new(wrapped: &str) -> Self {
Self {
wrapped: wrapped.to_string(),
emit_atomically: true,
buffer: Vec::new(),
}
}
pub fn emit(&mut self, s: &str) {
self.buffer.push(s.to_string());
}
pub fn emit_all(&mut self) {
self.buffer.clear();
}
pub fn has_buffered(&self) -> bool {
!self.buffer.is_empty()
}
pub fn buffered_count(&self) -> usize {
self.buffer.len()
}
}
#[derive(Debug, Clone)]
pub struct OSyncStream {
pub buffer: String,
pub target: String,
}
impl OSyncStream {
pub fn new(target: &str) -> Self {
Self {
buffer: String::new(),
target: target.to_string(),
}
}
pub fn write(&mut self, data: &str) {
self.buffer.push_str(data);
}
pub fn emit(&mut self) {
self.buffer.clear();
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AtomicRefAlignment {
Default,
Explicit(usize),
}
#[derive(Debug, Clone)]
pub struct AtomicRef {
pub value_type: String,
pub required_alignment: usize,
pub is_lock_free: bool,
pub memory_order: AtomicMemoryOrder,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AtomicMemoryOrder {
Relaxed,
Consume,
Acquire,
Release,
AcqRel,
SeqCst,
}
impl AtomicMemoryOrder {
pub fn as_llvm_str(&self) -> &'static str {
match self {
Self::Relaxed => "monotonic",
Self::Consume => "consume",
Self::Acquire => "acquire",
Self::Release => "release",
Self::AcqRel => "acq_rel",
Self::SeqCst => "seq_cst",
}
}
pub fn from_str(s: &str) -> Option<Self> {
match s {
"memory_order_relaxed" | "relaxed" => Some(Self::Relaxed),
"memory_order_consume" | "consume" => Some(Self::Consume),
"memory_order_acquire" | "acquire" => Some(Self::Acquire),
"memory_order_release" | "release" => Some(Self::Release),
"memory_order_acq_rel" | "acq_rel" => Some(Self::AcqRel),
"memory_order_seq_cst" | "seq_cst" => Some(Self::SeqCst),
_ => None,
}
}
}
impl fmt::Display for AtomicMemoryOrder {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Relaxed => write!(f, "memory_order_relaxed"),
Self::Consume => write!(f, "memory_order_consume"),
Self::Acquire => write!(f, "memory_order_acquire"),
Self::Release => write!(f, "memory_order_release"),
Self::AcqRel => write!(f, "memory_order_acq_rel"),
Self::SeqCst => write!(f, "memory_order_seq_cst"),
}
}
}
impl AtomicRef {
pub fn new(value_type: &str) -> Self {
let size = match value_type {
"int" | "unsigned" => 4,
"long" | "unsigned long" | "long long" | "unsigned long long" => 8,
"short" | "unsigned short" => 2,
"char" | "unsigned char" | "signed char" => 1,
_ => 4,
};
Self {
value_type: value_type.to_string(),
required_alignment: size,
is_lock_free: true,
memory_order: AtomicMemoryOrder::SeqCst,
}
}
pub fn validate_alignment(&self, actual_alignment: usize) -> Result<(), String> {
if actual_alignment < self.required_alignment {
return Err(format!(
"atomic_ref<{}> requires alignment >= {}, got {}",
self.value_type, self.required_alignment, actual_alignment
));
}
Ok(())
}
pub fn emit_load(&self, ptr: &str, order: AtomicMemoryOrder) -> String {
format!(
"load atomic {} {} {}, align {}",
self.value_type,
ptr,
order.as_llvm_str(),
self.required_alignment
)
}
pub fn emit_store(&self, ptr: &str, val: &str, order: AtomicMemoryOrder) -> String {
format!(
"store atomic {} {}, {} {}, align {}",
self.value_type,
val,
ptr,
order.as_llvm_str(),
self.required_alignment
)
}
pub fn emit_cas(&self, ptr: &str, expected: &str, desired: &str) -> String {
format!(
"cmpxchg atomic {} {}, {} {}, {} {} aquire, align {}",
self.value_type,
ptr,
self.value_type,
expected,
self.value_type,
desired,
self.required_alignment,
)
}
pub fn emit_fetch_add(&self, ptr: &str, val: &str) -> String {
format!(
"atomicrmw add {} {}, {} acq_rel, align {}",
self.value_type, ptr, val, self.required_alignment
)
}
}
#[derive(Debug, Clone)]
pub struct AtomicWaitOperations {
pub atomic_type: String,
}
impl AtomicWaitOperations {
pub fn new(atomic_type: &str) -> Self {
Self {
atomic_type: atomic_type.to_string(),
}
}
pub fn emit_wait(&self, ptr: &str, old_val: &str) -> String {
format!("call void @__cxa_atomic_wait(ptr {}, i64 {})", ptr, old_val)
}
pub fn emit_notify_one(&self, ptr: &str) -> String {
format!("call void @__cxa_atomic_notify_one(ptr {})", ptr)
}
pub fn emit_notify_all(&self, ptr: &str) -> String {
format!("call void @__cxa_atomic_notify_all(ptr {})", ptr)
}
pub fn supports_waiting(&self) -> bool {
true
}
}
#[derive(Debug, Clone)]
pub struct GeneratorDescriptor {
pub yield_type: String,
pub promise_type: String,
pub symmetric_transfer: bool,
}
impl GeneratorDescriptor {
pub fn new(yield_type: &str) -> Self {
Self {
yield_type: yield_type.to_string(),
promise_type: format!("std::generator<{}>::promise_type", yield_type),
symmetric_transfer: true,
}
}
pub fn emit_frame_alloca(&self) -> String {
"call ptr @llvm.coro.begin(ptr null, ptr %promise)".to_string()
}
pub fn emit_yield(&self, value: &str) -> String {
format!(
"call void @llvm.coro.suspend(ptr %frame, i1 true)\n store {} {}, ptr %frame.value",
self.yield_type, value
)
}
pub fn emit_return(&self) -> String {
"call i1 @llvm.coro.end(ptr %frame, i1 false)\n ret void".to_string()
}
}
#[derive(Debug, Clone)]
pub struct AsyncTaskDescriptor {
pub return_type: String,
pub awaitables: Vec<String>,
pub is_lazy: bool,
}
impl AsyncTaskDescriptor {
pub fn new(return_type: &str) -> Self {
Self {
return_type: return_type.to_string(),
awaitables: Vec::new(),
is_lazy: false,
}
}
pub fn with_awaitable(mut self, awaitable: &str) -> Self {
self.awaitables.push(awaitable.to_string());
self
}
pub fn emit_co_await(&self, expr: &str) -> String {
format!("call ptr @await_suspend(ptr %frame, ptr {})\n call void @llvm.coro.suspend(ptr %frame, i1 false)", expr)
}
pub fn emit_co_return(&self, value: &str) -> String {
format!("store {} {}, ptr %frame.result\n call i1 @llvm.coro.end(ptr %frame, i1 false)\n ret void",
self.return_type, value)
}
}
#[derive(Debug, Clone)]
pub struct CoroutineHandle {
pub frame_ptr: String,
pub done: bool,
}
impl CoroutineHandle {
pub fn new() -> Self {
Self {
frame_ptr: String::new(),
done: false,
}
}
pub fn resume(&mut self) {}
pub fn destroy(&self) {}
}
impl Default for CoroutineHandle {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OMPDirectiveKind {
Parallel,
ParallelFor,
ParallelForSimd,
For,
ForSimd,
Simd,
Single,
Master,
Critical,
Barrier,
Task,
TaskLoop,
Sections,
Section,
Ordered,
Atomic,
Target,
TargetParallelFor,
Teams,
Distribute,
DistributeParallelFor,
}
impl OMPDirectiveKind {
pub fn from_string(s: &str) -> Option<Self> {
match s.trim() {
"parallel" => Some(Self::Parallel),
"parallel for" => Some(Self::ParallelFor),
"parallel for simd" => Some(Self::ParallelForSimd),
"for" => Some(Self::For),
"for simd" => Some(Self::ForSimd),
"simd" => Some(Self::Simd),
"single" => Some(Self::Single),
"master" => Some(Self::Master),
"critical" => Some(Self::Critical),
"barrier" => Some(Self::Barrier),
"task" => Some(Self::Task),
"taskloop" => Some(Self::TaskLoop),
"sections" => Some(Self::Sections),
"section" => Some(Self::Section),
"ordered" => Some(Self::Ordered),
"atomic" => Some(Self::Atomic),
"target" => Some(Self::Target),
"target parallel for" => Some(Self::TargetParallelFor),
"teams" => Some(Self::Teams),
"distribute" => Some(Self::Distribute),
"distribute parallel for" => Some(Self::DistributeParallelFor),
_ => None,
}
}
pub fn runtime_function(&self) -> &'static str {
match self {
Self::Parallel => "__kmpc_fork_call",
Self::ParallelFor | Self::ParallelForSimd | Self::For | Self::ForSimd => {
"__kmpc_for_static_init_4"
}
Self::Simd => "__kmpc_simd",
Self::Barrier => "__kmpc_barrier",
Self::Critical => "__kmpc_critical",
Self::Single => "__kmpc_single",
Self::Master => "__kmpc_master",
Self::Atomic => "__kmpc_atomic",
Self::Ordered => "__kmpc_ordered",
Self::Task => "__kmpc_omp_task",
Self::TaskLoop => "__kmpc_taskloop",
Self::Sections => "__kmpc_sections_init",
Self::Section => "__kmpc_section",
Self::Target => "__tgt_target",
Self::TargetParallelFor => "__tgt_target_teams",
Self::Teams => "__kmpc_teams",
Self::Distribute | Self::DistributeParallelFor => "__kmpc_distribute_static_init_4",
}
}
}
#[derive(Debug, Clone)]
pub struct OMPConfig {
pub directive: OMPDirectiveKind,
pub clauses: Vec<OMPClause>,
pub num_threads: Option<usize>,
pub schedule: OMPSchedule,
}
#[derive(Debug, Clone)]
pub enum OMPClause {
Private(String),
Shared(String),
FirstPrivate(String),
LastPrivate(String),
Reduction(OMPReduction),
Nowait,
Ordered,
Collapse(usize),
If(String),
NumThreads(usize),
Default(OMPDefaultKind),
Copyin(String),
Schedule(OMPSchedule),
Map(OMPMapKind, String),
Depend(OMPDependKind, String),
Aligned(String, usize),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OMPDefaultKind {
Shared,
None,
FirstPrivate,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OMPMapKind {
To,
From,
ToFrom,
Alloc,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OMPDependKind {
In,
Out,
InOut,
}
#[derive(Debug, Clone)]
pub struct OMPReduction {
pub operator: String,
pub variables: Vec<String>,
}
impl OMPReduction {
pub fn new(op: &str, vars: Vec<String>) -> Self {
Self {
operator: op.to_string(),
variables: vars,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OMPScheduleKind {
Static,
Dynamic,
Guided,
Auto,
Runtime,
}
#[derive(Debug, Clone)]
pub struct OMPSchedule {
pub kind: OMPScheduleKind,
pub chunk_size: Option<usize>,
}
impl OMPSchedule {
pub fn new(kind: OMPScheduleKind) -> Self {
Self {
kind,
chunk_size: None,
}
}
pub fn with_chunk(kind: OMPScheduleKind, chunk: usize) -> Self {
Self {
kind,
chunk_size: Some(chunk),
}
}
}
pub struct OMPRuntime {
pub enabled: bool,
}
impl OMPRuntime {
pub fn new() -> Self {
Self { enabled: true }
}
pub fn emit_parallel_begin(&self, outlined_fn: &str, num_threads: Option<usize>) -> String {
let nt = num_threads
.map(|n| format!("i32 {}", n))
.unwrap_or_else(|| "i32 0".to_string());
format!(
"call void @__kmpc_fork_call(ptr @ident, i32 0, ptr @{})\n ; num_threads = {}",
outlined_fn, nt
)
}
pub fn emit_parallel_end(&self) -> String {
"call void @__kmpc_end_parallel(ptr @ident)".to_string()
}
pub fn emit_for_static_init(
&self,
_lower: i32,
_upper: i32,
stride: i32,
chunk: i32,
) -> String {
format!("call void @__kmpc_for_static_init_4(ptr @ident, i32 0, i32 0, ptr %is_last, ptr %lb, ptr %ub, ptr %st, i32 {}, i32 {})", chunk, stride)
}
pub fn emit_for_static_fini(&self) -> String {
"call void @__kmpc_for_static_fini(ptr @ident, i32 0)".to_string()
}
pub fn emit_barrier(&self) -> String {
"call void @__kmpc_barrier(ptr @ident, i32 0)".to_string()
}
pub fn emit_critical_enter(&self, name: &str) -> String {
format!(
"call void @__kmpc_critical(ptr @ident, i32 0, i64 8) ; {}",
name
)
}
pub fn emit_critical_exit(&self, name: &str) -> String {
format!(
"call void @__kmpc_end_critical(ptr @ident, i32 0, i64 8) ; {}",
name
)
}
}
impl Default for OMPRuntime {
fn default() -> Self {
Self::new()
}
}
pub struct OMPLowering {
pub runtime: OMPRuntime,
pub config: OMPConfig,
}
impl OMPLowering {
pub fn new(config: OMPConfig) -> Self {
Self {
runtime: OMPRuntime::new(),
config,
}
}
pub fn lower_parallel_for(&self, loop_body: &str) -> String {
let mut ir = String::new();
ir.push_str("; OpenMP parallel for lowered\n");
ir.push_str(
&self
.runtime
.emit_parallel_begin("__omp_outlined_fn", self.config.num_threads),
);
ir.push_str("\n; outlined fn body\n");
ir.push_str(&self.runtime.emit_for_static_init(0, 100, 1, 1));
ir.push_str("\n");
ir.push_str(loop_body);
ir.push_str("\n");
ir.push_str(&self.runtime.emit_for_static_fini());
ir.push_str("\n");
ir.push_str(&self.runtime.emit_parallel_end());
ir
}
pub fn lower_simd(&self, loop_body: &str) -> String {
format!("; OpenMP simd loop\n ; vectorized body:\n{}", loop_body)
}
}
#[derive(Debug, Clone)]
pub struct SimdType {
pub element_type: String,
pub width: usize,
pub native_type: String,
}
impl SimdType {
pub fn new(element_type: &str, width: usize) -> Self {
Self {
element_type: element_type.to_string(),
width,
native_type: format!("<{} x {}>", width, element_type),
}
}
pub fn emit_add(&self, a: &str, b: &str) -> String {
format!("add <{} x {}> {}, {}", self.width, self.element_type, a, b)
}
pub fn emit_mul(&self, a: &str, b: &str) -> String {
format!("mul <{} x {}> {}, {}", self.width, self.element_type, a, b)
}
pub fn emit_reduce_sum(&self, vec: &str) -> String {
format!(
"call {} @llvm.vector.reduce.add.v{}{}(<{} x {}> {})",
self.element_type, self.width, self.element_type, self.width, self.element_type, vec
)
}
pub fn emit_reduce_min(&self, vec: &str) -> String {
format!(
"call {} @llvm.vector.reduce.smin.v{}{}(<{} x {}> {})",
self.element_type, self.width, self.element_type, self.width, self.element_type, vec
)
}
pub fn emit_reduce_max(&self, vec: &str) -> String {
format!(
"call {} @llvm.vector.reduce.smax.v{}{}(<{} x {}> {})",
self.element_type, self.width, self.element_type, self.width, self.element_type, vec
)
}
pub fn emit_masked_load(&self, ptr: &str, mask: &str, passthru: &str) -> String {
format!(
"call <{} x {}> @llvm.masked.load.v{}{}(ptr {}, i32 1, <{} x i1> {}, <{} x {}> {})",
self.width,
self.element_type,
self.width,
self.element_type,
ptr,
self.width,
mask,
self.width,
self.element_type,
passthru
)
}
pub fn emit_masked_store(&self, val: &str, ptr: &str, mask: &str) -> String {
format!(
"call void @llvm.masked.store.v{}{}(<{} x {}> {}, ptr {}, i32 1, <{} x i1> {})",
self.width,
self.element_type,
self.width,
self.element_type,
val,
ptr,
self.width,
mask
)
}
pub fn emit_gather(&self, ptr: &str, mask: &str, passthru: &str) -> String {
format!("call <{} x {}> @llvm.masked.gather.v{}{}(<{} x ptr> {}, i32 1, <{} x i1> {}, <{} x {}> {})",
self.width, self.element_type, self.width, self.element_type,
self.width, ptr, self.width, mask, self.width, self.element_type, passthru)
}
pub fn emit_scatter(&self, val: &str, ptr: &str, mask: &str) -> String {
format!("call void @llvm.masked.scatter.v{}{}(<{} x {}> {}, <{} x ptr> {}, i32 1, <{} x i1> {})",
self.width, self.element_type, self.width, self.element_type, val,
self.width, ptr, self.width, mask)
}
}
#[derive(Debug, Clone)]
pub struct SimdRegister {
pub name: String,
pub bit_width: usize,
pub element_types: Vec<String>,
}
impl SimdRegister {
pub fn x86_sse() -> Self {
Self {
name: "__m128".to_string(),
bit_width: 128,
element_types: vec!["float".into(), "int".into(), "double".into()],
}
}
pub fn x86_avx() -> Self {
Self {
name: "__m256".to_string(),
bit_width: 256,
element_types: vec!["float".into(), "int".into(), "double".into()],
}
}
pub fn x86_avx512() -> Self {
Self {
name: "__m512".to_string(),
bit_width: 512,
element_types: vec!["float".into(), "int".into(), "double".into()],
}
}
pub fn neon() -> Self {
Self {
name: "float32x4_t".to_string(),
bit_width: 128,
element_types: vec!["float".into(), "int32".into()],
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExecutionPolicy {
Seq,
Par,
ParUnseq,
Unseq,
}
impl ExecutionPolicy {
pub fn as_str(&self) -> &'static str {
match self {
Self::Seq => "seq",
Self::Par => "par",
Self::ParUnseq => "par_unseq",
Self::Unseq => "unseq",
}
}
pub fn from_str(s: &str) -> Option<Self> {
match s {
"seq" | "sequenced_policy" => Some(Self::Seq),
"par" | "parallel_policy" => Some(Self::Par),
"par_unseq" | "parallel_unsequenced_policy" => Some(Self::ParUnseq),
"unseq" | "unsequenced_policy" => Some(Self::Unseq),
_ => None,
}
}
pub fn is_parallel(&self) -> bool {
matches!(self, Self::Par | Self::ParUnseq)
}
pub fn is_vectorized(&self) -> bool {
matches!(self, Self::ParUnseq | Self::Unseq)
}
}
#[derive(Debug, Clone)]
pub struct ParallelAlgorithm {
pub name: String,
pub policy: ExecutionPolicy,
pub element_type: String,
pub uses_thread_pool: bool,
pub can_vectorize: bool,
}
impl ParallelAlgorithm {
pub fn new(name: &str, policy: ExecutionPolicy, element_type: &str) -> Self {
Self {
name: name.to_string(),
policy,
element_type: element_type.to_string(),
uses_thread_pool: policy.is_parallel(),
can_vectorize: policy.is_vectorized(),
}
}
pub fn emit_parallel_sort(&self, begin: &str, end: &str) -> String {
if self.policy.is_parallel() {
format!(
"call void @__parallel_sort(ptr {}, ptr {}, i64 {})",
begin,
end,
self.element_type.len()
)
} else {
format!("call void @std_sort(ptr {}, ptr {})", begin, end)
}
}
pub fn emit_parallel_transform(
&self,
begin: &str,
end: &str,
out: &str,
op_name: &str,
) -> String {
if self.policy.is_parallel() {
format!(
"call void @__parallel_transform(ptr {}, ptr {}, ptr {}, ptr @{})",
begin, end, out, op_name
)
} else {
format!(
"call void @std_transform(ptr {}, ptr {}, ptr {}, ptr @{})",
begin, end, out, op_name
)
}
}
pub fn emit_parallel_reduce(
&self,
begin: &str,
end: &str,
init: &str,
op_name: &str,
) -> String {
if self.policy.is_parallel() {
format!(
"call {} @__parallel_reduce(ptr {}, ptr {}, {} {}, ptr @{})",
self.element_type, begin, end, self.element_type, init, op_name
)
} else {
format!(
"call {} @std_reduce(ptr {}, ptr {}, {} {}, ptr @{})",
self.element_type, begin, end, self.element_type, init, op_name
)
}
}
pub fn emit_parallel_for_each(&self, begin: &str, end: &str, fn_name: &str) -> String {
if self.policy.is_parallel() {
format!(
"call void @__parallel_for_each(ptr {}, ptr {}, ptr @{})",
begin, end, fn_name
)
} else {
format!(
"call void @std_for_each(ptr {}, ptr {}, ptr @{})",
begin, end, fn_name
)
}
}
}
#[derive(Debug, Clone)]
pub struct ThreadPool {
pub num_threads: usize,
pub tasks: Vec<ThreadPoolTask>,
}
#[derive(Debug, Clone)]
pub struct ThreadPoolTask {
pub id: usize,
pub function: String,
}
impl ThreadPool {
pub fn new(num_threads: usize) -> Self {
Self {
num_threads,
tasks: Vec::new(),
}
}
pub fn submit(&mut self, task: ThreadPoolTask) {
self.tasks.push(task);
}
pub fn emit_thread_pool(&self) -> String {
format!("call ptr @thread_pool_create(i64 {})", self.num_threads)
}
pub fn emit_submit(&self, task_id: usize, fn_name: &str) -> String {
format!(
"call void @thread_pool_submit(ptr %pool, i64 {}, ptr @{})",
task_id, fn_name
)
}
pub fn emit_wait(&self) -> String {
"call void @thread_pool_wait(ptr %pool)".to_string()
}
}
pub struct ConcurrencyRegistry {
enabled: HashMap<String, bool>,
}
impl ConcurrencyRegistry {
pub fn new() -> Self {
let mut enabled = HashMap::new();
for feature in &[
"jthread",
"stop_token",
"latch",
"barrier",
"counting_semaphore",
"binary_semaphore",
"syncbuf",
"atomic_ref",
"atomic_wait",
"coroutine",
"generator",
"openmp",
"simd",
"parallel_algorithms",
] {
enabled.insert(feature.to_string(), true);
}
Self { enabled }
}
pub fn is_enabled(&self, feature: &str) -> bool {
self.enabled.get(feature).copied().unwrap_or(false)
}
pub fn enable(&mut self, feature: &str) {
self.enabled.insert(feature.to_string(), true);
}
pub fn disable(&mut self, feature: &str) {
self.enabled.insert(feature.to_string(), false);
}
pub fn enabled_features(&self) -> Vec<String> {
self.enabled
.iter()
.filter_map(|(k, &v)| if v { Some(k.clone()) } else { None })
.collect()
}
pub fn feature_count(&self) -> usize {
self.enabled.len()
}
}
impl Default for ConcurrencyRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_jthread_creation() {
let jt = JThreadDescriptor::new("worker");
assert_eq!(jt.callable, "worker");
assert_eq!(jt.state, JThreadState::Unstarted);
}
#[test]
fn test_jthread_builder() {
let jt = JThreadBuilder::new("work")
.with_arg("42")
.with_stop_token()
.build();
assert_eq!(jt.callable, "work");
assert_eq!(jt.args.len(), 1);
assert!(jt.has_stop_token);
}
#[test]
fn test_jthread_emit_create() {
let jt = JThreadDescriptor::new("worker");
let ir = jt.emit_create();
assert!(ir.contains("pthread_create"));
assert!(ir.contains("worker"));
}
#[test]
fn test_jthread_emit_join() {
let jt = JThreadDescriptor::new("worker");
assert!(jt.emit_join().contains("pthread_join"));
}
#[test]
fn test_jthread_state_display() {
assert_eq!(JThreadState::Running.to_string(), "running");
assert_eq!(JThreadState::Joined.to_string(), "joined");
}
#[test]
fn test_stop_source_request_stop() {
let mut src = StopSource::new(1);
assert!(!src.stop_requested);
src.request_stop();
assert!(src.stop_requested);
}
#[test]
fn test_stop_source_get_token() {
let src = StopSource::new(42);
let token = src.get_token();
assert_eq!(token.source_id, 42);
assert!(!token.stop_requested());
assert!(token.stop_possible());
}
#[test]
fn test_stop_state_request() {
let state = StopState::new();
assert!(!state.stop_requested());
assert!(state.request_stop());
assert!(state.stop_requested());
assert!(!state.request_stop());
}
#[test]
fn test_stop_callback() {
let cb = StopCallback::new("on_stop");
assert_eq!(cb.function, "on_stop");
assert!(cb.execute_on_request);
}
#[test]
fn test_latch_count_down() {
let mut latch = Latch::new(3);
assert_eq!(latch.current, 3);
latch.count_down(1);
assert_eq!(latch.current, 2);
latch.count_down(2);
assert_eq!(latch.current, 0);
assert!(latch.try_wait());
}
#[test]
fn test_latch_try_wait() {
assert!(!Latch::new(2).try_wait());
}
#[test]
fn test_latch_arrive_and_wait() {
let mut l = Latch::new(1);
l.arrive_and_wait();
assert_eq!(l.current, 0);
}
#[test]
fn test_latch_max() {
assert_eq!(Latch::new(5).max_val(), 5);
}
#[test]
fn test_latch_emit_create() {
assert!(Latch::new(4).emit_create().contains("latch_create"));
}
#[test]
fn test_barrier_arrive() {
let mut b = Barrier::new(2);
let t = b.arrive();
assert_eq!(t.phase, 0);
assert_eq!(b.arrival_count, 1);
}
#[test]
fn test_barrier_new_phase() {
let mut b = Barrier::new(2);
b.arrive();
b.arrive();
assert_eq!(b.phase, 1);
assert_eq!(b.arrival_count, 0);
}
#[test]
fn test_barrier_with_completion() {
let b = Barrier::with_completion(3, "end_phase");
assert_eq!(b.completion_function, Some("end_phase".to_string()));
}
#[test]
fn test_barrier_emit_create() {
assert!(Barrier::new(3).emit_create().contains("barrier_create"));
}
#[test]
fn test_barrier_emit_arrive() {
assert!(Barrier::new(2).emit_arrive().contains("barrier_arrive"));
}
#[test]
fn test_counting_semaphore_acquire() {
let mut s = Semaphore::new_counting(3, 2);
assert_eq!(s.counter, 2);
s.acquire();
assert_eq!(s.counter, 1);
}
#[test]
fn test_binary_semaphore_try_acquire() {
let mut s = Semaphore::new_binary(1);
assert!(s.try_acquire());
assert!(!s.try_acquire());
}
#[test]
fn test_semaphore_release() {
let mut s = Semaphore::new_counting(3, 1);
s.release(2);
assert_eq!(s.counter, 3);
}
#[test]
fn test_semaphore_display() {
assert_eq!(
SemaphoreKind::Counting(5).to_string(),
"counting_semaphore<5>"
);
assert_eq!(SemaphoreKind::Binary.to_string(), "binary_semaphore");
}
#[test]
fn test_syncbuf_emit() {
let mut buf = SyncBuf::new("cout");
buf.emit("hello");
buf.emit(" world");
assert!(buf.has_buffered());
assert_eq!(buf.buffered_count(), 2);
}
#[test]
fn test_syncbuf_emit_all() {
let mut buf = SyncBuf::new("cout");
buf.emit("test");
buf.emit_all();
assert!(!buf.has_buffered());
}
#[test]
fn test_osyncstream_write() {
let mut s = OSyncStream::new("cout");
s.write("hello");
assert!(!s.buffer.is_empty());
}
#[test]
fn test_atomic_ref_new() {
let ar = AtomicRef::new("int");
assert_eq!(ar.value_type, "int");
assert_eq!(ar.required_alignment, 4);
}
#[test]
fn test_atomic_ref_validation() {
let ar = AtomicRef::new("int");
assert!(ar.validate_alignment(4).is_ok());
assert!(ar.validate_alignment(2).is_err());
}
#[test]
fn test_atomic_ref_emit_load() {
let ar = AtomicRef::new("int");
assert!(ar
.emit_load("ptr", AtomicMemoryOrder::Acquire)
.contains("acquire"));
}
#[test]
fn test_atomic_ref_emit_store() {
let ar = AtomicRef::new("int");
assert!(ar
.emit_store("ptr", "42", AtomicMemoryOrder::Release)
.contains("release"));
}
#[test]
fn test_atomic_ref_emit_cas() {
let ar = AtomicRef::new("int");
assert!(ar.emit_cas("ptr", "0", "1").contains("cmpxchg"));
}
#[test]
fn test_atomic_memory_order_from_str() {
assert_eq!(
AtomicMemoryOrder::from_str("acquire"),
Some(AtomicMemoryOrder::Acquire)
);
assert_eq!(AtomicMemoryOrder::from_str("bad"), None);
}
#[test]
fn test_atomic_memory_order_display() {
assert_eq!(
AtomicMemoryOrder::SeqCst.to_string(),
"memory_order_seq_cst"
);
}
#[test]
fn test_atomic_wait_operations() {
let ops = AtomicWaitOperations::new("int");
assert!(ops.emit_wait("ptr", "42").contains("__cxa_atomic_wait"));
assert!(ops
.emit_notify_one("ptr")
.contains("__cxa_atomic_notify_one"));
assert!(ops
.emit_notify_all("ptr")
.contains("__cxa_atomic_notify_all"));
}
#[test]
fn test_generator_creation() {
let g = GeneratorDescriptor::new("int");
assert_eq!(g.yield_type, "int");
assert!(g.symmetric_transfer);
}
#[test]
fn test_generator_emit_yield() {
let g = GeneratorDescriptor::new("int");
assert!(g.emit_yield("42").contains("llvm.coro.suspend"));
}
#[test]
fn test_generator_emit_return() {
let g = GeneratorDescriptor::new("int");
assert!(g.emit_return().contains("llvm.coro.end"));
}
#[test]
fn test_async_task() {
let task = AsyncTaskDescriptor::new("int").with_awaitable("future<int>");
assert_eq!(task.return_type, "int");
assert_eq!(task.awaitables.len(), 1);
}
#[test]
fn test_async_task_emit_co_await() {
let task = AsyncTaskDescriptor::new("void");
assert!(task.emit_co_await("%future").contains("await_suspend"));
}
#[test]
fn test_coroutine_handle_default() {
let h = CoroutineHandle::default();
assert!(!h.done);
}
#[test]
fn test_omp_directive_from_string() {
assert_eq!(
OMPDirectiveKind::from_string("parallel"),
Some(OMPDirectiveKind::Parallel)
);
assert_eq!(
OMPDirectiveKind::from_string("parallel for"),
Some(OMPDirectiveKind::ParallelFor)
);
assert_eq!(
OMPDirectiveKind::from_string("critical"),
Some(OMPDirectiveKind::Critical)
);
assert_eq!(OMPDirectiveKind::from_string("nonexistent"), None);
}
#[test]
fn test_omp_directive_runtime_function() {
assert_eq!(
OMPDirectiveKind::Parallel.runtime_function(),
"__kmpc_fork_call"
);
assert_eq!(
OMPDirectiveKind::Barrier.runtime_function(),
"__kmpc_barrier"
);
assert_eq!(OMPDirectiveKind::Target.runtime_function(), "__tgt_target");
}
#[test]
fn test_omp_runtime_emit_parallel() {
let rt = OMPRuntime::new();
let ir = rt.emit_parallel_begin("myfn", Some(4));
assert!(ir.contains("__kmpc_fork_call"));
assert!(ir.contains("myfn"));
}
#[test]
fn test_omp_runtime_emit_barrier() {
let rt = OMPRuntime::new();
assert!(rt.emit_barrier().contains("__kmpc_barrier"));
}
#[test]
fn test_omp_runtime_emit_critical() {
let rt = OMPRuntime::new();
assert!(rt.emit_critical_enter("lock1").contains("__kmpc_critical"));
assert!(rt
.emit_critical_exit("lock1")
.contains("__kmpc_end_critical"));
}
#[test]
fn test_omp_lowering_parallel_for() {
let config = OMPConfig {
directive: OMPDirectiveKind::ParallelFor,
clauses: vec![],
num_threads: Some(4),
schedule: OMPSchedule::new(OMPScheduleKind::Static),
};
let lower = OMPLowering::new(config);
let ir = lower.lower_parallel_for("call void @do_work()");
assert!(ir.contains("__kmpc_fork_call"));
assert!(ir.contains("do_work"));
}
#[test]
fn test_omp_reduction() {
let red = OMPReduction::new("+", vec!["x".into(), "y".into()]);
assert_eq!(red.operator, "+");
assert_eq!(red.variables.len(), 2);
}
#[test]
fn test_omp_schedule() {
let s = OMPSchedule::with_chunk(OMPScheduleKind::Dynamic, 16);
assert_eq!(s.kind, OMPScheduleKind::Dynamic);
assert_eq!(s.chunk_size, Some(16));
}
#[test]
fn test_simd_type_emit_add() {
let st = SimdType::new("float", 4);
assert!(st.emit_add("%a", "%b").contains("add <4 x float>"));
}
#[test]
fn test_simd_type_emit_reduce_sum() {
let st = SimdType::new("float", 4);
assert!(st.emit_reduce_sum("%v").contains("llvm.vector.reduce.add"));
}
#[test]
fn test_simd_type_emit_masked_load() {
let st = SimdType::new("float", 4);
let ir = st.emit_masked_load("ptr", "%m", "%p");
assert!(ir.contains("llvm.masked.load"));
}
#[test]
fn test_simd_type_emit_masked_store() {
let st = SimdType::new("int", 8);
let ir = st.emit_masked_store("%v", "ptr", "%m");
assert!(ir.contains("llvm.masked.store"));
}
#[test]
fn test_simd_type_emit_gather() {
let st = SimdType::new("float", 4);
assert!(st
.emit_gather("ptrs", "%m", "%p")
.contains("llvm.masked.gather"));
}
#[test]
fn test_simd_type_emit_scatter() {
let st = SimdType::new("float", 4);
assert!(st
.emit_scatter("%v", "ptrs", "%m")
.contains("llvm.masked.scatter"));
}
#[test]
fn test_simd_register_x86() {
let sse = SimdRegister::x86_sse();
assert_eq!(sse.name, "__m128");
assert_eq!(sse.bit_width, 128);
let avx = SimdRegister::x86_avx();
assert_eq!(avx.bit_width, 256);
let avx512 = SimdRegister::x86_avx512();
assert_eq!(avx512.bit_width, 512);
}
#[test]
fn test_simd_register_neon() {
let neon = SimdRegister::neon();
assert_eq!(neon.name, "float32x4_t");
}
#[test]
fn test_execution_policy_from_str() {
assert_eq!(ExecutionPolicy::from_str("par"), Some(ExecutionPolicy::Par));
assert_eq!(ExecutionPolicy::from_str("seq"), Some(ExecutionPolicy::Seq));
assert_eq!(
ExecutionPolicy::from_str("par_unseq"),
Some(ExecutionPolicy::ParUnseq)
);
}
#[test]
fn test_execution_policy_is_parallel() {
assert!(ExecutionPolicy::Par.is_parallel());
assert!(!ExecutionPolicy::Seq.is_parallel());
}
#[test]
fn test_execution_policy_is_vectorized() {
assert!(ExecutionPolicy::ParUnseq.is_vectorized());
assert!(!ExecutionPolicy::Par.is_vectorized());
}
#[test]
fn test_parallel_algorithm_sort() {
let alg = ParallelAlgorithm::new("sort", ExecutionPolicy::Par, "int");
assert!(alg
.emit_parallel_sort("beg", "end")
.contains("__parallel_sort"));
}
#[test]
fn test_parallel_algorithm_seq_sort() {
let alg = ParallelAlgorithm::new("sort", ExecutionPolicy::Seq, "int");
assert!(alg.emit_parallel_sort("beg", "end").contains("std_sort"));
}
#[test]
fn test_parallel_algorithm_transform() {
let alg = ParallelAlgorithm::new("transform", ExecutionPolicy::Par, "float");
assert!(alg
.emit_parallel_transform("b", "e", "o", "op")
.contains("__parallel_transform"));
}
#[test]
fn test_parallel_algorithm_reduce() {
let alg = ParallelAlgorithm::new("reduce", ExecutionPolicy::ParUnseq, "double");
let ir = alg.emit_parallel_reduce("b", "e", "0.0", "plus");
assert!(ir.contains("__parallel_reduce"));
}
#[test]
fn test_parallel_algorithm_for_each() {
let alg = ParallelAlgorithm::new("for_each", ExecutionPolicy::Par, "int");
assert!(alg
.emit_parallel_for_each("b", "e", "work")
.contains("__parallel_for_each"));
}
#[test]
fn test_thread_pool() {
let mut pool = ThreadPool::new(8);
pool.submit(ThreadPoolTask {
id: 1,
function: "task1".into(),
});
assert_eq!(pool.tasks.len(), 1);
assert!(pool.emit_thread_pool().contains("thread_pool_create"));
}
#[test]
fn test_thread_pool_emit() {
let pool = ThreadPool::new(4);
assert!(pool.emit_submit(0, "fn").contains("thread_pool_submit"));
assert!(pool.emit_wait().contains("thread_pool_wait"));
}
#[test]
fn test_concurrency_registry_new() {
let reg = ConcurrencyRegistry::new();
assert!(reg.is_enabled("jthread"));
assert!(reg.is_enabled("openmp"));
assert!(reg.is_enabled("simd"));
assert!(reg.feature_count() >= 10);
}
#[test]
fn test_concurrency_registry_enable_disable() {
let mut reg = ConcurrencyRegistry::new();
assert!(reg.is_enabled("latch"));
reg.disable("latch");
assert!(!reg.is_enabled("latch"));
reg.enable("latch");
assert!(reg.is_enabled("latch"));
}
#[test]
fn test_concurrency_registry_enabled_features() {
let mut reg = ConcurrencyRegistry::new();
reg.disable("openmp");
let enabled = reg.enabled_features();
assert!(!enabled.contains(&"openmp".to_string()));
assert!(enabled.contains(&"jthread".to_string()));
}
}