j2k-metal 0.7.2

// SPDX-License-Identifier: MIT OR Apache-2.0

#[cfg(test)]
use std::cell::Cell;
use std::sync::OnceLock;
use std::time::Instant;

use j2k_profile::{env_flag_from_env, ProfileStageMode, StageModeCache};
use metal::{CommandBufferRef, ComputeCommandEncoderRef};

mod direct;

pub(crate) use direct::{
    decode_profile_label, emit_metal_profile_row, MetalDirectProfileRow, MetalProfileFormat,
};

const CLASSIC_SELECTIVE_BYPASS_ENV: &str = "J2K_METAL_CLASSIC_SELECTIVE_BYPASS";
const METAL_PROFILE_STAGES_ENV: &str = "J2K_METAL_PROFILE_STAGES";
const METAL_PROFILE_SIGNPOSTS_ENV: &str = "J2K_METAL_PROFILE_SIGNPOSTS";
const METAL_PROFILE_DECODE_SPLIT_COMMANDS_ENV: &str = "J2K_METAL_PROFILE_DECODE_SPLIT_COMMANDS";
const METAL_PROFILE_COEFFICIENT_PREP_SPLIT_COMMANDS_ENV: &str =
    "J2K_METAL_PROFILE_COEFFICIENT_PREP_SPLIT_COMMANDS";
const METAL_PROFILE_CLASSIC_TIER1_DENSITY_ENV: &str = "J2K_METAL_PROFILE_CLASSIC_TIER1_DENSITY";
const METAL_PROFILE_CLASSIC_TIER1_RAW_PACK_ENV: &str = "J2K_METAL_PROFILE_CLASSIC_TIER1_RAW_PACK";
const METAL_PROFILE_CLASSIC_TIER1_ARITHMETIC_PACK_ENV: &str =
    "J2K_METAL_PROFILE_CLASSIC_TIER1_ARITHMETIC_PACK";
const METAL_PROFILE_CLASSIC_TIER1_SYMBOL_PLAN_ENV: &str =
    "J2K_METAL_PROFILE_CLASSIC_TIER1_SYMBOL_PLAN";
const METAL_PROFILE_CLASSIC_TIER1_PASS_PLAN_ENV: &str = "J2K_METAL_PROFILE_CLASSIC_TIER1_PASS_PLAN";
const METAL_PROFILE_CLASSIC_TIER1_TOKEN_EMIT_ENV: &str =
    "J2K_METAL_PROFILE_CLASSIC_TIER1_TOKEN_EMIT";
const METAL_PROFILE_CLASSIC_TIER1_SPLIT_TOKEN_EMIT_ENV: &str =
    "J2K_METAL_PROFILE_CLASSIC_TIER1_SPLIT_TOKEN_EMIT";
const METAL_PROFILE_CLASSIC_TIER1_TOKEN_PACK_ENV: &str =
    "J2K_METAL_PROFILE_CLASSIC_TIER1_TOKEN_PACK";
const CLASSIC_TIER1_GPU_TOKEN_PACK_ENV: &str = "J2K_METAL_CLASSIC_TIER1_GPU_TOKEN_PACK";
const CLASSIC_TIER1_SPLIT_GPU_TOKEN_PACK_ENV: &str = "J2K_METAL_CLASSIC_TIER1_SPLIT_GPU_TOKEN_PACK";
const CLASSIC_TIER1_SPLIT_MQ_BYTE_GPU_TOKEN_PACK_ENV: &str =
    "J2K_METAL_CLASSIC_TIER1_SPLIT_MQ_BYTE_GPU_TOKEN_PACK";

pub(crate) type HybridSignpostName = u32;

pub(crate) const SIGNPOST_DECODE_HYBRID_CPU_TIER1: HybridSignpostName = 1;
pub(crate) const SIGNPOST_DECODE_HYBRID_COEFFICIENT_UPLOAD: HybridSignpostName = 2;
pub(crate) const SIGNPOST_DECODE_HYBRID_COMMAND_WAIT: HybridSignpostName = 3;
pub(crate) const SIGNPOST_DECODE_HYBRID_IDWT_COMMAND_ENCODE: HybridSignpostName = 4;
pub(crate) const SIGNPOST_DECODE_HYBRID_STORE_COMMAND_ENCODE: HybridSignpostName = 5;
pub(crate) const SIGNPOST_DECODE_HYBRID_MCT_PACK_COMMAND_ENCODE: HybridSignpostName = 6;
pub(crate) const SIGNPOST_ENCODE_HYBRID_COMMAND_WAIT: HybridSignpostName = 7;
pub(crate) const SIGNPOST_ENCODE_HYBRID_RESULT_HARVEST: HybridSignpostName = 8;
pub(crate) const SIGNPOST_ENCODE_HYBRID_CLASSIC_TIER1_SETUP: HybridSignpostName = 9;
pub(crate) const SIGNPOST_ENCODE_HYBRID_CLASSIC_TIER1_COMMAND_ENCODE: HybridSignpostName = 10;
pub(crate) const SIGNPOST_ENCODE_HYBRID_CLASSIC_PACKET_PLAN: HybridSignpostName = 11;
pub(crate) const SIGNPOST_ENCODE_HYBRID_CLASSIC_PACKET_BUFFER_SETUP: HybridSignpostName = 12;
pub(crate) const SIGNPOST_ENCODE_HYBRID_CLASSIC_PACKETIZATION_COMMAND_ENCODE: HybridSignpostName =
    13;
pub(crate) const SIGNPOST_ENCODE_HYBRID_CLASSIC_PAYLOAD_COPY_COMMAND_ENCODE: HybridSignpostName =
    14;
pub(crate) const SIGNPOST_ENCODE_HYBRID_CLASSIC_CODESTREAM_ASSEMBLY_COMMAND_ENCODE:
    HybridSignpostName = 15;
pub(crate) const SIGNPOST_ENCODE_HYBRID_HT_TIER1_SETUP: HybridSignpostName = 16;
pub(crate) const SIGNPOST_ENCODE_HYBRID_HT_TIER1_COMMAND_ENCODE: HybridSignpostName = 17;
pub(crate) const SIGNPOST_ENCODE_HYBRID_HT_PACKET_PLAN: HybridSignpostName = 18;
pub(crate) const SIGNPOST_ENCODE_HYBRID_HT_PACKET_BUFFER_SETUP: HybridSignpostName = 19;
pub(crate) const SIGNPOST_ENCODE_HYBRID_HT_PACKET_BLOCK_PREP_COMMAND_ENCODE: HybridSignpostName =
    20;
pub(crate) const SIGNPOST_ENCODE_HYBRID_HT_PACKETIZATION_COMMAND_ENCODE: HybridSignpostName = 21;
pub(crate) const SIGNPOST_ENCODE_HYBRID_HT_PAYLOAD_COPY_COMMAND_ENCODE: HybridSignpostName = 22;
pub(crate) const SIGNPOST_ENCODE_HYBRID_HT_CODESTREAM_ASSEMBLY_COMMAND_ENCODE: HybridSignpostName =
    23;

#[cfg(test)]
std::thread_local! {
    static CLASSIC_GPU_TOKEN_PACK_ROUTE_OVERRIDE: Cell<Option<bool>> = const { Cell::new(None) };
    static METAL_PROFILE_STAGES_OVERRIDE: Cell<Option<ProfileStageMode>> = const { Cell::new(None) };
}

fn env_flag_enabled(name: &str) -> bool {
    env_flag_from_env(name)
}

pub(crate) fn classic_selective_bypass_disabled() -> bool {
    matches!(std::env::var(CLASSIC_SELECTIVE_BYPASS_ENV), Ok(value) if value == "0")
}

#[cfg(test)]
pub(crate) struct ClassicGpuTokenPackRouteOverrideGuard {
    previous: Option<bool>,
}

#[cfg(test)]
impl Drop for ClassicGpuTokenPackRouteOverrideGuard {
    fn drop(&mut self) {
        CLASSIC_GPU_TOKEN_PACK_ROUTE_OVERRIDE.with(|slot| slot.set(self.previous));
    }
}

#[cfg(test)]
pub(crate) fn force_classic_gpu_token_pack_route_for_test(
    enabled: bool,
) -> ClassicGpuTokenPackRouteOverrideGuard {
    let previous = CLASSIC_GPU_TOKEN_PACK_ROUTE_OVERRIDE.with(|slot| slot.replace(Some(enabled)));
    ClassicGpuTokenPackRouteOverrideGuard { previous }
}

#[cfg(test)]
pub(crate) struct MetalProfileStagesOverrideGuard {
    previous: Option<ProfileStageMode>,
}

#[cfg(test)]
impl Drop for MetalProfileStagesOverrideGuard {
    fn drop(&mut self) {
        METAL_PROFILE_STAGES_OVERRIDE.with(|slot| slot.set(self.previous));
    }
}

#[cfg(test)]
pub(crate) fn force_metal_profile_stages_for_test(
    enabled: bool,
) -> MetalProfileStagesOverrideGuard {
    force_metal_profile_stage_mode_for_test(if enabled {
        ProfileStageMode::Rows
    } else {
        ProfileStageMode::Disabled
    })
}

#[cfg(test)]
pub(crate) fn force_metal_profile_stage_mode_for_test(
    mode: ProfileStageMode,
) -> MetalProfileStagesOverrideGuard {
    let previous = METAL_PROFILE_STAGES_OVERRIDE.with(|slot| slot.replace(Some(mode)));
    MetalProfileStagesOverrideGuard { previous }
}

pub(crate) fn classic_tier1_gpu_token_pack_requested() -> bool {
    #[cfg(test)]
    if let Some(enabled) = CLASSIC_GPU_TOKEN_PACK_ROUTE_OVERRIDE.with(Cell::get) {
        return enabled;
    }
    env_flag_enabled(CLASSIC_TIER1_GPU_TOKEN_PACK_ENV)
}

pub(crate) fn classic_tier1_split_gpu_token_pack_requested() -> bool {
    env_flag_enabled(CLASSIC_TIER1_SPLIT_GPU_TOKEN_PACK_ENV)
}

fn classic_tier1_split_mq_byte_gpu_token_pack_setting() -> Option<bool> {
    match std::env::var(CLASSIC_TIER1_SPLIT_MQ_BYTE_GPU_TOKEN_PACK_ENV) {
        Ok(value) if value == "1" => Some(true),
        Ok(value) if value == "0" => Some(false),
        _ => None,
    }
}

pub(crate) fn classic_tier1_split_mq_byte_gpu_token_pack_requested() -> bool {
    classic_tier1_split_mq_byte_gpu_token_pack_setting() == Some(true)
}

pub(crate) fn classic_tier1_split_mq_byte_gpu_token_pack_disabled() -> bool {
    classic_tier1_split_mq_byte_gpu_token_pack_setting() == Some(false)
}

pub(crate) fn metal_profile_stage_mode() -> ProfileStageMode {
    static MODE: StageModeCache = StageModeCache::new();
    #[cfg(test)]
    if let Some(mode) = METAL_PROFILE_STAGES_OVERRIDE.with(Cell::get) {
        return mode;
    }
    MODE.mode_from_env(METAL_PROFILE_STAGES_ENV)
}

pub(crate) fn metal_profile_stages_enabled() -> bool {
    metal_profile_stage_mode() != ProfileStageMode::Disabled
}

pub(crate) fn elapsed_since_us(started: Instant) -> u128 {
    j2k_profile::elapsed_us(Some(started))
}

fn metal_profile_signposts_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    *ENABLED.get_or_init(|| env_flag_enabled(METAL_PROFILE_SIGNPOSTS_ENV))
}

pub(crate) fn metal_profile_decode_split_commands_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    metal_profile_stages_enabled()
        && *ENABLED.get_or_init(|| env_flag_enabled(METAL_PROFILE_DECODE_SPLIT_COMMANDS_ENV))
}

pub(crate) fn metal_profile_coefficient_prep_split_commands_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    metal_profile_stages_enabled()
        && *ENABLED
            .get_or_init(|| env_flag_enabled(METAL_PROFILE_COEFFICIENT_PREP_SPLIT_COMMANDS_ENV))
}

pub(crate) fn metal_profile_classic_tier1_density_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    metal_profile_stages_enabled()
        && *ENABLED.get_or_init(|| env_flag_enabled(METAL_PROFILE_CLASSIC_TIER1_DENSITY_ENV))
}

pub(crate) fn metal_profile_classic_tier1_raw_pack_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    metal_profile_stages_enabled()
        && *ENABLED.get_or_init(|| env_flag_enabled(METAL_PROFILE_CLASSIC_TIER1_RAW_PACK_ENV))
}

pub(crate) fn metal_profile_classic_tier1_arithmetic_pack_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    metal_profile_stages_enabled()
        && *ENABLED
            .get_or_init(|| env_flag_enabled(METAL_PROFILE_CLASSIC_TIER1_ARITHMETIC_PACK_ENV))
}

pub(crate) fn metal_profile_classic_tier1_symbol_plan_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    metal_profile_stages_enabled()
        && *ENABLED.get_or_init(|| {
            env_flag_enabled(METAL_PROFILE_CLASSIC_TIER1_SYMBOL_PLAN_ENV)
                || env_flag_enabled(METAL_PROFILE_CLASSIC_TIER1_PASS_PLAN_ENV)
        })
}

pub(crate) fn metal_profile_classic_tier1_pass_plan_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    metal_profile_stages_enabled()
        && *ENABLED.get_or_init(|| env_flag_enabled(METAL_PROFILE_CLASSIC_TIER1_PASS_PLAN_ENV))
}

pub(crate) fn metal_profile_classic_tier1_token_emit_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    metal_profile_stages_enabled()
        && *ENABLED.get_or_init(|| {
            env_flag_enabled(METAL_PROFILE_CLASSIC_TIER1_TOKEN_EMIT_ENV)
                || env_flag_enabled(METAL_PROFILE_CLASSIC_TIER1_TOKEN_PACK_ENV)
        })
}

pub(crate) fn metal_profile_classic_tier1_split_token_emit_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    metal_profile_stages_enabled()
        && *ENABLED
            .get_or_init(|| env_flag_enabled(METAL_PROFILE_CLASSIC_TIER1_SPLIT_TOKEN_EMIT_ENV))
}

pub(crate) fn metal_profile_classic_tier1_token_pack_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    metal_profile_stages_enabled()
        && *ENABLED.get_or_init(|| env_flag_enabled(METAL_PROFILE_CLASSIC_TIER1_TOKEN_PACK_ENV))
}

pub(crate) fn label_command_buffer(command_buffer: &CommandBufferRef, label: &str) {
    if metal_profile_stages_enabled() {
        command_buffer.set_label(label);
    }
}

type OsSignpostId = u64;

const OS_SIGNPOST_ID_NULL: OsSignpostId = 0;
const OS_SIGNPOST_ID_INVALID: OsSignpostId = OsSignpostId::MAX;

unsafe extern "C" {
    fn j2k_metal_signpost_begin(name: HybridSignpostName) -> OsSignpostId;
    fn j2k_metal_signpost_end(name: HybridSignpostName, id: OsSignpostId);
}

pub(crate) struct HybridStageSignpost {
    id: OsSignpostId,
    name: HybridSignpostName,
}

impl Drop for HybridStageSignpost {
    fn drop(&mut self) {
        // SAFETY: The shim accepts the signpost id returned by the matching begin call.
        unsafe {
            j2k_metal_signpost_end(self.name, self.id);
        }
    }
}

pub(crate) fn hybrid_stage_signpost(name: HybridSignpostName) -> Option<HybridStageSignpost> {
    if !metal_profile_signposts_enabled() {
        return None;
    }
    // SAFETY: The signpost shim is provided by this crate's build and has no Rust aliasing contract.
    let id = unsafe { j2k_metal_signpost_begin(name) };
    if id == OS_SIGNPOST_ID_NULL || id == OS_SIGNPOST_ID_INVALID {
        return None;
    }
    Some(HybridStageSignpost { id, name })
}

pub(crate) fn label_compute_encoder(encoder: &ComputeCommandEncoderRef, label: &str) {
    if metal_profile_stages_enabled() {
        encoder.set_label(label);
    }
}