cubecl-core 0.11.0

use cubecl_ir::AddressType;
use cubecl_runtime::server::CopyDescriptor;
use cubecl_zspace::{MetadataError, Shape, Strides, Tiling};

use crate::{
    self as cubecl,
    compute::{KernelBuilder, KernelLauncher},
    frontend::container::{buffer_len::expand_buffer_length_native, slice},
    prelude::*,
};

use super::Tensor;

#[derive(CubeType, CubeLaunch, Clone, Copy)]
#[expand(derive(Clone, Copy))]
pub struct TensorMeta {
    pub len: usize,
    pub rank: usize,
}

/// Argument to be used for [tensors](Tensor) passed as arguments to kernels.
#[derive(Debug)]
pub enum TensorArg {
    /// The tensor is passed with a tensor handle.
    Handle {
        /// The tensor handle.
        handle: TensorBinding,
    },
    /// The tensor is aliasing another input tensor.
    Alias {
        /// The position of the input tensor.
        input_pos: usize,
        strides: Strides,
        shape: Shape,
    },
}

/// Tensor representation with a reference to the [server handle](cubecl_runtime::server::Handle),
/// the strides and the shape.
pub struct TensorBinding {
    pub handle: cubecl_runtime::server::BufferBinding,
    pub strides: Strides,
    pub shape: Shape,
    /// How many fragments each logical dim is stored as; untiled by default.
    /// The shape and strides stay physical, as on [`Metadata`](cubecl_zspace::Metadata).
    pub tiling: Tiling,
}

impl Clone for TensorBinding {
    fn clone(&self) -> Self {
        Self {
            handle: self.handle.clone(),
            strides: self.strides.clone(),
            shape: self.shape.clone(),
            tiling: self.tiling,
        }
    }
}

impl TensorBinding {
    pub fn size(&self) -> usize {
        self.shape.iter().product()
    }

    /// Address type required to fully index this tensor handle, assuming scalar access.
    pub fn required_address_type(&self, elem_size: usize) -> AddressType {
        AddressType::from_len(self.handle.size() as usize / elem_size)
    }

    /// This binding with `tiling` labelling its physical dims.
    ///
    /// # Errors
    ///
    /// When `tiling` does not describe this rank: see [`Tiling::logical_rank`].
    pub fn with_tiling(mut self, tiling: Tiling) -> Result<Self, MetadataError> {
        tiling.logical_rank(self.shape.rank())?;
        self.tiling = tiling;
        Ok(self)
    }

    /// An alias reaches its tensor through another argument's dims, which carry
    /// no tiling of their own, so a tiled binding says so rather than aliasing
    /// as if it were plain.
    fn assert_untiled(&self, op: &str) {
        assert!(
            !self.tiling.is_tiled(),
            "TensorBinding::{op} on a storage-tiled binding is not supported: {:?}",
            self.tiling
        );
    }
}

impl core::fmt::Debug for TensorBinding {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        writeln!(
            f,
            "TensorHandleRef {{ strides: {:?}, shape: {:?}, tiling: {:?} }}",
            self.strides, self.shape, self.tiling
        )
    }
}

/// Compilation argument for a [tensor](Tensor).
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct TensorCompilationArg {
    pub meta: TensorMetaCompilationArg,
    pub buffer: BufferCompilationArg,
}

impl<C: CubePrimitive> LaunchArg for Tensor<C> {
    type RuntimeArg = TensorArg;
    type CompilationArg = TensorCompilationArg;

    fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
        let elem_size = launcher.with_scope(|scope| C::__expand_size(scope));
        let vector_size = launcher.with_scope(|scope| C::__expand_vector_size(scope));
        let len = arg.size() / vector_size;
        let meta_arg = TensorMetaLaunch::new(len, arg.shape().len());
        let buffer = match &arg {
            TensorArg::Handle { .. } => BufferCompilationArg { inplace: None },
            TensorArg::Alias { input_pos, .. } => BufferCompilationArg {
                inplace: Some(*input_pos),
            },
        };
        launcher.register_tensor(arg, elem_size);
        let meta = TensorMeta::register(meta_arg, launcher);
        TensorCompilationArg { meta, buffer }
    }

    fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> TensorExpand<C> {
        let buffer = match arg.buffer.inplace {
            Some(id) => builder.inplace(id),
            None => builder.tensor(C::__expand_as_type(&builder.scope)),
        };
        let meta = TensorMeta::expand(&arg.meta, builder);
        let scope = &builder.scope;
        let len = expand_buffer_length_native(scope, buffer);
        let buffer =
            slice::from_raw_parts::<C>(scope, buffer, 0usize.into_expand(scope), len.into());
        TensorExpand { meta, buffer }
    }
}

impl<C: CubePrimitive> LaunchArg for OwnedTensor<C> {
    type RuntimeArg = TensorArg;
    type CompilationArg = TensorCompilationArg;

    fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
        Tensor::<C>::register(arg, launcher)
    }

    fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> OwnedTensorExpand<C> {
        let tensor = Tensor::<C>::expand(arg, builder);
        OwnedTensorExpand {
            meta: tensor.meta,
            buffer: tensor.buffer.expand.into(),
        }
    }
}

impl TensorArg {
    /// Create a new tensor argument specified with its vectorization factor.
    ///
    /// # Safety
    ///
    /// If you provide wrong strides or shapes, it might create undefined behavior caused by
    /// out-of-bound reads and writes.
    pub unsafe fn from_raw_parts(
        handle: cubecl_runtime::server::Handle,
        strides: Strides,
        shape: Shape,
    ) -> Self {
        unsafe { Self::from_raw_parts_binding(handle.binding(), strides, shape) }
    }

    pub(crate) unsafe fn from_raw_parts_binding(
        handle: cubecl_runtime::server::BufferBinding,
        strides: Strides,
        shape: Shape,
    ) -> Self {
        unsafe {
            Self::Handle {
                handle: TensorBinding::from_raw_parts_binding(handle, strides, shape),
            }
        }
    }

    /// Create an alias argument.
    pub fn into_alias(self, position: usize) -> Self {
        match self {
            TensorArg::Handle { handle } => handle.into_alias(position),
            alias @ TensorArg::Alias { .. } => alias,
        }
    }

    pub fn size(&self) -> usize {
        match self {
            TensorArg::Handle { handle } => handle.size(),
            TensorArg::Alias { shape, .. } => shape.iter().product(),
        }
    }

    pub fn shape(&self) -> &[usize] {
        match self {
            TensorArg::Handle { handle } => &handle.shape,
            TensorArg::Alias { shape, .. } => shape,
        }
    }

    pub fn strides(&self) -> &[usize] {
        match self {
            TensorArg::Handle { handle } => &handle.strides,
            TensorArg::Alias { strides, .. } => strides,
        }
    }
}

impl TensorArg {
    pub fn into_buffer_arg(self) -> BufferArg {
        match self {
            TensorArg::Handle { handle } => {
                let handle = unsafe {
                    let size = handle.size();
                    BufferBinding::from_raw_parts_binding(handle.handle, size)
                };
                BufferArg::Handle { handle }
            }
            TensorArg::Alias {
                input_pos, shape, ..
            } => BufferArg::Alias {
                input_pos,
                length: [shape.iter().product()],
            },
        }
    }
}

impl TensorBinding {
    /// Convert the handle into a [tensor argument](TensorArg).
    pub fn into_tensor_arg(self) -> TensorArg {
        TensorArg::Handle { handle: self }
    }
    /// Convert the handle into a [tensor argument](TensorArg).
    pub fn into_alias(self, index: usize) -> TensorArg {
        self.assert_untiled("into_alias");
        TensorArg::Alias {
            input_pos: index,
            strides: self.strides,
            shape: self.shape,
        }
    }
    /// Convert the handle into a [tensor argument](TensorArg).
    pub fn as_alias(&self, index: usize) -> TensorArg {
        self.assert_untiled("as_alias");
        TensorArg::Alias {
            input_pos: index,
            strides: self.strides.clone(),
            shape: self.shape.clone(),
        }
    }
    /// Convert the handle into a [buffer argument](BufferArg).
    pub fn into_buffer_arg(self) -> BufferArg {
        unsafe { BufferArg::from_raw_parts_binding(self.handle, self.shape.iter().product()) }
    }

    /// Create a handle from raw parts.
    ///
    /// # Safety
    ///
    /// If you provide wrong strides or shapes, it might create undefined behavior caused by
    /// out-of-bounds reads and writes.
    pub unsafe fn from_raw_parts(
        handle: cubecl_runtime::server::Handle,
        strides: Strides,
        shape: Shape,
    ) -> Self {
        unsafe { Self::from_raw_parts_binding(handle.binding(), strides, shape) }
    }

    /// Create a handle from raw parts.
    ///
    /// # Safety
    ///
    /// If you provide wrong strides or shapes, it might create undefined behavior caused by
    /// out-of-bounds reads and writes.
    pub unsafe fn from_raw_parts_binding(
        handle: cubecl_runtime::server::BufferBinding,
        strides: Strides,
        shape: Shape,
    ) -> Self {
        Self {
            handle,
            strides,
            shape,
            tiling: Tiling::UNTILED,
        }
    }

    pub fn into_copy_descriptor(self, elem_size: usize) -> CopyDescriptor {
        CopyDescriptor {
            handle: self.handle,
            shape: self.shape,
            strides: self.strides,
            elem_size,
        }
    }
}