Function flatten_forward 

pub fn flatten_forward<T: Triton, D: Num, const BLOCK_B: i32, const BLOCK_N: i32>(
    input_ptr: T::Pointer<D>,
    output_ptr: T::Pointer<D>,
    B: i32,
    N: i32,
    stride_ib: i32,
    stride_in: i32,
)

Copy a [B, N] tensor with arbitrary input strides to a contiguous row-major [B, N] output.

This is the forward pass of the flatten operation. In a neural network pipeline it is used to materialise a contiguous copy of a potentially non-contiguous activation tensor (e.g. when transitioning from convolutional to fully-connected layers). When the input is already row-major contiguous (stride_ib = N, stride_in = 1) this is a simple memcpy; when the input uses a different layout (e.g. column-major: stride_ib = 1, stride_in = B) the kernel performs the necessary reordering so that downstream kernels can assume unit strides.

Grid: one flat 1D pid that encodes (pid_b, pid_n) = (pid / num_pid_n, pid % num_pid_n).