burn_cubecl_fusion/optim/reduce/

optimization.rs

use super::args::{
    FusedReduceInput, FusedReduceInputLaunch, FusedReduceOutput, FusedReduceOutputLaunch,
};
#[cfg(feature = "autotune")]
use super::tune::fused_reduce_autotune;
use crate::{
    CubeFusionHandle, FallbackOperation,
    engine::{
        codegen::ir::{
            FuseArg, FuseBlockConfig, FuseType, GlobalArgsLaunch, RefLayout,
            multi_block_variables_init,
        },
        launch::{
            FuseTraceLauncher,
            runner::{TraceRunner, Vectorization},
        },
        trace::{FuseTrace, TraceError, TuneOutput},
    },
    optim::{FusedOperation, elemwise::ElemwiseRunner, reduce::args::FusedReduceArgs},
};
use burn_backend::cubecl::{dtype_to_storage_type, elem_type_to_dtype};
use burn_fusion::stream::Context;
use burn_ir::ReduceDimOpIr;
use burn_std::DType;
use cubecl::{client::Client, prelude::*};
use cubek::reduce::{
    ReduceDtypes, ReduceError, VectorizationMode,
    components::instructions::ReduceOperationConfig,
    init_tensors,
    launch::{RoutineStrategy, reduce_kernel_virtual},
    output_vectorization_axis,
    routines::{
        ReduceBlueprint, ReduceLaunchSettings, ReduceProblem, ReduceVectorSettings, Routine,
        cube::CubeRoutine, plane::PlaneRoutine, unit::UnitRoutine,
    },
};
use serde::{Deserialize, Serialize};
use std::sync::Arc;

#[cfg(not(feature = "autotune"))]
use cubek::reduce::routines::{BlueprintStrategy, unit::UnitStrategy};

pub struct ReduceOptimization {
    pub(crate) info: Arc<ReduceOptimizationInfo>,
}

pub(crate) struct ReduceOptimizationInfo {
    pub(crate) trace: FuseTrace,
    trace_read_fallback: FuseTrace,
    trace_write_fallback: FuseTrace,
    pub(crate) client: Client,
    pub(crate) device: cubecl::Device,
    pub(crate) len: usize,
    pub(crate) len_read: usize,
    pub(crate) reduce: FusedReduce,
    settings: ReduceSettings,
}

impl ReduceOptimizationInfo {
    pub fn from_state(device: &cubecl::Device, state: ReduceOptimizationState) -> Self {
        let client = device.client();

        Self {
            trace: state.trace,
            trace_read_fallback: state.trace_read_fallback,
            trace_write_fallback: state.trace_write_fallback,
            client,
            device: device.clone(),
            len: state.len,
            len_read: state.len_read,
            reduce: state.reduce,
            settings: state.settings,
        }
    }
    pub fn to_state(&self) -> ReduceOptimizationState {
        ReduceOptimizationState {
            trace: self.trace.clone(),
            trace_read_fallback: self.trace_read_fallback.clone(),
            trace_write_fallback: self.trace_write_fallback.clone(),
            len: self.len,
            len_read: self.len_read,
            reduce: self.reduce.clone(),
            settings: self.settings,
        }
    }
}

#[derive(Serialize, Deserialize, Copy, Clone)]
pub enum ReduceSettings {
    Always,
    /// We only activate fuse-on-write when the reduction isn't on the last dimension, otherwise
    /// vectorization is impossible. Only [VectorizationMode::Perpendicular] supports vectorization.
    ///
    /// We could still fuse some output operations, but it would probably lead to worse performance.
    OnlyParallel,
    Never,
}

pub(crate) struct ReduceOptimizationTuneArg {
    pub(crate) info: Arc<ReduceOptimizationInfo>,
    pub(crate) fallback: Arc<Box<dyn FallbackOperation>>,
}

impl Clone for ReduceOptimizationTuneArg {
    fn clone(&self) -> Self {
        Self {
            info: self.info.clone(),
            fallback: self.fallback.clone(),
        }
    }
}

#[derive(Clone, Copy, Serialize, Deserialize, Debug)]
pub enum ReduceInstruction {
    ArgMax,
    ArgMin,
    Mean,
    Prod,
    Sum,
    Max,
    Min,
    MaxAbs,
    Any,
    All,
}

pub trait ReduceFallbackFn: Send + Sync {
    fn run(&self, context: &mut Context<CubeFusionHandle>);
}

#[derive(Serialize, Deserialize)]
pub struct ReduceOptimizationState {
    pub(crate) trace: FuseTrace,
    pub(crate) trace_read_fallback: FuseTrace,
    pub(crate) trace_write_fallback: FuseTrace,
    pub(crate) reduce: FusedReduce,
    pub(crate) len: usize,
    pub(crate) len_read: usize,
    pub(crate) settings: ReduceSettings,
}

impl core::fmt::Debug for ReduceOptimizationState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!(
            "{{ len_read: {}, len_total: {} }}",
            self.len_read, self.len
        ))
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct FusedReduce {
    pub(crate) input: FuseArg,
    pub(crate) output: FuseArg,
    pub(crate) acc: FuseType,
    pub(crate) axis: usize,
    pub(crate) op: ReduceDimOpIr,
    pub(crate) use_planes: bool,
    pub(crate) shared: bool,
    pub(crate) inst: ReduceInstruction,
}

#[derive(new)]
pub struct FusedReduceLaunch<'a> {
    reduce: &'a FusedReduce,
    strategy: RoutineStrategy,
}

#[derive(Debug)]
pub enum FusedReduceError {
    Reduce(ReduceError),
    InvalidSelection(Box<&'static str>),
    InvalidInput,
}

impl From<ReduceError> for FusedReduceError {
    fn from(value: ReduceError) -> Self {
        Self::Reduce(value)
    }
}

impl ReduceOptimizationTuneArg {
    pub fn execute_fused(
        &self,
        context: &mut Context<CubeFusionHandle>,
        strategy: RoutineStrategy,
    ) -> Result<TuneOutput, TraceError<FusedReduceError>> {
        let launch = FusedReduceLaunch::new(&self.info.reduce, strategy);
        let launcher = FuseTraceLauncher::new(&self.info.trace, &launch);
        launcher.launch(&self.info.client, &self.info.device, context)
    }

    pub fn execute_fallback(&self, context: &mut Context<CubeFusionHandle>) -> TuneOutput {
        let launcher = FuseTraceLauncher::new(&self.info.trace_read_fallback, &ElemwiseRunner);

        #[allow(unused_mut)] // It is used when `autotune-checks` is activated.
        let mut output_read = launcher
            .launch(&self.info.client, &self.info.device, context)
            .unwrap();

        self.fallback.run(context);

        #[cfg(feature = "autotune-checks")]
        if let TuneOutput::Checked { handles } = &mut output_read {
            let out_desc = context.tensors.get(&self.info.reduce.op.out.id).unwrap();
            let handle_out = context
                .handles
                .get_handle(&out_desc.id, &burn_ir::TensorStatus::ReadOnly);

            handles.insert(
                self.info.reduce.op.out.id,
                (out_desc.shape.clone(), handle_out.clone()),
            );
        }

        let launcher = FuseTraceLauncher::new(&self.info.trace_write_fallback, &ElemwiseRunner);

        let output_write = launcher
            .launch(&self.info.client, &self.info.device, context)
            .unwrap();

        output_read.merge(output_write)
    }
}

#[allow(clippy::too_many_arguments)]
impl ReduceOptimization {
    pub fn new(
        trace: FuseTrace,
        trace_read_fallback: FuseTrace,
        trace_write_fallback: FuseTrace,
        client: Client,
        device: cubecl::Device,
        len: usize,
        len_read: usize,
        reduce: FusedReduce,
        settings: ReduceSettings,
    ) -> Self {
        let info = ReduceOptimizationInfo {
            trace,
            trace_read_fallback,
            trace_write_fallback,
            client,
            device,
            len,
            len_read,
            reduce,
            settings,
        };

        Self {
            info: Arc::new(info),
        }
    }
    /// Execute the optimization.
    pub fn execute(
        &mut self,
        context: &mut Context<CubeFusionHandle>,
        fallback: impl FnOnce(usize) -> Box<dyn FallbackOperation>,
    ) {
        // The index of the fallback reduce is the number of ops fused as read.
        let fallback = fallback(self.info.len_read);
        let arg = ReduceOptimizationTuneArg {
            info: self.info.clone(),
            fallback: Arc::new(fallback),
        };

        // Same as `ReduceBroadcastedOptimization::execute`: an empty input leaves the fused
        // kernels either an axis they reject or nothing to launch over.
        let input = context.tensors.get(&self.info.reduce.op.input.id).unwrap();
        if input.shape.num_elements() == 0 {
            arg.execute_fallback(context);
            return;
        }

        #[cfg(feature = "autotune")]
        fused_reduce_autotune(arg, context);

        #[cfg(not(feature = "autotune"))]
        if arg
            .execute_fused(
                context,
                RoutineStrategy::Unit(BlueprintStrategy::Inferred(UnitStrategy)),
            )
            .is_err()
        {
            arg.execute_fallback(context);
        }
    }

    pub fn num_output_buffers(&self) -> usize {
        self.info.trace_read_fallback.resources.outputs.len()
    }

    pub fn to_state(&self) -> ReduceOptimizationState {
        ReduceOptimizationState {
            trace: self.info.trace.clone(),
            trace_read_fallback: self.info.trace_read_fallback.clone(),
            trace_write_fallback: self.info.trace_write_fallback.clone(),
            reduce: self.info.reduce.clone(),
            len: self.info.len,
            len_read: self.info.len_read,
            settings: self.info.settings,
        }
    }

    pub fn from_state(device: &cubecl::Device, state: ReduceOptimizationState) -> Self {
        let client = device.client();

        let info = ReduceOptimizationInfo {
            trace: state.trace,
            trace_read_fallback: state.trace_read_fallback,
            trace_write_fallback: state.trace_write_fallback,
            reduce: state.reduce,
            len: state.len,
            len_read: state.len_read,
            client,
            device: device.clone(),
            settings: state.settings,
        };

        Self {
            info: Arc::new(info),
        }
    }

    /// Returns the number of output buffers added by fusion.
    pub fn num_ops_fused(&self) -> usize {
        self.info.len
    }
}

// TODO: Implement better vectorization here.
impl Vectorization for FusedReduceLaunch<'_> {}

impl TraceRunner for FusedReduceLaunch<'_> {
    type Error = FusedReduceError;

    fn run<'a>(
        &'a self,
        client: &'a Client,
        inputs: GlobalArgsLaunch,
        outputs: GlobalArgsLaunch,
        configs: &'a [FuseBlockConfig],
    ) -> Result<(), FusedReduceError> {
        let [config_read, config_write] = [&configs[0], &configs[1]];
        // An output-concrete reference indexes the output arguments; shape and
        // strides must both be resolved against that list.
        let ref_args = match &config_read.ref_layout {
            RefLayout::Concrete(FuseArg::Output(..)) => &outputs,
            _ => &inputs,
        };
        let shape = ref_args.shape_ref(&config_read.ref_layout, config_read.rank);
        let ref_strides = ref_args.strides_ref(&config_read.ref_layout, config_read.rank);
        let reduce_count: usize = shape
            .iter()
            .enumerate()
            .map(|(i, s)| if i == self.reduce.axis { 1 } else { *s })
            .product();

        let vectorization_mode = match self.reduce.axis == config_read.rank - 1 {
            true => VectorizationMode::Parallel,
            false => VectorizationMode::Perpendicular,
        };
        let address_type = inputs
            .required_address_type()
            .max(outputs.required_address_type());

        let settings = ReduceVectorSettings {
            vectorization_mode,
            vector_size_input: config_read.width,
            vector_size_output: config_write.width,
            // Fused-reduce selection is cached per anchored key, so the
            // unchecked comptime fast paths are never stable here.
            unchecked_fast_paths: false,
            // The fused read runs the trace, which may write materialized
            // intermediates to global outputs.
            fuse_on_read: true,
        };
        let problem = ReduceProblem {
            reduce_len: shape[self.reduce.axis],
            reduce_count,
            axis: self.reduce.axis,
            dtypes: ReduceDtypes {
                input: dtype_to_storage_type(self.reduce.op.input.dtype),
                output: dtype_to_storage_type(self.reduce.op.out.dtype),
                accumulation: self.reduce.acc.into_elem(),
            },
            address_type,
            instruction: reduce_instruction2config(&self.reduce.inst),
        };

        let (blueprint, settings) = match self.strategy.clone() {
            RoutineStrategy::Unit(strategy) => {
                let routine = UnitRoutine;
                routine.prepare(client, problem, settings, strategy)?
            }
            RoutineStrategy::Plane(strategy) => {
                let routine = PlaneRoutine;
                routine.prepare(client, problem, settings, strategy)?
            }
            RoutineStrategy::Cube(strategy) => {
                let routine = CubeRoutine;
                routine.prepare(client, problem, settings, strategy)?
            }
        };

        let out_vec_axis =
            output_vectorization_axis(&ref_strides, self.reduce.axis, vectorization_mode);

        let kwargs = ReduceKwArgs {
            client,
            inputs,
            outputs,
            reduce_axis: self.reduce.axis,
            out_vec_axis,
            config_fuse_read: config_read.clone(),
            config_fuse_write: config_write.clone(),
            input: self.reduce.input.clone(),
            output: self.reduce.output.clone(),
            blueprint,
            settings,
        };
        let result = launch_reduce_mixed_precision(
            kwargs,
            self.reduce.inst,
            self.reduce.op.input.dtype,
            self.reduce.op.out.dtype,
            elem_type_to_dtype(self.reduce.acc.into_elem()),
        );

        match result {
            Ok(_) => Ok(()),
            Err(err) => Err(FusedReduceError::Reduce(ReduceError::Launch(err))),
        }
    }
}

struct ReduceKwArgs<'b> {
    client: &'b Client,
    inputs: GlobalArgsLaunch,
    outputs: GlobalArgsLaunch,
    reduce_axis: usize,
    out_vec_axis: usize,
    blueprint: ReduceBlueprint,
    settings: ReduceLaunchSettings,
    config_fuse_read: FuseBlockConfig,
    config_fuse_write: FuseBlockConfig,
    input: FuseArg,
    output: FuseArg,
}

fn launch_reduce_mixed_precision(
    kwargs: ReduceKwArgs<'_>,
    instruction: ReduceInstruction,
    dtype_input: DType,
    dtype_output: DType,
    dtype_acc: DType,
) -> Result<(), LaunchError> {
    let config = reduce_instruction2config(&instruction);
    launch_reduce(kwargs, config, dtype_input, dtype_output, dtype_acc)
}

pub(crate) fn reduce_instruction2config(instruction: &ReduceInstruction) -> ReduceOperationConfig {
    match instruction {
        ReduceInstruction::ArgMax => ReduceOperationConfig::ArgMax,
        ReduceInstruction::ArgMin => ReduceOperationConfig::ArgMin,
        ReduceInstruction::Prod => ReduceOperationConfig::Prod,
        ReduceInstruction::Mean => ReduceOperationConfig::Mean,
        ReduceInstruction::Sum => ReduceOperationConfig::Sum,
        ReduceInstruction::Max => ReduceOperationConfig::Max,
        ReduceInstruction::Min => ReduceOperationConfig::Min,
        ReduceInstruction::MaxAbs => ReduceOperationConfig::MaxAbs,
        ReduceInstruction::Any => ReduceOperationConfig::Any,
        ReduceInstruction::All => ReduceOperationConfig::All,
    }
}

fn launch_reduce(
    kwargs: ReduceKwArgs<'_>,
    inst: ReduceOperationConfig,
    dtype_input: DType,
    dtype_output: DType,
    dtype_acc: DType,
) -> Result<(), LaunchError> {
    unsafe {
        reduce_kernel_fused::launch_unchecked(
            kwargs.client,
            kwargs.settings.cube_count,
            kwargs.settings.cube_dim,
            kwargs.settings.address_type,
            kwargs.config_fuse_read.width,
            kwargs.config_fuse_write.width,
            FusedReduceInputLaunch::new(kwargs.inputs, kwargs.config_fuse_read, kwargs.input),
            FusedReduceOutputLaunch::new(kwargs.outputs, kwargs.config_fuse_write, kwargs.output),
            kwargs.reduce_axis,
            kwargs.out_vec_axis,
            kwargs.blueprint,
            inst,
            dtype_to_storage_type(dtype_input),
            dtype_to_storage_type(dtype_output),
            dtype_to_storage_type(dtype_acc),
        )
    };

    Ok(())
}

#[cube(launch_unchecked, address_type = "dynamic")]
pub fn reduce_kernel_fused<In: Numeric, SizeIn: Size, Out: Numeric, SizeOut: Size, Acc: Numeric>(
    input: &FusedReduceInput,
    output: &mut FusedReduceOutput,
    reduce_axis: usize,
    out_vec_axis: usize,
    #[comptime] blueprint: ReduceBlueprint,
    #[comptime] config: ReduceOperationConfig,
    #[define(In)] _input_dtype: ElemType,
    #[define(Out)] _output_dtype: ElemType,
    #[define(Acc)] _acc_dtype: ElemType,
) {
    multi_block_variables_init(&input.config, &mut output.global.variables);
    multi_block_variables_init(&output.config, &mut output.global.variables);

    let (input, mut output) =
        init_tensors::<FusedReduceArgs, In, SizeIn, Out, SizeOut>(input, output);

    reduce_kernel_virtual::<In, SizeIn, Out, SizeOut, Acc>(
        &input,
        &mut output,
        reduce_axis,
        out_vec_axis,
        blueprint,
        config,
    );
}

/// Name of the reduce fusion optimization.
pub const NAME: &str = "Reduce";

impl ReduceOptimizationInfo {
    /// The highest relative shape id across this reduce's trace and both of its fallbacks.
    pub(crate) fn max_relative_shape_id(&self) -> Option<usize> {
        [
            self.trace.max_relative_shape_id(),
            self.trace_read_fallback.max_relative_shape_id(),
            self.trace_write_fallback.max_relative_shape_id(),
        ]
        .into_iter()
        .flatten()
        .max()
    }
}

impl FusedOperation for ReduceOptimization {
    fn max_relative_shape_id(&self) -> Option<usize> {
        self.info.max_relative_shape_id()
    }

    const NAME: &'static str = self::NAME;
    type State = ReduceOptimizationState;

    fn num_ops_fused(&self) -> usize {
        Self::num_ops_fused(self)
    }

    fn run(
        &mut self,
        context: &mut Context<CubeFusionHandle>,
        fallback: &dyn Fn(usize) -> Box<dyn FallbackOperation>,
    ) {
        Self::execute(self, context, |index| fallback(index))
    }

    fn to_state(&self) -> Self::State {
        Self::to_state(self)
    }

    fn from_state(device: &cubecl::Device, state: Self::State) -> Self {
        Self::from_state(device, state)
    }
}