From f08385402c686ab93696252ed6b9d7171b08fcf6 Mon Sep 17 00:00:00 2001 From: Ryutaro Okada <1015ryu88@gmail.com> Date: Sat, 31 Jan 2026 02:11:10 +0100 Subject: [PATCH] Add fused reduction support Co-authored-by: Anastasiia Filippova --- mlx/backend/cpu/compiled.cpp | 25 ++ mlx/backend/cuda/CMakeLists.txt | 19 +- mlx/backend/cuda/jit_module.cpp | 53 ++- mlx/backend/cuda/reduce.cu | 38 +++ mlx/backend/cuda/reduce/all_reduce.cu | 56 +-- mlx/backend/cuda/reduce/all_reduce.cuh | 76 +++++ mlx/backend/cuda/reduce/fused_reduce.cu | 322 ++++++++++++++++++ mlx/backend/cuda/reduce/reduce.cuh | 9 + mlx/backend/cuda/reduce/reduce_ops.cuh | 7 + mlx/backend/cuda/reduce/reduce_utils.cuh | 13 +- mlx/backend/metal/CMakeLists.txt | 1 + mlx/backend/metal/jit/includes.h | 1 + mlx/backend/metal/kernels/fused_reduce.h | 5 + .../metal/kernels/reduction/reduce_all.h | 15 +- mlx/backend/metal/reduce.cpp | 191 ++++++++++- mlx/backend/no_cpu/compiled.cpp | 7 + mlx/backend/no_gpu/primitives.cpp | 1 + mlx/compile.cpp | 164 ++++++++- mlx/primitives.h | 29 ++ python/tests/test_compile.py | 21 ++ tests/compile_tests.cpp | 101 ++++-- 21 files changed, 1014 insertions(+), 140 deletions(-) create mode 100644 mlx/backend/cuda/reduce/all_reduce.cuh create mode 100644 mlx/backend/cuda/reduce/fused_reduce.cu create mode 100644 mlx/backend/metal/kernels/fused_reduce.h diff --git a/mlx/backend/cpu/compiled.cpp b/mlx/backend/cpu/compiled.cpp index f5b516e44b..05ea0b4aeb 100644 --- a/mlx/backend/cpu/compiled.cpp +++ b/mlx/backend/cpu/compiled.cpp @@ -1,6 +1,7 @@ // Copyright © 2023-2024 Apple Inc. #include +#include #include #include #include @@ -291,6 +292,30 @@ inline void build_kernel( os << "}" << std::endl; } +void CompiledReduce::eval_cpu( + const std::vector& inputs, + std::vector& outputs) { + auto& encoder = cpu::get_command_encoder(stream()); + assert(outputs.size() == 1); + const auto& compiled_tape = tape(); + std::vector prefix_tape( + compiled_tape.begin(), compiled_tape.end() - 1); + Compiled prefix( + stream(), + this->inputs(), + {prefix_tape.back()}, + prefix_tape, + constant_ids()); + array prefix_out( + prefix.output_shapes(inputs)[0], prefix_tape.back().dtype(), nullptr, {}); + std::vector prefix_outputs{std::move(prefix_out)}; + prefix.eval_cpu(inputs, prefix_outputs); + encoder.add_temporary(prefix_outputs[0]); + + auto& reduce = static_cast(compiled_tape.back().primitive()); + reduce.eval_cpu(prefix_outputs, outputs[0]); +} + void Compiled::eval_cpu( const std::vector& inputs, std::vector& outputs) { diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index a82c5ad6e9..dd9af46be0 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -46,6 +46,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cu ${CMAKE_CURRENT_SOURCE_DIR}/reduce/all_reduce.cu ${CMAKE_CURRENT_SOURCE_DIR}/reduce/col_reduce.cu + ${CMAKE_CURRENT_SOURCE_DIR}/reduce/fused_reduce.cu ${CMAKE_CURRENT_SOURCE_DIR}/reduce/init_reduce.cu ${CMAKE_CURRENT_SOURCE_DIR}/reduce/row_reduce.cu ${CMAKE_CURRENT_SOURCE_DIR}/rms_norm.cu @@ -103,7 +104,8 @@ file( GLOB MLX_JIT_SOURCES RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "${CMAKE_CURRENT_SOURCE_DIR}/device/*.h" - "${CMAKE_CURRENT_SOURCE_DIR}/device/*.cuh") + "${CMAKE_CURRENT_SOURCE_DIR}/device/*.cuh" + "${CMAKE_CURRENT_SOURCE_DIR}/reduce/*.cuh") string(JOIN ":" MLX_JIT_SOURCES_ARG ${MLX_JIT_SOURCES}) add_custom_command( OUTPUT gen/cuda_jit_sources.h @@ -235,13 +237,16 @@ install(DIRECTORY ${cccl_SOURCE_DIR}/include/cuda DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl) install(DIRECTORY ${cccl_SOURCE_DIR}/include/nv DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl) +# CUB and Thrust are needed to JIT compile unary -> reduction kernels. +install(DIRECTORY ${cccl_SOURCE_DIR}/include/cub + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl) +install(DIRECTORY ${cccl_SOURCE_DIR}/include/thrust + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl) -# The binary of C++ tests will not be installed so it can not find the CCCL -# headers, and we have to hard-code the path. -if(MLX_BUILD_TESTS) - target_compile_definitions(mlx_dirs - PRIVATE MLX_CCCL_DIR="${cccl_SOURCE_DIR}/include") -endif() +# JIT cannot find the CCCL headers when compiling unary -> reduction kernels, so +# we hard-code the path to the downloaded CCCL. +target_compile_definitions(mlx_dirs + PRIVATE MLX_CCCL_DIR="${cccl_SOURCE_DIR}/include") # Use fixed version of NVTX. FetchContent_Declare( diff --git a/mlx/backend/cuda/jit_module.cpp b/mlx/backend/cuda/jit_module.cpp index 0d107a8e18..961884f218 100644 --- a/mlx/backend/cuda/jit_module.cpp +++ b/mlx/backend/cuda/jit_module.cpp @@ -259,47 +259,32 @@ bool compiler_supports_device_sass(Device& device) { } #define INCLUDE_PREFIX "mlx/backend/cuda/device/" +#define REDUCE_PREFIX "mlx/backend/cuda/reduce/" constexpr const char* g_include_names[] = { - INCLUDE_PREFIX "atomic_ops.cuh", - INCLUDE_PREFIX "binary_ops.cuh", - INCLUDE_PREFIX "cast_op.cuh", - INCLUDE_PREFIX "config.h", - INCLUDE_PREFIX "complex.cuh", - INCLUDE_PREFIX "cute_dequant.cuh", - INCLUDE_PREFIX "fp16_math.cuh", - INCLUDE_PREFIX "gemm_sm70.cuh", - INCLUDE_PREFIX "hadamard.cuh", - INCLUDE_PREFIX "indexing.cuh", - INCLUDE_PREFIX "qmm_naive.cuh", - INCLUDE_PREFIX "qmm_sm80.cuh", - INCLUDE_PREFIX "qmm_sm90.cuh", - INCLUDE_PREFIX "scatter_ops.cuh", - INCLUDE_PREFIX "unary_ops.cuh", - INCLUDE_PREFIX "ternary_ops.cuh", - INCLUDE_PREFIX "utils.cuh", + INCLUDE_PREFIX "atomic_ops.cuh", INCLUDE_PREFIX "binary_ops.cuh", + INCLUDE_PREFIX "cast_op.cuh", INCLUDE_PREFIX "config.h", + INCLUDE_PREFIX "complex.cuh", INCLUDE_PREFIX "cute_dequant.cuh", + INCLUDE_PREFIX "fp16_math.cuh", INCLUDE_PREFIX "gemm_sm70.cuh", + INCLUDE_PREFIX "hadamard.cuh", INCLUDE_PREFIX "indexing.cuh", + INCLUDE_PREFIX "qmm_naive.cuh", INCLUDE_PREFIX "qmm_sm80.cuh", + INCLUDE_PREFIX "qmm_sm90.cuh", INCLUDE_PREFIX "scatter_ops.cuh", + INCLUDE_PREFIX "unary_ops.cuh", INCLUDE_PREFIX "ternary_ops.cuh", + INCLUDE_PREFIX "utils.cuh", REDUCE_PREFIX "all_reduce.cuh", + REDUCE_PREFIX "reduce_ops.cuh", REDUCE_PREFIX "reduce_utils.cuh", }; #undef INCLUDE_PREFIX +#undef REDUCE_PREFIX constexpr const char* g_headers[] = { - jit_source_atomic_ops, - jit_source_binary_ops, - jit_source_cast_op, - jit_source_config, - jit_source_complex, - jit_source_cute_dequant, - jit_source_fp16_math, - jit_source_gemm_sm70, - jit_source_hadamard, - jit_source_indexing, - jit_source_qmm_naive, - jit_source_qmm_sm80, - jit_source_qmm_sm90, - jit_source_scatter_ops, - jit_source_unary_ops, - jit_source_ternary_ops, - jit_source_utils, + jit_source_atomic_ops, jit_source_binary_ops, jit_source_cast_op, + jit_source_config, jit_source_complex, jit_source_cute_dequant, + jit_source_fp16_math, jit_source_gemm_sm70, jit_source_hadamard, + jit_source_indexing, jit_source_qmm_naive, jit_source_qmm_sm80, + jit_source_qmm_sm90, jit_source_scatter_ops, jit_source_unary_ops, + jit_source_ternary_ops, jit_source_utils, jit_source_all_reduce, + jit_source_reduce_ops, jit_source_reduce_utils, }; void compile( diff --git a/mlx/backend/cuda/reduce.cu b/mlx/backend/cuda/reduce.cu index 833fdbf669..4ef30d8dfe 100644 --- a/mlx/backend/cuda/reduce.cu +++ b/mlx/backend/cuda/reduce.cu @@ -78,4 +78,42 @@ void Reduce::eval_gpu(const std::vector& inputs, array& out) { throw std::runtime_error("No plan reached in reduce."); } +void CompiledReduce::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + assert(inputs.size() == 1); + assert(outputs.size() == 1); + + const auto& reduce = static_cast(tape().back().primitive()); + auto [reduce_type, axes] = reduce.state(); + array in = inputs[0]; + auto& out = outputs[0]; + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + + if (in.size() == 0) { + init_reduce(encoder, in, out, reduce_type); + return; + } + + auto plan = get_reduction_plan(in, axes); + bool broadcasted = false; + for (int i = 0, j = 0; i < in.ndim() && !broadcasted; i++) { + if (j < axes.size() && axes[j] == i) { + j++; + } else { + broadcasted = in.strides(i) == 0; + } + } + if (plan.type == GeneralReduce || broadcasted || !in.flags().contiguous) { + array in_copy = contiguous_copy_gpu(in, s); + encoder.add_temporary(in_copy); + in = in_copy; + plan = get_reduction_plan(in, axes); + } + + std::vector contiguous_inputs{std::move(in)}; + fused_reduce(encoder, *this, contiguous_inputs, out, axes, plan, s); +} + } // namespace mlx::core diff --git a/mlx/backend/cuda/reduce/all_reduce.cu b/mlx/backend/cuda/reduce/all_reduce.cu index 167a282afb..b226312dc8 100644 --- a/mlx/backend/cuda/reduce/all_reduce.cu +++ b/mlx/backend/cuda/reduce/all_reduce.cu @@ -1,64 +1,10 @@ // Copyright © 2025 Apple Inc. -#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/reduce/all_reduce.cuh" #include "mlx/backend/cuda/reduce/reduce.cuh" -#include -#include -#include - namespace mlx::core { -namespace cu { - -namespace cg = cooperative_groups; - -template -__global__ void all_reduce(T* in, U* out, size_t block_step, size_t size) { - // TODO: Process multiple "rows" in each thread - constexpr int M = 1; - - auto grid = cg::this_grid(); - auto block = cg::this_thread_block(); - auto warp = cg::tiled_partition(block); - - const U init = cu::ReduceInit::value(); - ReduceOp op; - - T vals[N]; - U accs[M]; - accs[0] = init; - - size_t start = grid.block_rank() * block_step; - size_t end = start + block_step; - size_t check = min(end, size); - - size_t i = start; - for (; i + block.size() * N <= check; i += block.size() * N) { - cub::LoadDirectBlockedVectorized(block.thread_rank(), in + i, vals); - for (int j = 0; j < N; j++) { - accs[0] = op(accs[0], cast_to(vals[j])); - } - } - - if (i < check) { - cub::LoadDirectBlocked( - block.thread_rank(), in + i, vals, check - i, cast_to(init)); - for (int i = 0; i < N; i++) { - accs[0] = op(accs[0], cast_to(vals[i])); - } - } - - __shared__ U shared_accumulators[32]; - block_reduce(block, warp, accs, shared_accumulators, op, init); - - if (block.thread_rank() == 0) { - out[grid.block_rank()] = accs[0]; - } -} - -} // namespace cu - void all_reduce( cu::CommandEncoder& encoder, const array& in, diff --git a/mlx/backend/cuda/reduce/all_reduce.cuh b/mlx/backend/cuda/reduce/all_reduce.cuh new file mode 100644 index 0000000000..e661af56e1 --- /dev/null +++ b/mlx/backend/cuda/reduce/all_reduce.cuh @@ -0,0 +1,76 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#include "mlx/backend/cuda/device/config.h" +#include "mlx/backend/cuda/reduce/reduce_ops.cuh" + +#include +#include +#include + +namespace mlx::core::cu { + +namespace cg = cooperative_groups; + +template +__device__ void all_reduce_impl( + T* in, + U* out, + size_t block_step, + size_t size, + PrefixOp prefix) { + // TODO: Process multiple "rows" in each thread + constexpr int M = 1; + + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + + const U init = ReduceInit::value(); + ReduceOp op; + + T vals[N]; + U accs[M]; + accs[0] = init; + + size_t start = grid.block_rank() * block_step; + size_t end = start + block_step; + size_t check = min(end, size); + + size_t i = start; + for (; i + block.size() * N <= check; i += block.size() * N) { + cub::LoadDirectBlockedVectorized(block.thread_rank(), in + i, vals); + for (int j = 0; j < N; j++) { + accs[0] = op(accs[0], cast_to(prefix(vals[j]))); + } + } + + if (i < check) { + cub::LoadDirectBlocked( + block.thread_rank(), in + i, vals, check - i, cast_to(init)); + for (int j = 0; j < N; j++) { + accs[0] = op(accs[0], cast_to(prefix(vals[j]))); + } + } + + __shared__ U shared_accumulators[32]; + block_reduce(block, warp, accs, shared_accumulators, op, init); + + if (block.thread_rank() == 0) { + out[grid.block_rank()] = accs[0]; + } +} + +template < + typename T, + typename U, + typename ReduceOp, + int N = 4, + typename PrefixOp = Identity> +__global__ void all_reduce(T* in, U* out, size_t block_step, size_t size) { + all_reduce_impl( + in, out, block_step, size, PrefixOp{}); +} + +} // namespace mlx::core::cu diff --git a/mlx/backend/cuda/reduce/fused_reduce.cu b/mlx/backend/cuda/reduce/fused_reduce.cu new file mode 100644 index 0000000000..1dd37a9724 --- /dev/null +++ b/mlx/backend/cuda/reduce/fused_reduce.cu @@ -0,0 +1,322 @@ +// Copyright © 2025 Apple Inc. + +#include "mlx/backend/common/compiled.h" +#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/jit_module.h" +#include "mlx/backend/cuda/reduce/all_reduce.cuh" +#include "mlx/backend/cuda/reduce/reduce.cuh" +#include "mlx/backend/cuda/utils.h" +#include "mlx/graph_utils.h" +#include "mlx/primitives.h" + +#include +#include +#include +#include + +namespace mlx::core { + +namespace cu { + +// Builder to generate prefix functor code that will be +// applied in reduction kernel +struct FusedReducePrefixBuilder { + std::string os; + const std::string& kernel_name; + const std::vector& inputs; + const std::vector& tape; + const std::function& is_constant; + + void build_prefix_struct() { + NodeNamer namer; + + std::unordered_set constant_set; + for (size_t i = 0; i < inputs.size(); ++i) { + if (is_constant(i)) { + constant_set.insert(inputs[i].id()); + } + } + + os += "struct " + kernel_name + "_Prefix {\n"; + + os += "\n template \n"; + os += " __device__ __forceinline__ T operator()(T val) const {\n"; + + // Read the first input (reduction input) from the passed value + // Find the non-constant input + for (size_t i = 0; i < inputs.size(); ++i) { + const auto& x = inputs[i]; + const std::string& xname = namer.get_name(x); + std::string type = dtype_to_cuda_type(x.dtype()); + if (is_constant(i)) { + std::ostringstream ss; + print_constant(ss, x); + os += fmt::format( + " {} tmp_{} = static_cast<{}>({}); // constant\n", + type, + xname, + type, + ss.str()); + } else { + os += fmt::format( + " {} tmp_{} = static_cast<{}>(val);\n", type, xname, type); + } + } + + for (const auto& x : tape) { + const std::string& xname = namer.get_name(x); + std::string type = dtype_to_cuda_type(x.dtype()); + std::string value; + if (is_static_cast(x.primitive())) { + value = fmt::format( + "static_cast<{}>(tmp_{})", type, namer.get_name(x.inputs()[0])); + } else { + value = x.primitive().name(); + value += "{}("; + for (size_t i = 0; i < x.inputs().size() - 1; ++i) { + value += fmt::format("tmp_{}, ", namer.get_name(x.inputs()[i])); + } + value += fmt::format("tmp_{})", namer.get_name(x.inputs().back())); + } + os += fmt::format(" {} tmp_{} = {};\n", type, xname, value); + } + + // Return the result (last tape item or first input if tape is empty) + if (!tape.empty()) { + os += fmt::format( + " return static_cast(tmp_{});\n", namer.get_name(tape.back())); + } else { + // Find the non-constant input to return + for (size_t i = 0; i < inputs.size(); ++i) { + if (!is_constant(i)) { + os += fmt::format( + " return static_cast(tmp_{});\n", + namer.get_name(inputs[i])); + break; + } + } + } + + os += " }\n"; + + os += "};\n\n"; + } +}; + +} // namespace cu + +constexpr const char* g_fused_reduce_includes = R"( +#include "mlx/backend/cuda/device/binary_ops.cuh" +#include "mlx/backend/cuda/device/ternary_ops.cuh" +#include "mlx/backend/cuda/device/unary_ops.cuh" +#include "mlx/backend/cuda/device/utils.cuh" +#include "mlx/backend/cuda/reduce/all_reduce.cuh" + +#include + +#define inf cuda::std::numeric_limits::infinity() +)"; + +std::string get_reduce_op_name(Reduce::ReduceType reduce_type) { + switch (reduce_type) { + case Reduce::ReduceType::And: + return "And"; + case Reduce::ReduceType::Or: + return "Or"; + case Reduce::ReduceType::Sum: + return "Sum"; + case Reduce::ReduceType::Prod: + return "Prod"; + case Reduce::ReduceType::Max: + return "Max"; + case Reduce::ReduceType::Min: + return "Min"; + default: + throw std::runtime_error("Unknown reduce type"); + } +} + +void fused_all_reduce( + cu::CommandEncoder& encoder, + const CompiledReduce& compiled, + const std::vector& inputs, + array& out, + const Stream& stream) { + nvtx3::scoped_range r("fused_all_reduce"); + + // Copied from all_reduce.cu + constexpr int N_READS = 8; + + const auto& prefix_inputs = compiled.inputs(); + const auto& tape = compiled.tape(); + std::vector prefix_tape(tape.begin(), tape.end() - 1); + const auto& prefix_constant_ids = compiled.constant_ids(); + const auto& reduce = static_cast(tape.back().primitive()); + + auto is_constant = [&](size_t i) -> bool { + return prefix_constant_ids.count(prefix_inputs[i].id()) > 0; + }; + + std::string kernel_name = compiled.lib_name() + "_fused_reduce"; + + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + + // Copied from all_reduce.cu + auto get_args = [](int size, int N) { + int threads = std::min(512, (size + N - 1) / N); + threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; + int reductions_per_step = threads * N; + size_t steps_needed = + (size + reductions_per_step - 1) / reductions_per_step; + + int blocks; + if (steps_needed < 32) { + blocks = 1; + } else if (steps_needed < 128) { + blocks = 32; + } else if (steps_needed < 512) { + blocks = 128; + } else if (steps_needed < 1024) { + blocks = 512; + } else { + blocks = 1024; + } + + size_t steps_per_block = (steps_needed + blocks - 1) / blocks; + size_t block_step = steps_per_block * reductions_per_step; + + return std::make_tuple(blocks, threads, block_step); + }; + + std::string reduce_op = get_reduce_op_name(reduce.state().first); + std::string in_type = dtype_to_cuda_type(prefix_inputs[0].dtype()); + std::string prefix_type = kernel_name + "_Prefix"; + + std::string full_kernel_name = + fmt::format("mlx::core::cu::{}_all_reduce", kernel_name); + + cu::JitModule& mod = cu::get_jit_module(stream.device, kernel_name, [&]() { + cu::FusedReducePrefixBuilder builder{ + g_fused_reduce_includes, + kernel_name, + prefix_inputs, + prefix_tape, + is_constant}; + builder.os += + "namespace mlx::core::cu {\n\n" + "namespace cg = cooperative_groups;\n\n"; + + // Generate the prefix struct + builder.build_prefix_struct(); + + builder.os += fmt::format( + "__global__ void {}_all_reduce({}* in, {}* out, size_t block_step, size_t size) {{\n" + " {} prefix{{}};\n" + " all_reduce_impl<{}, {}, {}, {}, {}>(in, out, block_step, size, prefix);\n" + "}}\n", + kernel_name, + in_type, + in_type, + prefix_type, + in_type, + in_type, + reduce_op, + N_READS, + prefix_type); + + builder.os += "\n} // namespace mlx::core::cu\n"; + + std::vector kernel_names; + kernel_names.push_back(full_kernel_name); + + return std::make_tuple( + false, std::move(builder.os), std::move(kernel_names)); + }); + + int blocks, threads; + size_t block_step; + + size_t insize = inputs[0].size(); + Reduce::ReduceType reduce_type = reduce.state().first; + + std::tie(blocks, threads, block_step) = get_args(insize, N_READS); + + encoder.set_input_array(inputs[0]); + + // If the reduction needs more than 1 block -- use fused kernel with + // fused prefix, then reduce intermediate to final output using build compiled + // all_reduce + if (blocks > 1) { + array intermediate({blocks}, out.dtype(), nullptr, {}); + intermediate.set_data(cu::malloc_async(intermediate.nbytes(), encoder)); + encoder.add_temporary(intermediate); + encoder.set_output_array(intermediate); + + // First pass: apply fused prefix and reduce to intermediate + cu::KernelArgs args; + args.append(inputs[0]); + args.append(intermediate); + args.append(block_step); + args.append(insize); + + auto kernel = mod.get_kernel(full_kernel_name); + encoder.add_kernel_node(kernel, blocks, threads, 0, args.args()); + + // Second pass: reduce intermediate to final output + size_t intermediate_size = intermediate.size(); + std::tie(blocks, threads, block_step) = + get_args(intermediate_size, N_READS); + encoder.set_input_array(intermediate); + encoder.set_output_array(out); + + dispatch_all_types(out.dtype(), [&](auto type_tag) { + dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) { + using OP = MLX_GET_TYPE(reduce_type_tag); + using T = cuda_type_t; + using U = typename cu::ReduceResult::type; + auto kernel2 = cu::all_reduce; + encoder.add_kernel_node( + kernel2, + blocks, + threads, + 0, + gpu_ptr(intermediate), + gpu_ptr(out), + block_step, + intermediate_size); + }); + }); + } else { + // Single block: direct reduction with fused prefix + encoder.set_output_array(out); + + cu::KernelArgs args; + args.append(inputs[0]); + args.append(out); + args.append(block_step); + args.append(insize); + + auto kernel = mod.get_kernel(full_kernel_name); + encoder.add_kernel_node(kernel, blocks, threads, 0, args.args()); + } +} + +void fused_reduce( + cu::CommandEncoder& encoder, + const CompiledReduce& compiled, + const std::vector& inputs, + array& out, + const std::vector& axes, + const ReductionPlan& plan, + const Stream& stream) { + if (plan.type == ContiguousAllReduce) { + fused_all_reduce(encoder, compiled, inputs, out, stream); + return; + } + + // TODO: Implement fused row_reduce and col_reduce + throw std::runtime_error( + "Fused reduce not yet implemented for this reduction type"); +} + +} // namespace mlx::core diff --git a/mlx/backend/cuda/reduce/reduce.cuh b/mlx/backend/cuda/reduce/reduce.cuh index 02e495594a..21ff2448bc 100644 --- a/mlx/backend/cuda/reduce/reduce.cuh +++ b/mlx/backend/cuda/reduce/reduce.cuh @@ -68,4 +68,13 @@ void init_reduce( array& out, Reduce::ReduceType reduce_type); +void fused_reduce( + cu::CommandEncoder& encoder, + const CompiledReduce& compiled, + const std::vector& inputs, + array& out, + const std::vector& axes, + const ReductionPlan& plan, + const Stream& stream); + } // namespace mlx::core diff --git a/mlx/backend/cuda/reduce/reduce_ops.cuh b/mlx/backend/cuda/reduce/reduce_ops.cuh index 6c6b1827ce..a8302dded4 100644 --- a/mlx/backend/cuda/reduce/reduce_ops.cuh +++ b/mlx/backend/cuda/reduce/reduce_ops.cuh @@ -9,6 +9,13 @@ namespace mlx::core::cu { +struct Identity { + template + __device__ __forceinline__ T operator()(T x) const { + return x; + } +}; + // Reduce ops. struct And { __device__ __forceinline__ bool operator()(bool a, bool b) { diff --git a/mlx/backend/cuda/reduce/reduce_utils.cuh b/mlx/backend/cuda/reduce/reduce_utils.cuh index fdbe723378..4122e32b42 100644 --- a/mlx/backend/cuda/reduce/reduce_utils.cuh +++ b/mlx/backend/cuda/reduce/reduce_utils.cuh @@ -2,15 +2,18 @@ #pragma once -#include - -#include "mlx/backend/common/utils.h" -#include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/device/utils.cuh" #include #include +// Host-only includes +#ifndef __CUDACC_RTC__ +#include +#include "mlx/backend/common/utils.h" +#include "mlx/backend/cuda/device.h" +#endif + namespace mlx::core { namespace cu { @@ -90,6 +93,7 @@ block_reduce(Block block, Warp warp, T (&vals)[N], T* smem, Op op, T init) { } // namespace cu +#ifndef __CUDACC_RTC__ inline void allocate_same_layout( array& out, const array& in, @@ -141,5 +145,6 @@ inline void allocate_same_layout( fl, allocator::free); } +#endif } // namespace mlx::core diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index ea4a995ade..3ae0d9a6fb 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -27,6 +27,7 @@ make_jit_source(unary_ops kernels/erf.h kernels/expm1f.h kernels/fp8.h) make_jit_source(binary_ops) make_jit_source(ternary_ops) make_jit_source(reduce_utils kernels/atomic.h kernels/reduction/ops.h) +make_jit_source(fused_reduce kernels/defines.h kernels/reduction/reduce_all.h) make_jit_source(indexing/scatter kernels/indexing/indexing.h) make_jit_source(indexing/masked_scatter) make_jit_source(indexing/gather kernels/indexing/indexing.h) diff --git a/mlx/backend/metal/jit/includes.h b/mlx/backend/metal/jit/includes.h index 4fb1be1110..fee0daf8ab 100644 --- a/mlx/backend/metal/jit/includes.h +++ b/mlx/backend/metal/jit/includes.h @@ -9,6 +9,7 @@ const char* binary_ops(); const char* unary_ops(); const char* ternary_ops(); const char* reduce_utils(); +const char* fused_reduce(); const char* gather(); const char* scatter(); const char* masked_scatter(); diff --git a/mlx/backend/metal/kernels/fused_reduce.h b/mlx/backend/metal/kernels/fused_reduce.h new file mode 100644 index 0000000000..3b3a8892e7 --- /dev/null +++ b/mlx/backend/metal/kernels/fused_reduce.h @@ -0,0 +1,5 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include "mlx/backend/metal/kernels/reduction/reduce_all.h" diff --git a/mlx/backend/metal/kernels/reduction/reduce_all.h b/mlx/backend/metal/kernels/reduction/reduce_all.h index 47ad63fbbd..3749872eb8 100644 --- a/mlx/backend/metal/kernels/reduction/reduce_all.h +++ b/mlx/backend/metal/kernels/reduction/reduce_all.h @@ -1,11 +1,19 @@ // Copyright © 2023-2024 Apple Inc. +template +struct Identity { + T operator()(T x) thread { + return x; + } +}; + template < typename T, typename U, typename Op, typename IdxT = int64_t, - int N_READS = REDUCE_N_READS> + int N_READS = REDUCE_N_READS, + typename Prefix = Identity> [[kernel]] void all_reduce( const device T* in [[buffer(0)]], device U* out [[buffer(1)]], @@ -18,6 +26,7 @@ template < uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]]) { Op op; + Prefix prefix; threadgroup U shared_vals[simd_size]; U total = Op::init; @@ -37,13 +46,13 @@ template < for (IdxT b = 0; b < blocks; b++) { for (int i = 0; i < N_READS; i++) { - total = op(cast_to(in[i]), total); + total = op(cast_to(prefix(in[i])), total); } in += lsize.x * N_READS; } if (extra > 0) { for (int i = 0; i < extra; i++) { - total = op(cast_to(in[i]), total); + total = op(cast_to(prefix(in[i])), total); } } diff --git a/mlx/backend/metal/reduce.cpp b/mlx/backend/metal/reduce.cpp index ac11ac6359..aa61f0d172 100644 --- a/mlx/backend/metal/reduce.cpp +++ b/mlx/backend/metal/reduce.cpp @@ -2,13 +2,17 @@ #include #include +#include +#include "mlx/backend/common/compiled.h" #include "mlx/backend/gpu/copy.h" #include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/jit/includes.h" #include "mlx/backend/metal/kernels.h" #include "mlx/backend/metal/kernels/defines.h" #include "mlx/backend/metal/reduce.h" #include "mlx/backend/metal/utils.h" +#include "mlx/graph_utils.h" #include "mlx/primitives.h" #include "mlx/utils.h" @@ -301,6 +305,94 @@ std::pair remap_reduce_types( return {in.dtype(), in.dtype()}; } +MTL::ComputePipelineState* get_fused_all_reduce_kernel( + metal::Device& d, + const CompiledReduce& compiled, + const array& in, + const array& out, + const std::string& op_name) { + const auto& prefix_inputs = compiled.inputs(); + const auto& tape = compiled.tape(); + const auto& constant_ids = compiled.constant_ids(); + + if (prefix_inputs.size() != 1 || tape.size() < 2) { + throw std::runtime_error( + "Fused Metal all-reduce currently requires one prefix input"); + } + + auto is_constant = [&](size_t i) { + return constant_ids.count(prefix_inputs[i].id()) > 0; + }; + + const std::string kernel_name = compiled.lib_name() + "_fused_all_reduce"; + const std::string prefix_name = kernel_name + "_prefix"; + + auto lib = d.get_library(kernel_name, [&]() { + std::string source = metal::utils(); + concatenate( + source, + metal::unary_ops(), + metal::reduce_utils(), + metal::fused_reduce()); + + NodeNamer source_namer; + source += "struct " + prefix_name + " {\n"; + source += " " + get_type_string(tape[tape.size() - 2].dtype()) + + " operator()(" + get_type_string(prefix_inputs[0].dtype()) + + " val) thread {\n"; + + for (size_t i = 0; i < prefix_inputs.size(); ++i) { + const auto& x = prefix_inputs[i]; + const auto& xname = source_namer.get_name(x); + const auto type = get_type_string(x.dtype()); + if (is_constant(i)) { + std::ostringstream value; + print_constant(value, x); + source += " " + type + " tmp_" + xname + " = static_cast<" + type + + ">(" + value.str() + ");\n"; + } else { + source += " " + type + " tmp_" + xname + " = static_cast<" + type + + ">(val);\n"; + } + } + + for (auto it = tape.begin(); it != tape.end() - 1; ++it) { + const auto& x = *it; + const auto& xname = source_namer.get_name(x); + source += " " + get_type_string(x.dtype()) + " tmp_" + xname + " = "; + if (is_static_cast(x.primitive())) { + source += "static_cast<" + get_type_string(x.dtype()) + ">("; + source += "tmp_" + source_namer.get_name(x.inputs()[0]) + ");\n"; + } else { + source += x.primitive().name(); + source += "()("; + for (size_t i = 0; i + 1 < x.inputs().size(); ++i) { + source += "tmp_" + source_namer.get_name(x.inputs()[i]) + ", "; + } + source += "tmp_" + source_namer.get_name(x.inputs().back()) + ");\n"; + } + } + + source += " return tmp_" + source_namer.get_name(tape[tape.size() - 2]) + + ";\n }\n};\n"; + + std::string op_type = op_name; + op_type[0] = std::toupper(op_type[0]); + const auto out_type = get_type_string(out.dtype()); + source += get_template_definition( + kernel_name, + "all_reduce", + get_type_string(in.dtype()), + out_type, + op_type + "<" + out_type + ">", + "int64_t", + REDUCE_N_READS, + prefix_name); + return source; + }); + return d.get_kernel(kernel_name, lib); +} + void init_reduce( array& out, const std::string& op_name, @@ -327,20 +419,16 @@ void init_reduce( compute_encoder.dispatch_threads(grid_dims, group_dims); } -void all_reduce_dispatch( +void all_reduce_dispatch_impl( const array& in, array& out, const std::string& op_name, + const Dtype& out_type, + MTL::ComputePipelineState* kernel, CommandEncoder& compute_encoder, metal::Device& d, const Stream& s) { - // Set the kernel - auto [in_type, out_type] = remap_reduce_types(in, op_name); const std::string func_name = "all_reduce"; - std::string kname = func_name; - concatenate(kname, "_", op_name, type_to_name(in_type)); - auto kernel = get_reduce_kernel( - d, kname, func_name, op_name, in_type, out_type, "int64_t"); compute_encoder.set_compute_pipeline_state(kernel); size_t in_size = in.size(); @@ -408,6 +496,36 @@ void all_reduce_dispatch( } } +void all_reduce_dispatch( + const array& in, + array& out, + const std::string& op_name, + CommandEncoder& compute_encoder, + metal::Device& d, + const Stream& s) { + auto [in_type, out_type] = remap_reduce_types(in, op_name); + const std::string func_name = "all_reduce"; + std::string kname = func_name; + concatenate(kname, "_", op_name, type_to_name(in_type)); + auto kernel = get_reduce_kernel( + d, kname, func_name, op_name, in_type, out_type, "int64_t"); + all_reduce_dispatch_impl( + in, out, op_name, out_type, kernel, compute_encoder, d, s); +} + +void fused_all_reduce_dispatch( + const array& in, + array& out, + const std::string& op_name, + const CompiledReduce& compiled, + CommandEncoder& compute_encoder, + metal::Device& d, + const Stream& s) { + auto kernel = get_fused_all_reduce_kernel(d, compiled, in, out, op_name); + all_reduce_dispatch_impl( + in, out, op_name, out.dtype(), kernel, compute_encoder, d, s); +} + void row_reduce_small( const array& in, array& out, @@ -1061,4 +1179,63 @@ void Reduce::eval_gpu(const std::vector& inputs, array& out) { } } +void CompiledReduce::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + assert(inputs.size() == 1); + assert(outputs.size() == 1); + + const auto& reduce = static_cast(tape().back().primitive()); + auto [reduce_type, axes] = reduce.state(); + array in = inputs[0]; + auto& out = outputs[0]; + + size_t min_bytes = std::max(out.nbytes(), 4ul); + out.set_data(allocator::malloc(min_bytes)); + + std::string op_name; + switch (reduce_type) { + case Reduce::And: + op_name = "and"; + break; + case Reduce::Or: + op_name = "or"; + break; + case Reduce::Sum: + op_name = "sum"; + break; + case Reduce::Prod: + op_name = "prod"; + break; + case Reduce::Min: + op_name = out.dtype() == bool_ ? "and" : "min"; + break; + case Reduce::Max: + op_name = out.dtype() == bool_ ? "or" : "max"; + break; + } + + auto& s = stream(); + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + if (in.size() == 0) { + init_reduce(out, op_name, compute_encoder, d, s); + return; + } + + auto plan = get_reduction_plan(in, axes); + if (plan.type == GeneralReduce) { + array in_copy = contiguous_copy_gpu(in, s); + compute_encoder.add_temporary(in_copy); + in = in_copy; + plan = get_reduction_plan(in, axes); + } + + if (plan.type != ContiguousAllReduce) { + throw std::runtime_error( + "Fused Metal reduce is only implemented for contiguous all-reduce"); + } + fused_all_reduce_dispatch(in, out, op_name, *this, compute_encoder, d, s); +} + } // namespace mlx::core diff --git a/mlx/backend/no_cpu/compiled.cpp b/mlx/backend/no_cpu/compiled.cpp index 09987e7d58..f22db9e47d 100644 --- a/mlx/backend/no_cpu/compiled.cpp +++ b/mlx/backend/no_cpu/compiled.cpp @@ -22,4 +22,11 @@ void Compiled::eval_cpu( "[Compiled::eval_cpu] CPU compilation not supported on the platform."); } +void CompiledReduce::eval_cpu( + const std::vector& inputs, + std::vector& outputs) { + throw std::runtime_error( + "[CompiledReduce::eval_cpu] CPU compilation not supported on the platform."); +} + } // namespace mlx::core diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index 7f60d0d83a..19c31f0c5a 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -75,6 +75,7 @@ NO_GPU(Broadcast) NO_GPU(BroadcastAxes) NO_GPU(Ceil) NO_GPU_MULTI(Compiled) +NO_GPU_MULTI(CompiledReduce) NO_GPU(Concatenate) NO_GPU(Conjugate) NO_GPU(Contiguous) diff --git a/mlx/compile.cpp b/mlx/compile.cpp index bb17c44962..fccb0ad54b 100644 --- a/mlx/compile.cpp +++ b/mlx/compile.cpp @@ -72,7 +72,8 @@ bool is_noop(const Primitive& p) { } bool is_reduction(const Primitive& p) { - return typeid(p) == typeid(Reduce) || typeid(p) == typeid(ArgReduce); + // TODO: Add fused-prefix support for ArgReduce. + return typeid(p) == typeid(Reduce); } bool is_fusable(const Primitive& p) { @@ -122,6 +123,13 @@ Compiled::Compiled( } // computation performed os << a.primitive().name(); + const auto& primitive = a.primitive(); + if (typeid(primitive) == typeid(Reduce)) { + const auto& reduce = static_cast(primitive); + for (auto axis : reduce.state().second) { + os << "A" << axis; + } + } // name of inputs to the function for (auto& inp : a.inputs()) { os << namer.get_name(inp); @@ -209,10 +217,54 @@ std::vector Compiled::output_shapes(const std::vector& inputs) { out_shape[i] = std::max(out_shape[i], in.shape()[i - dd]); } } - // All outputs have the same shape + // Elementwise outputs have the same shape. return std::vector(outputs_.size(), out_shape); } +CompiledReduce::CompiledReduce( + Stream stream, + std::vector inputs, + std::vector outputs, + std::vector tape, + std::unordered_set constant_ids) + : Compiled( + stream, + std::move(inputs), + std::move(outputs), + std::move(tape), + std::move(constant_ids)) { + if (this->tape().size() < 2) { + throw std::invalid_argument( + "[CompiledReduce] Expected an elementwise prefix and a reduction"); + } + const auto& terminal = this->tape().back().primitive(); + if (typeid(terminal) != typeid(Reduce)) { + throw std::invalid_argument( + "[CompiledReduce] Expected a terminal Reduce primitive"); + } +} + +std::vector CompiledReduce::output_shapes( + const std::vector& inputs) { + size_t nd = 0; + for (const auto& in : inputs) { + nd = std::max(nd, in.ndim()); + } + Shape out_shape(nd, 0); + for (const auto& in : inputs) { + auto dd = nd - in.ndim(); + for (auto i = dd; i < nd; ++i) { + out_shape[i] = std::max(out_shape[i], in.shape()[i - dd]); + } + } + + const auto& reduce = static_cast(tape().back().primitive()); + for (auto axis : reduce.state().second) { + out_shape[axis] = 1; + } + return {std::move(out_shape)}; +} + namespace detail { std::atomic& compile_mode() { @@ -829,6 +881,114 @@ void compile_fuse( if (global_cache.find(arr.id()) != global_cache.end()) { continue; } + // Fuse an elementwise prefix and a terminal all-reduce into one + // CompiledReduce primitive. + if (arr.has_primitive() && is_reduction(arr.primitive()) && + static_cast(arr.primitive()).state().second.size() == + arr.inputs()[0].ndim()) { + auto& reduction_input = arr.inputs()[0]; + Stream reduction_stream = arr.primitive().stream(); + const int max_prefix_depth = max_compile_depth - 1; // 1 for reduction + + std::vector prefix_tape; + std::vector prefix_inputs; + std::unordered_set visited; + + std::function collect_prefix; + collect_prefix = [&](const array& a, int depth) { + // Skip if already processed + if (visited.count(a.id())) { + return; + } + // Stop fusing if: + // depth limit exceeded + // non unary primitive + // does not have primitive + // stream mismatch + // is a function input or output + if (depth >= max_prefix_depth || !a.has_primitive() || + !is_unary(a.primitive()) || + a.primitive().stream() != reduction_stream || + input_ids.count(a.id()) || output_map.count(a.id())) { + prefix_inputs.push_back(a); + visited.insert(a.id()); + return; + } + // Check if the input is used multiple times + auto pit = parents_map.find(a.id()); + if (pit != parents_map.end() && pit->second.size() > 1) { + prefix_inputs.push_back(a); + visited.insert(a.id()); + return; + } + visited.insert(a.id()); + for (auto& in : a.inputs()) { + collect_prefix(in, depth + 1); + } + prefix_tape.push_back(a); + }; + + collect_prefix(reduction_input, 0); + + // If there are operations that we can fuse + if (!prefix_tape.empty()) { + std::unordered_set constant_ids; + for (auto& in : prefix_inputs) { + if (in.size() == 1 && !in.has_primitive() && + input_ids.find(in.id()) == input_ids.end()) { + constant_ids.insert(in.id()); + } + } + + std::vector fused_tape = prefix_tape; + fused_tape.push_back(arr); + std::unordered_set fused_ids; + for (const auto& a : fused_tape) { + fused_ids.insert(a.id()); + } + + auto compiled_output = array( + arr.shape(), + arr.dtype(), + std::make_shared( + reduction_stream, + prefix_inputs, + std::vector{arr}, + std::move(fused_tape), + std::move(constant_ids)), + prefix_inputs); + + new_tape.push_back(compiled_output); + + for (int j = 0; j < prefix_inputs.size(); ++j) { + auto& pairs = parents_map[prefix_inputs[j].id()]; + pairs.erase( + std::remove_if( + pairs.begin(), + pairs.end(), + [&](auto& p) { return fused_ids.count(p.first.id()); }), + pairs.end()); + pairs.push_back({compiled_output, j}); + } + + for (auto& p : prefix_tape) { + global_cache.insert(p.id()); + parents_map.erase(p.id()); + } + global_cache.insert(arr.id()); + + merge_one(compiled_output, arr, parents_map); + if (auto it = output_map.find(arr.id()); it != output_map.end()) { + it->second = compiled_output; + } + continue; + } + + // Keep reductions without a fusable prefix unchanged. + new_tape.push_back(arr); + global_cache.insert(arr.id()); + continue; + } // Two pass recursion: // First pass: diff --git a/mlx/primitives.h b/mlx/primitives.h index 3a3d0ba5e5..f1144f2863 100644 --- a/mlx/primitives.h +++ b/mlx/primitives.h @@ -659,6 +659,18 @@ class MLX_API Compiled : public Primitive { return kernel_lib_; } + const std::vector& inputs() const { + return inputs_; + } + + const std::vector& tape() const { + return tape_; + } + + const std::unordered_set& constant_ids() const { + return constant_ids_; + } + private: const std::vector inputs_; const std::vector outputs_; @@ -670,6 +682,23 @@ class MLX_API Compiled : public Primitive { std::string kernel_lib_; }; +class MLX_API CompiledReduce final : public Compiled { + public: + explicit CompiledReduce( + Stream stream, + std::vector inputs, + std::vector outputs, + std::vector tape, + std::unordered_set constant_ids); + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override; + void eval_gpu(const std::vector& inputs, std::vector& outputs) + override; + + std::vector output_shapes(const std::vector& inputs) override; +}; + class Concatenate : public UnaryPrimitive { public: explicit Concatenate(Stream stream, int axis) diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 1e1b20b05d..0504ae17b2 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -1616,6 +1616,27 @@ def p(x): expected = w[::-1, :, ::-1, :] + 1.0 self.assertTrue(mx.array_equal(p(w[::-1, :, ::-1, :]), expected)) + def test_compile_unary_reduction(self): + + x = mx.ones(shape=(4096, 4096)) + y = mx.ones(shape=(4096, 4096)) + + @mx.compile + def abs_max(x): + return mx.max(mx.abs(x)) + + # Metal all-reduce uses one pass up to 4096 elements and two passes + # above that threshold. Negative endpoints ensure that skipping abs + # changes the maximum by one. + for shape in ((32, 32), (128, 128)): + size = math.prod(shape) + x = mx.arange(-size // 2, size // 2, dtype=mx.float32).reshape(shape) + mx.eval(x) + + out = abs_max(x) + expected = mx.max(mx.abs(x)) + self.assertTrue(mx.array_equal(out, expected)) + def test_compile_abs_unsigned(self): # abs has to compile for the wider unsigned types too fun = lambda x: mx.abs(x) + 1 diff --git a/tests/compile_tests.cpp b/tests/compile_tests.cpp index 30c2f887ac..e84cf7504c 100644 --- a/tests/compile_tests.cpp +++ b/tests/compile_tests.cpp @@ -294,6 +294,12 @@ auto unary_fused_3(const std::vector& inputs) { return std::vector{exp(abs(negative(sum(inputs[0], true))))}; } +auto unary_fused_reduction_with_prefix_output( + const std::vector& inputs) { + auto x = abs(inputs[0]); + return std::vector{x, sum(exp(x), true)}; +} + TEST_CASE("test compile unary fused") { // NB: some of these tests are brittle and may need to be // updated if we change compile conditions @@ -327,13 +333,11 @@ TEST_CASE("test compile unary fused") { CHECK_EQ(out.size(), 1); auto& p = out[0].primitive(); - // NB: this test is brittle, will need to update - // it if we change compile conditions - CHECK_EQ(typeid(p), typeid(Reduce)); - auto cout = out[0].inputs()[0]; - auto& cp = cout.primitive(); - CHECK_EQ(typeid(cp), typeid(Compiled)); - CHECK_EQ(cout.inputs()[0].id(), x.id()); + CHECK_EQ(typeid(p), typeid(CompiledReduce)); + auto& compiled = static_cast(p); + auto& terminal = compiled.tape().back().primitive(); + CHECK_EQ(typeid(terminal), typeid(Reduce)); + CHECK_EQ(out[0].inputs()[0].id(), x.id()); } { @@ -350,6 +354,21 @@ TEST_CASE("test compile unary fused") { CHECK_EQ(sout.inputs()[0].id(), x.id()); } + { + auto cfun = compile(unary_fused_reduction_with_prefix_output); + auto x = array({1.0, -2.0}); + auto out = cfun({x}); + + auto& prefix_output = out[0].primitive(); + auto& reduction_output = out[1].primitive(); + CHECK_EQ(typeid(prefix_output), typeid(Abs)); + CHECK_EQ(typeid(reduction_output), typeid(CompiledReduce)); + CHECK_EQ(out[1].inputs()[0].id(), out[0].id()); + auto expected = unary_fused_reduction_with_prefix_output({x}); + CHECK(allclose(out[0], expected[0]).item()); + CHECK(allclose(out[1], expected[1]).item()); + } + // Is equivalent works { auto out1 = compile(unary_fused_1)({array(1.0)}); @@ -376,11 +395,6 @@ auto binary_fused_2(const std::vector& inputs) { return std::vector{x + inputs[0]}; } -// Binary into unary into un-compilable -auto binary_fused_3(const std::vector& inputs) { - return std::vector{sum(abs(inputs[0] + inputs[1]), true)}; -} - TEST_CASE("test compile binary fused") { { auto cfun = compile(binary_fused_0); @@ -419,22 +433,6 @@ TEST_CASE("test compile binary fused") { CHECK_EQ(out.inputs()[0].id(), x.id()); CHECK_EQ(out.inputs()[1].id(), y.id()); } - - { - auto cfun = compile(binary_fused_3); - auto x = array({1.0, 2.0}); - auto y = array({1.0, 2.0}); - auto out = cfun({x, y})[0]; - - auto& p = out.primitive(); - CHECK_EQ(typeid(p), typeid(Reduce)); - - auto cout = out.inputs()[0]; - auto& cp = cout.primitive(); - CHECK_EQ(typeid(cp), typeid(Compiled)); - CHECK_EQ(cout.inputs()[0].id(), x.id()); - CHECK_EQ(cout.inputs()[1].id(), y.id()); - } } auto gelu_1(const std::vector& inputs) { @@ -875,3 +873,50 @@ TEST_CASE("test compile throwing first trace does not poison cache") { REQUIRE_EQ(out.size(), 1); CHECK_EQ(out[0].item(), 3.0f); } + +TEST_CASE("test compile unary reduction one pass") { + auto fun = [](const std::vector& inputs) { + auto x = inputs[0]; + return std::vector{max(abs(x))}; + }; + auto in = reshape(arange(-8192, 8192, float32), {128, 128}); + eval(in); + auto expected = fun({in})[0]; + auto out = compile(fun)({in})[0]; + CHECK(array_equal(out, expected).item()); +} + +TEST_CASE("test compile unary reduction two passes") { + auto fun = [](const std::vector& inputs) { + auto x = inputs[0]; + return std::vector{max(abs(x))}; + }; + auto in = reshape(arange(-524288, 524288, float32), {1024, 1024}); + eval(in); + auto expected = fun({in})[0]; + auto out = compile(fun)({in})[0]; + CHECK(array_equal(out, expected).item()); +} + +TEST_CASE("test compile partial reduction") { + auto fun = [](const std::vector& inputs) { + return std::vector{max(abs(inputs[0]), 1, true)}; + }; + auto x = reshape(array({1.0f, -2.0f}), {1, 2}); + auto out = compile(fun)({x})[0]; + + auto& primitive = out.primitive(); + CHECK_EQ(typeid(primitive), typeid(Reduce)); + CHECK(array_equal(out, fun({x})[0]).item()); +} + +TEST_CASE("test compile unary+constant reduction two passes") { + auto fun = [](const std::vector& inputs) { + auto x = inputs[0]; + return std::vector{sum(abs(x) + 1.0f)}; + }; + auto in = ones({1024, 1024}); + auto expected = fun({in})[0]; + auto out = compile(fun)({ones({1024, 1024})})[0]; + CHECK(array_equal(out, expected).item()); +}