cubecl-wgpu 0.11.0

use cubecl_server::kernel::BufferIOAttr;
use cubecl_server::kernel::CompiledKernel;
use std::fmt::Display;

use cubecl_core::{
    Compiler, WgpuCompilationOptions,
    prelude::{CubeKernel, KernelDefinition},
    server::{LaunchError, ResourceLimitError},
};
#[cfg(feature = "msl")]
use cubecl_cpp::shared::MslComputeKernel;
use cubecl_environment::backtrace::BackTrace;
use cubecl_ir::DeviceProperties;
use cubecl_server::compiler::CompilationError;
use derive_more::derive::From;
use thiserror::Error;

#[cfg(feature = "spirv")]
use crate::ParamsTransfer;
use crate::{CompilerInfo, WgpuBackend, WgpuServer};

use super::wgsl;

#[cfg(feature = "msl")]
pub use cubecl_cpp::MslCompiler;
#[cfg(feature = "spirv")]
pub use cubecl_spirv::SpirvCompiler;
pub use wgsl::WgslCompiler;

/// Compiler that dispatches to the most appropriate shader language backend for the active
/// `wgpu` backend.
///
/// The variant is selected at runtime by [`WgpuCompiler::init`] based on the `wgpu::Backend`
/// in use and the enabled cargo features (`spirv`, `msl`).
#[allow(clippy::large_enum_variant)]
#[derive(Debug, Clone)]
pub enum AutoCompiler {
    /// WGSL backend, available on every `wgpu` backend.
    Wgsl(WgslCompiler),
    /// SPIR-V backend used on Vulkan when the device supports the required features.
    #[cfg(feature = "spirv")]
    SpirV(cubecl_spirv::SpirvCompiler),
    /// Metal Shading Language backend used on the Metal backend.
    #[cfg(feature = "msl")]
    Msl(MslCompiler),
}

/// Owned compiled kernel representation matching the variants of [`AutoCompiler`].
#[derive(From)]
#[allow(clippy::large_enum_variant)]
pub enum AutoRepresentation {
    /// WGSL compute shader source.
    Wgsl(wgsl::ComputeShader),
    /// Compiled SPIR-V kernel.
    #[cfg(feature = "spirv")]
    SpirV(cubecl_spirv::SpirvKernel),
    /// Compiled Metal Shading Language kernel.
    #[cfg(feature = "msl")]
    Msl(MslComputeKernel),
}

/// Borrowed counterpart of [`AutoRepresentation`], useful when only read access is needed.
#[derive(From, Clone, Copy)]
#[allow(clippy::large_enum_variant)]
pub enum AutoRepresentationRef<'a> {
    /// Borrowed WGSL compute shader.
    Wgsl(&'a wgsl::ComputeShader),
    /// Borrowed SPIR-V kernel.
    #[cfg(feature = "spirv")]
    SpirV(&'a cubecl_spirv::SpirvKernel),
    /// Borrowed Metal Shading Language kernel.
    #[cfg(feature = "msl")]
    Msl(&'a MslComputeKernel),
}

#[cfg(feature = "spirv")]
impl AutoRepresentation {
    /// Returns the SPIR-V kernel if this representation is the SPIR-V variant.
    pub fn as_spirv(&self) -> Option<&cubecl_spirv::SpirvKernel> {
        match self {
            AutoRepresentation::SpirV(repr) => Some(repr),
            _ => None,
        }
    }
}

#[cfg(feature = "msl")]
impl AutoRepresentation {
    /// Returns the MSL kernel if this representation is the MSL variant.
    pub fn as_msl(&self) -> Option<&MslComputeKernel> {
        match self {
            AutoRepresentation::Msl(repr) => Some(repr),
            _ => None,
        }
    }
}

impl AutoRepresentation {
    /// Borrow this representation as an [`AutoRepresentationRef`].
    pub fn as_ref(&self) -> AutoRepresentationRef<'_> {
        match self {
            AutoRepresentation::Wgsl(compute_shader) => AutoRepresentationRef::Wgsl(compute_shader),
            #[cfg(feature = "spirv")]
            AutoRepresentation::SpirV(spirv_kernel) => AutoRepresentationRef::SpirV(spirv_kernel),
            #[cfg(feature = "msl")]
            AutoRepresentation::Msl(compute_shader) => AutoRepresentationRef::Msl(compute_shader),
        }
    }
}

impl Display for AutoRepresentation {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AutoRepresentation::Wgsl(compute_shader) => compute_shader.fmt(f),
            #[cfg(feature = "spirv")]
            AutoRepresentation::SpirV(spirv_kernel) => spirv_kernel.fmt(f),
            #[cfg(feature = "msl")]
            AutoRepresentation::Msl(compute_shader) => compute_shader.fmt(f),
        }
    }
}

impl Compiler for AutoCompiler {
    type Representation = AutoRepresentation;

    type CompilationOptions = WgpuCompilationOptions;

    fn compile(
        &mut self,
        kernel: KernelDefinition,
        compilation_options: &Self::CompilationOptions,
    ) -> Result<Self::Representation, CompilationError> {
        let kernel = match self {
            AutoCompiler::Wgsl(wgsl_compiler) => {
                Compiler::compile(wgsl_compiler, kernel, compilation_options)?.into()
            }
            #[cfg(feature = "spirv")]
            AutoCompiler::SpirV(spirv_compiler) => {
                Compiler::compile(spirv_compiler, kernel, compilation_options)?.into()
            }
            #[cfg(feature = "msl")]
            AutoCompiler::Msl(msl_compiler) => {
                // override compilation options with cpp compiler options for metal
                use cubecl_cpp;
                let compilation_options = cubecl_cpp::shared::CompilationOptions::default();
                Compiler::compile(msl_compiler, kernel, &compilation_options)?.into()
            }
        };

        Ok(kernel)
    }

    fn buffer_io(repr: &Self::Representation) -> Option<Vec<BufferIOAttr>> {
        match repr {
            AutoRepresentation::Wgsl(shader) => WgslCompiler::buffer_io(shader),
            #[cfg(feature = "spirv")]
            AutoRepresentation::SpirV(kernel) => cubecl_spirv::SpirvCompiler::buffer_io(kernel),
            #[cfg(feature = "msl")]
            AutoRepresentation::Msl(kernel) => MslCompiler::buffer_io(kernel),
        }
    }

    fn extension(&self) -> &'static str {
        match self {
            AutoCompiler::Wgsl(_) => "wgsl",
            #[cfg(feature = "spirv")]
            AutoCompiler::SpirV(_) => "spv",
            #[cfg(feature = "msl")]
            AutoCompiler::Msl(_) => "msl",
        }
    }

    fn lang_tag(&self) -> &'static str {
        match self {
            AutoCompiler::Wgsl(compiler) => compiler.lang_tag(),
            #[cfg(feature = "spirv")]
            AutoCompiler::SpirV(compiler) => compiler.lang_tag(),
            #[cfg(feature = "msl")]
            AutoCompiler::Msl(compiler) => compiler.lang_tag(),
        }
    }
}

impl WgpuCompiler for AutoCompiler {
    #[cfg(feature = "msl")]
    fn validate_runtime(
        backend: wgpu::Backend,
        options: &WgpuCompilationOptions,
        requested_backend: WgpuBackend,
    ) -> Result<(), WgpuCompilerInitError> {
        if requested_backend == WgpuBackend::Metal {
            MslCompiler::validate_runtime(backend, options, requested_backend)?;
        }

        Ok(())
    }

    fn init(backend: wgpu::Backend, options: &WgpuCompilationOptions) -> Self {
        let _ = options; // Unused without `spirv` feature
        match backend {
            #[cfg(feature = "spirv")]
            wgpu::Backend::Vulkan if options.supports_vulkan_compiler => {
                AutoCompiler::SpirV(Default::default())
            }
            #[cfg(feature = "msl")]
            wgpu::Backend::Metal if options.supports_msl_compiler => {
                AutoCompiler::Msl(Default::default())
            }
            _ => AutoCompiler::Wgsl(Default::default()),
        }
    }

    fn compile_kernel(
        &mut self,
        server: &mut WgpuServer<AutoCompiler>,
        kernel: Box<dyn CubeKernel>,
        definition: KernelDefinition,
    ) -> Result<CompiledKernel<Self>, CompilationError> {
        match self {
            AutoCompiler::Wgsl(_) => {
                CompiledKernel::compile(&*kernel, definition, self, &server.compilation_options)
            }
            #[cfg(feature = "spirv")]
            AutoCompiler::SpirV(_) => crate::vulkan::compile(self, server, kernel, definition),
            #[cfg(feature = "msl")]
            AutoCompiler::Msl(_) => {
                CompiledKernel::compile(&*kernel, definition, self, &server.compilation_options)
            }
        }
    }

    fn validate_ir(
        &self,
        repr: &Option<Self::Representation>,
        props: &DeviceProperties,
    ) -> Result<(), LaunchError> {
        let shared_bytes = repr.as_ref().map(|repr| match repr {
            AutoRepresentation::Wgsl(repr) => repr.shared_memory_size,
            #[cfg(feature = "msl")]
            AutoRepresentation::Msl(repr) => repr.shared_memory_size,
            #[cfg(feature = "spirv")]
            AutoRepresentation::SpirV(repr) => repr.shared_size,
        });
        check_shared_memory(shared_bytes, props)
    }

    fn normalize_repr(
        &self,
        repr: Option<Self::Representation>,
    ) -> (CompilerInfo, Option<AutoRepresentation>) {
        let compiler_info = match &repr {
            #[cfg(feature = "spirv")]
            Some(AutoRepresentation::SpirV(repr)) => CompilerInfo::Vulkan {
                params_transfer: match repr.immediate_size {
                    Some(_) => ParamsTransfer::Immediate,
                    None => ParamsTransfer::Uniform,
                },
            },
            #[cfg(feature = "msl")]
            Some(AutoRepresentation::Msl(_)) => CompilerInfo::Metal,
            Some(AutoRepresentation::Wgsl(_)) => CompilerInfo::WGSL,
            None => CompilerInfo::None,
        };

        (compiler_info, repr)
    }
}

impl WgpuCompiler for WgslCompiler {
    fn init(_backend: wgpu::Backend, _options: &WgpuCompilationOptions) -> Self {
        Self
    }

    fn compile_kernel(
        &mut self,
        server: &mut WgpuServer<Self>,
        kernel: Box<dyn CubeKernel>,
        definition: KernelDefinition,
    ) -> Result<CompiledKernel<Self>, CompilationError> {
        CompiledKernel::compile(&*kernel, definition, self, &server.compilation_options)
    }

    fn validate_ir(
        &self,
        repr: &Option<Self::Representation>,
        props: &DeviceProperties,
    ) -> Result<(), LaunchError> {
        let shared_bytes = repr.as_ref().map(|repr| repr.shared_memory_size);
        check_shared_memory(shared_bytes, props)
    }

    fn normalize_repr(
        &self,
        repr: Option<Self::Representation>,
    ) -> (CompilerInfo, Option<AutoRepresentation>) {
        (CompilerInfo::WGSL, repr.map(|r| r.into()))
    }
}

#[cfg(feature = "msl")]
impl WgpuCompiler for MslCompiler {
    fn validate_runtime(
        backend: wgpu::Backend,
        options: &WgpuCompilationOptions,
        _requested_backend: WgpuBackend,
    ) -> Result<(), WgpuCompilerInitError> {
        if backend != wgpu::Backend::Metal {
            return Err(WgpuCompilerInitError::BackendMismatch {
                compiler: "MslCompiler",
                expected: wgpu::Backend::Metal,
                actual: backend,
            });
        }

        if !options.supports_msl_compiler {
            return Err(WgpuCompilerInitError::Unavailable {
                compiler: "MslCompiler",
                backend,
            });
        }

        Ok(())
    }

    fn init(_backend: wgpu::Backend, _options: &WgpuCompilationOptions) -> Self {
        Self::default()
    }

    fn compile_kernel(
        &mut self,
        _server: &mut WgpuServer<Self>,
        kernel: Box<dyn CubeKernel>,
        definition: KernelDefinition,
    ) -> Result<CompiledKernel<Self>, CompilationError> {
        // The MSL compiler uses its own CompilationOptions, not WgpuCompilationOptions.
        let compilation_options = cubecl_cpp::shared::CompilationOptions::default();
        CompiledKernel::compile(&*kernel, definition, self, &compilation_options)
    }

    fn validate_ir(
        &self,
        repr: &Option<Self::Representation>,
        props: &DeviceProperties,
    ) -> Result<(), LaunchError> {
        let shared_bytes = repr.as_ref().map(|repr| repr.shared_memory_size);
        check_shared_memory(shared_bytes, props)
    }

    fn normalize_repr(
        &self,
        repr: Option<Self::Representation>,
    ) -> (CompilerInfo, Option<AutoRepresentation>) {
        (CompilerInfo::Metal, repr.map(|r| r.into()))
    }
}

#[cfg(feature = "spirv")]
impl WgpuCompiler for cubecl_spirv::SpirvCompiler {
    fn init(_backend: wgpu::Backend, _options: &WgpuCompilationOptions) -> Self {
        Self
    }

    fn compile_kernel(
        &mut self,
        server: &mut WgpuServer<Self>,
        kernel: Box<dyn CubeKernel>,
        definition: KernelDefinition,
    ) -> Result<CompiledKernel<Self>, CompilationError> {
        crate::vulkan::compile(self, server, kernel, definition)
    }

    fn validate_ir(
        &self,
        repr: &Option<Self::Representation>,
        props: &DeviceProperties,
    ) -> Result<(), LaunchError> {
        let shared_bytes = repr.as_ref().map(|repr| repr.shared_size);
        check_shared_memory(shared_bytes, props)
    }

    fn normalize_repr(
        &self,
        repr: Option<Self::Representation>,
    ) -> (CompilerInfo, Option<AutoRepresentation>) {
        let params_transfer = match repr.as_ref().and_then(|r| r.immediate_size) {
            Some(_) => ParamsTransfer::Immediate,
            None => ParamsTransfer::Uniform,
        };
        (
            CompilerInfo::Vulkan { params_transfer },
            repr.map(|r| r.into()),
        )
    }
}

fn check_shared_memory(
    shared_bytes: Option<usize>,
    props: &DeviceProperties,
) -> Result<(), LaunchError> {
    let max_smem = props.hardware.max_shared_memory_size;
    if let Some(shared_bytes) = shared_bytes
        && shared_bytes > max_smem
    {
        return Err(ResourceLimitError::SharedMemory {
            requested: shared_bytes,
            max: max_smem,
            backtrace: BackTrace::capture(),
        }
        .into());
    }
    Ok(())
}

/// An explicitly selected wgpu compiler cannot serve the selected graphics backend or device.
#[non_exhaustive]
#[derive(Debug, Error, PartialEq, Eq)]
pub enum WgpuCompilerInitError {
    /// The compiler only emits shaders for a different graphics backend.
    #[error("`{compiler}` requires the {expected:?} wgpu backend, but {actual:?} was selected")]
    BackendMismatch {
        compiler: &'static str,
        expected: wgpu::Backend,
        actual: wgpu::Backend,
    },
    /// The selected device does not meet the compiler's requirements.
    #[error(
        "`{compiler}` was requested explicitly, but it is unavailable on the selected {backend:?} device; use `AutoCompiler` with an automatic device to allow WGSL fallback"
    )]
    Unavailable {
        compiler: &'static str,
        backend: wgpu::Backend,
    },
}

/// Extension trait implemented by every compiler usable with the `wgpu` runtime.
///
/// The base [`Compiler`] trait already exposes a `compile` method that turns a
/// [`KernelDefinition`] into a backend representation. [`WgpuCompiler`] sits one level
/// higher: it owns the wgpu-specific lifecycle around a [`CubeKernel`](cubecl_server::kernel::CubeKernel)
/// kernel — initializing the compiler for a given `wgpu::Backend`, compiling a kernel using
/// the server's [`WgpuCompilationOptions`], validating the resulting IR against the device,
/// and projecting the typed representation into the runtime-erased [`AutoRepresentation`].
pub trait WgpuCompiler: Compiler {
    /// Check that this compiler can serve the selected graphics backend and device.
    ///
    /// Explicit compilers should return an error instead of silently selecting another shader
    /// language. `requested_backend` preserves the unresolved selection; [`WgpuBackend::Auto`]
    /// allows automatic API selection. With `msl` enabled, [`AutoCompiler`] requires native MSL
    /// for an explicit Metal request; otherwise it permits WGSL fallback. Explicit compilers apply
    /// their own requirements.
    fn validate_runtime(
        _backend: wgpu::Backend,
        _options: &WgpuCompilationOptions,
        _requested_backend: WgpuBackend,
    ) -> Result<(), WgpuCompilerInitError> {
        Ok(())
    }

    /// Build the compiler instance appropriate for the given `wgpu` backend.
    ///
    /// `options` is consulted to decide between alternative implementations (for example, to
    /// opt into the SPIR-V compiler on Vulkan when the device advertises the required
    /// features).
    fn init(backend: wgpu::Backend, options: &WgpuCompilationOptions) -> Self;

    /// Validate that the compiled representation fits within the device's resource limits.
    ///
    /// Today this checks shared memory usage; additional checks may be added without
    /// breaking the contract.
    fn validate_ir(
        &self,
        repr: &Option<Self::Representation>,
        props: &DeviceProperties,
    ) -> Result<(), LaunchError>;

    /// Compile a runtime kernel into a [`CompiledKernel`] ready for pipeline creation.
    ///
    /// Distinct from [`Compiler::compile`], which only translates a [`KernelDefinition`].
    /// This entry point operates on a full server-level kernel and pulls compilation
    /// options from `server`, so its signature cannot collide with the base trait method.
    fn compile_kernel(
        &mut self,
        server: &mut WgpuServer<Self>,
        kernel: Box<dyn CubeKernel>,
        definition: KernelDefinition,
    ) -> Result<CompiledKernel<Self>, CompilationError>;

    /// Normalize the backend-specific representation into the [`AutoRepresentation`] shared
    /// by every wgpu compiler, and report the [`CompilerInfo`] derived from it.
    ///
    /// The [`CompilerInfo`] tells the server which parameter-passing strategy to use for
    /// the resulting pipeline.
    fn normalize_repr(
        &self,
        repr: Option<Self::Representation>,
    ) -> (CompilerInfo, Option<AutoRepresentation>);
}

#[cfg(all(test, feature = "msl"))]
mod tests {
    use super::*;

    #[test]
    fn auto_compiler_falls_back_to_wgsl_when_msl_is_unavailable() {
        let options = WgpuCompilationOptions::default();
        assert!(
            AutoCompiler::validate_runtime(wgpu::Backend::Metal, &options, WgpuBackend::Auto)
                .is_ok()
        );
        let compiler = AutoCompiler::init(wgpu::Backend::Metal, &options);

        assert!(matches!(compiler, AutoCompiler::Wgsl(_)));
    }

    #[test]
    fn explicit_msl_compiler_rejects_wgsl_fallback() {
        for requested_backend in [WgpuBackend::Auto, WgpuBackend::Metal] {
            let result = MslCompiler::validate_runtime(
                wgpu::Backend::Metal,
                &WgpuCompilationOptions::default(),
                requested_backend,
            );

            assert_eq!(
                result,
                Err(WgpuCompilerInitError::Unavailable {
                    compiler: "MslCompiler",
                    backend: wgpu::Backend::Metal,
                })
            );
        }
    }

    #[test]
    fn explicit_msl_compiler_rejects_non_metal_backend() {
        let options = WgpuCompilationOptions {
            supports_msl_compiler: true,
            ..Default::default()
        };
        let result =
            MslCompiler::validate_runtime(wgpu::Backend::Vulkan, &options, WgpuBackend::Auto);

        assert_eq!(
            result,
            Err(WgpuCompilerInitError::BackendMismatch {
                compiler: "MslCompiler",
                expected: wgpu::Backend::Metal,
                actual: wgpu::Backend::Vulkan,
            })
        );
    }

    #[test]
    fn explicit_msl_compiler_accepts_supported_metal_device() {
        let options = WgpuCompilationOptions {
            supports_msl_compiler: true,
            ..Default::default()
        };

        assert!(
            MslCompiler::validate_runtime(wgpu::Backend::Metal, &options, WgpuBackend::Auto)
                .is_ok()
        );
    }

    #[test]
    fn explicit_wgsl_on_metal_does_not_require_native_msl() {
        for supports_msl_compiler in [false, true] {
            let options = WgpuCompilationOptions {
                supports_msl_compiler,
                ..Default::default()
            };
            assert!(
                WgslCompiler::validate_runtime(wgpu::Backend::Metal, &options, WgpuBackend::Metal)
                    .is_ok()
            );
        }
    }

    #[test]
    fn auto_compiler_on_explicit_metal_requires_native_msl() {
        let options = WgpuCompilationOptions::default();
        assert_eq!(
            AutoCompiler::validate_runtime(wgpu::Backend::Metal, &options, WgpuBackend::Metal),
            Err(WgpuCompilerInitError::Unavailable {
                compiler: "MslCompiler",
                backend: wgpu::Backend::Metal,
            })
        );
        let supported = WgpuCompilationOptions {
            supports_msl_compiler: true,
            ..options
        };
        assert!(
            AutoCompiler::validate_runtime(wgpu::Backend::Metal, &supported, WgpuBackend::Metal)
                .is_ok()
        );
    }
}