From 7cc382f0a7d85f22e920c7d007e270d87e12cc0c Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 18:31:30 +0200 Subject: [PATCH 01/24] madspace: rename dtypes to int32/float64 and add float32 DataType becomes {dt_int32, dt_float32, dt_float64, batch_sizes}. float32 is supported in npy I/O, DLPack import/export, graph JSON and the Python enum. Tensor views now check that the element type matches the tensor dtype. --- .github/gpu_ci/madspace_checks.py | 4 +- madspace/generate_headers.py | 14 +- madspace/include/madspace/compgraphs/type.hpp | 62 ++++--- madspace/include/madspace/driver/tensor.hpp | 35 ++-- madspace/src/compgraphs/instruction.cpp | 57 ++++--- madspace/src/compgraphs/type.cpp | 39 +++-- madspace/src/cpu/device.cpp | 11 +- madspace/src/cpu/runtime.cpp | 69 ++++---- madspace/src/driver/context.cpp | 9 +- madspace/src/driver/discrete_optimizer.cpp | 2 +- madspace/src/driver/event_histograms.cpp | 57 ++++--- madspace/src/driver/io.cpp | 15 +- madspace/src/driver/systematics.cpp | 155 ++++++++++-------- madspace/src/driver/tensor.cpp | 13 +- madspace/src/gpu/device.cu | 14 +- madspace/src/gpu/runtime.cu | 142 +++++++++------- madspace/src/kernels/operations.hpp | 21 ++- .../src/phasespace/channel_weight_network.cpp | 6 +- madspace/src/phasespace/discrete_flow.cpp | 2 +- madspace/src/phasespace/discrete_sampler.cpp | 10 +- madspace/src/phasespace/flow.cpp | 2 +- madspace/src/phasespace/integrand.cpp | 10 +- madspace/src/phasespace/mlp.cpp | 12 +- madspace/src/phasespace/pdf.cpp | 30 ++-- madspace/src/phasespace/phasespace.cpp | 72 ++++---- madspace/src/phasespace/vegas.cpp | 10 +- madspace/src/python/function_runtime.cpp | 21 ++- madspace/src/python/madspace.cpp | 5 +- madspace/tests/test_histogram_op.py | 10 +- madspace/tests/test_simd_layout.py | 12 +- 30 files changed, 521 insertions(+), 400 deletions(-) diff --git a/.github/gpu_ci/madspace_checks.py b/.github/gpu_ci/madspace_checks.py index c2a557b652..0ebf7833ea 100644 --- a/.github/gpu_ci/madspace_checks.py +++ b/.github/gpu_ci/madspace_checks.py @@ -37,7 +37,7 @@ def to_numpy(tensor, torch): def check_histogram(torch, device, n=100000, n_bins=40): - out_type = ms.Type(ms.DataType.float, ms.BatchSize.one, [n_bins + 2]) + out_type = ms.Type(ms.DataType.float64, ms.BatchSize.one, [n_bins + 2]) fb = ms.FunctionBuilder( ms.NamedTypes([(name, ms.batch_float) for name in ('x', 'w', 'lo', 'hi')]), ms.NamedTypes([('values', out_type), ('square_values', out_type)]), @@ -62,7 +62,7 @@ def check_histogram(torch, device, n=100000, n_bins=40): def check_rank4_add(torch, device, n=4096): - rank4 = ms.Type(ms.DataType.float, ms.batch_size, [2, 3, 4]) + rank4 = ms.Type(ms.DataType.float64, ms.batch_size, [2, 3, 4]) fb = ms.FunctionBuilder(ms.NamedTypes([('x', rank4)]), ms.NamedTypes([('y', rank4)])) x = fb.input(0) fb.output(0, fb.add(x, x)) diff --git a/madspace/generate_headers.py b/madspace/generate_headers.py index d77b66ba10..16ccdec4b9 100644 --- a/madspace/generate_headers.py +++ b/madspace/generate_headers.py @@ -181,12 +181,16 @@ def instruction_set_python(commands, out_dir): f.write("}\n}\n") +# yaml shorthands for the default element types +DTYPE_NAMES = {"float": "float64", "int": "int32"} + + def format_type(data): type = data["type"] if type[0] == "size": - return f'{{DataType::dt_int, true, {{"{type[1]}"}}, true}}' + return f'{{DataType::dt_int32, true, {{"{type[1]}"}}, true}}' - dtype = f"DataType::dt_{type[0]}" + dtype = f"DataType::dt_{DTYPE_NAMES.get(type[0], type[0])}" single = len(type) > 1 and type[1] == "single" single_str = "true" if single else "false" shape = ", ".join( @@ -354,8 +358,10 @@ def runtime_backward_mixin(commands, device, out_dir): def is_single_input(arg): - return isinstance(arg.get("type"), list) and len(arg["type"]) > 1 and ( - arg["type"][1] == "single" + return ( + isinstance(arg.get("type"), list) + and len(arg["type"]) > 1 + and (arg["type"][1] == "single") ) diff --git a/madspace/include/madspace/compgraphs/type.hpp b/madspace/include/madspace/compgraphs/type.hpp index 3084f9dc0d..dcded6f22a 100644 --- a/madspace/include/madspace/compgraphs/type.hpp +++ b/madspace/include/madspace/compgraphs/type.hpp @@ -13,20 +13,44 @@ namespace madspace { /** * Element data type of a compute-graph @ref Value. * - * `dt_int` and `dt_float` are ordinary integer and floating-point tensors. + * `dt_int32`, `dt_float32` and `dt_float64` are ordinary integer and + * floating-point tensors. * `batch_sizes` marks a value that instead carries a list of batch sizes, used * to split and recombine batches in multi-channel functions. */ enum class DataType { - dt_int, ///< Integer tensor. - dt_float, ///< Floating-point tensor. + dt_int32, ///< 32-bit integer tensor. + dt_float32, ///< Single-precision floating-point tensor. + dt_float64, ///< Double-precision floating-point tensor. batch_sizes, ///< A list of batch sizes rather than tensor data. }; using me_int_t = int; template -concept ScalarType = std::same_as || std::same_as; +concept ScalarType = + std::same_as || std::same_as || std::same_as; + +/// Scalar types that can be stored in a literal @ref Value. +template +concept LiteralType = std::same_as || std::same_as; + +/// The `DataType` corresponding to the C++ element type `T`. +template +constexpr DataType dtype_of() { + if constexpr (std::is_same_v) { + return DataType::dt_int32; + } else if constexpr (std::is_same_v) { + return DataType::dt_float32; + } else { + return DataType::dt_float64; + } +} + +/// True for the floating-point data types `dt_float32` and `dt_float64`. +inline bool is_float(DataType dtype) { + return dtype == DataType::dt_float32 || dtype == DataType::dt_float64; +} /** * Symbolic size of the batch dimension of a @ref Value. @@ -138,19 +162,19 @@ inline bool operator!=(const Type& lhs, const Type& rhs) { using TypeVec = std::vector; -const Type single_float{DataType::dt_float, BatchSize::One{}, {}}; -const Type single_int{DataType::dt_int, BatchSize::One{}, {}}; +const Type single_float{DataType::dt_float64, BatchSize::One{}, {}}; +const Type single_int{DataType::dt_int32, BatchSize::One{}, {}}; inline Type single_float_array(int count) { - return {DataType::dt_float, BatchSize::one, {count}}; + return {DataType::dt_float64, BatchSize::one, {count}}; } inline Type single_int_array(int count) { - return {DataType::dt_int, BatchSize::one, {count}}; + return {DataType::dt_int32, BatchSize::one, {count}}; } inline Type single_float_array_2d(int count1, int count2) { - return {DataType::dt_float, BatchSize::one, {count1, count2}}; + return {DataType::dt_float64, BatchSize::one, {count1, count2}}; } inline Type single_int_array_2d(int count1, int count2) { - return {DataType::dt_int, BatchSize::one, {count1, count2}}; + return {DataType::dt_int32, BatchSize::one, {count1, count2}}; } const BatchSize batch_size = BatchSize("batch_size"); @@ -158,21 +182,21 @@ Type batch_size_array(int count); /// A `batch_sizes` `Type` splitting `batch_size` into `count` symbolic /// per-channel sizes that add up to the total. Type multichannel_batch_size(int count); -const Type batch_float{DataType::dt_float, batch_size, {}}; -const Type batch_int{DataType::dt_int, batch_size, {}}; -const Type batch_four_vec{DataType::dt_float, batch_size, {4}}; +const Type batch_float{DataType::dt_float64, batch_size, {}}; +const Type batch_int{DataType::dt_int32, batch_size, {}}; +const Type batch_four_vec{DataType::dt_float64, batch_size, {4}}; /// A `float` `Type` with a batch dimension and a static trailing shape of /// `count` elements. inline Type batch_float_array(int count) { - return {DataType::dt_float, batch_size, {count}}; + return {DataType::dt_float64, batch_size, {count}}; } inline Type batch_int_array(int count) { - return {DataType::dt_int, batch_size, {count}}; + return {DataType::dt_int32, batch_size, {count}}; } /// A `Type` of `count` four-vectors: batch dimension plus a trailing shape of /// `(count, 4)`. inline Type batch_four_vec_array(int count) { - return {DataType::dt_float, batch_size, {count, 4}}; + return {DataType::dt_float64, batch_size, {count, 4}}; } using TensorValue = std::tuple< @@ -207,7 +231,7 @@ struct Value { /// Scalar floating-point constant. Value(double value) : type(single_float), literal_value(value) {} - template + template Value(const std::vector>& values) : Value( [&] { @@ -231,10 +255,10 @@ struct Value { {static_cast(values.size()), static_cast(values.at(0).size())} ) {} - template + template Value(const std::vector& values, const std::vector& shape = {}) : type{ - std::is_same_v ? DataType::dt_int : DataType::dt_float, + std::is_same_v ? DataType::dt_int32 : DataType::dt_float64, BatchSize::one, shape.size() == 0 ? std::vector{static_cast(values.size())} : shape diff --git a/madspace/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp index 0223c9dc50..7bdc99a066 100644 --- a/madspace/include/madspace/driver/tensor.hpp +++ b/madspace/include/madspace/driver/tensor.hpp @@ -446,11 +446,7 @@ class Tensor { /// Single-value tensor holding `value`, allocated on `device`. template Tensor(T value, DevicePtr device, AllocHint hint = AllocHint::normal) : - impl(new TensorImpl{ - std::is_same_v ? DataType::dt_int : DataType::dt_float, - {1}, - device - }) { + impl(new TensorImpl{dtype_of(), {1}, device}) { auto size = init_stride(); allocate(size, *device, hint); device->memcpy(impl->data, &value, sizeof(value)); @@ -461,11 +457,7 @@ class Tensor { template Tensor(T value, const D& device, AllocHint hint = AllocHint::normal) : - impl(new TensorImpl{ - std::is_same_v ? DataType::dt_int : DataType::dt_float, - {1}, - device.device_ptr() - }) { + impl(new TensorImpl{dtype_of(), {1}, device.device_ptr()}) { auto size = init_stride(); allocate(size, device, hint); device.memcpy(impl->data, &value, sizeof(value)); @@ -479,8 +471,8 @@ class Tensor { impl(new TensorImpl{ std::visit( Overloaded{ - [](std::vector) { return DataType::dt_int; }, - [](std::vector) { return DataType::dt_float; }, + [](std::vector) { return DataType::dt_int32; }, + [](std::vector) { return DataType::dt_float64; }, }, std::get<1>(value) ), @@ -530,6 +522,7 @@ class Tensor { TensorView view() { check_impl(); check_view_dim(dim, impl->shape.size()); + check_view_dtype(); T* data = static_cast(impl->data); return TensorView(data, impl->stride.data(), impl->shape.data()); } @@ -539,6 +532,7 @@ class Tensor { const TensorView view() const { check_impl(); check_view_dim(dim, impl->shape.size()); + check_view_dtype(); T* data = static_cast(impl->data); return TensorView(data, impl->stride.data(), impl->shape.data()); } @@ -547,6 +541,7 @@ class Tensor { /// Like @ref view, with the leading `flatten_count` dimensions merged into one. PackedTensorView flat_view(std::size_t flatten_count) const { check_impl(); + check_view_dtype(); T* data = static_cast(impl->data); if (flatten_count <= 1) { check_view_dim(dim, impl->shape.size()); @@ -632,9 +627,11 @@ class Tensor { std::size_t dtype_size() const { check_impl(); switch (impl->dtype) { - case DataType::dt_int: + case DataType::dt_int32: return sizeof(me_int_t); - case DataType::dt_float: + case DataType::dt_float32: + return sizeof(float); + case DataType::dt_float64: return sizeof(double); case DataType::batch_sizes: return 0; @@ -920,6 +917,16 @@ class Tensor { } } + // reading a tensor through the wrong element type gives garbage without an error + template + void check_view_dtype() const { + if constexpr (ScalarType) { + if (impl->dtype != dtype_of() && impl->dtype != DataType::batch_sizes) { + throw std::invalid_argument("view element type does not match dtype"); + } + } + } + template void allocate(std::size_t size, const D& device, AllocHint hint) { auto [data, parent] = device.allocate(size, hint); diff --git a/madspace/src/compgraphs/instruction.cpp b/madspace/src/compgraphs/instruction.cpp index b05d0486ae..f608bfc0d9 100644 --- a/madspace/src/compgraphs/instruction.cpp +++ b/madspace/src/compgraphs/instruction.cpp @@ -138,7 +138,7 @@ me_int_t Instruction::int_literal_arg( ) const { auto& arg_val = args.at(index); auto& arg_type = arg_val.type; - if (arg_type.dtype != DataType::dt_int || arg_type.batch_size != BatchSize::one || + if (arg_type.dtype != DataType::dt_int32 || arg_type.batch_size != BatchSize::one || arg_type.shape.size() != 0 || !std::holds_alternative(arg_val.literal_value)) { throw std::invalid_argument( @@ -488,7 +488,7 @@ TypeVec BatchSplitInstruction::signature(const ValueVec& args) const { TypeVec BatchSplitByIndexInstruction::signature(const ValueVec& args) const { check_arg_count(args, 2); auto& indices_type = args.at(0).type; - if (indices_type.dtype != DataType::dt_int || indices_type.shape.size() != 0) { + if (indices_type.dtype != DataType::dt_int32 || indices_type.shape.size() != 0) { throw std::invalid_argument( std::format("{}, argument 1: expected batch of integers", name()) ); @@ -503,7 +503,7 @@ TypeVec BatchSplitByIndexInstruction::signature(const ValueVec& args) const { auto last_batch_size = indices_type.batch_size; for (int i = 0; i < size; ++i) { auto batch_size = i == size - 1 ? last_batch_size : BatchSize(); - output_types.push_back({DataType::dt_int, batch_size, {}}); + output_types.push_back({DataType::dt_int32, batch_size, {}}); last_batch_size = last_batch_size - batch_size; } return output_types; @@ -547,7 +547,8 @@ TypeVec BatchMergeByIndexInstruction::signature(const ValueVec& args) const { ) ); } - if (indices_type.dtype != DataType::dt_int || indices_type.shape.size() != 0) { + if (indices_type.dtype != DataType::dt_int32 || + indices_type.shape.size() != 0) { throw std::invalid_argument( std::format( "{}, argument {}: expected batch of integers", name(), i + 2 @@ -679,7 +680,7 @@ TypeVec OffsetIndicesInstruction::signature(const ValueVec& args) const { batch_sizes_out.type.batch_size_list.end(), BatchSize::zero ); - return {{DataType::dt_int, total_batch_size, {}}}; + return {{DataType::dt_int32, total_batch_size, {}}}; } TypeVec FullInstruction::signature(const ValueVec& args) const { @@ -749,7 +750,7 @@ TypeVec AcceptNormInstruction::signature(const ValueVec& args) const { } ++i; } - return {{DataType::dt_float, BatchSize::one, {}}}; + return {{DataType::dt_float64, BatchSize::one, {}}}; } TypeVec RqsReshapeInstruction::signature(const ValueVec& args) const { @@ -757,7 +758,7 @@ TypeVec RqsReshapeInstruction::signature(const ValueVec& args) const { int bin_count = int_literal_arg(args, 1); auto& input_type = args.at(0).type; - if (input_type.dtype != DataType::dt_float || input_type.shape.size() != 1 || + if (input_type.dtype != DataType::dt_float64 || input_type.shape.size() != 1 || input_type.shape.at(0) % (3 * bin_count + 1) != 0) { throw std::invalid_argument( std::format( @@ -769,28 +770,28 @@ TypeVec RqsReshapeInstruction::signature(const ValueVec& args) const { int dim = input_type.shape.at(0) / (3 * bin_count + 1); return { - {DataType::dt_float, input_type.batch_size, {dim, bin_count}}, - {DataType::dt_float, input_type.batch_size, {dim, bin_count}}, - {DataType::dt_float, input_type.batch_size, {dim, bin_count + 1}}, + {DataType::dt_float64, input_type.batch_size, {dim, bin_count}}, + {DataType::dt_float64, input_type.batch_size, {dim, bin_count}}, + {DataType::dt_float64, input_type.batch_size, {dim, bin_count + 1}}, }; } TypeVec NonzeroInstruction::signature(const ValueVec& args) const { check_arg_count(args, 1); auto& input_type = args.at(0).type; - if (input_type.dtype != DataType::dt_float || input_type.shape.size() != 0) { + if (input_type.dtype != DataType::dt_float64 || input_type.shape.size() != 0) { throw std::invalid_argument( std::format("{}, argument 1: expected batch of floats", name()) ); } - return {{DataType::dt_int, BatchSize(), {}}}; + return {{DataType::dt_int32, BatchSize(), {}}}; } TypeVec BatchGatherInstruction::signature(const ValueVec& args) const { check_arg_count(args, 2); auto& indices_type = args.at(0).type; auto& values_type = args.at(1).type; - if (indices_type.dtype != DataType::dt_int || indices_type.shape.size() != 0) { + if (indices_type.dtype != DataType::dt_int32 || indices_type.shape.size() != 0) { throw std::invalid_argument( std::format("{}, argument 1: expected batch of integers", name()) ); @@ -808,13 +809,13 @@ TypeVec BatchScatterInstruction::signature(const ValueVec& args) const { auto& indices_type = args.at(0).type; auto& target_type = args.at(1).type; auto& source_type = args.at(2).type; - if (indices_type.dtype != DataType::dt_int || indices_type.shape.size() != 0) { + if (indices_type.dtype != DataType::dt_int32 || indices_type.shape.size() != 0) { throw std::invalid_argument( std::format("{}, argument 1: expected batch of integers", name()) ); } - if (target_type.dtype != DataType::dt_float && - target_type.dtype != DataType::dt_int) { + if (target_type.dtype != DataType::dt_float64 && + target_type.dtype != DataType::dt_int32) { throw std::invalid_argument( std::format("{}, argument 2: expected data type float or int", name()) ); @@ -841,7 +842,7 @@ TypeVec RandomInstruction::signature(const ValueVec& args) const { ); } int count = int_literal_arg(args, 1); - return {{DataType::dt_float, batch_size_type.batch_size_list.at(0), {count}}}; + return {{DataType::dt_float64, batch_size_type.batch_size_list.at(0), {count}}}; } TypeVec RandomIntInstruction::signature(const ValueVec& args) const { @@ -860,19 +861,19 @@ TypeVec RandomIntInstruction::signature(const ValueVec& args) const { std::format("{}, argument 2: expected single batch size", name()) ); } - return {{DataType::dt_int, batch_size_type.batch_size_list.at(0), {}}}; + return {{DataType::dt_int32, batch_size_type.batch_size_list.at(0), {}}}; } TypeVec UnweightInstruction::signature(const ValueVec& args) const { check_arg_count(args, 2); auto& weights_type = args.at(0).type; auto& max_weight_type = args.at(1).type; - if (weights_type.dtype != DataType::dt_float || weights_type.shape.size() != 0) { + if (weights_type.dtype != DataType::dt_float64 || weights_type.shape.size() != 0) { throw std::invalid_argument( std::format("{}, argument 1: expected batch of floats", name()) ); } - if (max_weight_type.dtype != DataType::dt_float || + if (max_weight_type.dtype != DataType::dt_float64 || max_weight_type.batch_size != BatchSize::one || max_weight_type.shape.size() != 0) { throw std::invalid_argument( @@ -882,7 +883,8 @@ TypeVec UnweightInstruction::signature(const ValueVec& args) const { BatchSize out_batch_size; return { - {DataType::dt_int, out_batch_size, {}}, {DataType::dt_float, out_batch_size, {}} + {DataType::dt_int32, out_batch_size, {}}, + {DataType::dt_float64, out_batch_size, {}} }; } @@ -925,7 +927,7 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { } switch (input_key) { case UMAMI_IN_MOMENTA: - if (input_type.dtype != DataType::dt_float || + if (input_type.dtype != DataType::dt_float64 || input_type.shape.size() != 2 || input_type.shape.at(1) != 4) { throw std::invalid_argument( std::format( @@ -939,7 +941,7 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { case UMAMI_IN_RANDOM_COLOR: case UMAMI_IN_RANDOM_HELICITY: case UMAMI_IN_RANDOM_DIAGRAM: - if (input_type.dtype != DataType::dt_float || + if (input_type.dtype != DataType::dt_float64 || input_type.shape.size() != 0) { throw std::invalid_argument( std::format( @@ -952,7 +954,8 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { case UMAMI_IN_HELICITY_INDEX: case UMAMI_IN_DIAGRAM_INDEX: case UMAMI_IN_CHANNEL_INDEX: - if (input_type.dtype != DataType::dt_int || input_type.shape.size() != 0) { + if (input_type.dtype != DataType::dt_int32 || + input_type.shape.size() != 0) { throw std::invalid_argument( std::format( "matrix_element, argument {}: expected batch of integers", i + 1 @@ -979,17 +982,17 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { std::vector shape; switch (output_key) { case UMAMI_OUT_MATRIX_ELEMENT: - dtype = DataType::dt_float; + dtype = DataType::dt_float64; shape = {}; break; case UMAMI_OUT_DIAGRAM_AMP2: - dtype = DataType::dt_float; + dtype = DataType::dt_float64; shape = {output_size}; break; case UMAMI_OUT_COLOR_INDEX: case UMAMI_OUT_HELICITY_INDEX: case UMAMI_OUT_DIAGRAM_INDEX: - dtype = DataType::dt_int; + dtype = DataType::dt_int32; shape = {}; break; default: diff --git a/madspace/src/compgraphs/type.cpp b/madspace/src/compgraphs/type.cpp index c82dd051ef..17fd02acbb 100644 --- a/madspace/src/compgraphs/type.cpp +++ b/madspace/src/compgraphs/type.cpp @@ -86,11 +86,14 @@ BatchSize BatchSize::add(const BatchSize& other, int factor) const { std::ostream& madspace::operator<<(std::ostream& out, const DataType& dtype) { switch (dtype) { - case DataType::dt_float: - out << "float"; + case DataType::dt_float32: + out << "float32"; break; - case DataType::dt_int: - out << "int"; + case DataType::dt_float64: + out << "float64"; + break; + case DataType::dt_int32: + out << "int32"; break; case DataType::batch_sizes: out << "batch_sizes"; @@ -237,11 +240,14 @@ void madspace::to_json(json& j, const Value& value) { void madspace::to_json(json& j, const DataType& dtype) { switch (dtype) { - case DataType::dt_int: - j = "int"; + case DataType::dt_int32: + j = "int32"; + break; + case DataType::dt_float32: + j = "float32"; break; - case DataType::dt_float: - j = "float"; + case DataType::dt_float64: + j = "float64"; break; case DataType::batch_sizes: j = "batch_sizes"; @@ -280,14 +286,15 @@ void madspace::from_json(const json& j, Value& value) { auto shape = j.at("shape").get>(); auto j_data = j.at("data"); switch (dtype) { - case DataType::dt_int: + case DataType::dt_int32: if (shape.size() == 0) { value = j_data.get(); } else { value = Value(j_data.get>(), shape); } break; - case DataType::dt_float: + case DataType::dt_float32: + case DataType::dt_float64: if (shape.size() == 0) { value = unescape_special_values(j_data); } else { @@ -297,6 +304,8 @@ void madspace::from_json(const json& j, Value& value) { } value = Value(tensor, shape); } + // float32 literals are stored as double and only converted at runtime + value.type.dtype = dtype; break; case DataType::batch_sizes: throw std::invalid_argument("invalid data type"); @@ -305,10 +314,12 @@ void madspace::from_json(const json& j, Value& value) { void madspace::from_json(const json& j, DataType& dtype) { auto str = j.get(); - if (str == "int") { - dtype = DataType::dt_int; - } else if (str == "float") { - dtype = DataType::dt_float; + if (str == "int32") { + dtype = DataType::dt_int32; + } else if (str == "float32") { + dtype = DataType::dt_float32; + } else if (str == "float64") { + dtype = DataType::dt_float64; } else if (str == "batch_sizes") { dtype = DataType::batch_sizes; } else { diff --git a/madspace/src/cpu/device.cpp b/madspace/src/cpu/device.cpp index 649fce51ed..e3ef27b902 100644 --- a/madspace/src/cpu/device.cpp +++ b/madspace/src/cpu/device.cpp @@ -10,14 +10,15 @@ namespace { template void tensor_copy_impl(const Tensor& source, Tensor& target, const D& device) { - if (source.dtype() == DataType::dt_float && target.dtype() == DataType::dt_float) { + if (source.dtype() == DataType::dt_float64 && + target.dtype() == DataType::dt_float64) { tensor_foreach_dynamic_single< kernel_copy, kernel_copy, 1, 1>({&source}, {&target}, target.size(0), device); - } else if (source.dtype() == DataType::dt_int && - target.dtype() == DataType::dt_int) { + } else if (source.dtype() == DataType::dt_int32 && + target.dtype() == DataType::dt_int32) { tensor_foreach_dynamic_single< kernel_copy_int, kernel_copy_int, @@ -30,13 +31,13 @@ void tensor_copy_impl(const Tensor& source, Tensor& target, const D& device) { template void tensor_zero_impl(Tensor& tensor, const D& device) { - if (tensor.dtype() == DataType::dt_float) { + if (tensor.dtype() == DataType::dt_float64) { tensor_foreach_dynamic_single< kernel_zero, kernel_zero, 1, 1>({&tensor}, {&tensor}, tensor.size(0), device); - } else if (tensor.dtype() == DataType::dt_int) { + } else if (tensor.dtype() == DataType::dt_int32) { tensor_foreach_dynamic_single< kernel_zero_int, kernel_zero_int, diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 65a71dfa63..6530e39de6 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -182,7 +182,7 @@ void op_matmul( std::size_t batch_size = input.size(0); std::size_t dims_in = input.size(1); std::size_t dims_out = weight.size(1); - output = Tensor(DataType::dt_float, {batch_size, dims_out}, device); + output = Tensor(DataType::dt_float64, {batch_size, dims_out}, device); output.copy_from(bias_orig, device); if (batch_size == 0) { return; @@ -220,17 +220,17 @@ void backward_op_matmul( if (!input_grad) { input_grad = - Tensor(DataType::dt_float, input.shape(), device, AllocHint::local_grad); + Tensor(DataType::dt_float64, input.shape(), device, AllocHint::local_grad); // input_grad.zero(device); } if (!weight_grad) { weight_grad = - Tensor(DataType::dt_float, weight.shape(), device, AllocHint::local_grad); + Tensor(DataType::dt_float64, weight.shape(), device, AllocHint::local_grad); // weight_grad.zero(device); } if (!bias_grad) { bias_grad = - Tensor(DataType::dt_float, {1, dims_out}, device, AllocHint::local_grad); + Tensor(DataType::dt_float64, {1, dims_out}, device, AllocHint::local_grad); // bias_grad.zero(device); } if (batch_size == 0) { @@ -287,7 +287,7 @@ void op_nonzero( auto& input = locals[instruction.input_indices[0]]; auto batch_size = input.size(0); auto& output = locals[instruction.output_indices[0]]; - Tensor output_tmp(DataType::dt_int, {batch_size}); + Tensor output_tmp(DataType::dt_int32, {batch_size}); auto input_view_flat = input.flat_view(0); auto output_view_flat = output_tmp.flat_view(0); device.submit( @@ -332,13 +332,13 @@ void batch_gather_impl( ) { Sizes out_shape = values.shape(); out_shape[0] = indices.size(0); - if (values.dtype() == DataType::dt_float) { - selection = Tensor(DataType::dt_float, out_shape, device); + if (values.dtype() == DataType::dt_float64) { + selection = Tensor(DataType::dt_float64, out_shape, device); batch_gather_impl_body, dim, double>( indices, values, selection, device ); - } else if (values.dtype() == DataType::dt_int) { - selection = Tensor(DataType::dt_int, out_shape, device); + } else if (values.dtype() == DataType::dt_int32) { + selection = Tensor(DataType::dt_int32, out_shape, device); batch_gather_impl_body, dim, me_int_t>( indices, values, selection, device ); @@ -396,11 +396,11 @@ template void batch_scatter_impl( Tensor& indices, Tensor& source, Tensor& output, const D& device ) { - if (source.dtype() == DataType::dt_float) { + if (source.dtype() == DataType::dt_float64) { batch_scatter_impl_body, dim, double>( indices, source, output, device ); - } else if (source.dtype() == DataType::dt_int) { + } else if (source.dtype() == DataType::dt_int32) { batch_scatter_impl_body, dim, me_int_t>( indices, source, output, device ); @@ -462,7 +462,7 @@ void op_batch_split_by_index( views.reserve(count); for (auto [size, output_index] : zip(sizes, instruction.output_indices)) { auto& output = locals[output_index]; - output = Tensor(DataType::dt_int, {size}, device); + output = Tensor(DataType::dt_int32, {size}, device); views.push_back(output.view()); } std::fill(sizes.begin(), sizes.end(), 0); @@ -509,7 +509,7 @@ void batch_reduce_sum_mean_impl( auto& input = locals[instruction.input_indices[0]]; auto& output = locals[instruction.output_indices[0]]; std::size_t batch_size = input.size(0); - output = Tensor(DataType::dt_float, {keepdim ? batch_size : 1}, device); + output = Tensor(DataType::dt_float64, {keepdim ? batch_size : 1}, device); device.sync_barrier(); auto input_view_flat = input.flat_view(0); @@ -553,7 +553,7 @@ void batch_reduce_sum_mean_backward_impl( auto& output_grad = local_grads[instruction.output_indices[0]]; if (!input_grad) { input_grad = - Tensor(DataType::dt_float, input.shape(), device, AllocHint::local_grad); + Tensor(DataType::dt_float64, input.shape(), device, AllocHint::local_grad); // input_grad.zero(device); } device.sync_barrier(); @@ -674,7 +674,7 @@ void op_offset_indices( auto& sizes_out = locals[instruction.input_indices[1]].batch_sizes(); std::size_t total_size = std::accumulate(sizes_out.begin(), sizes_out.end(), 0); auto& output = locals[instruction.output_indices[0]]; - output = Tensor(DataType::dt_int, {total_size}, device); + output = Tensor(DataType::dt_int32, {total_size}, device); auto flat_view = output.flat_view(0); std::size_t sum_offset = 0, sum_out = 0; @@ -709,7 +709,7 @@ void op_quantile( return; } - output = Tensor(DataType::dt_float, {1}, device); + output = Tensor(DataType::dt_float64, {1}, device); device.foreach ( input.size(0), [&](std::size_t count, std::size_t offset) { @@ -738,7 +738,7 @@ void op_random( auto batch_size = locals[instruction.input_indices[0]].batch_sizes()[0]; auto& output = locals[instruction.output_indices[0]]; auto dim = instruction.output_shapes[0][0]; - output = Tensor(DataType::dt_float, {batch_size, dim}, device); + output = Tensor(DataType::dt_float64, {batch_size, dim}, device); auto flat_view = output.flat_view(2); auto& runtime = instruction.runtime; device.foreach ( @@ -760,7 +760,7 @@ void op_random_int( auto batch_size = locals[instruction.input_indices[0]].batch_sizes()[0]; auto max_val = locals[instruction.input_indices[1]].batch_sizes()[0]; auto& output = locals[instruction.output_indices[0]]; - output = Tensor(DataType::dt_float, {batch_size}, device); + output = Tensor(DataType::dt_float64, {batch_size}, device); auto flat_view = output.flat_view(0); auto& runtime = instruction.runtime; device.foreach ( @@ -788,8 +788,8 @@ void op_unweight( auto& uw_weights = locals[instruction.output_indices[1]]; auto batch_size = weights.size(0); - Tensor indices_tmp(DataType::dt_int, {batch_size}); - Tensor uw_weights_tmp(DataType::dt_float, {batch_size}); + Tensor indices_tmp(DataType::dt_int32, {batch_size}); + Tensor uw_weights_tmp(DataType::dt_float64, {batch_size}); auto weights_view_flat = weights.flat_view(0); auto max_weight_view_flat = max_weight.flat_view(0); @@ -846,8 +846,8 @@ void op_histogram( Sizes shape(out_shape.size() + 1); shape[0] = 1; std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1); - values = Tensor(DataType::dt_float, shape, device); - square_values = Tensor(DataType::dt_float, shape, device); + values = Tensor(DataType::dt_float64, shape, device); + square_values = Tensor(DataType::dt_float64, shape, device); device.sync_barrier(); auto weights_view_flat = weights.flat_view(0); @@ -915,8 +915,8 @@ void op_vegas_histogram( Sizes shape(out_shape.size() + 1); shape[0] = 1; std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1); - values = Tensor(DataType::dt_float, shape, device); - counts = Tensor(DataType::dt_int, shape, device); + values = Tensor(DataType::dt_float64, shape, device); + counts = Tensor(DataType::dt_int32, shape, device); device.sync_barrier(); auto input_view_flat = input.flat_view(0); @@ -974,8 +974,8 @@ void op_discrete_histogram( Sizes shape(out_shape.size() + 1); shape[0] = 1; std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1); - values = Tensor(DataType::dt_float, shape, device); - counts = Tensor(DataType::dt_int, shape, device); + values = Tensor(DataType::dt_float64, shape, device); + counts = Tensor(DataType::dt_int32, shape, device); std::size_t batch_size = locals[instruction.batch_size_index].size(0); @@ -1132,8 +1132,9 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur ); for (auto& in : instr.inputs) { auto& local = function.locals().at(in.local_index); - bool is_literal = !std::holds_alternative(local.literal_value); - bool is_int_intermediate = local.type.dtype == DataType::dt_int && + bool is_literal = + !std::holds_alternative(local.literal_value); + bool is_int_intermediate = local.type.dtype == DataType::dt_int32 && static_cast(in.local_index) >= function.inputs().size() && !global_locals.contains(in.local_index); _instructions.back().discard_input_grad.push_back( @@ -1373,7 +1374,7 @@ std::pair CpuRuntime::run_backward_single( } Tensor all_global_grads( - DataType::dt_float, {_grad_global_total_size}, AllocHint::global_grad + DataType::dt_float64, {_grad_global_total_size}, AllocHint::global_grad ); // all_global_grads.zero(); TensorVec global_grads = all_global_grads.split_and_reshape(_grad_global_shapes); @@ -1392,9 +1393,9 @@ std::pair CpuRuntime::run_backward_single( for (auto [output_index, output_dtype] : zip(instr.output_indices, instr.output_dtypes)) { auto& grad = local_grads[output_index]; - if (!grad && output_dtype == DataType::dt_float) { + if (!grad && output_dtype == DataType::dt_float64) { grad = Tensor( - DataType::dt_float, + DataType::dt_float64, locals[output_index].shape(), device, AllocHint::local_grad @@ -1567,7 +1568,7 @@ std::pair CpuRuntime::run_backward_concurrent( } Tensor all_global_grads( - DataType::dt_float, {_grad_global_total_size}, AllocHint::global_grad + DataType::dt_float64, {_grad_global_total_size}, AllocHint::global_grad ); // all_global_grads.zero(); TensorVec global_grads = all_global_grads.split_and_reshape(_grad_global_shapes); @@ -1606,9 +1607,9 @@ std::pair CpuRuntime::run_backward_concurrent( for (auto [output_index, output_dtype] : zip(instr.output_indices, instr.output_dtypes)) { auto& grad = local_grads[output_index]; - if (!grad && output_dtype == DataType::dt_float) { + if (!grad && output_dtype == DataType::dt_float64) { grad = Tensor( - DataType::dt_float, + DataType::dt_float64, locals[output_index].shape(), device, AllocHint::local_grad diff --git a/madspace/src/driver/context.cpp b/madspace/src/driver/context.cpp index f988b8f731..d3ceff92f6 100644 --- a/madspace/src/driver/context.cpp +++ b/madspace/src/driver/context.cpp @@ -12,7 +12,9 @@ using namespace madspace; using json = nlohmann::json; namespace { -UmamiStatus umami_key_query_not_implemented(bool const**, int*) { return UMAMI_ERROR_NOT_IMPLEMENTED; } +UmamiStatus umami_key_query_not_implemented(bool const**, int*) { + return UMAMI_ERROR_NOT_IMPLEMENTED; +} thread_local std::optional current_caller_stream; } // namespace @@ -160,8 +162,7 @@ void MatrixElementApi::throw_error(const std::string& message) const { ); } -const MatrixElementApi& -Context::load_matrix_element( +const MatrixElementApi& Context::load_matrix_element( const std::string& file, const std::string& param_card, const std::unordered_map& parameters @@ -364,7 +365,7 @@ Tensor Context::cached_tensor(std::size_t size) { } } std::size_t word_count = (size + 7) / 8; - Tensor new_tensor(DataType::dt_float, {word_count}, device()); + Tensor new_tensor(DataType::dt_float64, {word_count}, device()); tensors.insert(tensors.begin() + size_pos, new_tensor); return new_tensor; } diff --git a/madspace/src/driver/discrete_optimizer.cpp b/madspace/src/driver/discrete_optimizer.cpp index 2e12fa3a30..14f2116400 100644 --- a/madspace/src/driver/discrete_optimizer.cpp +++ b/madspace/src/driver/discrete_optimizer.cpp @@ -35,7 +35,7 @@ void DiscreteOptimizer::optimize() { //_sample_count = next_sample_count; for (auto [prob_name, data_item] : zip(_prob_names, _data)) { auto& [weight_sums, counts] = data_item; - Tensor prob(DataType::dt_float, _contexts.at(0)->global(prob_name).shape()); + Tensor prob(DataType::dt_float64, _contexts.at(0)->global(prob_name).shape()); auto prob_view = prob.view()[0]; double norm = 0.; diff --git a/madspace/src/driver/event_histograms.cpp b/madspace/src/driver/event_histograms.cpp index a00376ee51..6af627b235 100644 --- a/madspace/src/driver/event_histograms.cpp +++ b/madspace/src/driver/event_histograms.cpp @@ -27,8 +27,9 @@ EventHistograms::EventHistograms( ) : _specs(specs), _reference_weight( - std::isfinite(reference_weight) && reference_weight != 0. ? reference_weight - : 1. + std::isfinite(reference_weight) && reference_weight != 0. + ? reference_weight + : 1. ) { if (!context) { throw std::invalid_argument("EventHistograms requires a context"); @@ -59,15 +60,19 @@ EventHistograms::EventHistograms( "EventHistograms: one observable per histogram expected" ); } - _runtimes.push_back(RuntimeData{ - .runtime = build_runtime(obs->values.function(), context, false), - .particle_count = obs->particle_count, - }); + _runtimes.push_back( + RuntimeData{ + .runtime = build_runtime(obs->values.function(), context, false), + .particle_count = obs->particle_count, + } + ); } } void EventHistograms::fill( - EventBuffer& buffer, const std::vector& syst_weights, std::size_t weight_count + EventBuffer& buffer, + const std::vector& syst_weights, + std::size_t weight_count ) { std::size_t count = buffer.event_count(); if (count == 0) { @@ -93,7 +98,7 @@ void EventHistograms::fill( for (auto& [subproc, indices] : groups) { auto& runtime_data = _runtimes.at(subproc).value(); std::size_t pc = runtime_data.particle_count; - Tensor momenta(DataType::dt_float, {indices.size(), pc, 4}); + Tensor momenta(DataType::dt_float64, {indices.size(), pc, 4}); auto view = momenta.view(); for (std::size_t n = 0; n < indices.size(); ++n) { for (std::size_t j = 0; j < pc; ++j) { @@ -111,7 +116,7 @@ void EventHistograms::fill( } for (std::size_t o = 0; o < _specs.size(); ++o) { if (!_observable_index.at(o)) { - continue; // filled from the event weight below + continue; // filled from the event weight below } Tensor out = outputs.at(_observable_index.at(o).value()).cpu().contiguous(); if (out.shape().size() == 1) { @@ -192,21 +197,22 @@ json EventHistograms::to_json(const SystematicsCalculator* systematics) const { return result; } double n = static_cast(_event_count); - auto normalised = [&](const std::vector& sums, - const std::vector& squares, - std::vector& values, - std::vector& errors) { - values.resize(sums.size()); - errors.resize(sums.size()); - for (std::size_t b = 0; b < sums.size(); ++b) { - // the event weights are "average" normalised: the cross section in a - // bin is the mean of (weight x indicator), its error the standard - // error of that mean - values[b] = sums[b] / n; - double var = (squares[b] - sums[b] * sums[b] / n) / n; - errors[b] = var > 0. ? std::sqrt(var / n) : 0.; - } - }; + auto normalised = + [&](const std::vector& sums, + const std::vector& squares, + std::vector& values, + std::vector& errors) { + values.resize(sums.size()); + errors.resize(sums.size()); + for (std::size_t b = 0; b < sums.size(); ++b) { + // the event weights are "average" normalised: the cross section in a + // bin is the mean of (weight x indicator), its error the standard + // error of that mean + values[b] = sums[b] / n; + double var = (squares[b] - sums[b] * sums[b] / n) / n; + errors[b] = var > 0. ? std::sqrt(var / n) : 0.; + } + }; std::vector ids; std::vector scale_indices; std::vector pdf_groups; @@ -217,7 +223,8 @@ json EventHistograms::to_json(const SystematicsCalculator* systematics) const { } for (std::size_t o = 0; o < _specs.size(); ++o) { auto& spec = _specs[o]; - std::vector> values(_weight_count + 1), errors(_weight_count + 1); + std::vector> values(_weight_count + 1), + errors(_weight_count + 1); for (std::size_t c = 0; c <= _weight_count; ++c) { normalised(_sums[o][c], _square_sums[o][c], values[c], errors[c]); } diff --git a/madspace/src/driver/io.cpp b/madspace/src/driver/io.cpp index 833dfaee9e..c80c72c6d4 100644 --- a/madspace/src/driver/io.cpp +++ b/madspace/src/driver/io.cpp @@ -15,10 +15,13 @@ namespace { std::string dtype_to_str(DataType dtype) { switch (dtype) { - case DataType::dt_float: + case DataType::dt_float32: + return "function(), _context, false), @@ -269,22 +275,27 @@ SystematicsCalculator::SystematicsCalculator( std::map alpha_s_indices; for (auto& spec : _config.pdf_members) { if (!_config.has_pdf) { - _warnings.push_back(std::format( - "PDF variation {} member {} ignored: the beams have no PDF", - spec.set_name, - spec.member - )); + _warnings.push_back( + std::format( + "PDF variation {} member {} ignored: the beams have no PDF", + spec.set_name, + spec.member + ) + ); continue; } bool same_set = spec.set_lhaid == _config.nominal_lhaid; if (!same_set && !_mur_supported) { - _warnings.push_back(std::format( - "PDF variation {} member {} ignored: the events have no uniform " - "alpha_s power and no matrix element to re-evaluate, so the alpha_s " - "of a different PDF set cannot be applied", - spec.set_name, - spec.member - )); + _warnings.push_back( + std::format( + "PDF variation {} member {} ignored: the events have no uniform " + "alpha_s power and no matrix element to re-evaluate, so the " + "alpha_s " + "of a different PDF set cannot be applied", + spec.set_name, + spec.member + ) + ); continue; } std::size_t alpha_s_index; @@ -367,11 +378,10 @@ void SystematicsCalculator::build_variations() { // PDF variations: when members of the nominal set are requested, the group // starts with the nominal member itself so that the group is complete - bool has_nominal_set_members = std::any_of( - _members.begin(), - _members.end(), - [&](auto& spec) { return spec.set_lhaid == _config.nominal_lhaid; } - ); + bool has_nominal_set_members = + std::any_of(_members.begin(), _members.end(), [&](auto& spec) { + return spec.set_lhaid == _config.nominal_lhaid; + }); if (has_nominal_set_members) { _variations.push_back({id++, 1., 1., -1, -1}); } @@ -498,9 +508,9 @@ std::vector SystematicsCalculator::evaluate_pdf( if (count == 0) { return values; } - Tensor x_tensor(DataType::dt_float, {count}); - Tensor q_tensor(DataType::dt_float, {count}); - Tensor slot_tensor(DataType::dt_int, {count}); + Tensor x_tensor(DataType::dt_float64, {count}); + Tensor q_tensor(DataType::dt_float64, {count}); + Tensor slot_tensor(DataType::dt_int32, {count}); auto x_view = x_tensor.view(); auto q_view = q_tensor.view(); auto slot_view = slot_tensor.view(); @@ -530,7 +540,7 @@ std::vector SystematicsCalculator::evaluate_alpha_s( if (count == 0) { return values; } - Tensor q_tensor(DataType::dt_float, {count}); + Tensor q_tensor(DataType::dt_float64, {count}); auto q_view = q_tensor.view(); for (std::size_t i = 0; i < count; ++i) { q_view[i] = q[i]; @@ -548,8 +558,9 @@ std::vector SystematicsCalculator::evaluate_alpha_s( return values; } -std::vector> -SystematicsCalculator::event_momenta(EventBuffer& buffer, std::size_t event_index) const { +std::vector> SystematicsCalculator::event_momenta( + EventBuffer& buffer, std::size_t event_index +) const { std::vector> momenta; for (std::size_t j = 0; j < buffer.particle_count(); ++j) { auto particle = buffer.particle(event_index, j); @@ -571,9 +582,9 @@ std::vector SystematicsCalculator::matrix_elements( auto& me_data = _matrix_elements.at(subproc).value(); std::size_t count = indices.size(); std::size_t pc = me_data.particle_count; - Tensor momenta(DataType::dt_float, {count, pc, 4}); - Tensor alpha_s_tensor(DataType::dt_float, {count}); - Tensor flavor(DataType::dt_int, {count}); + Tensor momenta(DataType::dt_float64, {count, pc, 4}); + Tensor alpha_s_tensor(DataType::dt_float64, {count}); + Tensor flavor(DataType::dt_int32, {count}); auto mom_view = momenta.view(); auto alpha_view = alpha_s_tensor.view(); auto flavor_view = flavor.view(); @@ -606,8 +617,9 @@ std::vector SystematicsCalculator::matrix_elements( return values; } -LOReweightInfo -SystematicsCalculator::reweight_info(EventBuffer& buffer, std::size_t event_index) const { +LOReweightInfo SystematicsCalculator::reweight_info( + EventBuffer& buffer, std::size_t event_index +) const { auto& layout = buffer.layout(); bool has_beam1 = layout_has(layout, "fact_scale1"); bool has_beam2 = layout_has(layout, "fact_scale2"); @@ -727,41 +739,42 @@ void SystematicsCalculator::compute( // (nominal grid: every scale variation; member grids: the nominal scales) std::vector> r_pdf(count, std::vector(var_count, 1.)); if (use_pdf) { - auto evaluate_variations = [&](const PdfEvaluator& evaluator, - const std::vector& var_indices) { - std::vector x, q; - std::vector slots; - for (auto& in : inputs) { - for (std::size_t k : var_indices) { - auto& var = _variations.at(k); - if (has_beam1) { - x.push_back(in.x1); - q.push_back(var.muf * scale_of(in, var.dyn, in.muf1)); - slots.push_back(in.slot1); - } - if (has_beam2) { - x.push_back(in.x2); - q.push_back(var.muf * scale_of(in, var.dyn, in.muf2)); - slots.push_back(in.slot2); + auto evaluate_variations = + [&](const PdfEvaluator& evaluator, + const std::vector& var_indices) { + std::vector x, q; + std::vector slots; + for (auto& in : inputs) { + for (std::size_t k : var_indices) { + auto& var = _variations.at(k); + if (has_beam1) { + x.push_back(in.x1); + q.push_back(var.muf * scale_of(in, var.dyn, in.muf1)); + slots.push_back(in.slot1); + } + if (has_beam2) { + x.push_back(in.x2); + q.push_back(var.muf * scale_of(in, var.dyn, in.muf2)); + slots.push_back(in.slot2); + } } } - } - auto values = evaluate_pdf(evaluator, x, q, slots); - for (std::size_t i = 0, n = 0; i < count; ++i) { - for (std::size_t k : var_indices) { - double product = 1.; - if (has_beam1) { - product *= values[n++]; - } - if (has_beam2) { - product *= values[n++]; + auto values = evaluate_pdf(evaluator, x, q, slots); + for (std::size_t i = 0, n = 0; i < count; ++i) { + for (std::size_t k : var_indices) { + double product = 1.; + if (has_beam1) { + product *= values[n++]; + } + if (has_beam2) { + product *= values[n++]; + } + r_pdf[i][k] = inputs[i].nominal_product == 0. + ? 0. + : product / inputs[i].nominal_product; } - r_pdf[i][k] = inputs[i].nominal_product == 0. - ? 0. - : product / inputs[i].nominal_product; } - } - }; + }; std::vector nominal_grid_vars; std::vector> member_vars(_member_pdfs.size()); for (std::size_t k = 0; k < var_count; ++k) { @@ -842,7 +855,8 @@ void SystematicsCalculator::compute( if (in.use_me || in.qcd_power == 0) { continue; } - ratio[i] = in.alpha_s == 0. ? 0. : std::pow(varied[i] / in.alpha_s, in.qcd_power); + ratio[i] = + in.alpha_s == 0. ? 0. : std::pow(varied[i] / in.alpha_s, in.qcd_power); } for (auto& [subproc, indices] : me_events) { std::vector alpha_s(indices.size()); @@ -1000,8 +1014,7 @@ json SystematicsCalculator::summary() const { json result{ {"variations", variations}, {"nominal", - {{"pdf_set", _config.nominal_set_name}, - {"pdf_lhaid", _config.nominal_lhaid}}}, + {{"pdf_set", _config.nominal_set_name}, {"pdf_lhaid", _config.nominal_lhaid}}}, {"event_count", _event_count}, {"warnings", _warnings}, }; diff --git a/madspace/src/driver/tensor.cpp b/madspace/src/driver/tensor.cpp index 6842b1686d..c048bfbd2f 100644 --- a/madspace/src/driver/tensor.cpp +++ b/madspace/src/driver/tensor.cpp @@ -276,7 +276,7 @@ Tensor Tensor::simd_padded_literal( auto [dtype, buffer] = std::visit( [&](const std::vector& vec) { auto dtype = - std::is_same_v ? DataType::dt_int : DataType::dt_float; + std::is_same_v ? DataType::dt_int32 : DataType::dt_float64; std::vector padded(count * lanes); for (std::size_t j = 0; j < count; ++j) { std::fill_n(padded.begin() + j * lanes, lanes, vec[j]); @@ -294,7 +294,16 @@ Tensor Tensor::simd_padded_literal( stride_prod *= shape[i]; } Tensor ret(new TensorImpl{ - dtype, shape, device, buffer.data(), false, std::nullopt, buffer.impl, 1, stride, 1 + dtype, + shape, + device, + buffer.data(), + false, + std::nullopt, + buffer.impl, + 1, + stride, + 1 }); ret.impl->simd_padded = true; // non-negative integer scalars double as batch sizes, see the scalar constructor diff --git a/madspace/src/gpu/device.cu b/madspace/src/gpu/device.cu index 745e22f611..40e1a5070f 100644 --- a/madspace/src/gpu/device.cu +++ b/madspace/src/gpu/device.cu @@ -147,7 +147,8 @@ MemPool::MemPool( pool.parent_tensor = parent_tensor; pool.capacity = parent_tensor.byte_size(); } else { - pool.parent_tensor = Tensor(DataType::dt_float, {word_count}, async_device); + pool.parent_tensor = + Tensor(DataType::dt_float64, {word_count}, async_device); pool.capacity = word_count * 8; } pool.needed_size = word_count * 8; @@ -307,12 +308,13 @@ void AsyncGpuDevice::memcpy(void* to, void* from, std::size_t size) const { } void AsyncGpuDevice::tensor_copy(const Tensor& source, Tensor& target) const { - if (source.dtype() == DataType::dt_float && target.dtype() == DataType::dt_float) { + if (source.dtype() == DataType::dt_float64 && + target.dtype() == DataType::dt_float64) { tensor_foreach_dynamic, 1, 1>( {&source}, {&target}, target.size(0), *this ); - } else if (source.dtype() == DataType::dt_int && - target.dtype() == DataType::dt_int) { + } else if (source.dtype() == DataType::dt_int32 && + target.dtype() == DataType::dt_int32) { tensor_foreach_dynamic, 1, 1>( {&source}, {&target}, target.size(0), *this ); @@ -322,7 +324,7 @@ void AsyncGpuDevice::tensor_copy(const Tensor& source, Tensor& target) const { } void AsyncGpuDevice::tensor_zero(Tensor& tensor) const { - if (tensor.dtype() == DataType::dt_float) { + if (tensor.dtype() == DataType::dt_float64) { if (tensor.is_contiguous()) { check_error(gpuMemsetAsync(tensor.data(), 0, tensor.byte_size(), _stream)); } else { @@ -330,7 +332,7 @@ void AsyncGpuDevice::tensor_zero(Tensor& tensor) const { {&tensor}, {&tensor}, tensor.size(0), *this ); } - } else if (tensor.dtype() == DataType::dt_int) { + } else if (tensor.dtype() == DataType::dt_int32) { if (tensor.is_contiguous()) { check_error(gpuMemsetAsync(tensor.data(), 0, tensor.byte_size(), _stream)); } else { diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 8e9a06a218..7e94980367 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -40,8 +40,8 @@ void op_matrix_element( std::size_t output_count = locals[instruction.input_indices[2]].index_value(); TensorVec contiguous_inputs(input_count); std::vector input_keys(input_count); - // one more output: the stream (UMAMI_OUT_GPU_STREAM), so that the matrix element runs - // in stream order with the rest of the function, not on the default stream + // one more output: the stream (UMAMI_OUT_GPU_STREAM), so that the matrix element + // runs in stream order with the rest of the function, not on the default stream std::vector output_keys(output_count + 1); std::vector input_ptrs(input_count), output_ptrs(output_count + 1); for (std::size_t i = 0; i < input_count; ++i) { @@ -141,7 +141,7 @@ void op_matmul( std::size_t dims_in = input.size(1); std::size_t dims_out = weight.size(1); output = Tensor( - DataType::dt_float, + DataType::dt_float64, {batch_size, dims_out}, device, instruction.output_alloc_hints[0] @@ -205,7 +205,7 @@ void backward_op_matmul( if (!input_grad) { input_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, input.shape(), device, instruction.input_grad_alloc_hints[0] @@ -214,7 +214,7 @@ void backward_op_matmul( } if (!weight_grad) { weight_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, weight.shape(), device, instruction.input_grad_alloc_hints[1] @@ -223,7 +223,7 @@ void backward_op_matmul( } if (!bias_grad) { bias_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, {1, dims_out}, device, instruction.input_grad_alloc_hints[2] @@ -279,7 +279,7 @@ void backward_op_matmul( )); // compute bias_grad += sum_i output_grad_ij - Tensor ones(DataType::dt_float, {batch_size}, device, AllocHint::temporary); + Tensor ones(DataType::dt_float64, {batch_size}, device, AllocHint::temporary); launch_kernel( kernel_one, batch_size, device.stream(), batch_size, ones.view() ); @@ -326,9 +326,9 @@ void op_nonzero( auto& input = locals[instruction.input_indices[0]]; auto batch_size = input.size(0); auto& output = locals[instruction.output_indices[0]]; - Tensor indices_tmp(DataType::dt_int, {batch_size}, device, AllocHint::temporary); + Tensor indices_tmp(DataType::dt_int32, {batch_size}, device, AllocHint::temporary); Tensor output_tmp( - DataType::dt_int, {batch_size}, device, instruction.output_alloc_hints[0] + DataType::dt_int32, {batch_size}, device, instruction.output_alloc_hints[0] ); launch_kernel( kernel_nonzero, @@ -397,8 +397,8 @@ void batch_gather_impl( Sizes out_shape = values.shape(); out_shape[0] = batch_size; - if (values.dtype() == DataType::dt_float) { - selection = Tensor(DataType::dt_float, out_shape, device, hint); + if (values.dtype() == DataType::dt_float64) { + selection = Tensor(DataType::dt_float64, out_shape, device, hint); launch_kernel( batch_gather_kernel, batch_size, @@ -408,8 +408,8 @@ void batch_gather_impl( values.view(), selection.view() ); - } else if (values.dtype() == DataType::dt_int) { - selection = Tensor(DataType::dt_int, out_shape, device, hint); + } else if (values.dtype() == DataType::dt_int32) { + selection = Tensor(DataType::dt_int32, out_shape, device, hint); launch_kernel( batch_gather_kernel_int, batch_size, @@ -483,7 +483,7 @@ template void batch_scatter_impl( Tensor& indices, Tensor& source, Tensor& output, const AsyncGpuDevice& device ) { - if (source.dtype() == DataType::dt_float) { + if (source.dtype() == DataType::dt_float64) { auto batch_size = indices.size(0); launch_kernel( batch_scatter_kernel, @@ -494,7 +494,7 @@ void batch_scatter_impl( source.view(), output.view() ); - } else if (source.dtype() == DataType::dt_int) { + } else if (source.dtype() == DataType::dt_int32) { auto batch_size = indices.size(0); launch_kernel( batch_scatter_kernel_int, @@ -558,14 +558,14 @@ void op_batch_split_by_index( if (batch_size == 0) { for (std::size_t k = 0; k < count; ++k) { locals[instruction.output_indices[k]] = Tensor( - DataType::dt_int, {0}, device, instruction.output_alloc_hints[k] + DataType::dt_int32, {0}, device, instruction.output_alloc_hints[k] ); } return; } // per-bucket counts, needed on the host to size and split the outputs below - Tensor sizes(DataType::dt_int, {count}, device, AllocHint::temporary); + Tensor sizes(DataType::dt_int32, {count}, device, AllocHint::temporary); { std::size_t temp_storage_bytes = 0; check_error( @@ -582,7 +582,7 @@ void op_batch_split_by_index( ) ); Tensor temp( - DataType::dt_float, + DataType::dt_float64, {(temp_storage_bytes + 7) / 8}, device, AllocHint::temporary @@ -605,15 +605,15 @@ void op_batch_split_by_index( // stable sort of the original positions by bucket index: afterwards each // bucket's positions are contiguous, in their original relative order - Tensor values_in(DataType::dt_int, {batch_size}, device, AllocHint::temporary); + Tensor values_in(DataType::dt_int32, {batch_size}, device, AllocHint::temporary); auto values_in_ptr = thrust::device_pointer_cast(static_cast(values_in.data())); thrust::sequence( thrust_par.on(device.stream()), values_in_ptr, values_in_ptr + batch_size ); - Tensor keys_out(DataType::dt_int, {batch_size}, device, AllocHint::temporary); + Tensor keys_out(DataType::dt_int32, {batch_size}, device, AllocHint::temporary); Tensor values_out( - DataType::dt_int, {batch_size}, device, instruction.output_alloc_hints[0] + DataType::dt_int32, {batch_size}, device, instruction.output_alloc_hints[0] ); { std::size_t temp_storage_bytes = 0; @@ -632,7 +632,7 @@ void op_batch_split_by_index( ) ); Tensor temp( - DataType::dt_float, + DataType::dt_float64, {(temp_storage_bytes + 7) / 8}, device, AllocHint::temporary @@ -718,7 +718,7 @@ void batch_reduce_sum_mean_impl( AllocHint hint = instruction.output_alloc_hints[0]; std::size_t batch_size = input.size(0); - Tensor out(DataType::dt_float, {1}, device, hint); + Tensor out(DataType::dt_float64, {1}, device, hint); std::size_t temp_storage_bytes = 0; check_error( cub::DeviceReduce::Sum( @@ -731,7 +731,10 @@ void batch_reduce_sum_mean_impl( ) ); Tensor temp( - DataType::dt_float, {(temp_storage_bytes + 7) / 8}, device, AllocHint::temporary + DataType::dt_float64, + {(temp_storage_bytes + 7) / 8}, + device, + AllocHint::temporary ); check_error( cub::DeviceReduce::Sum( @@ -774,7 +777,7 @@ void batch_reduce_sum_mean_backward_impl( auto& input_grad = local_grads[instruction.input_indices[0]]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, input.shape(), device, instruction.input_grad_alloc_hints[0] @@ -782,7 +785,7 @@ void batch_reduce_sum_mean_backward_impl( // input_grad.zero(device); } - Tensor grad(DataType::dt_float, {1}, device, AllocHint::temporary); + Tensor grad(DataType::dt_float64, {1}, device, AllocHint::temporary); if (keepdim) { auto output_grad = local_grads[instruction.output_indices[0]].contiguous(device); @@ -799,7 +802,7 @@ void batch_reduce_sum_mean_backward_impl( ) ); Tensor temp( - DataType::dt_float, + DataType::dt_float64, {(temp_storage_bytes + 7) / 8}, device, AllocHint::temporary @@ -912,7 +915,7 @@ void op_offset_indices( std::size_t total_size = std::accumulate(sizes_out.begin(), sizes_out.end(), 0); auto& output = locals[instruction.output_indices[0]]; output = Tensor( - DataType::dt_int, {total_size}, device, instruction.output_alloc_hints[0] + DataType::dt_int32, {total_size}, device, instruction.output_alloc_hints[0] ); std::size_t sum_offset = 0, sum_out = 0; for (auto [size_offset, size_out] : zip(sizes_offset, sizes_out)) { @@ -961,7 +964,7 @@ void op_quantile( return; } - Tensor tmp(DataType::dt_float, {batch_size}, device, AllocHint::temporary); + Tensor tmp(DataType::dt_float64, {batch_size}, device, AllocHint::temporary); tmp.copy_from(input, device); std::size_t temp_storage_bytes; check_error( @@ -975,7 +978,10 @@ void op_quantile( ) ); Tensor tmp_sort( - DataType::dt_float, {(temp_storage_bytes + 7) / 8}, device, AllocHint::temporary + DataType::dt_float64, + {(temp_storage_bytes + 7) / 8}, + device, + AllocHint::temporary ); check_error( cub::DeviceMergeSort::SortKeys( @@ -989,7 +995,8 @@ void op_quantile( ); tmp_sort.reset(device); - output = Tensor(DataType::dt_float, {1}, device, instruction.output_alloc_hints[0]); + output = + Tensor(DataType::dt_float64, {1}, device, instruction.output_alloc_hints[0]); launch_kernel( kernel_gather_quantile, 1, @@ -1022,7 +1029,10 @@ void op_random( auto& output = locals[instruction.output_indices[0]]; auto dim = instruction.output_shapes[0][0]; output = Tensor( - DataType::dt_float, {batch_size, dim}, device, instruction.output_alloc_hints[0] + DataType::dt_float64, + {batch_size, dim}, + device, + instruction.output_alloc_hints[0] ); launch_rng_kernel( kernel_random, @@ -1058,7 +1068,7 @@ void op_random_int( auto max_val = locals[instruction.input_indices[1]].batch_sizes()[0]; auto& output = locals[instruction.output_indices[0]]; output = Tensor( - DataType::dt_int, {batch_size}, device, instruction.output_alloc_hints[0] + DataType::dt_int32, {batch_size}, device, instruction.output_alloc_hints[0] ); launch_rng_kernel( kernel_random_int, @@ -1106,16 +1116,18 @@ void op_unweight( gpuStream_t stream = device.stream(); if (batch_size == 0) { - indices = - Tensor(DataType::dt_float, {0}, device, instruction.output_alloc_hints[0]); - uw_weights = - Tensor(DataType::dt_float, {0}, device, instruction.output_alloc_hints[1]); + indices = Tensor( + DataType::dt_float64, {0}, device, instruction.output_alloc_hints[0] + ); + uw_weights = Tensor( + DataType::dt_float64, {0}, device, instruction.output_alloc_hints[1] + ); return; } - Tensor indices_tmp(DataType::dt_int, {batch_size}, device, AllocHint::temporary); + Tensor indices_tmp(DataType::dt_int32, {batch_size}, device, AllocHint::temporary); Tensor uw_weights_tmp( - DataType::dt_float, {batch_size}, device, AllocHint::temporary + DataType::dt_float64, {batch_size}, device, AllocHint::temporary ); launch_rng_kernel( kernel_unweight, @@ -1129,7 +1141,7 @@ void op_unweight( ); Tensor indices_compacted( - DataType::dt_int, {batch_size}, device, instruction.output_alloc_hints[0] + DataType::dt_int32, {batch_size}, device, instruction.output_alloc_hints[0] ); auto ptr_all = thrust::device_pointer_cast(static_cast(indices_tmp.data())); @@ -1145,8 +1157,9 @@ void op_unweight( std::size_t count = ptr_compacted_end - ptr_compacted; indices = indices_compacted.slice(0, 0, count); - uw_weights = - Tensor(DataType::dt_float, {count}, device, instruction.output_alloc_hints[1]); + uw_weights = Tensor( + DataType::dt_float64, {count}, device, instruction.output_alloc_hints[1] + ); auto ptr_all_weights = thrust::device_pointer_cast(static_cast(uw_weights_tmp.data())); auto ptr_uw_weights = @@ -1179,7 +1192,7 @@ void histogram_common( ) { auto policy = thrust_par.on(device.stream()); Tensor reduce_tmp( - DataType::dt_float, {n_dims * n_bins}, device, AllocHint::temporary + DataType::dt_float64, {n_dims * n_bins}, device, AllocHint::temporary ); auto indices_ptr = thrust::device_pointer_cast(static_cast(indices_tmp.data())); @@ -1252,8 +1265,9 @@ void op_vegas_histogram( shape[0] = 1; std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1); values = - Tensor(DataType::dt_float, shape, device, instruction.output_alloc_hints[0]); - counts = Tensor(DataType::dt_int, shape, device, instruction.output_alloc_hints[1]); + Tensor(DataType::dt_float64, shape, device, instruction.output_alloc_hints[0]); + counts = + Tensor(DataType::dt_int32, shape, device, instruction.output_alloc_hints[1]); std::size_t batch_size = locals[instruction.batch_size_index].size(0); std::size_t n_dims = input.size(1); @@ -1261,10 +1275,10 @@ void op_vegas_histogram( std::size_t padded_size = batch_size + n_bins; Tensor indices_tmp( - DataType::dt_int, {padded_size, n_dims}, device, AllocHint::temporary + DataType::dt_int32, {padded_size, n_dims}, device, AllocHint::temporary ); Tensor weights_tmp( - DataType::dt_float, {padded_size, n_dims}, device, AllocHint::temporary + DataType::dt_float64, {padded_size, n_dims}, device, AllocHint::temporary ); launch_kernel( kernel_prepare_vegas_hist, @@ -1315,15 +1329,18 @@ void op_discrete_histogram( shape[0] = 1; std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1); values = - Tensor(DataType::dt_float, shape, device, instruction.output_alloc_hints[0]); - counts = Tensor(DataType::dt_int, shape, device, instruction.output_alloc_hints[1]); + Tensor(DataType::dt_float64, shape, device, instruction.output_alloc_hints[0]); + counts = + Tensor(DataType::dt_int32, shape, device, instruction.output_alloc_hints[1]); std::size_t batch_size = locals[instruction.batch_size_index].size(0); std::size_t n_opts = values.size(1); std::size_t padded_size = batch_size + n_opts; - Tensor indices_tmp(DataType::dt_int, {padded_size}, device, AllocHint::temporary); - Tensor weights_tmp(DataType::dt_float, {padded_size}, device, AllocHint::temporary); + Tensor indices_tmp(DataType::dt_int32, {padded_size}, device, AllocHint::temporary); + Tensor weights_tmp( + DataType::dt_float64, {padded_size}, device, AllocHint::temporary + ); launch_kernel( kernel_prepare_discrete_hist, padded_size, @@ -1396,20 +1413,23 @@ void op_histogram( shape[0] = 1; std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1); values = - Tensor(DataType::dt_float, shape, device, instruction.output_alloc_hints[0]); + Tensor(DataType::dt_float64, shape, device, instruction.output_alloc_hints[0]); square_values = - Tensor(DataType::dt_float, shape, device, instruction.output_alloc_hints[1]); + Tensor(DataType::dt_float64, shape, device, instruction.output_alloc_hints[1]); std::size_t batch_size = locals[instruction.batch_size_index].size(0); std::size_t n_bins = values.size(1) - 2; std::size_t padded_size = batch_size + n_bins + 2; - Tensor indices_tmp(DataType::dt_int, {padded_size}, device, AllocHint::temporary); - Tensor weights_tmp(DataType::dt_float, {padded_size}, device, AllocHint::temporary); + Tensor indices_tmp(DataType::dt_int32, {padded_size}, device, AllocHint::temporary); + Tensor weights_tmp( + DataType::dt_float64, {padded_size}, device, AllocHint::temporary + ); Tensor square_weights_tmp( - DataType::dt_float, {padded_size}, device, AllocHint::temporary + DataType::dt_float64, {padded_size}, device, AllocHint::temporary ); - // the keys reduce_by_key writes: the n_bins bins, the underflow and the overflow bin - Tensor reduce_tmp(DataType::dt_float, {n_bins + 2}, device, AllocHint::temporary); + // the keys reduce_by_key writes: the n_bins bins, the underflow and the overflow + // bin + Tensor reduce_tmp(DataType::dt_float64, {n_bins + 2}, device, AllocHint::temporary); launch_kernel( kernel_prepare_hist, @@ -2139,7 +2159,7 @@ std::pair GpuRuntime::run_backward( } } Tensor all_global_grads( - DataType::dt_float, + DataType::dt_float64, {_grad_global_total_size}, init_device, AllocHint::global_grad @@ -2166,9 +2186,9 @@ std::pair GpuRuntime::run_backward( for (auto [output_index, output_dtype] : zip(instr.output_indices, instr.output_dtypes)) { auto& grad = local_grads[output_index]; - if (!grad && output_dtype == DataType::dt_float) { + if (!grad && output_dtype == DataType::dt_float64) { grad = Tensor( - DataType::dt_float, + DataType::dt_float64, locals[output_index].shape(), device, AllocHint::local_grad diff --git a/madspace/src/kernels/operations.hpp b/madspace/src/kernels/operations.hpp index b129ef6e8f..1bd937cdbc 100644 --- a/madspace/src/kernels/operations.hpp +++ b/madspace/src/kernels/operations.hpp @@ -81,13 +81,12 @@ void backward_batch_foreach( if (!input_grad || input_grad.size(0) != batch_size) { Sizes shape = input.shape(); shape[0] = batch_size; - input_grad = - Tensor(input.dtype(), shape, device, AllocHint::temporary); + input_grad = Tensor(input.dtype(), shape, device, AllocHint::temporary); } } else if (!input_grad) { auto& input = locals[input_index]; - /*if (input.size(0) != batch_size && input.dtype() == DataType::dt_float) { - throw std::runtime_error("backward not possible for broadcasting args"); + /*if (input.size(0) != batch_size && input.dtype() == DataType::dt_float64) + { throw std::runtime_error("backward not possible for broadcasting args"); }*/ input_grad = Tensor( input.dtype(), @@ -135,7 +134,7 @@ void backward_op_stack( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, locals[input_index].shape(), device, instruction.input_grad_alloc_hints[i] @@ -186,7 +185,7 @@ void backward_op_unstack( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] @@ -216,7 +215,7 @@ void backward_op_pop( auto input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] @@ -269,7 +268,7 @@ void backward_op_batch_cat( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, locals[input_index].shape(), device, instruction.input_grad_alloc_hints[i] @@ -305,7 +304,7 @@ void backward_op_batch_split( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] @@ -357,7 +356,7 @@ void backward_op_cat( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, locals[input_index].shape(), device, instruction.input_grad_alloc_hints[i] @@ -471,7 +470,7 @@ void backward_op_rqs_reshape( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float, + DataType::dt_float64, locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] diff --git a/madspace/src/phasespace/channel_weight_network.cpp b/madspace/src/phasespace/channel_weight_network.cpp index 343468299f..abb461d094 100644 --- a/madspace/src/phasespace/channel_weight_network.cpp +++ b/madspace/src/phasespace/channel_weight_network.cpp @@ -57,7 +57,7 @@ NamedVector ChannelWeightNetwork::build_function_impl( FunctionBuilder& fb, const NamedVector& args ) const { auto mask = - fb.global(_mask_name, DataType::dt_float, {static_cast(_channel_count)}); + fb.global(_mask_name, DataType::dt_float64, {static_cast(_channel_count)}); // auto net_input = // _preprocessing.build_function(fb, {args["momenta"], args["x1"], args["x2"]}); // auto net_output = _mlp.build_function(fb, net_input.values()).at(0); @@ -73,14 +73,14 @@ void ChannelWeightNetwork::initialize_globals( ) const { _mlp.initialize_globals(context, seed); - context->define_global(_mask_name, DataType::dt_float, {_channel_count}); + context->define_global(_mask_name, DataType::dt_float64, {_channel_count}); bool is_cpu = context->device() == cpu_device(); auto mask_global = context->global(_mask_name); Tensor mask; if (is_cpu) { mask = mask_global; } else { - mask = Tensor(DataType::dt_float, mask_global.shape()); + mask = Tensor(DataType::dt_float64, mask_global.shape()); } auto mask_view = mask.view()[0]; for (std::size_t i = 0; i < mask_view.size(); ++i) { diff --git a/madspace/src/phasespace/discrete_flow.cpp b/madspace/src/phasespace/discrete_flow.cpp index 5f283fcb19..ffd896b1ff 100644 --- a/madspace/src/phasespace/discrete_flow.cpp +++ b/madspace/src/phasespace/discrete_flow.cpp @@ -112,7 +112,7 @@ Mapping::Result DiscreteFlow::build_transform( if (dim_index == 0 && _first_prob_name) { probs = fb.global( _first_prob_name.value(), - DataType::dt_float, + DataType::dt_float64, {static_cast(option_count)} ); } else { diff --git a/madspace/src/phasespace/discrete_sampler.cpp b/madspace/src/phasespace/discrete_sampler.cpp index 39c6db5cf7..26c475cb26 100644 --- a/madspace/src/phasespace/discrete_sampler.cpp +++ b/madspace/src/phasespace/discrete_sampler.cpp @@ -111,8 +111,9 @@ Mapping::Result DiscreteSampler::build_transform( std::size_t condition_index = 0; for (auto [input, prob_name, option_count, has_prior] : zip(inputs, _prob_names, _option_counts, _dim_has_prior)) { - auto probs = - fb.global(prob_name, DataType::dt_float, {static_cast(option_count)}); + auto probs = fb.global( + prob_name, DataType::dt_float64, {static_cast(option_count)} + ); if (has_prior) { probs = fb.mul(probs, conditions.at(condition_index)); ++condition_index; @@ -139,8 +140,9 @@ void madspace::initialize_uniform_probs( ContextPtr context, const std::string& name, std::size_t option_count ) { bool is_cpu = context->device() == cpu_device(); - auto prob_global = context->define_global(name, DataType::dt_float, {option_count}); - auto prob = is_cpu ? prob_global : Tensor(DataType::dt_float, {1, option_count}); + auto prob_global = + context->define_global(name, DataType::dt_float64, {option_count}); + auto prob = is_cpu ? prob_global : Tensor(DataType::dt_float64, {1, option_count}); auto prob_view = prob.view()[0]; auto prob_value = 1. / option_count; for (std::size_t i = 0; i < prob_view.size(); ++i) { diff --git a/madspace/src/phasespace/flow.cpp b/madspace/src/phasespace/flow.cpp index 1f353fcf1e..208c59cabb 100644 --- a/madspace/src/phasespace/flow.cpp +++ b/madspace/src/phasespace/flow.cpp @@ -46,7 +46,7 @@ void vegas_init( auto grid_tensor = context->global(grid_name).cpu(); bool is_cpu = false; // context->device() == cpu_device(); Tensor bias_tensor = - is_cpu ? bias_global : Tensor(DataType::dt_float, bias_global.shape()); + is_cpu ? bias_global : Tensor(DataType::dt_float64, bias_global.shape()); auto bias_view = bias_tensor.view()[0]; auto grid_view = grid_tensor.view()[0]; std::size_t n_dims = grid_tensor.size(1); diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index ed70552e38..3db3d36f7c 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -269,13 +269,13 @@ NamedVector Integrand::build_function_impl( } NamedVector Integrand::compute_channel_part_ret_types() const { - auto acc_float = Type(DataType::dt_float, acc_batch_size, {}); - auto acc_int = Type(DataType::dt_int, acc_batch_size, {}); + auto acc_float = Type(DataType::dt_float64, acc_batch_size, {}); + auto acc_int = Type(DataType::dt_int32, acc_batch_size, {}); auto acc_float_array = [](int n) { - return Type(DataType::dt_float, acc_batch_size, {n}); + return Type(DataType::dt_float64, acc_batch_size, {n}); }; auto acc_four_vec_array = [](int n) { - return Type(DataType::dt_float, acc_batch_size, {n, 4}); + return Type(DataType::dt_float64, acc_batch_size, {n, 4}); }; bool has_multi_flavor = _pid_options.size() > 1; @@ -299,7 +299,7 @@ NamedVector Integrand::compute_channel_part_ret_types() const { } // outputs after cuts - ret.push_back("indices_acc", Type(DataType::dt_int, acc_batch_size, {})); + ret.push_back("indices_acc", Type(DataType::dt_int32, acc_batch_size, {})); ret.push_back("momenta_acc", acc_four_vec_array(particle_count)); if (_has_mirror) { if (!_madnis_training) { diff --git a/madspace/src/phasespace/mlp.cpp b/madspace/src/phasespace/mlp.cpp index 4de6a6537f..80cb554a9c 100644 --- a/madspace/src/phasespace/mlp.cpp +++ b/madspace/src/phasespace/mlp.cpp @@ -19,12 +19,12 @@ Value build_layer( ) { auto weight = fb.global( prefixed_name(prefix, std::format("layer{}.weight", layer_index)), - DataType::dt_float, + DataType::dt_float64, {output_dim, input_dim} ); auto bias = fb.global( prefixed_name(prefix, std::format("layer{}.bias", layer_index)), - DataType::dt_float, + DataType::dt_float64, {output_dim} ); auto linear_out = fb.matmul(input, weight, bias); @@ -59,18 +59,18 @@ void initialize_layer( prefixed_name(prefix, std::format("layer{}.weight", layer_index)); auto bias_name = prefixed_name(prefix, std::format("layer{}.bias", layer_index)); auto weight_tensor_global = context->define_global( - weight_name, DataType::dt_float, {output_dim, input_dim}, true + weight_name, DataType::dt_float64, {output_dim, input_dim}, true ); auto bias_tensor_global = - context->define_global(bias_name, DataType::dt_float, {output_dim}, true); + context->define_global(bias_name, DataType::dt_float64, {output_dim}, true); Tensor weight_tensor, bias_tensor; bool is_cpu = context->device() == cpu_device(); if (is_cpu) { weight_tensor = weight_tensor_global; bias_tensor = bias_tensor_global; } else { - weight_tensor = Tensor(DataType::dt_float, {1, output_dim, input_dim}); - bias_tensor = Tensor(DataType::dt_float, {1, output_dim}); + weight_tensor = Tensor(DataType::dt_float64, {1, output_dim, input_dim}); + bias_tensor = Tensor(DataType::dt_float64, {1, output_dim}); } auto weight_view = weight_tensor.view()[0]; diff --git a/madspace/src/phasespace/pdf.cpp b/madspace/src/phasespace/pdf.cpp index 0a437f341e..63bdc3f8e3 100644 --- a/madspace/src/phasespace/pdf.cpp +++ b/madspace/src/phasespace/pdf.cpp @@ -320,14 +320,14 @@ std::vector PdfGrid::logq2_shape(bool batch_dim) const { void PdfGrid::initialize_globals(ContextPtr context, const std::string& prefix) const { auto logx_tensor_global = context->define_global( - prefixed_name(prefix, "pdf_logx"), DataType::dt_float, logx_shape() + prefixed_name(prefix, "pdf_logx"), DataType::dt_float64, logx_shape() ); auto q2_tensor_global = context->define_global( - prefixed_name(prefix, "pdf_logq2"), DataType::dt_float, logq2_shape() + prefixed_name(prefix, "pdf_logq2"), DataType::dt_float64, logq2_shape() ); auto coeffs_tensor_global = context->define_global( prefixed_name(prefix, "pdf_coefficients"), - DataType::dt_float, + DataType::dt_float64, coefficients_shape() ); bool is_cpu = context->device() == cpu_device(); @@ -337,9 +337,9 @@ void PdfGrid::initialize_globals(ContextPtr context, const std::string& prefix) q2_tensor = q2_tensor_global; coeffs_tensor = coeffs_tensor_global; } else { - logx_tensor = Tensor(DataType::dt_float, logx_shape(true)); - q2_tensor = Tensor(DataType::dt_float, logq2_shape(true)); - coeffs_tensor = Tensor(DataType::dt_float, coefficients_shape(true)); + logx_tensor = Tensor(DataType::dt_float64, logx_shape(true)); + q2_tensor = Tensor(DataType::dt_float64, logq2_shape(true)); + coeffs_tensor = Tensor(DataType::dt_float64, coefficients_shape(true)); } initialize_logx(logx_tensor); initialize_logq2(q2_tensor); @@ -391,17 +391,17 @@ NamedVector PartonDensity::build_function_impl( auto q = args.at(1); auto grid_logx = fb.global( prefixed_name(_prefix, "pdf_logx"), - DataType::dt_float, + DataType::dt_float64, {_logx_shape.begin(), _logx_shape.end()} ); auto grid_logq2 = fb.global( prefixed_name(_prefix, "pdf_logq2"), - DataType::dt_float, + DataType::dt_float64, {_logq2_shape.begin(), _logq2_shape.end()} ); auto grid_coeffs = fb.global( prefixed_name(_prefix, "pdf_coefficients"), - DataType::dt_float, + DataType::dt_float64, {_coeffs_shape.begin(), _coeffs_shape.end()} ); if (_dynamic_pid) { @@ -546,11 +546,11 @@ void AlphaSGrid::initialize_globals( ContextPtr context, const std::string& prefix ) const { auto q2_tensor_global = context->define_global( - prefixed_name(prefix, "alpha_s_logq2"), DataType::dt_float, logq2_shape() + prefixed_name(prefix, "alpha_s_logq2"), DataType::dt_float64, logq2_shape() ); auto coeffs_tensor_global = context->define_global( prefixed_name(prefix, "alpha_s_coefficients"), - DataType::dt_float, + DataType::dt_float64, coefficients_shape() ); bool is_cpu = context->device() == cpu_device(); @@ -559,8 +559,8 @@ void AlphaSGrid::initialize_globals( q2_tensor = q2_tensor_global; coeffs_tensor = coeffs_tensor_global; } else { - q2_tensor = Tensor(DataType::dt_float, logq2_shape(true)); - coeffs_tensor = Tensor(DataType::dt_float, coefficients_shape(true)); + q2_tensor = Tensor(DataType::dt_float64, logq2_shape(true)); + coeffs_tensor = Tensor(DataType::dt_float64, coefficients_shape(true)); } initialize_logq2(q2_tensor); initialize_coefficients(coeffs_tensor); @@ -584,12 +584,12 @@ NamedVector RunningCoupling::build_function_impl( auto q = args.at(0); auto grid_logq2 = fb.global( prefixed_name(_prefix, "alpha_s_logq2"), - DataType::dt_float, + DataType::dt_float64, {_logq2_shape.begin(), _logq2_shape.end()} ); auto grid_coeffs = fb.global( prefixed_name(_prefix, "alpha_s_coefficients"), - DataType::dt_float, + DataType::dt_float64, {_coeffs_shape.begin(), _coeffs_shape.end()} ); return {{"alpha_s", fb.interpolate_alpha_s(q, grid_logq2, grid_coeffs)}}; diff --git a/madspace/src/phasespace/phasespace.cpp b/madspace/src/phasespace/phasespace.cpp index 41202b82b8..09b31f79d7 100644 --- a/madspace/src/phasespace/phasespace.cpp +++ b/madspace/src/phasespace/phasespace.cpp @@ -74,9 +74,8 @@ void update_mass_min_max( // Masses of the external momenta in the order they are handed to boost_beam, // i.e. after the channel permutation. A position whose mass differs between // permutations gets -1, and boost_beam reads that mass off the momentum. -std::vector lab_masses( - const Topology& topology, const nested_vector2& permutations -) { +std::vector +lab_masses(const Topology& topology, const nested_vector2& permutations) { std::vector masses = topology.incoming_masses(); const auto& out = topology.outgoing_masses(); masses.insert(masses.end(), out.begin(), out.end()); @@ -275,7 +274,7 @@ PhaseSpaceMapping::PhaseSpaceMapping( if (nd > 0) { in.push_back( "discrete", - Type{DataType::dt_int, batch_size, {static_cast(nd)}} + Type{DataType::dt_int32, batch_size, {static_cast(nd)}} ); } return in; @@ -291,9 +290,9 @@ PhaseSpaceMapping::PhaseSpaceMapping( : NamedVector{} ), _topology(topology), - _cuts(cuts.value_or(Cuts( - topology.outgoing_masses().size() + topology.incoming_masses().size() - ))), + _cuts(cuts.value_or( + Cuts(topology.outgoing_masses().size() + topology.incoming_masses().size()) + )), _pi_factors( std::pow(2 * PI, 4 - 3 * static_cast(topology.outgoing_masses().size())) ), @@ -361,15 +360,18 @@ PhaseSpaceMapping::PhaseSpaceMapping( auto at_or = [](const std::vector& v, std::size_t i, double none) { return i < v.size() ? v.at(i) : none; }; - auto pair_at = [](const std::vector>& m, - std::size_t i, - std::size_t j) { - return i < m.size() && j < m.at(i).size() ? m.at(i).at(j) : 0.; - }; + auto pair_at = + [](const std::vector>& m, std::size_t i, std::size_t j) { + return i < m.size() && j < m.at(i).size() ? m.at(i).at(j) : 0.; + }; constexpr double inf = std::numeric_limits::infinity(); std::vector topo_pt_min(n_out, inf), topo_eta_max(n_out, 0.); - std::vector> topo_m_inv_min(n_out, std::vector(n_out, inf)); - std::vector> topo_dr_min(n_out, std::vector(n_out, inf)); + std::vector> topo_m_inv_min( + n_out, std::vector(n_out, inf) + ); + std::vector> topo_dr_min( + n_out, std::vector(n_out, inf) + ); for (const auto& position : event_position) { for (std::size_t t = 0; t < n_out; ++t) { std::size_t i = position.at(t); @@ -407,11 +409,8 @@ PhaseSpaceMapping::PhaseSpaceMapping( constexpr std::size_t no_leaf = static_cast(-1); std::vector> node_leaves(_topology.decays().size()); { - std::vector decay_to_outgoing( - _topology.decays().size(), no_leaf - ); - for (std::size_t out_pos = 0; - out_pos < _topology.outgoing_indices().size(); + std::vector decay_to_outgoing(_topology.decays().size(), no_leaf); + for (std::size_t out_pos = 0; out_pos < _topology.outgoing_indices().size(); ++out_pos) { decay_to_outgoing.at(_topology.outgoing_indices().at(out_pos)) = out_pos; } @@ -578,9 +577,8 @@ PhaseSpaceMapping::PhaseSpaceMapping( // Two sources of such a floor: a cut on sqrt(s_hat) itself, and the pair // cuts, every one of which bounds the total exactly as it bounds a // propagator above. - double sqrt_s_hat_min = std::max( - read_cuts ? _cuts.sqrt_s_min() : 0., pair_floor(node_leaves.at(0)) - ); + double sqrt_s_hat_min = + std::max(read_cuts ? _cuts.sqrt_s_min() : 0., pair_floor(node_leaves.at(0))); // Only the luminosity mapping samples the root virtuality, and the root's // e_min is read nowhere else: build_forward_impl gives the root its s_min // only when it samples it, and the other propagators bound themselves by @@ -658,8 +656,7 @@ PhaseSpaceMapping::PhaseSpaceMapping( if (!is_com_decay || _map_luminosity) { double mass = decay.width == 0. ? 0. : decay.mass; double width = decay.width; - info.invariant = - Invariant(invariant_power, mass, width, decay.flat_window); + info.invariant = Invariant(invariant_power, mass, width, decay.flat_window); } } for (std::size_t index : _topology.decay_integration_order()) { @@ -695,8 +692,8 @@ PhaseSpaceMapping::PhaseSpaceMapping( ); } // Reorder the per-pair cut matrices (indexed by topology outgoing - // position) into the child order in which masses/pt are handed to the chain, - // mirroring the pt_min reordering above. Composite (non-leaf) + // position) into the child order in which masses/pt are handed to the + // chain, mirroring the pt_min reordering above. Composite (non-leaf) // children carry no pairwise cut. const auto& out_idx = topology.outgoing_indices(); const auto& child_indices = topology.decays().at(0).child_indices; @@ -1010,7 +1007,9 @@ Mapping::Result PhaseSpaceMapping::build_forward_impl( // boost into correct frame and apply cuts auto p_ext_lab = _map_luminosity - ? fb.boost_beam(p_ext_stack, Value(lab_masses(_topology, _permutations)), x1, x2) + ? fb.boost_beam( + p_ext_stack, Value(lab_masses(_topology, _permutations)), x1, x2 + ) : p_ext_stack; dets.push_back(_cuts.build_function(fb, {p_ext_lab}).at(0)); auto ps_weight = fb.cut_unphysical(fb.product(dets), p_ext_lab, x1, x2); @@ -1023,11 +1022,11 @@ Mapping::Result PhaseSpaceMapping::build_inverse_impl( const NamedVector& conditions ) const { Value p_ext_lab = inputs.at(0), x1 = inputs.at(1), x2 = inputs.at(2); - Value p_ext_stack = - _map_luminosity ? fb.boost_beam_inverse( - p_ext_lab, Value(lab_masses(_topology, _permutations)), x1, x2 - ) - : p_ext_lab; + Value p_ext_stack = _map_luminosity + ? fb.boost_beam_inverse( + p_ext_lab, Value(lab_masses(_topology, _permutations)), x1, x2 + ) + : p_ext_lab; // permute momenta if permutations are given if (_permutations.size() > 1) { @@ -1048,12 +1047,11 @@ Mapping::Result PhaseSpaceMapping::build_inverse_impl( std::vector decay_data( _topology.decays().begin(), _topology.decays().end() ); - for (auto [decay_index, mass, momentum] : - zip(_topology.outgoing_indices(), + for (auto [decay_index, mass, momentum] : zip( + _topology.outgoing_indices(), _topology.outgoing_masses(), - std::span( - p_ext.begin() + _topology.incoming_masses().size(), p_ext.end() - ))) { + std::span(p_ext.begin() + _topology.incoming_masses().size(), p_ext.end()) + )) { auto& data = decay_data.at(decay_index); data.mass = mass; data.mass2 = mass * mass; diff --git a/madspace/src/phasespace/vegas.cpp b/madspace/src/phasespace/vegas.cpp index 0d5f804cb4..4f3e03e456 100644 --- a/madspace/src/phasespace/vegas.cpp +++ b/madspace/src/phasespace/vegas.cpp @@ -39,7 +39,7 @@ Mapping::Result VegasMapping::build_forward_impl( ) const { auto grid = fb.global( _grid_name, - DataType::dt_float, + DataType::dt_float64, {static_cast(_dimension), static_cast(_bin_count) + 1} ); auto [output, dets] = fb.vegas_forward(inputs.at(0), grid); @@ -53,7 +53,7 @@ Mapping::Result VegasMapping::build_inverse_impl( ) const { auto grid = fb.global( _grid_name, - DataType::dt_float, + DataType::dt_float64, {static_cast(_dimension), static_cast(_bin_count) + 1} ); auto [output, dets] = fb.vegas_inverse(inputs.at(0), grid); @@ -62,7 +62,7 @@ Mapping::Result VegasMapping::build_inverse_impl( void VegasMapping::initialize_globals(ContextPtr context) const { context->define_global( - _grid_name, DataType::dt_float, {_dimension, _bin_count + 1} + _grid_name, DataType::dt_float64, {_dimension, _bin_count + 1} ); initialize_vegas_grid(context, _grid_name); } @@ -71,14 +71,14 @@ void madspace::initialize_vegas_grid(ContextPtr context, const std::string& grid bool is_cpu = context->device() == cpu_device(); auto grid_global = context->global(grid_name); if (grid_global.shape().size() != 3 || grid_global.size(0) != 1 || - grid_global.size(2) < 2 || grid_global.dtype() != DataType::dt_float) { + grid_global.size(2) < 2 || grid_global.dtype() != DataType::dt_float64) { throw std::runtime_error("Invalid vegas grid type"); } Tensor grid; if (is_cpu) { grid = grid_global; } else { - grid = Tensor(DataType::dt_float, grid_global.shape()); + grid = Tensor(DataType::dt_float64, grid_global.shape()); } auto grid_view = grid.view()[0]; diff --git a/madspace/src/python/function_runtime.cpp b/madspace/src/python/function_runtime.cpp index 161c4e5ec5..2abc56fdeb 100644 --- a/madspace/src/python/function_runtime.cpp +++ b/madspace/src/python/function_runtime.cpp @@ -260,10 +260,13 @@ py::object madspace_py::tensor_to_dlpack( } else { DLDataType dtype; switch (tensor.dtype()) { - case DataType::dt_float: + case DataType::dt_float32: + dtype = {kDLFloat, 32, 1}; + break; + case DataType::dt_float64: dtype = {kDLFloat, 64, 1}; break; - case DataType::dt_int: + case DataType::dt_int32: dtype = {kDLInt, 32, 1}; break; default: @@ -410,18 +413,20 @@ Tensor madspace_py::dlpack_to_tensor( bool is_batch_sizes = expected_type ? expected_type->dtype == DataType::batch_sizes : false; DataType dtype; - if (dl_tensor->dtype.code == kDLFloat && dl_tensor->dtype.bits == 64 && + if (dl_tensor->dtype.code == kDLFloat && + (dl_tensor->dtype.bits == 32 || dl_tensor->dtype.bits == 64) && dl_tensor->dtype.lanes == 1) { - dtype = DataType::dt_float; - if (expected_type && expected_type->dtype != DataType::dt_float) { + dtype = + dl_tensor->dtype.bits == 32 ? DataType::dt_float32 : DataType::dt_float64; + if (expected_type && expected_type->dtype != dtype) { throw std::invalid_argument( std::format("Argument {}: got unexpected dtype", arg_index + 1) ); } } else if (dl_tensor->dtype.code == kDLInt && dl_tensor->dtype.bits == 32 && dl_tensor->dtype.lanes == 1) { - dtype = DataType::dt_int; - if (expected_type && expected_type->dtype != DataType::dt_int && + dtype = DataType::dt_int32; + if (expected_type && expected_type->dtype != DataType::dt_int32 && !is_batch_sizes) { throw std::invalid_argument( std::format("Argument {}: got unexpected dtype", arg_index + 1) @@ -439,7 +444,7 @@ Tensor madspace_py::dlpack_to_tensor( } else { throw std::invalid_argument( std::format( - "Argument {}: input dtype must be 64-bit float or 32-bit int", + "Argument {}: input dtype must be 32/64-bit float or 32-bit int", arg_index + 1 ) ); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 7d97fb22a6..743c21ead5 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -148,8 +148,9 @@ PYBIND11_MODULE(_madspace_py, m) { m, "DataType", { - {"int", DataType::dt_int}, - {"float", DataType::dt_float}, + {"int32", DataType::dt_int32}, + {"float32", DataType::dt_float32}, + {"float64", DataType::dt_float64}, {"batch_sizes", DataType::batch_sizes}, }, "", diff --git a/madspace/tests/test_histogram_op.py b/madspace/tests/test_histogram_op.py index 3cc026c3c1..67bdc9efb5 100644 --- a/madspace/tests/test_histogram_op.py +++ b/madspace/tests/test_histogram_op.py @@ -14,7 +14,7 @@ def run_histogram(x, w, lo, hi): - out_type = ms.Type(ms.DataType.float, ms.BatchSize.one, [N_BINS + 2]) + out_type = ms.Type(ms.DataType.float64, ms.BatchSize.one, [N_BINS + 2]) fb = ms.FunctionBuilder( ms.NamedTypes([(name, ms.batch_float) for name in ("x", "w", "lo", "hi")]), ms.NamedTypes([("values", out_type), ("square_values", out_type)]), @@ -34,13 +34,17 @@ def test_histogram_underflow_bins_overflow(): # underflow: -1; bins of width 0.2: 0.1, 0.3, 0.5, -, 0.95; overflow: 1.5 and 2 expected = np.array([1.0, 2.0, 7.0, 3.0, 0.0, 4.0, 5.0 + 6.0]) np.testing.assert_allclose(values, expected) - np.testing.assert_allclose(square_values, [1.0, 4.0, 49.0, 9.0, 0.0, 16.0, 25.0 + 36.0]) + np.testing.assert_allclose( + square_values, [1.0, 4.0, 49.0, 9.0, 0.0, 16.0, 25.0 + 36.0] + ) def test_histogram_range_per_event(): """min and max are per-event inputs: the same value falls in different bins.""" x = np.full(3, 0.5) w = np.ones(3) - values, _ = run_histogram(x, w, np.array([0.0, 0.4, 1.0]), np.array([1.0, 1.4, 2.0])) + values, _ = run_histogram( + x, w, np.array([0.0, 0.4, 1.0]), np.array([1.0, 1.4, 2.0]) + ) # 0.5 in [0, 1): bin 2 (index 3); in [0.4, 1.4): bin 0 (index 1); below [1, 2): underflow np.testing.assert_allclose(values, [1.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0]) diff --git a/madspace/tests/test_simd_layout.py b/madspace/tests/test_simd_layout.py index 247ab81be8..e84948a30e 100644 --- a/madspace/tests/test_simd_layout.py +++ b/madspace/tests/test_simd_layout.py @@ -87,7 +87,9 @@ def body(fb, x): assert stacked.literal_value == ([3], [1.0, 2.0, 3.0]) return fb.mul(x, stacked) - func = build([("x", ms.batch_float_array(3))], [("y", ms.batch_float_array(3))], body) + func = build( + [("x", ms.batch_float_array(3))], [("y", ms.batch_float_array(3))], body + ) assert instruction_names(func) == ["mul"] @@ -245,7 +247,7 @@ def test_global_probability_gradient_is_summed_over_events(): """The gradient of a broadcast (global) probability vector is the sum of the contributions of all events.""" ctx = ms.Context() - ctx.define_global("test_simd_layout_probs", ms.DataType.float, [4], True) + ctx.define_global("test_simd_layout_probs", ms.DataType.float64, [4], True) ctx.get_global("test_simd_layout_probs").torch()[:] = torch.tensor( [[0.1, 0.4, 0.3, 0.2]] ) @@ -253,7 +255,7 @@ def test_global_probability_gradient_is_summed_over_events(): ms.NamedTypes([("index", ms.batch_int)]), ms.NamedTypes([("r", ms.batch_float), ("det", ms.batch_float)]), ) - probs = fb.get_global("test_simd_layout_probs", ms.DataType.float, [4]) + probs = fb.get_global("test_simd_layout_probs", ms.DataType.float64, [4]) r, det = fb.sample_discrete_probs_inverse(fb.input(0), probs) fb.output(0, r) fb.output(1, det) @@ -286,13 +288,13 @@ def test_broadcast_gradient_through_stack(): """stack of a broadcast (global) and a batched value: the gradient of the global accumulates the whole batch.""" ctx = ms.Context() - ctx.define_global("test_simd_layout_scalar", ms.DataType.float, [], True) + ctx.define_global("test_simd_layout_scalar", ms.DataType.float64, [], True) ctx.get_global("test_simd_layout_scalar").torch()[:] = 1.5 fb = ms.FunctionBuilder( ms.NamedTypes([("x", ms.batch_float)]), ms.NamedTypes([("y", ms.batch_float_array(2))]), ) - g = fb.get_global("test_simd_layout_scalar", ms.DataType.float, []) + g = fb.get_global("test_simd_layout_scalar", ms.DataType.float64, []) fb.output(0, fb.stack([g, fb.input(0)])) module = FunctionModule(fb.function(), ctx) param = module.global_params["test_simd_layout_scalar"] From 1933c24a43da9572d5f13cb5f014ff1ce80b7487 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 18:42:08 +0200 Subject: [PATCH 02/24] madspace: float32/float64 variants of NN instructions Instructions with a 'precision' field in instruction_set.yaml are expanded into _float32 and _float64, each with its own opcode. The FunctionBuilder method dispatches on the dtype of its arguments. Kernel policies are templated on the float type (CpuTypes32, GpuTypes32). - matmul and the activations are available in both precisions (sgemm/dgemm) - new instructions convert_float32 and convert_float64 - device copy/zero/add and gradient allocation handle float32 - activation backward kernels now accumulate into the input gradient - fix the shape check of single (non-batch) inputs from Python --- madspace/generate_headers.py | 110 ++++++++-- .../include/madspace/compgraphs/function.hpp | 4 + madspace/instruction_set.yaml | 36 ++++ madspace/src/compgraphs/function.cpp | 30 ++- madspace/src/compgraphs/type.cpp | 4 +- madspace/src/cpu/device.cpp | 64 ++++-- madspace/src/cpu/kernel_definitions.hpp | 22 +- madspace/src/cpu/runtime.cpp | 196 +++++++++++++----- madspace/src/gpu/device.cu | 74 ++++--- madspace/src/gpu/gpu_abstraction.cuh | 4 + madspace/src/gpu/kernel_definitions.cuh | 25 ++- madspace/src/gpu/runtime.cu | 191 +++++++++-------- madspace/src/kernels/nn.hpp | 113 +++++++--- madspace/src/python/function_runtime.cpp | 6 +- madspace/tests/test_float32.py | 124 +++++++++++ 15 files changed, 757 insertions(+), 246 deletions(-) create mode 100644 madspace/tests/test_float32.py diff --git a/madspace/generate_headers.py b/madspace/generate_headers.py index 16ccdec4b9..e379291264 100644 --- a/madspace/generate_headers.py +++ b/madspace/generate_headers.py @@ -25,6 +25,7 @@ def main(): commands = {} for sec in data: commands.update({key: value for key, value in sec.items() if key != "title"}) + commands = expand_precisions(commands) for i, cmd in enumerate(commands.values()): cmd["opcode"] = i sections = [ @@ -45,6 +46,51 @@ def main(): kernel_checks(commands, out_dir) +# C++ element type and kernel policy suffix for each float precision +PRECISIONS = {"float32": ("float", "32"), "float64": ("double", "")} + + +def expand_precisions(commands): + """Split each instruction with a `precision` list into one instruction per + precision, named `_`. In these, `float` in a type means the + instruction's precision. The expanded entries keep the original name in + `base_name` and the kernel policy suffix in `types_suffix`.""" + expanded = {} + for name, cmd in commands.items(): + precisions = cmd.get("precision") + if precisions is None: + cmd["base_name"] = name + cmd["types_suffix"] = "" + cmd["dispatch"] = None + expanded[name] = cmd + continue + for precision in precisions: + ctype, suffix = PRECISIONS[precision] + sub_cmd = dict(cmd) + for key in ("inputs", "outputs"): + if sub_cmd[key] == "any": + raise ValueError(f"{name}: precision needs explicit {key}") + sub_cmd[key] = [ + { + **arg, + "type": ( + [precision, *arg["type"][1:]] + if arg.get("type", [None])[0] == "float" + else arg.get("type") + ), + } + for arg in cmd[key] + ] + sub_cmd["base_name"] = name + sub_cmd["types_suffix"] = suffix + sub_cmd["ctype"] = ctype + sub_cmd["doc_cmd"] = cmd + # the builder method is generated once, for the first precision + sub_cmd["dispatch"] = precisions if precision == precisions[0] else [] + expanded[f"{name}_{precision}"] = sub_cmd + return expanded + + def write_autogen(f): f.write( "// This file was automatically generated based on instruction_set.yaml\n" @@ -118,14 +164,23 @@ def function_builder_mixin(commands, out_dir): write_autogen(f) first = True for name, cmd in commands.items(): + if cmd["dispatch"] == []: + continue if first: first = False else: f.write("\n") - f.write(doc_comment(name, cmd)) + base_name = cmd["base_name"] + f.write(doc_comment(base_name, cmd.get("doc_cmd", cmd))) if cmd["inputs"] == "any": parameters = "ValueVec args" instruction_call = f'instruction("{name}", args)' + elif cmd["dispatch"]: + parameters = ", ".join(f"Value {arg['name']}" for arg in cmd["inputs"]) + arguments = ", ".join(arg["name"] for arg in cmd["inputs"]) + instruction_call = ( + f'precision_instruction("{base_name}", {{{arguments}}})' + ) else: parameters = ", ".join(f"Value {arg['name']}" for arg in cmd["inputs"]) arguments = ", ".join(arg["name"] for arg in cmd["inputs"]) @@ -152,7 +207,7 @@ def function_builder_mixin(commands, out_dir): f" return {{{return_array}}};" ) - f.write(f"{return_type} {name}({parameters}) {{\n{func_body}\n}}\n") + f.write(f"{return_type} {base_name}({parameters}) {{\n{func_body}\n}}\n") def instruction_set_python(commands, out_dir): @@ -170,6 +225,9 @@ def instruction_set_python(commands, out_dir): ) for name, cmd in commands.items(): + if cmd["dispatch"] == []: + continue + name = cmd["base_name"] f.write(f' fb.def("{name}", &FunctionBuilder::{name}') if cmd["inputs"] == "any": f.write(', py::arg("args")') @@ -259,22 +317,28 @@ def runtime_mixin(commands, device, out_dir): for name, cmd in commands.items(): opcode = cmd["opcode"] + base_name = cmd["base_name"] + suffix = cmd["types_suffix"] if cmd.get("custom_op", False): - func = f"op_{name}" + func = f"op_{base_name}" + if "ctype" in cmd: + func += f"<{cmd['ctype']}>" else: n_inputs = len(cmd["inputs"]) n_outputs = len(cmd["outputs"]) dims = cmd.get("dims", 1) vectorized = cmd.get("vectorized", True) + kernel_name = f"kernel_{base_name}" if device == "cpu": - if vectorized: - kernel = f"kernel_{name}, kernel_{name}" - else: - kernel = f"kernel_{name}, kernel_{name}" + simd = "Simd" if vectorized else "Cpu" + kernel = ( + f"{kernel_name}, " + f"{kernel_name}<{simd}Types{suffix}>" + ) device_arg = ", DeviceType" elif device == "gpu": - kernel = f"kernel_{name}" + kernel = f"{kernel_name}" device_arg = "" foreach_func = ( f"tensor_foreach_dynamic<{kernel}, {n_inputs}, {n_outputs}{device_arg}>" @@ -297,10 +361,15 @@ def runtime_backward_mixin(commands, device, out_dir): if not cmd.get("differentiable", False): continue opcode = cmd["opcode"] + base_name = cmd["base_name"] + suffix = cmd["types_suffix"] if cmd.get("custom_op", False): + func = f"backward_op_{base_name}" + if "ctype" in cmd: + func += f"<{cmd['ctype']}>" f.write( f"case {opcode}:\n" - f" backward_op_{name}(instr, locals, local_grads, device);\n" + f" {func}(instr, locals, local_grads, device);\n" f" break;\n" ) else: @@ -322,20 +391,16 @@ def runtime_backward_mixin(commands, device, out_dir): out_stored_str = ",".join(str(i) for i in out_stored) vectorized = cmd.get("vectorized", True) + kernel_name = f"backward_kernel_{base_name}" if device == "cpu": - if vectorized: - kernel = ( - f"backward_kernel_{name}, " - f"backward_kernel_{name}" - ) - else: - kernel = ( - f"backward_kernel_{name}, " - f"backward_kernel_{name}" - ) + simd = "Simd" if vectorized else "Cpu" + kernel = ( + f"{kernel_name}, " + f"{kernel_name}<{simd}Types{suffix}>" + ) device_arg = ", DeviceType" elif device == "gpu": - kernel = f"backward_kernel_{name}" + kernel = f"{kernel_name}" device_arg = "" dims = cmd.get("dims", 1) @@ -375,7 +440,8 @@ def kernel_checks(commands, out_dir): def check(kernel, index, single, what): f.write( f"static_assert(\n" - f" single_params)>::value[{index}] == " + f" single_params)>" + f"::value[{index}] == " f"{'true' if single else 'false'},\n" f' "{kernel}: parameter {index} ({what}) must ' f'{"" if single else "not "}be single"\n' @@ -385,6 +451,8 @@ def check(kernel, index, single, what): for name, cmd in commands.items(): if cmd.get("custom_op", False) or not cmd.get("vectorized", True): continue + suffix = cmd["types_suffix"] + name = cmd["base_name"] inputs = cmd["inputs"] if is_single_input(inputs[0]): raise ValueError( diff --git a/madspace/include/madspace/compgraphs/function.hpp b/madspace/include/madspace/compgraphs/function.hpp index 42c0a48dd2..9f373f6d81 100644 --- a/madspace/include/madspace/compgraphs/function.hpp +++ b/madspace/include/madspace/compgraphs/function.hpp @@ -128,6 +128,10 @@ class FunctionBuilder { ValueVec instruction(const std::string& name, const ValueVec& args); /// Append @p instruction with arguments @p args; returns its outputs. ValueVec instruction(InstructionPtr instruction, const ValueVec& args); + /// Calls the float32 or float64 variant `_float32` / `_float64` of + /// an instruction, chosen by the first non-literal float argument. Float64 + /// literals are converted for a float32 call. + ValueVec precision_instruction(const std::string& name, const ValueVec& args); /// Index of the execution stream new instructions are added to. std::size_t current_stream() const { return _current_stream; } /// Select the execution stream new instructions are added to. diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 7621b2ff91..8d0378d397 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -2869,7 +2869,36 @@ interpolate_alpha_s: --- title: Normalizing flows and MLPs +convert_float32: + inputs: + - name: in1 + type: [float64, ...] + desc: "double-precision tensor" + outputs: + - name: out + type: [float32, ...] + desc: "the input rounded to single precision" + dims: 0 + vectorized: False + differentiable: True + desc: "Converts a tensor to single precision. Used at the boundaries of single-precision @ref MLP networks." + +convert_float64: + inputs: + - name: in1 + type: [float32, ...] + desc: "single-precision tensor" + outputs: + - name: out + type: [float64, ...] + desc: "the input in double precision" + dims: 0 + vectorized: False + differentiable: True + desc: "Converts a tensor to double precision. Used at the boundaries of single-precision @ref MLP networks." + matmul: + precision: [float32, float64] inputs: - name: x type: [float, n] @@ -2891,6 +2920,7 @@ matmul: desc: "Affine layer y = weight x + bias. Building block of @ref MLP." relu: + precision: [float32, float64] inputs: - name: in1 type: [float, ...] @@ -2904,6 +2934,7 @@ relu: desc: "Element-wise ReLU activation. Building block of @ref MLP." leaky_relu: + precision: [float32, float64] inputs: - name: in1 type: [float, ...] @@ -2917,6 +2948,7 @@ leaky_relu: desc: "Element-wise leaky ReLU activation. Building block of @ref MLP." elu: + precision: [float32, float64] inputs: - name: in1 type: [float, ...] @@ -2930,6 +2962,7 @@ elu: desc: "Element-wise ELU activation. Building block of @ref MLP." gelu: + precision: [float32, float64] inputs: - name: in1 type: [float, ...] @@ -2943,6 +2976,7 @@ gelu: desc: "Element-wise GELU activation. Building block of @ref MLP." sigmoid: + precision: [float32, float64] inputs: - name: in1 type: [float, ...] @@ -2957,6 +2991,7 @@ sigmoid: desc: "Element-wise sigmoid activation. Building block of @ref MLP." scaled_tanh: + precision: [float32, float64] inputs: - name: in1 type: [float, ...] @@ -2971,6 +3006,7 @@ scaled_tanh: desc: "Element-wise scaled hyperbolic tangent activation. Building block of @ref MLP." softplus: + precision: [float32, float64] inputs: - name: in1 type: [float, ...] diff --git a/madspace/src/compgraphs/function.cpp b/madspace/src/compgraphs/function.cpp index bc2c307d66..ecaa86994e 100644 --- a/madspace/src/compgraphs/function.cpp +++ b/madspace/src/compgraphs/function.cpp @@ -263,6 +263,32 @@ ValueVec FunctionBuilder::instruction(const std::string& name, const ValueVec& a return instruction(find_instr->second.get(), args); } +ValueVec +FunctionBuilder::precision_instruction(const std::string& name, const ValueVec& args) { + auto is_literal = [](const Value& arg) { + return !std::holds_alternative(arg.literal_value); + }; + DataType dtype = DataType::dt_float64; + for (auto& arg : args) { + if (!is_literal(arg) && is_float(arg.type.dtype)) { + dtype = arg.type.dtype; + break; + } + } + if (dtype == DataType::dt_float64) { + return instruction(name + "_float64", args); + } + ValueVec converted_args; + for (auto& arg : args) { + converted_args.push_back( + is_literal(arg) && arg.type.dtype == DataType::dt_float64 + ? convert_float32(arg) + : arg + ); + } + return instruction(name + "_float32", converted_args); +} + ValueVec FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) { // check argument types and determine if all arguments are constant @@ -315,8 +341,8 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) { return arg.local_index == -1 && std::holds_alternative(arg.literal_value); }; - bool all_float_literals = args.size() > 0 && - std::all_of(args.begin(), args.end(), is_float_literal); + bool all_float_literals = + args.size() > 0 && std::all_of(args.begin(), args.end(), is_float_literal); bool all_int_literals = args.size() > 0 && std::all_of(args.begin(), args.end(), is_int_literal); if (all_float_literals) { diff --git a/madspace/src/compgraphs/type.cpp b/madspace/src/compgraphs/type.cpp index 17fd02acbb..1762d9811d 100644 --- a/madspace/src/compgraphs/type.cpp +++ b/madspace/src/compgraphs/type.cpp @@ -293,7 +293,6 @@ void madspace::from_json(const json& j, Value& value) { value = Value(j_data.get>(), shape); } break; - case DataType::dt_float32: case DataType::dt_float64: if (shape.size() == 0) { value = unescape_special_values(j_data); @@ -304,9 +303,8 @@ void madspace::from_json(const json& j, Value& value) { } value = Value(tensor, shape); } - // float32 literals are stored as double and only converted at runtime - value.type.dtype = dtype; break; + case DataType::dt_float32: case DataType::batch_sizes: throw std::invalid_argument("invalid data type"); } diff --git a/madspace/src/cpu/device.cpp b/madspace/src/cpu/device.cpp index e3ef27b902..bdae82c913 100644 --- a/madspace/src/cpu/device.cpp +++ b/madspace/src/cpu/device.cpp @@ -10,40 +10,61 @@ namespace { template void tensor_copy_impl(const Tensor& source, Tensor& target, const D& device) { - if (source.dtype() == DataType::dt_float64 && - target.dtype() == DataType::dt_float64) { + if (source.dtype() != target.dtype()) { + throw std::runtime_error("invalid dtype in copy"); + } + switch (source.dtype()) { + case DataType::dt_float64: tensor_foreach_dynamic_single< kernel_copy, kernel_copy, 1, 1>({&source}, {&target}, target.size(0), device); - } else if (source.dtype() == DataType::dt_int32 && - target.dtype() == DataType::dt_int32) { + break; + case DataType::dt_float32: + tensor_foreach_dynamic_single< + kernel_copy, + kernel_copy, + 1, + 1>({&source}, {&target}, target.size(0), device); + break; + case DataType::dt_int32: tensor_foreach_dynamic_single< kernel_copy_int, kernel_copy_int, 1, 1>({&source}, {&target}, target.size(0), device); - } else { + break; + default: throw std::runtime_error("invalid dtype in copy"); } } template void tensor_zero_impl(Tensor& tensor, const D& device) { - if (tensor.dtype() == DataType::dt_float64) { + switch (tensor.dtype()) { + case DataType::dt_float64: tensor_foreach_dynamic_single< kernel_zero, kernel_zero, 1, 1>({&tensor}, {&tensor}, tensor.size(0), device); - } else if (tensor.dtype() == DataType::dt_int32) { + break; + case DataType::dt_float32: + tensor_foreach_dynamic_single< + kernel_zero, + kernel_zero, + 1, + 1>({&tensor}, {&tensor}, tensor.size(0), device); + break; + case DataType::dt_int32: tensor_foreach_dynamic_single< kernel_zero_int, kernel_zero_int, 1, 1>({&tensor}, {&tensor}, tensor.size(0), device); - } else { + break; + default: throw std::runtime_error("invalid dtype in zero"); } } @@ -52,11 +73,28 @@ template void tensor_add_impl(const Tensor& source, Tensor& target, const D& device) { // a broadcast target (e.g. the gradient of a literal or global) accumulates // the whole source batch - tensor_foreach_dynamic_single< - kernel_add_inplace, - kernel_add_inplace, - 1, - 1>({&source}, {&target}, std::max(source.size(0), target.size(0)), device); + std::size_t batch_size = std::max(source.size(0), target.size(0)); + if (source.dtype() != target.dtype()) { + throw std::runtime_error("invalid dtype in add"); + } + switch (source.dtype()) { + case DataType::dt_float64: + tensor_foreach_dynamic_single< + kernel_add_inplace, + kernel_add_inplace, + 1, + 1>({&source}, {&target}, batch_size, device); + break; + case DataType::dt_float32: + tensor_foreach_dynamic_single< + kernel_add_inplace, + kernel_add_inplace, + 1, + 1>({&source}, {&target}, batch_size, device); + break; + default: + throw std::runtime_error("invalid dtype in add"); + } } } // namespace diff --git a/madspace/src/cpu/kernel_definitions.hpp b/madspace/src/cpu/kernel_definitions.hpp index 79916ef220..6e5f3bbc4c 100644 --- a/madspace/src/cpu/kernel_definitions.hpp +++ b/madspace/src/cpu/kernel_definitions.hpp @@ -10,32 +10,43 @@ namespace madspace { namespace kernels { -struct CpuTypes { +template +struct CpuTypesT { template - using FIn = const TensorView; + using FIn = const TensorView; template using IIn = const TensorView; template - using FOut = TensorView; + using FOut = TensorView; template using IOut = TensorView; template - using FSIn = const TensorView; + using FSIn = const TensorView; template using ISIn = const TensorView; - using FVal = double; + using FVal = F; using IVal = me_int_t; using BVal = bool; + // the same policy with a different float type, for conversion kernels + using F32 = CpuTypesT; + using F64 = CpuTypesT; }; +using CpuTypes = CpuTypesT; +using CpuTypes32 = CpuTypesT; inline double where(bool condition, double val_true, double val_false) { return condition ? val_true : val_false; } +inline float where(bool condition, float val_true, float val_false) { + return condition ? val_true : val_false; +} inline me_int_t where(bool condition, me_int_t val_true, me_int_t val_false) { return condition ? val_true : val_false; } inline double min(double arg1, double arg2) { return arg1 < arg2 ? arg1 : arg2; } inline double max(double arg1, double arg2) { return arg1 > arg2 ? arg1 : arg2; } +inline float min(float arg1, float arg2) { return arg1 < arg2 ? arg1 : arg2; } +inline float max(float arg1, float arg2) { return arg1 > arg2 ? arg1 : arg2; } inline std::size_t single_index(me_int_t arg) { return arg; } using std::acos; @@ -82,6 +93,7 @@ struct SimdTypes { #else // USE_SIMD using SimdTypes = CpuTypes; #endif // USE_SIMD +using SimdTypes32 = CpuTypes32; } // namespace kernels } // namespace madspace diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 6530e39de6..d83406bb22 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -29,6 +29,22 @@ extern "C" void dgemm_( int* ldc ); +extern "C" void sgemm_( + char* transa, + char* transb, + int* m, + int* n, + int* k, + float* alpha, + float* a, + int* lda, + float* b, + int* ldb, + float* beta, + float* c, + int* ldc +); + extern "C" void dgemv_( char* trans, int* m, @@ -43,6 +59,68 @@ extern "C" void dgemv_( int* incy ); +extern "C" void sgemv_( + char* trans, + int* m, + int* n, + float* alpha, + float* a, + int* lda, + float* x, + int* incx, + float* beta, + float* y, + int* incy +); + +namespace { + +template +void gemm( + char transa, + char transb, + int m, + int n, + int k, + F alpha, + F* a, + int lda, + F* b, + int ldb, + F beta, + F* c, + int ldc +) { + if constexpr (std::is_same_v) { + sgemm_(&transa, &transb, &m, &n, &k, &alpha, a, &lda, b, &ldb, &beta, c, &ldc); + } else { + dgemm_(&transa, &transb, &m, &n, &k, &alpha, a, &lda, b, &ldb, &beta, c, &ldc); + } +} + +template +void gemv( + char trans, + int m, + int n, + F alpha, + F* a, + int lda, + F* x, + int incx, + F beta, + F* y, + int incy +) { + if constexpr (std::is_same_v) { + sgemv_(&trans, &m, &n, &alpha, a, &lda, x, &incx, &beta, y, &incy); + } else { + dgemv_(&trans, &m, &n, &alpha, a, &lda, x, &incx, &beta, y, &incy); + } +} + +} // namespace + using namespace madspace; using namespace madspace::cpu; using namespace madspace::kernels; @@ -170,7 +248,7 @@ void op_matrix_element( ); } -template +template void op_matmul( const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device ) { @@ -182,7 +260,7 @@ void op_matmul( std::size_t batch_size = input.size(0); std::size_t dims_in = input.size(1); std::size_t dims_out = weight.size(1); - output = Tensor(DataType::dt_float64, {batch_size, dims_out}, device); + output = Tensor(dtype_of(), {batch_size, dims_out}, device); output.copy_from(bias_orig, device); if (batch_size == 0) { return; @@ -190,18 +268,25 @@ void op_matmul( device.sync_barrier(); device.submit([=]() mutable { - char transa = 'N', transb = 'T'; - int m = batch_size, n = dims_out, k = dims_in; - double alpha = 1., beta = 1.; - int lda = batch_size, ldb = dims_out, ldc = batch_size; - double* a = static_cast(input.data()); - double* b = static_cast(weight.data()); - double* c = static_cast(output.data()); - dgemm_(&transa, &transb, &m, &n, &k, &alpha, a, &lda, b, &ldb, &beta, c, &ldc); + gemm( + 'N', + 'T', + batch_size, + dims_out, + dims_in, + 1., + static_cast(input.data()), + batch_size, + static_cast(weight.data()), + dims_out, + 1., + static_cast(output.data()), + batch_size + ); }); } -template +template void backward_op_matmul( const CpuRuntime::Instruction& instruction, TensorVec& locals, @@ -220,18 +305,14 @@ void backward_op_matmul( if (!input_grad) { input_grad = - Tensor(DataType::dt_float64, input.shape(), device, AllocHint::local_grad); - // input_grad.zero(device); + Tensor(dtype_of(), input.shape(), device, AllocHint::local_grad); } if (!weight_grad) { weight_grad = - Tensor(DataType::dt_float64, weight.shape(), device, AllocHint::local_grad); - // weight_grad.zero(device); + Tensor(dtype_of(), weight.shape(), device, AllocHint::local_grad); } if (!bias_grad) { - bias_grad = - Tensor(DataType::dt_float64, {1, dims_out}, device, AllocHint::local_grad); - // bias_grad.zero(device); + bias_grad = Tensor(dtype_of(), {1, dims_out}, device, AllocHint::local_grad); } if (batch_size == 0) { return; @@ -240,42 +321,59 @@ void backward_op_matmul( // compute input_grad += output_grad * weight device.submit([=]() mutable { - char transa = 'N', transb = 'N'; - int m = batch_size, n = dims_in, k = dims_out; - double alpha = 1., beta = 1.; - int lda = batch_size, ldb = dims_out, ldc = batch_size; - double* a = static_cast(output_grad.data()); - ; - double* b = static_cast(weight.data()); - double* c = static_cast(input_grad.data()); - ; - dgemm_(&transa, &transb, &m, &n, &k, &alpha, a, &lda, b, &ldb, &beta, c, &ldc); + gemm( + 'N', + 'N', + batch_size, + dims_in, + dims_out, + 1., + static_cast(output_grad.data()), + batch_size, + static_cast(weight.data()), + dims_out, + 1., + static_cast(input_grad.data()), + batch_size + ); }); // compute weight_grad += output_grad.T * input device.submit([=]() mutable { - char transa = 'T', transb = 'N'; - int m = dims_out, n = dims_in, k = batch_size; - double alpha = 1., beta = 1.; - int lda = batch_size, ldb = batch_size, ldc = dims_out; - double* a = static_cast(output_grad.data()); - double* b = static_cast(input.data()); - double* c = static_cast(weight_grad.data()); - dgemm_(&transa, &transb, &m, &n, &k, &alpha, a, &lda, b, &ldb, &beta, c, &ldc); + gemm( + 'T', + 'N', + dims_out, + dims_in, + batch_size, + 1., + static_cast(output_grad.data()), + batch_size, + static_cast(input.data()), + batch_size, + 1., + static_cast(weight_grad.data()), + dims_out + ); }); // compute bias_grad += sum_i output_grad_ij device.submit([=]() mutable { // TODO: we should probably do this differently... - std::vector ones(batch_size, 1.); - char trans = 'T'; - int m = batch_size, n = dims_out; - double alpha = 1., beta = 1.; - int lda = batch_size, incx = 1, incy = 1; - double* a = static_cast(output_grad.data()); - double* x = ones.data(); - double* y = static_cast(bias_grad.data()); - dgemv_(&trans, &m, &n, &alpha, a, &lda, x, &incx, &beta, y, &incy); + std::vector ones(batch_size, 1.); + gemv( + 'T', + batch_size, + dims_out, + 1., + static_cast(output_grad.data()), + batch_size, + ones.data(), + 1, + 1., + static_cast(bias_grad.data()), + 1 + ); }); } @@ -1393,9 +1491,9 @@ std::pair CpuRuntime::run_backward_single( for (auto [output_index, output_dtype] : zip(instr.output_indices, instr.output_dtypes)) { auto& grad = local_grads[output_index]; - if (!grad && output_dtype == DataType::dt_float64) { + if (!grad && is_float(output_dtype)) { grad = Tensor( - DataType::dt_float64, + output_dtype, locals[output_index].shape(), device, AllocHint::local_grad @@ -1607,9 +1705,9 @@ std::pair CpuRuntime::run_backward_concurrent( for (auto [output_index, output_dtype] : zip(instr.output_indices, instr.output_dtypes)) { auto& grad = local_grads[output_index]; - if (!grad && output_dtype == DataType::dt_float64) { + if (!grad && is_float(output_dtype)) { grad = Tensor( - DataType::dt_float64, + output_dtype, locals[output_index].shape(), device, AllocHint::local_grad diff --git a/madspace/src/gpu/device.cu b/madspace/src/gpu/device.cu index 40e1a5070f..d20239177b 100644 --- a/madspace/src/gpu/device.cu +++ b/madspace/src/gpu/device.cu @@ -308,47 +308,71 @@ void AsyncGpuDevice::memcpy(void* to, void* from, std::size_t size) const { } void AsyncGpuDevice::tensor_copy(const Tensor& source, Tensor& target) const { - if (source.dtype() == DataType::dt_float64 && - target.dtype() == DataType::dt_float64) { + if (source.dtype() != target.dtype()) { + throw std::runtime_error("invalid dtype in copy"); + } + switch (source.dtype()) { + case DataType::dt_float64: tensor_foreach_dynamic, 1, 1>( {&source}, {&target}, target.size(0), *this ); - } else if (source.dtype() == DataType::dt_int32 && - target.dtype() == DataType::dt_int32) { + break; + case DataType::dt_float32: + tensor_foreach_dynamic, 1, 1>( + {&source}, {&target}, target.size(0), *this + ); + break; + case DataType::dt_int32: tensor_foreach_dynamic, 1, 1>( {&source}, {&target}, target.size(0), *this ); - } else { + break; + default: throw std::runtime_error("invalid dtype in copy"); } } void AsyncGpuDevice::tensor_zero(Tensor& tensor) const { - if (tensor.dtype() == DataType::dt_float64) { - if (tensor.is_contiguous()) { - check_error(gpuMemsetAsync(tensor.data(), 0, tensor.byte_size(), _stream)); - } else { - tensor_foreach_dynamic, 1, 1>( - {&tensor}, {&tensor}, tensor.size(0), *this - ); - } - } else if (tensor.dtype() == DataType::dt_int32) { - if (tensor.is_contiguous()) { - check_error(gpuMemsetAsync(tensor.data(), 0, tensor.byte_size(), _stream)); - } else { - tensor_foreach_dynamic, 1, 1>( - {&tensor}, {&tensor}, tensor.size(0), *this - ); - } - } else { + if (tensor.dtype() == DataType::batch_sizes) { throw std::runtime_error("invalid dtype in zero"); } + if (tensor.is_contiguous()) { + check_error(gpuMemsetAsync(tensor.data(), 0, tensor.byte_size(), _stream)); + return; + } + switch (tensor.dtype()) { + case DataType::dt_float64: + tensor_foreach_dynamic, 1, 1>( + {&tensor}, {&tensor}, tensor.size(0), *this + ); + break; + case DataType::dt_float32: + tensor_foreach_dynamic, 1, 1>( + {&tensor}, {&tensor}, tensor.size(0), *this + ); + break; + default: + tensor_foreach_dynamic, 1, 1>( + {&tensor}, {&tensor}, tensor.size(0), *this + ); + break; + } } void AsyncGpuDevice::tensor_add(const Tensor& source, Tensor& target) const { - tensor_foreach_dynamic, 1, 1>( - {&source}, {&target}, target.size(0), *this - ); + if (source.dtype() == DataType::dt_float32 && + target.dtype() == DataType::dt_float32) { + tensor_foreach_dynamic, 1, 1>( + {&source}, {&target}, target.size(0), *this + ); + } else if (source.dtype() == DataType::dt_float64 && + target.dtype() == DataType::dt_float64) { + tensor_foreach_dynamic, 1, 1>( + {&source}, {&target}, target.size(0), *this + ); + } else { + throw std::runtime_error("invalid dtype in add"); + } } void AsyncGpuDevice::tensor_cpu(const Tensor& source, Tensor& target) const { diff --git a/madspace/src/gpu/gpu_abstraction.cuh b/madspace/src/gpu/gpu_abstraction.cuh index 5c62d5f700..5d71f2c89e 100644 --- a/madspace/src/gpu/gpu_abstraction.cuh +++ b/madspace/src/gpu/gpu_abstraction.cuh @@ -45,6 +45,8 @@ #define gpublasSetStream cublasSetStream #define gpublasDgemm cublasDgemm #define gpublasDgemv cublasDgemv +#define gpublasSgemm cublasSgemm +#define gpublasSgemv cublasSgemv #define GPUBLAS_OP_N CUBLAS_OP_N #define GPUBLAS_OP_T CUBLAS_OP_T @@ -95,6 +97,8 @@ #define gpublasSetStream rocblas_set_stream #define gpublasDgemm rocblas_dgemm #define gpublasDgemv rocblas_dgemv +#define gpublasSgemm rocblas_sgemm +#define gpublasSgemv rocblas_sgemv #define GPUBLAS_OP_N rocblas_operation_none #define GPUBLAS_OP_T rocblas_operation_transpose diff --git a/madspace/src/gpu/kernel_definitions.cuh b/madspace/src/gpu/kernel_definitions.cuh index a4f8c1c03b..721c02bddd 100644 --- a/madspace/src/gpu/kernel_definitions.cuh +++ b/madspace/src/gpu/kernel_definitions.cuh @@ -8,28 +8,37 @@ namespace madspace { namespace kernels { -struct GpuTypes { +template +struct GpuTypesT { template - using FIn = const gpu::GpuTensorView; + using FIn = const gpu::GpuTensorView; template using IIn = const gpu::GpuTensorView; template - using FOut = gpu::GpuTensorView; + using FOut = gpu::GpuTensorView; template using IOut = gpu::GpuTensorView; // single inputs are broadcast through a zero batch stride on the GPU template - using FSIn = const gpu::GpuTensorView; + using FSIn = const gpu::GpuTensorView; template using ISIn = const gpu::GpuTensorView; - using FVal = double; + using FVal = F; using IVal = me_int_t; using BVal = bool; + // the same policy with a different float type, for conversion kernels + using F32 = GpuTypesT; + using F64 = GpuTypesT; }; +using GpuTypes = GpuTypesT; +using GpuTypes32 = GpuTypesT; inline __device__ double where(bool condition, double val_true, double val_false) { return condition ? val_true : val_false; } +inline __device__ float where(bool condition, float val_true, float val_false) { + return condition ? val_true : val_false; +} inline __device__ me_int_t where(bool condition, me_int_t val_true, me_int_t val_false) { return condition ? val_true : val_false; @@ -40,6 +49,12 @@ inline __device__ double min(double arg1, double arg2) { inline __device__ double max(double arg1, double arg2) { return arg1 > arg2 ? arg1 : arg2; } +inline __device__ float min(float arg1, float arg2) { + return arg1 < arg2 ? arg1 : arg2; +} +inline __device__ float max(float arg1, float arg2) { + return arg1 > arg2 ? arg1 : arg2; +} inline __device__ std::size_t single_index(me_int_t arg) { return arg; } using ::acos; diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 7e94980367..7cbaa49faf 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -125,6 +125,25 @@ void op_matrix_element( } } +template +auto gpublas_gemm(Args... args) { + if constexpr (std::is_same_v) { + return gpublasSgemm(args...); + } else { + return gpublasDgemm(args...); + } +} + +template +auto gpublas_gemv(Args... args) { + if constexpr (std::is_same_v) { + return gpublasSgemv(args...); + } else { + return gpublasDgemv(args...); + } +} + +template void op_matmul( const GpuRuntime::Instruction& instruction, TensorVec& locals, @@ -141,10 +160,7 @@ void op_matmul( std::size_t dims_in = input.size(1); std::size_t dims_out = weight.size(1); output = Tensor( - DataType::dt_float64, - {batch_size, dims_out}, - device, - instruction.output_alloc_hints[0] + dtype_of(), {batch_size, dims_out}, device, instruction.output_alloc_hints[0] ); output.copy_from(bias, device); if (batch_size == 0) { @@ -153,36 +169,39 @@ void op_matmul( gpublasHandle_t handle = instruction.runtime.gpublas_handle(); check_error(gpublasSetStream(handle, device.stream())); - double alpha = 1., beta = 1.; - check_error(gpublasDgemm( - handle, - GPUBLAS_OP_N, - GPUBLAS_OP_T, - batch_size, - dims_out, - dims_in, - &alpha, - static_cast(input.data()), - batch_size, - static_cast(weight.data()), - dims_out, - &beta, - static_cast(output.data()), - batch_size - )); + F alpha = 1., beta = 1.; + check_error( + gpublas_gemm( + handle, + GPUBLAS_OP_N, + GPUBLAS_OP_T, + batch_size, + dims_out, + dims_in, + &alpha, + static_cast(input.data()), + batch_size, + static_cast(weight.data()), + dims_out, + &beta, + static_cast(output.data()), + batch_size + ) + ); input.reset(device); weight.reset(device); bias.reset(device); } -__global__ void -kernel_one(std::size_t batch_size, GpuTensorView output) { +template +__global__ void kernel_one(std::size_t batch_size, GpuTensorView output) { me_int_t i = blockDim.x * blockIdx.x + threadIdx.x; if (i < batch_size) { output[i] = 1.; } } +template void backward_op_matmul( const GpuRuntime::Instruction& instruction, TensorVec& locals, @@ -205,30 +224,18 @@ void backward_op_matmul( if (!input_grad) { input_grad = Tensor( - DataType::dt_float64, - input.shape(), - device, - instruction.input_grad_alloc_hints[0] + dtype_of(), input.shape(), device, instruction.input_grad_alloc_hints[0] ); - // input_grad.zero(device); } if (!weight_grad) { weight_grad = Tensor( - DataType::dt_float64, - weight.shape(), - device, - instruction.input_grad_alloc_hints[1] + dtype_of(), weight.shape(), device, instruction.input_grad_alloc_hints[1] ); - // weight_grad.zero(device); } if (!bias_grad) { bias_grad = Tensor( - DataType::dt_float64, - {1, dims_out}, - device, - instruction.input_grad_alloc_hints[2] + dtype_of(), {1, dims_out}, device, instruction.input_grad_alloc_hints[2] ); - // bias_grad.zero(device); } if (batch_size == 0) { input.reset(device); @@ -237,66 +244,72 @@ void backward_op_matmul( return; } - double alpha = 1., beta = 1.; + F alpha = 1., beta = 1.; gpublasHandle_t handle = instruction.runtime.gpublas_handle(); gpuStream_t stream = device.stream(); check_error(gpublasSetStream(handle, stream)); // compute input_grad += output_grad * weight - check_error(gpublasDgemm( - handle, - GPUBLAS_OP_N, - GPUBLAS_OP_N, - batch_size, - dims_in, - dims_out, - &alpha, - static_cast(output_grad.data()), - batch_size, - static_cast(weight.data()), - dims_out, - &beta, - static_cast(input_grad.data()), - batch_size - )); + check_error( + gpublas_gemm( + handle, + GPUBLAS_OP_N, + GPUBLAS_OP_N, + batch_size, + dims_in, + dims_out, + &alpha, + static_cast(output_grad.data()), + batch_size, + static_cast(weight.data()), + dims_out, + &beta, + static_cast(input_grad.data()), + batch_size + ) + ); // compute weight_grad += output_grad.T * input - check_error(gpublasDgemm( - handle, - GPUBLAS_OP_T, - GPUBLAS_OP_N, - dims_out, - dims_in, - batch_size, - &alpha, - static_cast(output_grad.data()), - batch_size, - static_cast(input.data()), - batch_size, - &beta, - static_cast(weight_grad.data()), - dims_out - )); + check_error( + gpublas_gemm( + handle, + GPUBLAS_OP_T, + GPUBLAS_OP_N, + dims_out, + dims_in, + batch_size, + &alpha, + static_cast(output_grad.data()), + batch_size, + static_cast(input.data()), + batch_size, + &beta, + static_cast(weight_grad.data()), + dims_out + ) + ); // compute bias_grad += sum_i output_grad_ij - Tensor ones(DataType::dt_float64, {batch_size}, device, AllocHint::temporary); + Tensor ones(dtype_of(), {batch_size}, device, AllocHint::temporary); launch_kernel( - kernel_one, batch_size, device.stream(), batch_size, ones.view() + kernel_one, batch_size, device.stream(), batch_size, ones.view() + ); + check_error( + gpublas_gemv( + handle, + GPUBLAS_OP_T, + batch_size, + dims_out, + &alpha, + static_cast(output_grad.data()), + batch_size, + static_cast(ones.data()), + 1, + &beta, + static_cast(bias_grad.data()), + 1 + ) ); - check_error(gpublasDgemv( - handle, - GPUBLAS_OP_T, - batch_size, - dims_out, - &alpha, - static_cast(output_grad.data()), - batch_size, - static_cast(ones.data()), - 1, - &beta, - static_cast(bias_grad.data()), - 1 - )); ones.reset(device); input.reset(device); weight.reset(device); @@ -2186,9 +2199,9 @@ std::pair GpuRuntime::run_backward( for (auto [output_index, output_dtype] : zip(instr.output_indices, instr.output_dtypes)) { auto& grad = local_grads[output_index]; - if (!grad && output_dtype == DataType::dt_float64) { + if (!grad && is_float(output_dtype)) { grad = Tensor( - DataType::dt_float64, + output_dtype, locals[output_index].shape(), device, AllocHint::local_grad diff --git a/madspace/src/kernels/nn.hpp b/madspace/src/kernels/nn.hpp index 176329a5e8..4cef7f7067 100644 --- a/madspace/src/kernels/nn.hpp +++ b/madspace/src/kernels/nn.hpp @@ -8,102 +8,153 @@ namespace kernels { // Kernels +// The activations are instantiated in float32 and float64. Literals are cast to +// FVal so that float32 code does not get promoted to double. + template KERNELSPEC void kernel_relu(FIn input, FOut output) { - FVal x(input); - output = where(x < 0., 0., x); + using F = FVal; + F x(input); + output = where(x < F(0.), F(0.), x); } template KERNELSPEC void backward_kernel_relu(FIn input, FIn output_grad, FOut input_grad) { - FVal x(input), g(output_grad); - input_grad = where(x < 0., 0., g); + using F = FVal; + F x(input), g(output_grad); + input_grad += where(x < F(0.), F(0.), g); } template KERNELSPEC void kernel_leaky_relu(FIn input, FOut output) { - FVal x(input); - output = where(x < 0., x * 0.01, x); + using F = FVal; + F x(input); + output = where(x < F(0.), x * F(0.01), x); } template KERNELSPEC void backward_kernel_leaky_relu( FIn input, FIn output_grad, FOut input_grad ) { - FVal x(input), g(output_grad); - input_grad = where(x < 0., g * 0.01, g); + using F = FVal; + F x(input), g(output_grad); + input_grad += where(x < F(0.), g * F(0.01), g); } template KERNELSPEC void kernel_elu(FIn input, FOut output) { - FVal x(input); - output = where(x < 0., expm1(x), x); + using F = FVal; + F x(input); + output = where(x < F(0.), expm1(x), x); } template KERNELSPEC void backward_kernel_elu(FIn input, FIn output_grad, FOut input_grad) { - FVal x(input), g(output_grad); - input_grad = where(x < 0., g * exp(x), g); + using F = FVal; + F x(input), g(output_grad); + input_grad += where(x < F(0.), g * exp(x), g); } template KERNELSPEC void kernel_gelu(FIn input, FOut output) { - output = 0.5 * input * (1. + erf(SQRT_HALF * input)); + using F = FVal; + F x(input); + output = F(0.5) * x * (F(1.) + erf(F(SQRT_HALF) * x)); } template KERNELSPEC void backward_kernel_gelu(FIn input, FIn output_grad, FOut input_grad) { - FVal x(input), g(output_grad); - auto cdf = 0.5 * (1. + erf(SQRT_HALF * x)); - auto pdf = 0.5 * SQRT_HALF * TWO_DIV_SQRT_PI * exp(-0.5 * x * x); - input_grad = g * (cdf + x * pdf); + using F = FVal; + F x(input), g(output_grad); + F cdf = F(0.5) * (F(1.) + erf(F(SQRT_HALF) * x)); + F pdf = F(0.5 * SQRT_HALF * TWO_DIV_SQRT_PI) * exp(F(-0.5) * x * x); + input_grad += g * (cdf + x * pdf); } template KERNELSPEC void kernel_sigmoid(FIn input, FOut output) { - FVal x = input; - output = 1. / (1. + exp(-x)); + using F = FVal; + F x(input); + output = F(1.) / (F(1.) + exp(-x)); } template KERNELSPEC void backward_kernel_sigmoid( FIn output, FIn output_grad, FOut input_grad ) { - FVal y(output), g(output_grad); - input_grad = g * y * (1. - y); + using F = FVal; + F y(output), g(output_grad); + input_grad += g * y * (F(1.) - y); } template KERNELSPEC void kernel_softplus(FIn input, FOut output) { - FVal x(input); - output = where(x < 20., log1p(exp(x)), x); + using F = FVal; + F x(input); + output = where(x < F(20.), log1p(exp(x)), x); } template KERNELSPEC void backward_kernel_softplus( FIn input, FIn output_grad, FOut input_grad ) { - FVal x(input), g(output_grad); - auto z = exp(x); - input_grad = where(x < 20., g * z / (1. + z), g); + using F = FVal; + F x(input), g(output_grad); + F z = exp(x); + input_grad += where(x < F(20.), g * z / (F(1.) + z), g); } template KERNELSPEC void kernel_scaled_tanh(FIn input, FOut output) { - FVal x(input); - output = 4. * tanh(x / 4.); + using F = FVal; + F x(input); + output = F(4.) * tanh(x / F(4.)); } template KERNELSPEC void backward_kernel_scaled_tanh( FIn output, FIn output_grad, FOut input_grad ) { - FVal g(output_grad); - FVal y_scaled = output / 4.; - input_grad = g * (1. - y_scaled) * (1. + y_scaled); + using F = FVal; + F g(output_grad), y(output); + F y_scaled = y / F(4.); + input_grad += g * (F(1.) - y_scaled) * (F(1.) + y_scaled); +} + +// Precision conversions. T is the policy of either precision; T::F32 and T::F64 +// select the two sides. + +template +KERNELSPEC void +kernel_convert_float32(FIn input, FOut output) { + output = FVal(FVal(input)); +} + +template +KERNELSPEC void backward_kernel_convert_float32( + FIn input, + FIn output_grad, + FOut input_grad +) { + input_grad += FVal(FVal(output_grad)); +} + +template +KERNELSPEC void +kernel_convert_float64(FIn input, FOut output) { + output = FVal(FVal(input)); +} + +template +KERNELSPEC void backward_kernel_convert_float64( + FIn input, + FIn output_grad, + FOut input_grad +) { + input_grad += FVal(FVal(output_grad)); } template diff --git a/madspace/src/python/function_runtime.cpp b/madspace/src/python/function_runtime.cpp index 2abc56fdeb..0d54307ef3 100644 --- a/madspace/src/python/function_runtime.cpp +++ b/madspace/src/python/function_runtime.cpp @@ -537,7 +537,7 @@ Tensor madspace_py::dlpack_to_tensor( std::format( "Argument {}: wrong input dimension. Expected {}, got {}", arg_index + 1, - expected_type->shape.size() + 1, + expected_type->shape.size() + is_batch, dl_tensor->ndim ) ); @@ -563,7 +563,7 @@ Tensor madspace_py::dlpack_to_tensor( } if (expected_type) { for (int j = is_batch; j < dl_tensor->ndim; ++j) { - if (shape.at(j) != expected_type->shape.at(j - is_batch)) { + if (dl_tensor->shape[j] != expected_type->shape.at(j - is_batch)) { throw std::invalid_argument( std::format( "Argument {}, dimension {}: shape mismatch. Expected {}, " @@ -571,7 +571,7 @@ Tensor madspace_py::dlpack_to_tensor( arg_index + 1, j, expected_type->shape.at(j - is_batch), - shape.at(j) + dl_tensor->shape[j] ) ); } diff --git a/madspace/tests/test_float32.py b/madspace/tests/test_float32.py new file mode 100644 index 0000000000..331c2aef70 --- /dev/null +++ b/madspace/tests/test_float32.py @@ -0,0 +1,124 @@ +"""Single-precision NN instructions: forward and backward against torch.""" + +import pytest +import torch + +import madspace as ms +from madspace.torch import FunctionModule + +torch.manual_seed(1234) + +N = 103 +DIMS_IN = 7 +DIMS_OUT = 5 + +ACTIVATIONS = { + "relu": torch.relu, + "leaky_relu": lambda x: torch.nn.functional.leaky_relu(x, 0.01), + "elu": torch.nn.functional.elu, + "gelu": torch.nn.functional.gelu, + "sigmoid": torch.sigmoid, + "softplus": torch.nn.functional.softplus, + "scaled_tanh": lambda x: 4 * torch.tanh(x / 4), +} + +DTYPES = { + "float32": (ms.DataType.float32, torch.float32, 1e-5), + "float64": (ms.DataType.float64, torch.float64, 1e-12), +} + + +def batch_type(dtype, shape): + return ms.Type(dtype, ms.batch_size, shape) + + +@pytest.mark.parametrize("precision", DTYPES.keys()) +@pytest.mark.parametrize("name", ACTIVATIONS.keys()) +def test_activation(name, precision): + ms_dtype, torch_dtype, tol = DTYPES[precision] + io_type = batch_type(ms_dtype, [DIMS_IN]) + fb = ms.FunctionBuilder( + ms.NamedTypes([("in", io_type)]), ms.NamedTypes([("out", io_type)]) + ) + fb.output(0, getattr(fb, name)(fb.input(0))) + module = FunctionModule(fb.function()) + + x = (3 * torch.randn((N, DIMS_IN), dtype=torch_dtype)).requires_grad_() + out = module(x) + assert out.dtype == torch_dtype + out_grad = torch.randn_like(out) + (out * out_grad).sum().backward() + + x_ref = x.detach().clone().requires_grad_() + out_ref = ACTIVATIONS[name](x_ref) + (out_ref * out_grad).sum().backward() + torch.testing.assert_close(out, out_ref, rtol=tol, atol=tol) + torch.testing.assert_close(x.grad, x_ref.grad, rtol=tol, atol=tol) + + +@pytest.mark.parametrize("precision", DTYPES.keys()) +def test_matmul(precision): + ms_dtype, torch_dtype, tol = DTYPES[precision] + fb = ms.FunctionBuilder( + ms.NamedTypes( + [ + ("x", batch_type(ms_dtype, [DIMS_IN])), + ("w", ms.Type(ms_dtype, ms.BatchSize.one, [DIMS_OUT, DIMS_IN])), + ("b", ms.Type(ms_dtype, ms.BatchSize.one, [DIMS_OUT])), + ] + ), + ms.NamedTypes([("y", batch_type(ms_dtype, [DIMS_OUT]))]), + ) + fb.output(0, fb.matmul(fb.input(0), fb.input(1), fb.input(2))) + module = FunctionModule(fb.function()) + + x = torch.randn((N, DIMS_IN), dtype=torch_dtype, requires_grad=True) + # single inputs are passed without batch dimension + w = torch.randn((DIMS_OUT, DIMS_IN), dtype=torch_dtype, requires_grad=True) + b = torch.randn((DIMS_OUT,), dtype=torch_dtype, requires_grad=True) + y = module(x, w, b) + assert y.dtype == torch_dtype + y_grad = torch.randn_like(y) + (y * y_grad).sum().backward() + + refs = [t.detach().clone().requires_grad_() for t in (x, w, b)] + y_ref = refs[0] @ refs[1].T + refs[2] + (y_ref * y_grad).sum().backward() + torch.testing.assert_close(y, y_ref, rtol=tol, atol=tol) + for t, t_ref in zip((x, w, b), refs): + torch.testing.assert_close(t.grad, t_ref.grad, rtol=tol, atol=tol) + + +def test_convert_roundtrip(): + in_type = batch_type(ms.DataType.float64, [3]) + fb = ms.FunctionBuilder( + ms.NamedTypes([("in", in_type)]), + ms.NamedTypes( + [("single", batch_type(ms.DataType.float32, [3])), ("out", in_type)] + ), + ) + single = fb.convert_float32(fb.input(0)) + fb.output(0, single) + fb.output(1, fb.convert_float64(fb.relu(single))) + module = FunctionModule(fb.function()) + + x = torch.randn((N, 3), dtype=torch.float64, requires_grad=True) + single_out, out = module(x) + assert single_out.dtype == torch.float32 + assert out.dtype == torch.float64 + torch.testing.assert_close(single_out, x.detach().float()) + torch.testing.assert_close(out, torch.relu(x.detach().float()).double()) + out.sum().backward() + torch.testing.assert_close(x.grad, (x.detach() > 0).double()) + + +def test_literal_converted_for_float32(): + io_type = batch_type(ms.DataType.float32, [1]) + fb = ms.FunctionBuilder( + ms.NamedTypes([("in", io_type)]), ms.NamedTypes([("out", io_type)]) + ) + # float64 literals are converted when the instruction dispatches to float32 + fb.output(0, fb.matmul(fb.input(0), ms.Value([2.0], [1, 1]), ms.Value([1.0]))) + module = FunctionModule(fb.function()) + x = torch.randn((N, 1), dtype=torch.float32) + torch.testing.assert_close(module(x), 2 * x + 1) From d6f1fe573ffdb8fb2d86ff8eca2631ac2348e25f Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 18:57:41 +0200 Subject: [PATCH 03/24] madspace: float32 SIMD vectors with twice the lanes Every SIMD backend gets a float32 family FVec32/IVec32/BVec32 with twice the lanes of the float64 vectors and int32 indices and masks of the same width (NEON 4, AVX2/AVX512Y 8, AVX512 16 lanes), used by SimdTypes32. - vector types carry their lane count, which the views and the foreach loop use instead of the global simd_vec_size - loads and gathers dispatch on a VecTag, as int loads exist in both families - jobs are aligned to and literals padded to the largest lane count --- madspace/include/madspace/driver/tensor.hpp | 2 +- madspace/src/cpu/device.hpp | 9 +- madspace/src/cpu/kernel_definitions.hpp | 20 +- madspace/src/cpu/runtime.cpp | 2 +- madspace/src/cpu/simd.hpp | 9 + madspace/src/cpu/simd_arm.hpp | 217 +++++++++++++++- madspace/src/cpu/simd_x86_256.hpp | 268 +++++++++++++++++++- madspace/src/cpu/simd_x86_256y.hpp | 256 ++++++++++++++++++- madspace/src/cpu/simd_x86_512.hpp | 259 ++++++++++++++++++- madspace/src/cpu/tensor.hpp | 58 +++-- 10 files changed, 1040 insertions(+), 60 deletions(-) diff --git a/madspace/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp index 7bdc99a066..39961f36ff 100644 --- a/madspace/include/madspace/driver/tensor.hpp +++ b/madspace/include/madspace/driver/tensor.hpp @@ -803,7 +803,7 @@ class Tensor { /// Maximum number of stored elements (value elements times lanes) of a /// literal for which @ref simd_padded_literal pads the storage. - static constexpr std::size_t max_simd_padded_size = 256; + static constexpr std::size_t max_simd_padded_size = 512; /** * Literal constant for the CPU SIMD kernels. diff --git a/madspace/src/cpu/device.hpp b/madspace/src/cpu/device.hpp index c6de22e360..65ea999413 100644 --- a/madspace/src/cpu/device.hpp +++ b/madspace/src/cpu/device.hpp @@ -105,15 +105,18 @@ class AsyncCpuDevice : public CpuDevice { return {1, batch_size}; } - std::size_t batch_size_vec = (batch_size + simd_vec_size - 1) / simd_vec_size; - std::size_t min_batch_size = 64 / simd_vec_size; + // jobs are aligned to the largest vector size, so that only the end of the + // batch is left to the scalar kernels for every precision + std::size_t batch_size_vec = + (batch_size + simd_vec_size_max - 1) / simd_vec_size_max; + std::size_t min_batch_size = std::max(64 / simd_vec_size_max, 1); std::size_t thread_count = _thread_pool.thread_count(); // std::size_t job_count = (batch_size + min_batch_size - 1) / min_batch_size; std::size_t job_count = batch_size_vec < thread_count * min_batch_size ? (batch_size_vec + min_batch_size - 1) / min_batch_size : thread_count; std::size_t job_size = - (batch_size_vec + job_count - 1) / job_count * simd_vec_size; + (batch_size_vec + job_count - 1) / job_count * simd_vec_size_max; // correct rounding errors to vector size job_count = (batch_size + job_size - 1) / job_size; return {job_count, job_size}; diff --git a/madspace/src/cpu/kernel_definitions.hpp b/madspace/src/cpu/kernel_definitions.hpp index 6e5f3bbc4c..bbdc47a5f2 100644 --- a/madspace/src/cpu/kernel_definitions.hpp +++ b/madspace/src/cpu/kernel_definitions.hpp @@ -90,10 +90,28 @@ struct SimdTypes { using IVal = IVec; using BVal = BVec; }; + +struct SimdTypes32 { + template + using FIn = const cpu::VectorizedTensorView; + template + using IIn = const cpu::VectorizedTensorView; + template + using FOut = cpu::VectorizedTensorView; + template + using IOut = cpu::VectorizedTensorView; + template + using FSIn = const cpu::SingleVectorizedTensorView; + template + using ISIn = const cpu::SingleVectorizedTensorView; + using FVal = FVec32; + using IVal = IVec32; + using BVal = BVec32; +}; #else // USE_SIMD using SimdTypes = CpuTypes; -#endif // USE_SIMD using SimdTypes32 = CpuTypes32; +#endif // USE_SIMD } // namespace kernels } // namespace madspace diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index d83406bb22..407051b821 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -1307,7 +1307,7 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur // padded so that the vectorized kernels can load it like a // batched tensor, see Tensor::simd_padded_literal _locals_init[local.local_index] = Tensor::simd_padded_literal( - LiteralValue(val), simd_vec_size, cpu_device() + LiteralValue(val), simd_vec_size_max, cpu_device() ); }, [](std::monostate val) {} diff --git a/madspace/src/cpu/simd.hpp b/madspace/src/cpu/simd.hpp index 7f10c491db..6fae8a4392 100644 --- a/madspace/src/cpu/simd.hpp +++ b/madspace/src/cpu/simd.hpp @@ -1,5 +1,10 @@ #pragma once +// Selects the vector type for loads and gathers, which can not be overloaded on +// the return type alone (e.g. int loads into the float64 or float32 family) +template +struct VecTag {}; + #ifdef USE_SIMD #ifdef USE_SIMD_NEON @@ -21,5 +26,9 @@ #else // USE_SIMD constexpr int simd_vec_size = 1; +constexpr int simd_vec_size_f32 = 1; #endif // USE_SIMD + +// the largest lane count, used to align jobs and pad literals +constexpr int simd_vec_size_max = simd_vec_size_f32; diff --git a/madspace/src/cpu/simd_arm.hpp b/madspace/src/cpu/simd_arm.hpp index 279ba6c138..d47e6f88ac 100644 --- a/madspace/src/cpu/simd_arm.hpp +++ b/madspace/src/cpu/simd_arm.hpp @@ -3,8 +3,11 @@ #include constexpr int simd_vec_size = 2; +// float32 vectors have twice the lanes of the float64 ones +constexpr int simd_vec_size_f32 = 4; struct FVec { + static constexpr int lanes = simd_vec_size; FVec() = default; FVec(float64x2_t _v) : v(_v) {}; FVec(double _v) : v(vdupq_n_f64(_v)) {}; @@ -18,6 +21,7 @@ struct FVec { }; struct IVec { + static constexpr int lanes = simd_vec_size; IVec() = default; IVec(int64x2_t _v) : v(_v) {}; IVec(int _v) : v(vdupq_n_s64(_v)) {}; @@ -31,6 +35,7 @@ struct IVec { }; struct BVec { + static constexpr int lanes = simd_vec_size; BVec() = default; BVec(uint64x2_t _v) : v(_v) {}; BVec(bool _v) : v(vdupq_n_u64(_v ? ~0ULL : 0)) {}; @@ -39,7 +44,11 @@ struct BVec { }; inline FVec vgather( - double* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices + VecTag, + double* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec indices ) { FVec ret; IVec strided_indices = indices * IVec(index_stride); @@ -51,7 +60,11 @@ inline FVec vgather( } inline IVec vgather( - int* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices + VecTag, + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec indices ) { IVec ret; IVec strided_indices = indices * IVec(index_stride); @@ -63,14 +76,18 @@ inline IVec vgather( } // contiguous loads and stores of the lanes of one batch element -inline FVec vload(double* base_ptr) { return vld1q_f64(base_ptr); } +inline FVec vload(VecTag, double* base_ptr) { return vld1q_f64(base_ptr); } -inline IVec vload(int* base_ptr) { return vmovl_s32(vld1_s32(base_ptr)); } +inline IVec vload(VecTag, int* base_ptr) { return vmovl_s32(vld1_s32(base_ptr)); } // the same value in every lane -inline FVec vload_broadcast(double* base_ptr) { return vld1q_dup_f64(base_ptr); } +inline FVec vload_broadcast(VecTag, double* base_ptr) { + return vld1q_dup_f64(base_ptr); +} -inline IVec vload_broadcast(int* base_ptr) { return vdupq_n_s64(*base_ptr); } +inline IVec vload_broadcast(VecTag, int* base_ptr) { + return vdupq_n_s64(*base_ptr); +} inline void vscatter( double* base_ptr, @@ -100,7 +117,9 @@ inline void vscatter( inline void vstore(double* base_ptr, FVec values) { vst1q_f64(base_ptr, values); } -inline void vstore(int* base_ptr, IVec values) { vst1_s32(base_ptr, vmovn_s64(values)); } +inline void vstore(int* base_ptr, IVec values) { + vst1_s32(base_ptr, vmovn_s64(values)); +} inline FVec where(BVec arg1, FVec arg2, FVec arg3) { return vbslq_f64(arg1, arg2, arg3); @@ -164,3 +183,187 @@ inline FVec log1p(FVec arg1) { return Sleef_log1pd2_u10(arg1); } inline FVec expm1(FVec arg1) { return Sleef_expm1d2_u10(arg1); } inline FVec erf(FVec arg1) { return Sleef_erfd2_u10(arg1); } inline FVec fma(FVec arg1, FVec arg2, FVec arg3) { return vfmaq_f64(arg3, arg1, arg2); } + +// float32 family: four lanes, with int32 indices and masks of the same width + +struct IVec32; + +struct FVec32 { + static constexpr int lanes = simd_vec_size_f32; + FVec32() = default; + FVec32(float32x4_t _v) : v(_v) {}; + FVec32(double _v) : v(vdupq_n_f32(static_cast(_v))) {}; + explicit FVec32(int32x4_t _v) : v(vcvtq_f32_s32(_v)) {}; + operator float32x4_t() { return v; } + FVec32 operator+=(FVec32 _v) { + v = vaddq_f32(v, _v); + return v; + } + float32x4_t v; +}; + +struct IVec32 { + static constexpr int lanes = simd_vec_size_f32; + IVec32() = default; + IVec32(int32x4_t _v) : v(_v) {}; + IVec32(int _v) : v(vdupq_n_s32(_v)) {}; + explicit IVec32(float32x4_t _v) : v(vcvtq_s32_f32(_v)) {}; + operator int32x4_t() { return v; } + IVec32 operator+=(IVec32 _v) { + v = vaddq_s32(v, _v); + return v; + } + int32x4_t v; +}; + +struct BVec32 { + static constexpr int lanes = simd_vec_size_f32; + BVec32() = default; + BVec32(uint32x4_t _v) : v(_v) {}; + BVec32(bool _v) : v(vdupq_n_u32(_v ? ~0U : 0)) {}; + operator uint32x4_t() { return v; } + uint32x4_t v; +}; + +inline IVec32 operator*(IVec32 arg1, IVec32 arg2) { return vmulq_s32(arg1, arg2); } + +inline FVec32 vgather( + VecTag, + float* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices +) { + int32x4_t idx = indices * IVec32(index_stride); + float values[4]; + for (int i = 0; i < 4; ++i) { + values[i] = base_ptr[idx[i] + batch_stride * i]; + } + return vld1q_f32(values); +} + +inline IVec32 vgather( + VecTag, + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices +) { + int32x4_t idx = indices * IVec32(index_stride); + int values[4]; + for (int i = 0; i < 4; ++i) { + values[i] = base_ptr[idx[i] + batch_stride * i]; + } + return vld1q_s32(values); +} + +inline FVec32 vload(VecTag, float* base_ptr) { return vld1q_f32(base_ptr); } + +inline IVec32 vload(VecTag, int* base_ptr) { return vld1q_s32(base_ptr); } + +inline FVec32 vload_broadcast(VecTag, float* base_ptr) { + return vld1q_dup_f32(base_ptr); +} + +inline IVec32 vload_broadcast(VecTag, int* base_ptr) { + return vld1q_dup_s32(base_ptr); +} + +inline void vscatter( + float* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices, + FVec32 values +) { + int32x4_t idx = indices * IVec32(index_stride); + float32x4_t vals = values; + for (int i = 0; i < 4; ++i) { + base_ptr[idx[i] + batch_stride * i] = vals[i]; + } +} + +inline void vscatter( + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices, + IVec32 values +) { + int32x4_t idx = indices * IVec32(index_stride); + int32x4_t vals = values; + for (int i = 0; i < 4; ++i) { + base_ptr[idx[i] + batch_stride * i] = vals[i]; + } +} + +inline void vstore(float* base_ptr, FVec32 values) { vst1q_f32(base_ptr, values); } + +inline void vstore(int* base_ptr, IVec32 values) { vst1q_s32(base_ptr, values); } + +inline FVec32 where(BVec32 arg1, FVec32 arg2, FVec32 arg3) { + return vbslq_f32(arg1, arg2, arg3); +} +inline IVec32 where(BVec32 arg1, IVec32 arg2, IVec32 arg3) { + return vbslq_s32(arg1, arg2, arg3); +} +inline std::size_t single_index(IVec32 arg) { return vgetq_lane_s32(arg, 0); } + +inline BVec32 operator==(FVec32 arg1, FVec32 arg2) { return vceqq_f32(arg1, arg2); } +inline BVec32 operator!=(FVec32 arg1, FVec32 arg2) { + return vmvnq_u32(vceqq_f32(arg1, arg2)); +} +inline BVec32 operator>(FVec32 arg1, FVec32 arg2) { return vcgtq_f32(arg1, arg2); } +inline BVec32 operator<(FVec32 arg1, FVec32 arg2) { return vcltq_f32(arg1, arg2); } +inline BVec32 operator>=(FVec32 arg1, FVec32 arg2) { return vcgeq_f32(arg1, arg2); } +inline BVec32 operator<=(FVec32 arg1, FVec32 arg2) { return vcleq_f32(arg1, arg2); } + +inline BVec32 operator==(IVec32 arg1, IVec32 arg2) { return vceqq_s32(arg1, arg2); } +inline BVec32 operator!=(IVec32 arg1, IVec32 arg2) { + return vmvnq_u32(vceqq_s32(arg1, arg2)); +} +inline BVec32 operator>(IVec32 arg1, IVec32 arg2) { return vcgtq_s32(arg1, arg2); } +inline BVec32 operator<(IVec32 arg1, IVec32 arg2) { return vcltq_s32(arg1, arg2); } +inline BVec32 operator>=(IVec32 arg1, IVec32 arg2) { return vcgeq_s32(arg1, arg2); } +inline BVec32 operator<=(IVec32 arg1, IVec32 arg2) { return vcleq_s32(arg1, arg2); } + +inline BVec32 operator&(BVec32 arg1, BVec32 arg2) { return vandq_u32(arg1, arg2); } +inline BVec32 operator|(BVec32 arg1, BVec32 arg2) { return vorrq_u32(arg1, arg2); } +inline BVec32 operator!(BVec32 arg1) { return vmvnq_u32(arg1); } + +inline FVec32 min(FVec32 arg1, FVec32 arg2) { return where(arg1 < arg2, arg1, arg2); } +inline FVec32 max(FVec32 arg1, FVec32 arg2) { return where(arg1 > arg2, arg1, arg2); } + +inline FVec32 operator-(FVec32 arg1) { return vnegq_f32(arg1); } +inline FVec32 operator+(FVec32 arg1, FVec32 arg2) { return vaddq_f32(arg1, arg2); } +inline FVec32 operator-(FVec32 arg1, FVec32 arg2) { return vsubq_f32(arg1, arg2); } +inline FVec32 operator*(FVec32 arg1, FVec32 arg2) { return vmulq_f32(arg1, arg2); } +inline FVec32 operator/(FVec32 arg1, FVec32 arg2) { return vdivq_f32(arg1, arg2); } +inline IVec32 operator-(IVec32 arg1) { return vnegq_s32(arg1); } +inline IVec32 operator+(IVec32 arg1, IVec32 arg2) { return vaddq_s32(arg1, arg2); } +inline IVec32 operator-(IVec32 arg1, IVec32 arg2) { return vsubq_s32(arg1, arg2); } + +inline BVec32 isnan(FVec32 arg) { return arg != arg; } + +inline FVec32 sqrt(FVec32 arg1) { return Sleef_sqrtf4_u05(arg1); } +inline FVec32 sin(FVec32 arg1) { return Sleef_sinf4_u10(arg1); } +inline FVec32 cos(FVec32 arg1) { return Sleef_cosf4_u10(arg1); } +inline FVec32 sinh(FVec32 arg1) { return Sleef_sinhf4_u10(arg1); } +inline FVec32 asinh(FVec32 arg1) { return Sleef_asinhf4_u10(arg1); } +inline FVec32 cosh(FVec32 arg1) { return Sleef_coshf4_u10(arg1); } +inline FVec32 tanh(FVec32 arg1) { return Sleef_tanhf4_u10(arg1); } +inline FVec32 atan2(FVec32 arg1, FVec32 arg2) { return Sleef_atan2f4_u10(arg1, arg2); } +inline FVec32 pow(FVec32 arg1, FVec32 arg2) { return Sleef_powf4_u10(arg1, arg2); } +inline FVec32 fabs(FVec32 arg1) { return Sleef_fabsf4(arg1); } +inline FVec32 log(FVec32 arg1) { return Sleef_logf4_u10(arg1); } +inline FVec32 tan(FVec32 arg1) { return Sleef_tanf4_u10(arg1); } +inline FVec32 atan(FVec32 arg1) { return Sleef_atanf4_u10(arg1); } +inline FVec32 acos(FVec32 arg1) { return Sleef_acosf4_u10(arg1); } +inline FVec32 atanh(FVec32 arg1) { return Sleef_atanhf4_u10(arg1); } +inline FVec32 exp(FVec32 arg1) { return Sleef_expf4_u10(arg1); } +inline FVec32 log1p(FVec32 arg1) { return Sleef_log1pf4_u10(arg1); } +inline FVec32 expm1(FVec32 arg1) { return Sleef_expm1f4_u10(arg1); } +inline FVec32 erf(FVec32 arg1) { return Sleef_erff4_u10(arg1); } +inline FVec32 fma(FVec32 arg1, FVec32 arg2, FVec32 arg3) { + return vfmaq_f32(arg3, arg1, arg2); +} diff --git a/madspace/src/cpu/simd_x86_256.hpp b/madspace/src/cpu/simd_x86_256.hpp index eb770f1158..9cc0e3a0b7 100644 --- a/madspace/src/cpu/simd_x86_256.hpp +++ b/madspace/src/cpu/simd_x86_256.hpp @@ -3,6 +3,8 @@ #include constexpr int simd_vec_size = 4; +// float32 vectors have twice the lanes of the float64 ones +constexpr int simd_vec_size_f32 = 8; inline __m128i truncate_i64_to_i32(__m256i arg) { __m256 x = _mm256_castsi256_ps(arg); @@ -12,6 +14,7 @@ inline __m128i truncate_i64_to_i32(__m256i arg) { } struct FVec { + static constexpr int lanes = simd_vec_size; FVec() = default; FVec(__m256d _v) : v(_v) {}; FVec(double _v) : v(_mm256_set1_pd(_v)) {}; @@ -25,6 +28,7 @@ struct FVec { }; struct IVec { + static constexpr int lanes = simd_vec_size; IVec() = default; IVec(__m256i _v) : v(_v) {}; IVec(int _v) : v(_mm256_set1_epi64x(_v)) {}; @@ -38,6 +42,7 @@ struct IVec { }; struct BVec { + static constexpr int lanes = simd_vec_size; BVec() = default; BVec(__m256d _v) : v(_v) {}; BVec(__m256i _v) : v(_mm256_castsi256_pd(_v)) {}; @@ -60,7 +65,11 @@ mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) { } inline FVec vgather( - double* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices + VecTag, + double* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec indices ) { return _mm256_i32gather_pd( base_ptr, mem_indices(batch_stride, index_stride, indices), 8 @@ -68,7 +77,11 @@ inline FVec vgather( } inline IVec vgather( - int* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices + VecTag, + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec indices ) { return _mm256_cvtepi32_epi64(_mm_i32gather_epi32( base_ptr, mem_indices(batch_stride, index_stride, indices), 4 @@ -76,16 +89,20 @@ inline IVec vgather( } // contiguous loads and stores of the lanes of one batch element -inline FVec vload(double* base_ptr) { return _mm256_loadu_pd(base_ptr); } +inline FVec vload(VecTag, double* base_ptr) { return _mm256_loadu_pd(base_ptr); } -inline IVec vload(int* base_ptr) { +inline IVec vload(VecTag, int* base_ptr) { return _mm256_cvtepi32_epi64(_mm_loadu_si128(reinterpret_cast<__m128i*>(base_ptr))); } // the same value in every lane -inline FVec vload_broadcast(double* base_ptr) { return _mm256_broadcast_sd(base_ptr); } +inline FVec vload_broadcast(VecTag, double* base_ptr) { + return _mm256_broadcast_sd(base_ptr); +} -inline IVec vload_broadcast(int* base_ptr) { return _mm256_set1_epi64x(*base_ptr); } +inline IVec vload_broadcast(VecTag, int* base_ptr) { + return _mm256_set1_epi64x(*base_ptr); +} inline void vscatter( double* base_ptr, @@ -130,7 +147,9 @@ inline void vscatter( } } -inline void vstore(double* base_ptr, FVec values) { _mm256_storeu_pd(base_ptr, values); } +inline void vstore(double* base_ptr, FVec values) { + _mm256_storeu_pd(base_ptr, values); +} inline void vstore(int* base_ptr, IVec values) { _mm_storeu_si128(reinterpret_cast<__m128i*>(base_ptr), truncate_i64_to_i32(values)); @@ -213,3 +232,238 @@ inline FVec erf(FVec arg1) { return Sleef_erfd4_u10avx2(arg1); } inline FVec fma(FVec arg1, FVec arg2, FVec arg3) { return _mm256_fmadd_pd(arg1, arg2, arg3); } + +// float32 family: eight lanes, with int32 indices and masks of the same width + +struct FVec32 { + static constexpr int lanes = simd_vec_size_f32; + FVec32() = default; + FVec32(__m256 _v) : v(_v) {}; + FVec32(double _v) : v(_mm256_set1_ps(static_cast(_v))) {}; + explicit FVec32(__m256i _v) : v(_mm256_cvtepi32_ps(_v)) {}; + operator __m256() { return v; } + FVec32 operator+=(FVec32 _v) { + v = _mm256_add_ps(v, _v); + return v; + } + __m256 v; +}; + +struct IVec32 { + static constexpr int lanes = simd_vec_size_f32; + IVec32() = default; + IVec32(__m256i _v) : v(_v) {}; + IVec32(int _v) : v(_mm256_set1_epi32(_v)) {}; + explicit IVec32(__m256 _v) : v(_mm256_cvttps_epi32(_v)) {}; + operator __m256i() { return v; } + IVec32 operator+=(IVec32 _v) { + v = _mm256_add_epi32(v, _v); + return v; + } + __m256i v; +}; + +struct BVec32 { + static constexpr int lanes = simd_vec_size_f32; + BVec32() = default; + BVec32(__m256 _v) : v(_v) {}; + BVec32(__m256i _v) : v(_mm256_castsi256_ps(_v)) {}; + BVec32(bool _v) : v(_mm256_castsi256_ps(_mm256_set1_epi32(_v ? -1 : 0))) {}; + operator __m256() { return v; } + operator __m256i() { return _mm256_castps_si256(v); } + __m256 v; +}; + +inline __m256i +mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) { + return _mm256_add_epi32( + _mm256_mullo_epi32(indices, _mm256_set1_epi32(index_stride)), + _mm256_mullo_epi32( + _mm256_set1_epi32(batch_stride), _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0) + ) + ); +} + +inline FVec32 vgather( + VecTag, + float* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices +) { + return _mm256_i32gather_ps( + base_ptr, mem_indices32(batch_stride, index_stride, indices), 4 + ); +} + +inline IVec32 vgather( + VecTag, + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices +) { + return _mm256_i32gather_epi32( + base_ptr, mem_indices32(batch_stride, index_stride, indices), 4 + ); +} + +inline FVec32 vload(VecTag, float* base_ptr) { + return _mm256_loadu_ps(base_ptr); +} + +inline IVec32 vload(VecTag, int* base_ptr) { + return _mm256_loadu_si256(reinterpret_cast<__m256i*>(base_ptr)); +} + +inline FVec32 vload_broadcast(VecTag, float* base_ptr) { + return _mm256_broadcast_ss(base_ptr); +} + +inline IVec32 vload_broadcast(VecTag, int* base_ptr) { + return _mm256_set1_epi32(*base_ptr); +} + +inline void vscatter( + float* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices, + FVec32 values +) { + float values_buf[8]; + int indices_buf[8]; + _mm256_storeu_ps(values_buf, values); + _mm256_storeu_si256( + reinterpret_cast<__m256i*>(indices_buf), + mem_indices32(batch_stride, index_stride, indices) + ); + for (int i = 0; i < 8; ++i) { + base_ptr[indices_buf[i]] = values_buf[i]; + } +} + +inline void vscatter( + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices, + IVec32 values +) { + int values_buf[8]; + int indices_buf[8]; + _mm256_storeu_si256(reinterpret_cast<__m256i*>(values_buf), values); + _mm256_storeu_si256( + reinterpret_cast<__m256i*>(indices_buf), + mem_indices32(batch_stride, index_stride, indices) + ); + for (int i = 0; i < 8; ++i) { + base_ptr[indices_buf[i]] = values_buf[i]; + } +} + +inline void vstore(float* base_ptr, FVec32 values) { + _mm256_storeu_ps(base_ptr, values); +} + +inline void vstore(int* base_ptr, IVec32 values) { + _mm256_storeu_si256(reinterpret_cast<__m256i*>(base_ptr), values); +} + +inline FVec32 where(BVec32 arg1, FVec32 arg2, FVec32 arg3) { + return _mm256_blendv_ps(arg3, arg2, arg1); +} +inline IVec32 where(BVec32 arg1, IVec32 arg2, IVec32 arg3) { + return _mm256_blendv_epi8(arg3, arg2, arg1); +} +inline std::size_t single_index(IVec32 arg) { + return _mm_cvtsi128_si32(_mm256_castsi256_si128(arg)); +} +inline FVec32 min(FVec32 arg1, FVec32 arg2) { return _mm256_min_ps(arg1, arg2); } +inline FVec32 max(FVec32 arg1, FVec32 arg2) { return _mm256_max_ps(arg1, arg2); } + +inline BVec32 operator==(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps(arg1, arg2, _CMP_EQ_OQ); +} +inline BVec32 operator!=(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps(arg1, arg2, _CMP_NEQ_UQ); +} +inline BVec32 operator>(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps(arg1, arg2, _CMP_GT_OQ); +} +inline BVec32 operator<(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps(arg1, arg2, _CMP_LT_OQ); +} +inline BVec32 operator>=(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps(arg1, arg2, _CMP_GE_OQ); +} +inline BVec32 operator<=(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps(arg1, arg2, _CMP_LE_OQ); +} + +inline BVec32 operator&(BVec32 arg1, BVec32 arg2) { + return _mm256_and_si256(arg1, arg2); +} +inline BVec32 operator|(BVec32 arg1, BVec32 arg2) { + return _mm256_or_si256(arg1, arg2); +} +inline BVec32 operator!(BVec32 arg1) { + return _mm256_xor_si256(arg1, _mm256_cmpeq_epi32(arg1, arg1)); +} + +inline BVec32 operator==(IVec32 arg1, IVec32 arg2) { + return _mm256_cmpeq_epi32(arg1, arg2); +} +inline BVec32 operator!=(IVec32 arg1, IVec32 arg2) { return !(arg1 == arg2); } +inline BVec32 operator>(IVec32 arg1, IVec32 arg2) { + return _mm256_cmpgt_epi32(arg1, arg2); +} +inline BVec32 operator>=(IVec32 arg1, IVec32 arg2) { + return (arg1 > arg2) | (arg1 == arg2); +} +inline BVec32 operator<(IVec32 arg1, IVec32 arg2) { return !(arg1 >= arg2); } +inline BVec32 operator<=(IVec32 arg1, IVec32 arg2) { return !(arg1 > arg2); } + +inline FVec32 operator-(FVec32 arg1) { + return _mm256_sub_ps(_mm256_set1_ps(0.f), arg1); +} +inline FVec32 operator+(FVec32 arg1, FVec32 arg2) { return _mm256_add_ps(arg1, arg2); } +inline FVec32 operator-(FVec32 arg1, FVec32 arg2) { return _mm256_sub_ps(arg1, arg2); } +inline FVec32 operator*(FVec32 arg1, FVec32 arg2) { return _mm256_mul_ps(arg1, arg2); } +inline FVec32 operator/(FVec32 arg1, FVec32 arg2) { return _mm256_div_ps(arg1, arg2); } +inline IVec32 operator-(IVec32 arg1) { + return _mm256_sub_epi32(_mm256_setzero_si256(), arg1); +} +inline IVec32 operator+(IVec32 arg1, IVec32 arg2) { + return _mm256_add_epi32(arg1, arg2); +} +inline IVec32 operator-(IVec32 arg1, IVec32 arg2) { + return _mm256_sub_epi32(arg1, arg2); +} + +inline BVec32 isnan(FVec32 arg) { return arg != arg; } + +inline FVec32 sqrt(FVec32 arg1) { return Sleef_sqrtf8_u05avx2(arg1); } +inline FVec32 sin(FVec32 arg1) { return Sleef_sinf8_u10avx2(arg1); } +inline FVec32 cos(FVec32 arg1) { return Sleef_cosf8_u10avx2(arg1); } +inline FVec32 sinh(FVec32 arg1) { return Sleef_sinhf8_u10avx2(arg1); } +inline FVec32 asinh(FVec32 arg1) { return Sleef_asinhf8_u10avx2(arg1); } +inline FVec32 cosh(FVec32 arg1) { return Sleef_coshf8_u10avx2(arg1); } +inline FVec32 tanh(FVec32 arg1) { return Sleef_tanhf8_u10avx2(arg1); } +inline FVec32 atan2(FVec32 arg1, FVec32 arg2) { + return Sleef_atan2f8_u10avx2(arg1, arg2); +} +inline FVec32 pow(FVec32 arg1, FVec32 arg2) { return Sleef_powf8_u10avx2(arg1, arg2); } +inline FVec32 fabs(FVec32 arg1) { return Sleef_fabsf8_avx2(arg1); } +inline FVec32 log(FVec32 arg1) { return Sleef_logf8_u10avx2(arg1); } +inline FVec32 tan(FVec32 arg1) { return Sleef_tanf8_u10avx2(arg1); } +inline FVec32 atan(FVec32 arg1) { return Sleef_atanf8_u10avx2(arg1); } +inline FVec32 acos(FVec32 arg1) { return Sleef_acosf8_u10avx2(arg1); } +inline FVec32 atanh(FVec32 arg1) { return Sleef_atanhf8_u10avx2(arg1); } +inline FVec32 exp(FVec32 arg1) { return Sleef_expf8_u10avx2(arg1); } +inline FVec32 log1p(FVec32 arg1) { return Sleef_log1pf8_u10avx2(arg1); } +inline FVec32 expm1(FVec32 arg1) { return Sleef_expm1f8_u10avx2(arg1); } +inline FVec32 erf(FVec32 arg1) { return Sleef_erff8_u10avx2(arg1); } +inline FVec32 fma(FVec32 arg1, FVec32 arg2, FVec32 arg3) { + return _mm256_fmadd_ps(arg1, arg2, arg3); +} diff --git a/madspace/src/cpu/simd_x86_256y.hpp b/madspace/src/cpu/simd_x86_256y.hpp index bdf193a1c6..4d54bc55d0 100644 --- a/madspace/src/cpu/simd_x86_256y.hpp +++ b/madspace/src/cpu/simd_x86_256y.hpp @@ -4,8 +4,11 @@ #include constexpr int simd_vec_size = 4; +// float32 vectors have twice the lanes of the float64 ones +constexpr int simd_vec_size_f32 = 8; struct FVec { + static constexpr int lanes = simd_vec_size; FVec() = default; FVec(__m256d _v) : v(_v) {}; FVec(double _v) : v(_mm256_set1_pd(_v)) {}; @@ -19,6 +22,7 @@ struct FVec { }; struct IVec { + static constexpr int lanes = simd_vec_size; IVec() = default; IVec(__m256i _v) : v(_v) {}; IVec(int _v) : v(_mm256_set1_epi64x(_v)) {}; @@ -32,6 +36,7 @@ struct IVec { }; struct BVec { + static constexpr int lanes = simd_vec_size; BVec() = default; BVec(bool _v) : v(_v ? 0x0F : 0) {}; BVec(__mmask8 _v) : v(_v) {}; @@ -54,7 +59,11 @@ mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) { } inline FVec vgather( - double* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices + VecTag, + double* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec indices ) { return _mm256_i64gather_pd( base_ptr, mem_indices(batch_stride, index_stride, indices), 8 @@ -62,7 +71,11 @@ inline FVec vgather( } inline IVec vgather( - int* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices + VecTag, + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec indices ) { return _mm256_cvtepi32_epi64(_mm256_i64gather_epi32( base_ptr, mem_indices(batch_stride, index_stride, indices), 4 @@ -70,16 +83,20 @@ inline IVec vgather( } // contiguous loads and stores of the lanes of one batch element -inline FVec vload(double* base_ptr) { return _mm256_loadu_pd(base_ptr); } +inline FVec vload(VecTag, double* base_ptr) { return _mm256_loadu_pd(base_ptr); } -inline IVec vload(int* base_ptr) { +inline IVec vload(VecTag, int* base_ptr) { return _mm256_cvtepi32_epi64(_mm_loadu_si128(reinterpret_cast<__m128i*>(base_ptr))); } // the same value in every lane -inline FVec vload_broadcast(double* base_ptr) { return _mm256_broadcast_sd(base_ptr); } +inline FVec vload_broadcast(VecTag, double* base_ptr) { + return _mm256_broadcast_sd(base_ptr); +} -inline IVec vload_broadcast(int* base_ptr) { return _mm256_set1_epi64x(*base_ptr); } +inline IVec vload_broadcast(VecTag, int* base_ptr) { + return _mm256_set1_epi64x(*base_ptr); +} inline void vscatter( double* base_ptr, @@ -108,7 +125,9 @@ inline void vscatter( ); } -inline void vstore(double* base_ptr, FVec values) { _mm256_storeu_pd(base_ptr, values); } +inline void vstore(double* base_ptr, FVec values) { + _mm256_storeu_pd(base_ptr, values); +} inline void vstore(int* base_ptr, IVec values) { _mm_storeu_si128( @@ -209,3 +228,226 @@ inline FVec erf(FVec arg1) { return Sleef_erfd4_u10avx2(arg1); } inline FVec fma(FVec arg1, FVec arg2, FVec arg3) { return _mm256_fmadd_pd(arg1, arg2, arg3); } + +// float32 family: eight lanes, with int32 indices and masks of the same width + +struct FVec32 { + static constexpr int lanes = simd_vec_size_f32; + FVec32() = default; + FVec32(__m256 _v) : v(_v) {}; + FVec32(double _v) : v(_mm256_set1_ps(static_cast(_v))) {}; + explicit FVec32(__m256i _v) : v(_mm256_cvtepi32_ps(_v)) {}; + operator __m256() { return v; } + FVec32 operator+=(FVec32 _v) { + v = _mm256_add_ps(v, _v); + return v; + } + __m256 v; +}; + +struct IVec32 { + static constexpr int lanes = simd_vec_size_f32; + IVec32() = default; + IVec32(__m256i _v) : v(_v) {}; + IVec32(int _v) : v(_mm256_set1_epi32(_v)) {}; + explicit IVec32(__m256 _v) : v(_mm256_cvttps_epi32(_v)) {}; + operator __m256i() { return v; } + IVec32 operator+=(IVec32 _v) { + v = _mm256_add_epi32(v, _v); + return v; + } + __m256i v; +}; + +struct BVec32 { + static constexpr int lanes = simd_vec_size_f32; + BVec32() = default; + BVec32(bool _v) : v(_v ? 0xFF : 0) {}; + BVec32(__mmask8 _v) : v(_v) {}; + operator __mmask8() { return v; } + __mmask8 v; +}; + +inline __m256i +mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) { + return _mm256_add_epi32( + _mm256_mullo_epi32(indices, _mm256_set1_epi32(index_stride)), + _mm256_mullo_epi32( + _mm256_set1_epi32(batch_stride), _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0) + ) + ); +} + +inline FVec32 vgather( + VecTag, + float* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices +) { + return _mm256_i32gather_ps( + base_ptr, mem_indices32(batch_stride, index_stride, indices), 4 + ); +} + +inline IVec32 vgather( + VecTag, + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices +) { + return _mm256_i32gather_epi32( + base_ptr, mem_indices32(batch_stride, index_stride, indices), 4 + ); +} + +inline FVec32 vload(VecTag, float* base_ptr) { + return _mm256_loadu_ps(base_ptr); +} + +inline IVec32 vload(VecTag, int* base_ptr) { + return _mm256_loadu_si256(reinterpret_cast<__m256i*>(base_ptr)); +} + +inline FVec32 vload_broadcast(VecTag, float* base_ptr) { + return _mm256_broadcast_ss(base_ptr); +} + +inline IVec32 vload_broadcast(VecTag, int* base_ptr) { + return _mm256_set1_epi32(*base_ptr); +} + +inline void vscatter( + float* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices, + FVec32 values +) { + _mm256_i32scatter_ps( + base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4 + ); +} + +inline void vscatter( + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices, + IVec32 values +) { + _mm256_i32scatter_epi32( + base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4 + ); +} + +inline void vstore(float* base_ptr, FVec32 values) { + _mm256_storeu_ps(base_ptr, values); +} + +inline void vstore(int* base_ptr, IVec32 values) { + _mm256_storeu_si256(reinterpret_cast<__m256i*>(base_ptr), values); +} + +inline FVec32 where(BVec32 arg1, FVec32 arg2, FVec32 arg3) { + return _mm256_mask_blend_ps(arg1, arg3, arg2); +} +inline IVec32 where(BVec32 arg1, IVec32 arg2, IVec32 arg3) { + return _mm256_mask_blend_epi32(arg1, arg3, arg2); +} +inline std::size_t single_index(IVec32 arg) { + return _mm_cvtsi128_si32(_mm256_castsi256_si128(arg)); +} +inline FVec32 min(FVec32 arg1, FVec32 arg2) { return _mm256_min_ps(arg1, arg2); } +inline FVec32 max(FVec32 arg1, FVec32 arg2) { return _mm256_max_ps(arg1, arg2); } + +inline BVec32 operator==(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps_mask(arg1, arg2, _CMP_EQ_OQ); +} +inline BVec32 operator!=(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps_mask(arg1, arg2, _CMP_NEQ_UQ); +} +inline BVec32 operator>(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps_mask(arg1, arg2, _CMP_GT_OQ); +} +inline BVec32 operator<(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps_mask(arg1, arg2, _CMP_LT_OQ); +} +inline BVec32 operator>=(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps_mask(arg1, arg2, _CMP_GE_OQ); +} +inline BVec32 operator<=(FVec32 arg1, FVec32 arg2) { + return _mm256_cmp_ps_mask(arg1, arg2, _CMP_LE_OQ); +} + +inline BVec32 operator&(BVec32 arg1, BVec32 arg2) { + return static_cast<__mmask8>(arg1.v & arg2.v); +} +inline BVec32 operator|(BVec32 arg1, BVec32 arg2) { + return static_cast<__mmask8>(arg1.v | arg2.v); +} +inline BVec32 operator!(BVec32 arg1) { return static_cast<__mmask8>(~arg1.v); } + +inline BVec32 operator==(IVec32 arg1, IVec32 arg2) { + return _mm256_cmpeq_epi32_mask(arg1, arg2); +} +inline BVec32 operator!=(IVec32 arg1, IVec32 arg2) { + return _mm256_cmpneq_epi32_mask(arg1, arg2); +} +inline BVec32 operator>(IVec32 arg1, IVec32 arg2) { + return _mm256_cmpgt_epi32_mask(arg1, arg2); +} +inline BVec32 operator>=(IVec32 arg1, IVec32 arg2) { + return _mm256_cmpge_epi32_mask(arg1, arg2); +} +inline BVec32 operator<(IVec32 arg1, IVec32 arg2) { + return _mm256_cmplt_epi32_mask(arg1, arg2); +} +inline BVec32 operator<=(IVec32 arg1, IVec32 arg2) { + return _mm256_cmple_epi32_mask(arg1, arg2); +} + +inline FVec32 operator-(FVec32 arg1) { + return _mm256_sub_ps(_mm256_set1_ps(0.f), arg1); +} +inline FVec32 operator+(FVec32 arg1, FVec32 arg2) { return _mm256_add_ps(arg1, arg2); } +inline FVec32 operator-(FVec32 arg1, FVec32 arg2) { return _mm256_sub_ps(arg1, arg2); } +inline FVec32 operator*(FVec32 arg1, FVec32 arg2) { return _mm256_mul_ps(arg1, arg2); } +inline FVec32 operator/(FVec32 arg1, FVec32 arg2) { return _mm256_div_ps(arg1, arg2); } +inline IVec32 operator-(IVec32 arg1) { + return _mm256_sub_epi32(_mm256_setzero_si256(), arg1); +} +inline IVec32 operator+(IVec32 arg1, IVec32 arg2) { + return _mm256_add_epi32(arg1, arg2); +} +inline IVec32 operator-(IVec32 arg1, IVec32 arg2) { + return _mm256_sub_epi32(arg1, arg2); +} + +inline BVec32 isnan(FVec32 arg) { return arg != arg; } + +inline FVec32 sqrt(FVec32 arg1) { return Sleef_sqrtf8_u05avx2(arg1); } +inline FVec32 sin(FVec32 arg1) { return Sleef_sinf8_u10avx2(arg1); } +inline FVec32 cos(FVec32 arg1) { return Sleef_cosf8_u10avx2(arg1); } +inline FVec32 sinh(FVec32 arg1) { return Sleef_sinhf8_u10avx2(arg1); } +inline FVec32 asinh(FVec32 arg1) { return Sleef_asinhf8_u10avx2(arg1); } +inline FVec32 cosh(FVec32 arg1) { return Sleef_coshf8_u10avx2(arg1); } +inline FVec32 tanh(FVec32 arg1) { return Sleef_tanhf8_u10avx2(arg1); } +inline FVec32 atan2(FVec32 arg1, FVec32 arg2) { + return Sleef_atan2f8_u10avx2(arg1, arg2); +} +inline FVec32 pow(FVec32 arg1, FVec32 arg2) { return Sleef_powf8_u10avx2(arg1, arg2); } +inline FVec32 fabs(FVec32 arg1) { return Sleef_fabsf8_avx2(arg1); } +inline FVec32 log(FVec32 arg1) { return Sleef_logf8_u10avx2(arg1); } +inline FVec32 tan(FVec32 arg1) { return Sleef_tanf8_u10avx2(arg1); } +inline FVec32 atan(FVec32 arg1) { return Sleef_atanf8_u10avx2(arg1); } +inline FVec32 acos(FVec32 arg1) { return Sleef_acosf8_u10avx2(arg1); } +inline FVec32 atanh(FVec32 arg1) { return Sleef_atanhf8_u10avx2(arg1); } +inline FVec32 exp(FVec32 arg1) { return Sleef_expf8_u10avx2(arg1); } +inline FVec32 log1p(FVec32 arg1) { return Sleef_log1pf8_u10avx2(arg1); } +inline FVec32 expm1(FVec32 arg1) { return Sleef_expm1f8_u10avx2(arg1); } +inline FVec32 erf(FVec32 arg1) { return Sleef_erff8_u10avx2(arg1); } +inline FVec32 fma(FVec32 arg1, FVec32 arg2, FVec32 arg3) { + return _mm256_fmadd_ps(arg1, arg2, arg3); +} diff --git a/madspace/src/cpu/simd_x86_512.hpp b/madspace/src/cpu/simd_x86_512.hpp index c143fb69a9..0a80a261dc 100644 --- a/madspace/src/cpu/simd_x86_512.hpp +++ b/madspace/src/cpu/simd_x86_512.hpp @@ -3,8 +3,11 @@ #include constexpr int simd_vec_size = 8; +// float32 vectors have twice the lanes of the float64 ones +constexpr int simd_vec_size_f32 = 16; struct FVec { + static constexpr int lanes = simd_vec_size; FVec() = default; FVec(__m512d _v) : v(_v) {}; FVec(double _v) : v(_mm512_set1_pd(_v)) {}; @@ -18,6 +21,7 @@ struct FVec { }; struct IVec { + static constexpr int lanes = simd_vec_size; IVec() = default; IVec(__m512i _v) : v(_v) {}; IVec(int _v) : v(_mm512_set1_epi64(_v)) {}; @@ -31,6 +35,7 @@ struct IVec { }; struct BVec { + static constexpr int lanes = simd_vec_size; BVec() = default; BVec(bool _v) : v(_v ? 0xFF : 0) {}; BVec(__mmask8 _v) : v(_v) {}; @@ -55,7 +60,11 @@ mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) { } inline FVec vgather( - double* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices + VecTag, + double* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec indices ) { return _mm512_i32gather_pd( mem_indices(batch_stride, index_stride, indices), base_ptr, 8 @@ -63,7 +72,11 @@ inline FVec vgather( } inline IVec vgather( - int* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices + VecTag, + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec indices ) { return _mm512_cvtepi32_epi64(_mm256_i32gather_epi32( base_ptr, mem_indices(batch_stride, index_stride, indices), 4 @@ -71,18 +84,22 @@ inline IVec vgather( } // contiguous loads and stores of the lanes of one batch element -inline FVec vload(double* base_ptr) { return _mm512_loadu_pd(base_ptr); } +inline FVec vload(VecTag, double* base_ptr) { return _mm512_loadu_pd(base_ptr); } -inline IVec vload(int* base_ptr) { +inline IVec vload(VecTag, int* base_ptr) { return _mm512_cvtepi32_epi64( _mm256_loadu_si256(reinterpret_cast<__m256i*>(base_ptr)) ); } // the same value in every lane -inline FVec vload_broadcast(double* base_ptr) { return _mm512_set1_pd(*base_ptr); } +inline FVec vload_broadcast(VecTag, double* base_ptr) { + return _mm512_set1_pd(*base_ptr); +} -inline IVec vload_broadcast(int* base_ptr) { return _mm512_set1_epi64(*base_ptr); } +inline IVec vload_broadcast(VecTag, int* base_ptr) { + return _mm512_set1_epi64(*base_ptr); +} inline void vscatter( double* base_ptr, @@ -119,7 +136,9 @@ inline void vscatter( } } -inline void vstore(double* base_ptr, FVec values) { _mm512_storeu_pd(base_ptr, values); } +inline void vstore(double* base_ptr, FVec values) { + _mm512_storeu_pd(base_ptr, values); +} inline void vstore(int* base_ptr, IVec values) { _mm256_storeu_si256( @@ -220,3 +239,229 @@ inline FVec erf(FVec arg1) { return Sleef_erfd8_u10avx512f(arg1); } inline FVec fma(FVec arg1, FVec arg2, FVec arg3) { return _mm512_fmadd_pd(arg1, arg2, arg3); } + +// float32 family: sixteen lanes, with int32 indices and masks of the same width + +struct FVec32 { + static constexpr int lanes = simd_vec_size_f32; + FVec32() = default; + FVec32(__m512 _v) : v(_v) {}; + FVec32(double _v) : v(_mm512_set1_ps(static_cast(_v))) {}; + explicit FVec32(__m512i _v) : v(_mm512_cvtepi32_ps(_v)) {}; + operator __m512() { return v; } + FVec32 operator+=(FVec32 _v) { + v = _mm512_add_ps(v, _v); + return v; + } + __m512 v; +}; + +struct IVec32 { + static constexpr int lanes = simd_vec_size_f32; + IVec32() = default; + IVec32(__m512i _v) : v(_v) {}; + IVec32(int _v) : v(_mm512_set1_epi32(_v)) {}; + explicit IVec32(__m512 _v) : v(_mm512_cvttps_epi32(_v)) {}; + operator __m512i() { return v; } + IVec32 operator+=(IVec32 _v) { + v = _mm512_add_epi32(v, _v); + return v; + } + __m512i v; +}; + +struct BVec32 { + static constexpr int lanes = simd_vec_size_f32; + BVec32() = default; + BVec32(bool _v) : v(_v ? 0xFFFF : 0) {}; + BVec32(__mmask16 _v) : v(_v) {}; + operator __mmask16() { return v; } + __mmask16 v; +}; + +inline __m512i +mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) { + return _mm512_add_epi32( + _mm512_mullo_epi32(indices, _mm512_set1_epi32(index_stride)), + _mm512_mullo_epi32( + _mm512_set1_epi32(batch_stride), + _mm512_set_epi32(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0) + ) + ); +} + +inline FVec32 vgather( + VecTag, + float* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices +) { + return _mm512_i32gather_ps( + mem_indices32(batch_stride, index_stride, indices), base_ptr, 4 + ); +} + +inline IVec32 vgather( + VecTag, + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices +) { + return _mm512_i32gather_epi32( + mem_indices32(batch_stride, index_stride, indices), base_ptr, 4 + ); +} + +inline FVec32 vload(VecTag, float* base_ptr) { + return _mm512_loadu_ps(base_ptr); +} + +inline IVec32 vload(VecTag, int* base_ptr) { + return _mm512_loadu_si512(base_ptr); +} + +inline FVec32 vload_broadcast(VecTag, float* base_ptr) { + return _mm512_set1_ps(*base_ptr); +} + +inline IVec32 vload_broadcast(VecTag, int* base_ptr) { + return _mm512_set1_epi32(*base_ptr); +} + +inline void vscatter( + float* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices, + FVec32 values +) { + _mm512_i32scatter_ps( + base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4 + ); +} + +inline void vscatter( + int* base_ptr, + std::size_t batch_stride, + std::size_t index_stride, + IVec32 indices, + IVec32 values +) { + _mm512_i32scatter_epi32( + base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4 + ); +} + +inline void vstore(float* base_ptr, FVec32 values) { + _mm512_storeu_ps(base_ptr, values); +} + +inline void vstore(int* base_ptr, IVec32 values) { + _mm512_storeu_si512(base_ptr, values); +} + +inline FVec32 where(BVec32 arg1, FVec32 arg2, FVec32 arg3) { + return _mm512_mask_blend_ps(arg1, arg3, arg2); +} +inline IVec32 where(BVec32 arg1, IVec32 arg2, IVec32 arg3) { + return _mm512_mask_blend_epi32(arg1, arg3, arg2); +} +inline std::size_t single_index(IVec32 arg) { + return _mm_cvtsi128_si32(_mm512_castsi512_si128(arg)); +} +inline FVec32 min(FVec32 arg1, FVec32 arg2) { return _mm512_min_ps(arg1, arg2); } +inline FVec32 max(FVec32 arg1, FVec32 arg2) { return _mm512_max_ps(arg1, arg2); } + +inline BVec32 operator==(FVec32 arg1, FVec32 arg2) { + return _mm512_cmp_ps_mask(arg1, arg2, _CMP_EQ_OQ); +} +inline BVec32 operator!=(FVec32 arg1, FVec32 arg2) { + return _mm512_cmp_ps_mask(arg1, arg2, _CMP_NEQ_UQ); +} +inline BVec32 operator>(FVec32 arg1, FVec32 arg2) { + return _mm512_cmp_ps_mask(arg1, arg2, _CMP_GT_OQ); +} +inline BVec32 operator<(FVec32 arg1, FVec32 arg2) { + return _mm512_cmp_ps_mask(arg1, arg2, _CMP_LT_OQ); +} +inline BVec32 operator>=(FVec32 arg1, FVec32 arg2) { + return _mm512_cmp_ps_mask(arg1, arg2, _CMP_GE_OQ); +} +inline BVec32 operator<=(FVec32 arg1, FVec32 arg2) { + return _mm512_cmp_ps_mask(arg1, arg2, _CMP_LE_OQ); +} + +inline BVec32 operator&(BVec32 arg1, BVec32 arg2) { + return static_cast<__mmask16>(arg1.v & arg2.v); +} +inline BVec32 operator|(BVec32 arg1, BVec32 arg2) { + return static_cast<__mmask16>(arg1.v | arg2.v); +} +inline BVec32 operator!(BVec32 arg1) { return static_cast<__mmask16>(~arg1.v); } + +inline BVec32 operator==(IVec32 arg1, IVec32 arg2) { + return _mm512_cmpeq_epi32_mask(arg1, arg2); +} +inline BVec32 operator!=(IVec32 arg1, IVec32 arg2) { + return _mm512_cmpneq_epi32_mask(arg1, arg2); +} +inline BVec32 operator>(IVec32 arg1, IVec32 arg2) { + return _mm512_cmpgt_epi32_mask(arg1, arg2); +} +inline BVec32 operator>=(IVec32 arg1, IVec32 arg2) { + return _mm512_cmpge_epi32_mask(arg1, arg2); +} +inline BVec32 operator<(IVec32 arg1, IVec32 arg2) { + return _mm512_cmplt_epi32_mask(arg1, arg2); +} +inline BVec32 operator<=(IVec32 arg1, IVec32 arg2) { + return _mm512_cmple_epi32_mask(arg1, arg2); +} + +inline FVec32 operator-(FVec32 arg1) { + return _mm512_sub_ps(_mm512_set1_ps(0.f), arg1); +} +inline FVec32 operator+(FVec32 arg1, FVec32 arg2) { return _mm512_add_ps(arg1, arg2); } +inline FVec32 operator-(FVec32 arg1, FVec32 arg2) { return _mm512_sub_ps(arg1, arg2); } +inline FVec32 operator*(FVec32 arg1, FVec32 arg2) { return _mm512_mul_ps(arg1, arg2); } +inline FVec32 operator/(FVec32 arg1, FVec32 arg2) { return _mm512_div_ps(arg1, arg2); } +inline IVec32 operator-(IVec32 arg1) { + return _mm512_sub_epi32(_mm512_setzero_si512(), arg1); +} +inline IVec32 operator+(IVec32 arg1, IVec32 arg2) { + return _mm512_add_epi32(arg1, arg2); +} +inline IVec32 operator-(IVec32 arg1, IVec32 arg2) { + return _mm512_sub_epi32(arg1, arg2); +} + +inline BVec32 isnan(FVec32 arg) { return arg != arg; } + +inline FVec32 sqrt(FVec32 arg1) { return Sleef_sqrtf16_u05avx512f(arg1); } +inline FVec32 sin(FVec32 arg1) { return Sleef_sinf16_u10avx512f(arg1); } +inline FVec32 cos(FVec32 arg1) { return Sleef_cosf16_u10avx512f(arg1); } +inline FVec32 sinh(FVec32 arg1) { return Sleef_sinhf16_u10avx512f(arg1); } +inline FVec32 asinh(FVec32 arg1) { return Sleef_asinhf16_u10avx512f(arg1); } +inline FVec32 cosh(FVec32 arg1) { return Sleef_coshf16_u10avx512f(arg1); } +inline FVec32 tanh(FVec32 arg1) { return Sleef_tanhf16_u10avx512f(arg1); } +inline FVec32 atan2(FVec32 arg1, FVec32 arg2) { + return Sleef_atan2f16_u10avx512f(arg1, arg2); +} +inline FVec32 pow(FVec32 arg1, FVec32 arg2) { + return Sleef_powf16_u10avx512f(arg1, arg2); +} +inline FVec32 fabs(FVec32 arg1) { return Sleef_fabsf16_avx512f(arg1); } +inline FVec32 log(FVec32 arg1) { return Sleef_logf16_u10avx512f(arg1); } +inline FVec32 tan(FVec32 arg1) { return Sleef_tanf16_u10avx512f(arg1); } +inline FVec32 atan(FVec32 arg1) { return Sleef_atanf16_u10avx512f(arg1); } +inline FVec32 acos(FVec32 arg1) { return Sleef_acosf16_u10avx512f(arg1); } +inline FVec32 atanh(FVec32 arg1) { return Sleef_atanhf16_u10avx512f(arg1); } +inline FVec32 exp(FVec32 arg1) { return Sleef_expf16_u10avx512f(arg1); } +inline FVec32 log1p(FVec32 arg1) { return Sleef_log1pf16_u10avx512f(arg1); } +inline FVec32 expm1(FVec32 arg1) { return Sleef_expm1f16_u10avx512f(arg1); } +inline FVec32 erf(FVec32 arg1) { return Sleef_erff16_u10avx512f(arg1); } +inline FVec32 fma(FVec32 arg1, FVec32 arg2, FVec32 arg3) { + return _mm512_fmadd_ps(arg1, arg2, arg3); +} diff --git a/madspace/src/cpu/tensor.hpp b/madspace/src/cpu/tensor.hpp index 0eeba4ba86..00d8345550 100644 --- a/madspace/src/cpu/tensor.hpp +++ b/madspace/src/cpu/tensor.hpp @@ -12,12 +12,12 @@ namespace madspace { namespace cpu { -// View onto a tensor for the vectorized kernels. The simd_vec_size lanes of a +// View onto a tensor for the vectorized kernels. The V::lanes lanes of a // vector are consecutive batch entries, so they are always contiguous in memory // (batch stride one) and are read and written with plain vector loads and stores. // Moving to the next vector advances the data pointer by vec_step elements: -// simd_vec_size for batched tensors, zero for literals that are stored padded to -// simd_vec_size identical lanes. Tensors that do not have this layout are handled +// V::lanes for batched tensors, zero for literals that are stored padded to +// simd_vec_size_max identical lanes. Tensors that do not have this layout are handled // by the scalar kernels instead (see tensor_foreach_impl). template class VectorizedTensorView { @@ -65,9 +65,7 @@ class VectorizedTensorView { T* ptr = &_data[offset(first_index)]; int i = 1; ((ptr = &ptr[index * _stride[i++]]), ...); - return { - ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _vec_step - }; + return {ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _vec_step}; } template @@ -78,15 +76,13 @@ class VectorizedTensorView { T* ptr = &_data[offset(first_index)]; int i = 1; ((ptr = &ptr[index * _stride[i++]]), ...); - return { - ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _vec_step - }; + return {ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _vec_step}; } operator V() const requires(_dim == 0) { - return vload(_data); + return vload(VecTag{}, _data); } V operator=(V value) @@ -99,7 +95,7 @@ class VectorizedTensorView { V operator+=(V value) requires(_dim == 0) { - V new_value = vload(_data) + value; + V new_value = vload(VecTag{}, _data) + value; vstore(_data, new_value); return new_value; } @@ -109,7 +105,7 @@ class VectorizedTensorView { std::size_t size(std::size_t index = 0) const { if (is_batch && index == 0) { - return _shape[0] / simd_vec_size; + return _shape[0] / V::lanes; } else { return _shape[index]; } @@ -119,14 +115,14 @@ class VectorizedTensorView { V gather(IVec indices) const requires(_dim == 1) { - return vgather(_data, 1, _stride[0], indices); + return vgather(VecTag{}, _data, 1, _stride[0], indices); } template void scatter_add(IVec indices, V values) requires(_dim == 1) { - V old_values = vgather(_data, 1, _stride[0], indices); + V old_values = vgather(VecTag{}, _data, 1, _stride[0], indices); vscatter(_data, 1, _stride[0], indices, old_values + values); } @@ -183,7 +179,7 @@ class SingleVectorizedTensorView { operator V() const requires(_dim == 0) { - return vload_broadcast(_data); + return vload_broadcast(VecTag{}, _data); } // the value as a plain scalar, e.g. for loop bounds or indices @@ -199,7 +195,7 @@ class SingleVectorizedTensorView { V gather(IVec indices) const requires(_dim == 1) { - return vgather(_data, 0, _stride[0], indices); + return vgather(VecTag{}, _data, 0, _stride[0], indices); } private: @@ -277,7 +273,7 @@ struct get_views { // return the tuple of VectorizedTensorViews where the type is extracted // from the signature of F. vec_steps holds the step between vectors for every -// argument (simd_vec_size, or zero for padded literals) +// argument (the lane count, or zero for padded literals) template struct get_vectorized_views; template @@ -336,6 +332,15 @@ struct first_param { static constexpr int dim = std::tuple_element_t<0, std::tuple>::dim; }; +// lane count of a vectorized kernel, taken from the vector type of its first view +template +struct vector_lanes; +template +struct vector_lanes { + static constexpr int value = std::remove_cvref_t< + std::tuple_element_t<0, std::tuple>>::VType::lanes; +}; + template inline void nested_for(std::size_t batch_size, std::size_t batch_offset, V... views) { std::size_t end_index = batch_offset + batch_size; @@ -437,7 +442,7 @@ inline void tensor_foreach_impl( // The vectorized kernel loads and stores the lanes of a vector with plain // contiguous memory accesses. It is only used if every batched argument has - // batch stride one (or is a literal padded to simd_vec_size lanes); single + // batch stride one (or is a literal padded to simd_vec_size_max lanes); single // inputs are read with broadcast loads and can have any layout. Otherwise // the whole batch is handled by the scalar kernel. constexpr std::size_t n_args = n_in + n_out + sizeof...(S); @@ -454,7 +459,7 @@ inline void tensor_foreach_impl( if constexpr (!ScalarType && !Single::value[I]) { auto& view = std::get(flat_views); if (view.shape[0] == batch_size && view.stride[0] == 1) { - vec_steps[I] = simd_vec_size; + vec_steps[I] = vector_lanes::value; } else if (view.simd_padded) { vec_steps[I] = 0; } else { @@ -469,10 +474,9 @@ inline void tensor_foreach_impl( device.foreach ( batch_size, - [flat_views, - vectorize, - vec_steps, - scalar_args...](std::size_t count, std::size_t offset) mutable { + [flat_views, vectorize, vec_steps, scalar_args...]( + std::size_t count, std::size_t offset + ) mutable { auto views = std::apply(get_views(), flat_views); std::size_t scalar_offset = offset; if constexpr (!std:: @@ -486,8 +490,10 @@ inline void tensor_foreach_impl( }, views ); - std::size_t vec_count = count / simd_vec_size; - std::size_t vec_offset = offset / simd_vec_size; + constexpr std::size_t lanes = + vector_lanes::value; + std::size_t vec_count = count / lanes; + std::size_t vec_offset = offset / lanes; std::apply( [vec_count, vec_offset](auto&&... args) { nested_for( @@ -496,7 +502,7 @@ inline void tensor_foreach_impl( }, vectorized_views ); - scalar_offset += vec_count * simd_vec_size; + scalar_offset += vec_count * lanes; } } std::size_t scalar_count = offset + count - scalar_offset; From c33c5fc14b5e7089e5b295b5c80f6dbc2e7e8460 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 19:16:10 +0200 Subject: [PATCH 04/24] madspace: train float32 parameters - the contiguous global gradient buffer takes the dtype of the trainable globals, which must all share one dtype - Adam step for float32 parameters - the gradient clipper computes the norm in float64 for float32 gradients - Tensor::cast for a CPU copy with converted element type --- .../madspace/driver/adam_optimizer.hpp | 3 +- madspace/include/madspace/driver/tensor.hpp | 3 ++ madspace/src/cpu/device.cpp | 51 +++++++++++++------ madspace/src/cpu/runtime.cpp | 11 +++- madspace/src/cpu/runtime.hpp | 1 + madspace/src/driver/adam_optimizer.cpp | 28 ++++++---- madspace/src/driver/tensor.cpp | 39 ++++++++++++++ madspace/src/gpu/device.cu | 39 +++++++++----- madspace/src/gpu/runtime.cu | 9 +++- madspace/src/gpu/runtime.hpp | 1 + 10 files changed, 144 insertions(+), 41 deletions(-) diff --git a/madspace/include/madspace/driver/adam_optimizer.hpp b/madspace/include/madspace/driver/adam_optimizer.hpp index d276f76610..4e78d61441 100644 --- a/madspace/include/madspace/driver/adam_optimizer.hpp +++ b/madspace/include/madspace/driver/adam_optimizer.hpp @@ -11,7 +11,8 @@ namespace madspace { /// `grad_clip_threshold`. class GradientClipper : public FunctionGenerator { public: - GradientClipper(); + /// @param dtype Float type of the gradients; the norm is computed in float64. + GradientClipper(DataType dtype = DataType::dt_float64); private: NamedVector build_function_impl( diff --git a/madspace/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp index 39961f36ff..56eeb042b6 100644 --- a/madspace/include/madspace/driver/tensor.hpp +++ b/madspace/include/madspace/driver/tensor.hpp @@ -700,6 +700,9 @@ class Tensor { } /// A copy of this tensor on the CPU, or `*this` if it is already there. Tensor cpu() const { return cpu(*impl->device); } + /// A contiguous CPU copy with the element type converted to `dtype`, or the + /// CPU tensor itself if it already has that type. + Tensor cast(DataType dtype) const; template /// Sets every element to zero. diff --git a/madspace/src/cpu/device.cpp b/madspace/src/cpu/device.cpp index bdae82c913..9c4d8f3185 100644 --- a/madspace/src/cpu/device.cpp +++ b/madspace/src/cpu/device.cpp @@ -121,22 +121,41 @@ void CpuDevice::adam_step( double bias_corr2_sqrt, double weight_decay ) const { - tensor_foreach_dynamic< - kernel_adam_step, - kernel_adam_step, - 1, - 3>( - {&gradient}, - {¶meter, &exp_avg, &exp_avg_sq}, - gradient.size(0), - *this, - step_size, - beta1, - beta2, - eps, - bias_corr2_sqrt, - weight_decay - ); + if (parameter.dtype() == DataType::dt_float32) { + tensor_foreach_dynamic< + kernel_adam_step, + kernel_adam_step, + 1, + 3>( + {&gradient}, + {¶meter, &exp_avg, &exp_avg_sq}, + gradient.size(0), + *this, + static_cast(step_size), + static_cast(beta1), + static_cast(beta2), + static_cast(eps), + static_cast(bias_corr2_sqrt), + static_cast(weight_decay) + ); + } else { + tensor_foreach_dynamic< + kernel_adam_step, + kernel_adam_step, + 1, + 3>( + {&gradient}, + {¶meter, &exp_avg, &exp_avg_sq}, + gradient.size(0), + *this, + step_size, + beta1, + beta2, + eps, + bias_corr2_sqrt, + weight_decay + ); + } } void AsyncCpuDevice::tensor_copy(const Tensor& source, Tensor& target) const { diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 407051b821..b7b48ffdab 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -1295,6 +1295,13 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur if (context->global_requires_grad(name)) { _requires_grad_init.at(value.local_index) = true; _grad_global_indices.push_back(value.local_index); + // the gradients of all globals are stored in one contiguous tensor + if (_grad_global_total_size > 0 && global.dtype() != _grad_global_dtype) { + throw std::invalid_argument( + "all globals that require a gradient must have the same dtype" + ); + } + _grad_global_dtype = global.dtype(); _grad_global_shapes.push_back(global.shape()); _grad_global_total_size += global.shape().product(); } @@ -1472,7 +1479,7 @@ std::pair CpuRuntime::run_backward_single( } Tensor all_global_grads( - DataType::dt_float64, {_grad_global_total_size}, AllocHint::global_grad + _grad_global_dtype, {_grad_global_total_size}, AllocHint::global_grad ); // all_global_grads.zero(); TensorVec global_grads = all_global_grads.split_and_reshape(_grad_global_shapes); @@ -1666,7 +1673,7 @@ std::pair CpuRuntime::run_backward_concurrent( } Tensor all_global_grads( - DataType::dt_float64, {_grad_global_total_size}, AllocHint::global_grad + _grad_global_dtype, {_grad_global_total_size}, AllocHint::global_grad ); // all_global_grads.zero(); TensorVec global_grads = all_global_grads.split_and_reshape(_grad_global_shapes); diff --git a/madspace/src/cpu/runtime.hpp b/madspace/src/cpu/runtime.hpp index 8324539c85..25406843bc 100644 --- a/madspace/src/cpu/runtime.hpp +++ b/madspace/src/cpu/runtime.hpp @@ -90,6 +90,7 @@ class CpuRuntime : public Runtime { SizeVec _grad_global_indices; std::vector _grad_global_shapes; std::size_t _grad_global_total_size; + DataType _grad_global_dtype = DataType::dt_float64; ContextPtr _context; ThreadResource& _rand_gens; bool _concurrent; diff --git a/madspace/src/driver/adam_optimizer.cpp b/madspace/src/driver/adam_optimizer.cpp index 7e05fe7430..cf55a256fd 100644 --- a/madspace/src/driver/adam_optimizer.cpp +++ b/madspace/src/driver/adam_optimizer.cpp @@ -4,20 +4,25 @@ using namespace madspace; -GradientClipper::GradientClipper() : +GradientClipper::GradientClipper(DataType dtype) : FunctionGenerator( "Unweighter", - {{"gradients_in", batch_float}, {"threshold", single_float}}, - {{"gradients_out", batch_float}, {"gradient_norm", single_float}} + {{"gradients_in", Type(dtype, batch_size, {})}, {"threshold", single_float}}, + {{"gradients_out", Type(dtype, batch_size, {})}, + {"gradient_norm", single_float}} ) {} NamedVector GradientClipper::build_function_impl( FunctionBuilder& fb, const NamedVector& args ) const { - Value grads_in = args.at(0); + bool single = args.at(0).type.dtype == DataType::dt_float32; + Value grads_in = single ? fb.convert_float64(args.at(0)) : args.at(0); Value grad_norm = fb.sqrt(fb.batch_reduce_sum(fb.square(grads_in))); Value factor = fb.min(fb.div(args.at(1), grad_norm), 1.0); Value grads_out = fb.mul(grads_in, factor); + if (single) { + grads_out = fb.convert_float32(grads_out); + } return { {"gradients_out", grads_out}, {"gradient_norm", grad_norm}, @@ -46,11 +51,6 @@ AdamOptimizer::AdamOptimizer( _eps(eps), _grad_clip_threshold(grad_clip_threshold), _weight_decay(weight_decay), - _grad_clipper( - grad_clip_threshold > 0. - ? build_runtime(GradientClipper().function(), context, false) - : nullptr - ), _one(1.0, context->device()), _loss_mean(std::numeric_limits::quiet_NaN()) { DevicePtr device = context->device(); @@ -61,6 +61,11 @@ AdamOptimizer::AdamOptimizer( } _parameter = context->reallocate_globals_contiguously(_param_names); _runtime = build_runtime(function, context); + if (grad_clip_threshold > 0.) { + _grad_clipper = build_runtime( + GradientClipper(_parameter.dtype()).function(), context, false + ); + } _exp_avg = Tensor(_parameter.dtype(), _parameter.shape(), _parameter.device()); _exp_avg.zero(); _exp_avg_sq = Tensor(_parameter.dtype(), _parameter.shape(), _parameter.device()); @@ -141,6 +146,11 @@ void AdamOptimizer::replace_function(const Function& function) { _runtime.reset(); _parameter = _context->reallocate_globals_contiguously(_param_names); _runtime = build_runtime(function, _context); + if (_grad_clipper) { + _grad_clipper = build_runtime( + GradientClipper(_parameter.dtype()).function(), _context, false + ); + } Tensor old_ea = _exp_avg, old_eas = _exp_avg_sq; _exp_avg = Tensor(_parameter.dtype(), _parameter.shape(), _parameter.device()); _exp_avg_sq = Tensor(_parameter.dtype(), _parameter.shape(), _parameter.device()); diff --git a/madspace/src/driver/tensor.cpp b/madspace/src/driver/tensor.cpp index c048bfbd2f..e954cab7d6 100644 --- a/madspace/src/driver/tensor.cpp +++ b/madspace/src/driver/tensor.cpp @@ -1,5 +1,6 @@ #include "madspace/driver/tensor.hpp" +#include #include using namespace madspace; @@ -231,6 +232,44 @@ std::size_t Tensor::init_stride() { return stride_prod * dtype_size(); } +Tensor Tensor::cast(DataType dtype) const { + Tensor source = cpu().contiguous(); + if (source.dtype() == dtype) { + return source; + } + Tensor target(dtype, source.shape()); + std::size_t count = source.shape().product(); + auto convert = [&](From* from) { + switch (dtype) { + case DataType::dt_int32: + std::copy_n(from, count, static_cast(target.data())); + break; + case DataType::dt_float32: + std::copy_n(from, count, static_cast(target.data())); + break; + case DataType::dt_float64: + std::copy_n(from, count, static_cast(target.data())); + break; + default: + throw std::invalid_argument("invalid dtype in cast"); + } + }; + switch (source.dtype()) { + case DataType::dt_int32: + convert(static_cast(source.data())); + break; + case DataType::dt_float32: + convert(static_cast(source.data())); + break; + case DataType::dt_float64: + convert(static_cast(source.data())); + break; + default: + throw std::invalid_argument("invalid dtype in cast"); + } + return target; +} + Tensor Tensor::simd_padded_literal( const LiteralValue& value, std::size_t lanes, DevicePtr device ) { diff --git a/madspace/src/gpu/device.cu b/madspace/src/gpu/device.cu index d20239177b..0d37e1a8a1 100644 --- a/madspace/src/gpu/device.cu +++ b/madspace/src/gpu/device.cu @@ -107,18 +107,33 @@ void GpuDevice::adam_step( wait_for_tensor(exp_avg); wait_for_tensor(exp_avg_sq); AsyncGpuDevice device(*this, gpuStreamPerThread, 0); - tensor_foreach_dynamic, 1, 3>( - {&gradient}, - {¶meter, &exp_avg, &exp_avg_sq}, - gradient.size(0), - device, - step_size, - beta1, - beta2, - eps, - bias_corr2_sqrt, - weight_decay - ); + if (parameter.dtype() == DataType::dt_float32) { + tensor_foreach_dynamic, 1, 3>( + {&gradient}, + {¶meter, &exp_avg, &exp_avg_sq}, + gradient.size(0), + device, + static_cast(step_size), + static_cast(beta1), + static_cast(beta2), + static_cast(eps), + static_cast(bias_corr2_sqrt), + static_cast(weight_decay) + ); + } else { + tensor_foreach_dynamic, 1, 3>( + {&gradient}, + {¶meter, &exp_avg, &exp_avg_sq}, + gradient.size(0), + device, + step_size, + beta1, + beta2, + eps, + bias_corr2_sqrt, + weight_decay + ); + } check_error(gpuStreamSynchronize(gpuStreamPerThread)); } diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 7cbaa49faf..968671c40c 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -1825,6 +1825,13 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) : if (context->global_requires_grad(name)) { _requires_grad_init.at(value.local_index) = true; _grad_global_indices.push_back(value.local_index); + // the gradients of all globals are stored in one contiguous tensor + if (_grad_global_total_size > 0 && global.dtype() != _grad_global_dtype) { + throw std::invalid_argument( + "all globals that require a gradient must have the same dtype" + ); + } + _grad_global_dtype = global.dtype(); _grad_global_shapes.push_back(global.shape()); _grad_global_total_size += global.shape().product(); } @@ -2172,7 +2179,7 @@ std::pair GpuRuntime::run_backward( } } Tensor all_global_grads( - DataType::dt_float64, + _grad_global_dtype, {_grad_global_total_size}, init_device, AllocHint::global_grad diff --git a/madspace/src/gpu/runtime.hpp b/madspace/src/gpu/runtime.hpp index 212e1f479a..15b85a0880 100644 --- a/madspace/src/gpu/runtime.hpp +++ b/madspace/src/gpu/runtime.hpp @@ -91,6 +91,7 @@ class GpuRuntime : public Runtime { SizeVec _grad_global_indices; std::vector _grad_global_shapes; std::size_t _grad_global_total_size; + DataType _grad_global_dtype = DataType::dt_float64; ContextPtr _context; ThreadResource> _streams; ThreadResource> _events; From 818d0e517c2e97c00c48816cf270f22fd6372274 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 19:16:22 +0200 Subject: [PATCH 05/24] madspace: precision argument for MLP and the networks using it MLP takes a precision (float32 or float64) for its weights and evaluation. Its input and output stay float64, with conversions at the boundary, so softmaxes, splines and Jacobians around it are unchanged. Flow, DiscreteFlow and ChannelWeightNetwork forward it as subnet_precision / precision. Tests for float32 MLP training against torch, float32 flows and saving and loading float32 globals. --- .../phasespace/channel_weight_network.hpp | 4 +- .../madspace/phasespace/discrete_flow.hpp | 4 +- madspace/include/madspace/phasespace/flow.hpp | 4 +- madspace/include/madspace/phasespace/mlp.hpp | 12 +++- .../src/phasespace/channel_weight_network.cpp | 6 +- madspace/src/phasespace/discrete_flow.cpp | 6 +- madspace/src/phasespace/flow.cpp | 11 ++-- madspace/src/phasespace/mlp.cpp | 63 +++++++++--------- madspace/src/python/madspace.cpp | 17 +++-- madspace/tests/test_float32.py | 63 +++++++++++++++--- madspace/tests/test_mlp.py | 66 ++++++++++++------- 11 files changed, 175 insertions(+), 81 deletions(-) diff --git a/madspace/include/madspace/phasespace/channel_weight_network.hpp b/madspace/include/madspace/phasespace/channel_weight_network.hpp index 23b14cafda..71c597709d 100644 --- a/madspace/include/madspace/phasespace/channel_weight_network.hpp +++ b/madspace/include/madspace/phasespace/channel_weight_network.hpp @@ -77,6 +77,7 @@ class ChannelWeightNetwork : public FunctionGenerator { * @param prefix Prefix for the trainable global names. * @param include_preprocessing If true, apply @ref MomentumPreprocessing * to the input first. + * @param precision Float type of the @ref MLP. */ ChannelWeightNetwork( std::size_t channel_count, @@ -85,7 +86,8 @@ class ChannelWeightNetwork : public FunctionGenerator { std::size_t layers = 3, MLP::Activation activation = MLP::leaky_relu, const std::string& prefix = "", - bool include_preprocessing = true + bool include_preprocessing = true, + DataType precision = DataType::dt_float64 ); /// The underlying network. diff --git a/madspace/include/madspace/phasespace/discrete_flow.hpp b/madspace/include/madspace/phasespace/discrete_flow.hpp index a4affa9bde..ec990ede60 100644 --- a/madspace/include/madspace/phasespace/discrete_flow.hpp +++ b/madspace/include/madspace/phasespace/discrete_flow.hpp @@ -45,6 +45,7 @@ class DiscreteFlow : public Mapping { * @param subnet_hidden_dim Hidden width of the @ref MLP subnetworks. * @param subnet_layers Number of layers in the subnetworks. * @param subnet_activation Activation of the subnetworks. + * @param subnet_precision Float type of the subnetworks (see @ref MLP). */ DiscreteFlow( const std::vector& option_counts, @@ -53,7 +54,8 @@ class DiscreteFlow : public Mapping { std::size_t condition_dim = 0, std::size_t subnet_hidden_dim = 32, std::size_t subnet_layers = 3, - MLP::Activation subnet_activation = MLP::leaky_relu + MLP::Activation subnet_activation = MLP::leaky_relu, + DataType subnet_precision = DataType::dt_float64 ); /// Number of categories per discrete dimension. const std::vector& option_counts() const { return _option_counts; } diff --git a/madspace/include/madspace/phasespace/flow.hpp b/madspace/include/madspace/phasespace/flow.hpp index 94eeaa3f82..4b267896a2 100644 --- a/madspace/include/madspace/phasespace/flow.hpp +++ b/madspace/include/madspace/phasespace/flow.hpp @@ -54,6 +54,7 @@ class Flow : public Mapping { * @param subnet_activation Activation of the subnetworks. * @param invert_spline Whether the spline is applied in the inverse * direction. + * @param subnet_precision Float type of the subnetworks (see @ref MLP). */ Flow( std::size_t input_dim, @@ -63,7 +64,8 @@ class Flow : public Mapping { std::size_t subnet_hidden_dim = 32, std::size_t subnet_layers = 3, MLP::Activation subnet_activation = MLP::leaky_relu, - bool invert_spline = true + bool invert_spline = true, + DataType subnet_precision = DataType::dt_float64 ); /// Number of transformed dimensions. std::size_t input_dim() const { return _input_dim; } diff --git a/madspace/include/madspace/phasespace/mlp.hpp b/madspace/include/madspace/phasespace/mlp.hpp index 75d34f0bcb..6a63f77f16 100644 --- a/madspace/include/madspace/phasespace/mlp.hpp +++ b/madspace/include/madspace/phasespace/mlp.hpp @@ -15,7 +15,9 @@ namespace madspace { * between them and a linear output layer. Used as the subnetwork of the * @ref Flow and @ref DiscreteFlow adaptive samplers [1]. The weights and biases * are compute-graph globals named with @p prefix. Call - * `initialize_globals(context)` once before use. + * `initialize_globals(context)` once before use. With a single-precision + * @p precision, the weights and the network evaluation are in float32, while the + * input and output stay in float64. * * `batch` is the leading batch dimension. * @@ -40,18 +42,23 @@ class MLP : public FunctionGenerator { * @param layers Number of layers. * @param activation Activation between the hidden layers. * @param prefix Prefix for the trainable global names. + * @param precision Float type of the weights and the evaluation, + * `dt_float32` or `dt_float64`. */ MLP(std::size_t input_dim, std::size_t output_dim, std::size_t hidden_dim = 32, std::size_t layers = 3, Activation activation = leaky_relu, - const std::string& prefix = ""); + const std::string& prefix = "", + DataType precision = DataType::dt_float64); /// Input width. std::size_t input_dim() const { return _input_dim; } /// Output width. std::size_t output_dim() const { return _output_dim; } + /// Float type of the weights and the evaluation. + DataType precision() const { return _precision; } /** * Register the network's trainable parameters on @p context. * @param context Context receiving the globals. @@ -79,6 +86,7 @@ class MLP : public FunctionGenerator { std::size_t _layers; Activation _activation; std::string _prefix; + DataType _precision; }; } // namespace madspace diff --git a/madspace/src/phasespace/channel_weight_network.cpp b/madspace/src/phasespace/channel_weight_network.cpp index abb461d094..f07e3d7cf9 100644 --- a/madspace/src/phasespace/channel_weight_network.cpp +++ b/madspace/src/phasespace/channel_weight_network.cpp @@ -33,7 +33,8 @@ ChannelWeightNetwork::ChannelWeightNetwork( std::size_t layers, MLP::Activation activation, const std::string& prefix, - bool include_preprocessing + bool include_preprocessing, + DataType precision ) : FunctionGenerator( "ChannelWeightNetwork", @@ -48,7 +49,8 @@ ChannelWeightNetwork::ChannelWeightNetwork( hidden_dim, layers, activation, - prefix + prefix, + precision ), _channel_count(channel_count), _mask_name(prefixed_name(prefix, "active_channels_mask")) {} diff --git a/madspace/src/phasespace/discrete_flow.cpp b/madspace/src/phasespace/discrete_flow.cpp index ffd896b1ff..bf5c899294 100644 --- a/madspace/src/phasespace/discrete_flow.cpp +++ b/madspace/src/phasespace/discrete_flow.cpp @@ -12,7 +12,8 @@ DiscreteFlow::DiscreteFlow( std::size_t condition_dim, std::size_t subnet_hidden_dim, std::size_t subnet_layers, - MLP::Activation subnet_activation + MLP::Activation subnet_activation, + DataType subnet_precision ) : Mapping( "DiscreteFlow", @@ -57,7 +58,8 @@ DiscreteFlow::DiscreteFlow( subnet_hidden_dim, subnet_layers, subnet_activation, - prefixed_name(prefix, std::format("subnet{}_", dim_index + 1)) + prefixed_name(prefix, std::format("subnet{}_", dim_index + 1)), + subnet_precision ); } else { _first_prob_name = prefixed_name(prefix, "prob0"); diff --git a/madspace/src/phasespace/flow.cpp b/madspace/src/phasespace/flow.cpp index 208c59cabb..ae7a15060f 100644 --- a/madspace/src/phasespace/flow.cpp +++ b/madspace/src/phasespace/flow.cpp @@ -142,7 +142,7 @@ void vegas_init( } if (!is_cpu) { - bias_global.copy_from(bias_tensor); + bias_global.copy_from(bias_tensor.cast(bias_global.dtype())); } } @@ -156,7 +156,8 @@ Flow::Flow( std::size_t subnet_hidden_dim, std::size_t subnet_layers, MLP::Activation subnet_activation, - bool invert_spline + bool invert_spline, + DataType subnet_precision ) : Mapping( "Flow", @@ -198,13 +199,15 @@ Flow::Flow( subnet_hidden_dim, subnet_layers, subnet_activation, - prefixed_name(prefix, std::format("subnet{}a", block_index + 1))), + prefixed_name(prefix, std::format("subnet{}a", block_index + 1)), + subnet_precision), MLP(indices1.size() + condition_dim, indices2.size() * (3 * bin_count + 1), subnet_hidden_dim, subnet_layers, subnet_activation, - prefixed_name(prefix, std::format("subnet{}b", block_index + 1))), + prefixed_name(prefix, std::format("subnet{}b", block_index + 1)), + subnet_precision), indices1, indices2} ); diff --git a/madspace/src/phasespace/mlp.cpp b/madspace/src/phasespace/mlp.cpp index 80cb554a9c..8da85ee264 100644 --- a/madspace/src/phasespace/mlp.cpp +++ b/madspace/src/phasespace/mlp.cpp @@ -15,16 +15,17 @@ Value build_layer( int output_dim, MLP::Activation activation, const std::string& prefix, - int layer_index + int layer_index, + DataType precision ) { auto weight = fb.global( prefixed_name(prefix, std::format("layer{}.weight", layer_index)), - DataType::dt_float64, + precision, {output_dim, input_dim} ); auto bias = fb.global( prefixed_name(prefix, std::format("layer{}.bias", layer_index)), - DataType::dt_float64, + precision, {output_dim} ); auto linear_out = fb.matmul(input, weight, bias); @@ -53,25 +54,19 @@ void initialize_layer( const std::string& prefix, int layer_index, std::optional seed, - bool zeros + bool zeros, + DataType precision ) { auto weight_name = prefixed_name(prefix, std::format("layer{}.weight", layer_index)); auto bias_name = prefixed_name(prefix, std::format("layer{}.bias", layer_index)); - auto weight_tensor_global = context->define_global( - weight_name, DataType::dt_float64, {output_dim, input_dim}, true - ); + auto weight_tensor_global = + context->define_global(weight_name, precision, {output_dim, input_dim}, true); auto bias_tensor_global = - context->define_global(bias_name, DataType::dt_float64, {output_dim}, true); - Tensor weight_tensor, bias_tensor; - bool is_cpu = context->device() == cpu_device(); - if (is_cpu) { - weight_tensor = weight_tensor_global; - bias_tensor = bias_tensor_global; - } else { - weight_tensor = Tensor(DataType::dt_float64, {1, output_dim, input_dim}); - bias_tensor = Tensor(DataType::dt_float64, {1, output_dim}); - } + context->define_global(bias_name, precision, {output_dim}, true); + // filled in double precision on the CPU, then converted and copied + Tensor weight_tensor(DataType::dt_float64, {1, output_dim, input_dim}); + Tensor bias_tensor(DataType::dt_float64, {1, output_dim}); auto weight_view = weight_tensor.view()[0]; auto bias_view = bias_tensor.view()[0]; @@ -105,10 +100,8 @@ void initialize_layer( } } - if (!is_cpu) { - weight_tensor_global.copy_from(weight_tensor); - bias_tensor_global.copy_from(bias_tensor); - } + weight_tensor_global.copy_from(weight_tensor.cast(precision)); + bias_tensor_global.copy_from(bias_tensor.cast(precision)); } } // namespace @@ -119,7 +112,8 @@ MLP::MLP( std::size_t hidden_dim, std::size_t layers, Activation activation, - const std::string& prefix + const std::string& prefix, + DataType precision ) : FunctionGenerator( "MLP", @@ -131,7 +125,11 @@ MLP::MLP( _hidden_dim(hidden_dim), _layers(layers), _activation(activation), - _prefix(prefix) { + _prefix(prefix), + _precision(precision) { + if (!is_float(precision)) { + throw std::invalid_argument("MLP precision must be a float type"); + } if (input_dim == 0) { throw std::invalid_argument("MLP input dimension cannot be 0"); } @@ -142,15 +140,16 @@ MLP::MLP( NamedVector MLP::build_function_impl(FunctionBuilder& fb, const NamedVector& args) const { + bool single = _precision == DataType::dt_float32; std::size_t dim = _input_dim; - Value x = args.at(0); + Value x = single ? fb.convert_float32(args.at(0)) : args.at(0); for (std::size_t i = 1; i < _layers; ++i) { - x = build_layer(fb, x, dim, _hidden_dim, _activation, _prefix, i); + x = build_layer(fb, x, dim, _hidden_dim, _activation, _prefix, i, _precision); dim = _hidden_dim; } - return { - {"output", build_layer(fb, x, dim, _output_dim, MLP::linear, _prefix, _layers)} - }; + Value output = + build_layer(fb, x, dim, _output_dim, MLP::linear, _prefix, _layers, _precision); + return {{"output", single ? fb.convert_float64(output) : output}}; } void MLP::initialize_globals( @@ -158,10 +157,14 @@ void MLP::initialize_globals( ) const { std::size_t dim = _input_dim; for (std::size_t i = 1; i < _layers; ++i) { - initialize_layer(context, dim, _hidden_dim, _prefix, i, seed, false); + initialize_layer( + context, dim, _hidden_dim, _prefix, i, seed, false, _precision + ); dim = _hidden_dim; } - initialize_layer(context, dim, _output_dim, _prefix, _layers, seed, true); + initialize_layer( + context, dim, _output_dim, _prefix, _layers, seed, true, _precision + ); } std::vector MLP::global_names() const { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 743c21ead5..4a9ca4bf78 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1451,17 +1451,20 @@ PYBIND11_MODULE(_madspace_py, m) { std::size_t, std::size_t, MLP::Activation, - const std::string&>(), + const std::string&, + DataType>(), py::arg("input_dim"), py::arg("output_dim"), py::arg("hidden_dim") = 32, py::arg("layers") = 3, py::arg("activation") = MLP::leaky_relu, py::arg("prefix") = "", + py::arg("precision") = DataType::dt_float64, pydoc::doc("MLP::MLP") ) .def("input_dim", &MLP::input_dim, pydoc::doc("MLP::input_dim")) .def("output_dim", &MLP::output_dim, pydoc::doc("MLP::output_dim")) + .def("precision", &MLP::precision, pydoc::doc("MLP::precision")) .def( "initialize_globals", &MLP::initialize_globals, @@ -1486,7 +1489,8 @@ PYBIND11_MODULE(_madspace_py, m) { std::size_t, std::size_t, MLP::Activation, - bool>(), + bool, + DataType>(), py::arg("input_dim"), py::arg("condition_dim") = 0, py::arg("prefix") = "", @@ -1495,6 +1499,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("subnet_layers") = 3, py::arg("subnet_activation") = MLP::leaky_relu, py::arg("invert_spline") = true, + py::arg("subnet_precision") = DataType::dt_float64, pydoc::doc("Flow::Flow") ) .def("input_dim", &Flow::input_dim, pydoc::doc("Flow::input_dim")) @@ -1573,7 +1578,8 @@ PYBIND11_MODULE(_madspace_py, m) { std::size_t, MLP::Activation, const std::string&, - bool>(), + bool, + DataType>(), py::arg("channel_count"), py::arg("particle_count"), py::arg("hidden_dim") = 32, @@ -1581,6 +1587,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("activation") = MLP::leaky_relu, py::arg("prefix") = "", py::arg("include_preprocessing") = true, + py::arg("precision") = DataType::dt_float64, pydoc::doc("ChannelWeightNetwork::ChannelWeightNetwork") ) .def("mlp", &ChannelWeightNetwork::mlp, pydoc::doc("ChannelWeightNetwork::mlp")) @@ -1650,7 +1657,8 @@ PYBIND11_MODULE(_madspace_py, m) { std::size_t, std::size_t, std::size_t, - MLP::Activation>(), + MLP::Activation, + DataType>(), py::arg("option_counts"), py::arg("prefix") = "", py::arg("dims_with_prior") = std::vector{}, @@ -1658,6 +1666,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("subnet_hidden_dim") = 32, py::arg("subnet_layers") = 3, py::arg("subnet_activation") = MLP::leaky_relu, + py::arg("subnet_precision") = DataType::dt_float64, pydoc::doc("DiscreteFlow::DiscreteFlow") ) .def( diff --git a/madspace/tests/test_float32.py b/madspace/tests/test_float32.py index 331c2aef70..7b0e0e42e8 100644 --- a/madspace/tests/test_float32.py +++ b/madspace/tests/test_float32.py @@ -6,7 +6,8 @@ import madspace as ms from madspace.torch import FunctionModule -torch.manual_seed(1234) +# a local generator, so that the global RNG state seen by other tests is unchanged +GEN = torch.Generator().manual_seed(1234) N = 103 DIMS_IN = 7 @@ -43,10 +44,12 @@ def test_activation(name, precision): fb.output(0, getattr(fb, name)(fb.input(0))) module = FunctionModule(fb.function()) - x = (3 * torch.randn((N, DIMS_IN), dtype=torch_dtype)).requires_grad_() + x = ( + 3 * torch.randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN) + ).requires_grad_() out = module(x) assert out.dtype == torch_dtype - out_grad = torch.randn_like(out) + out_grad = torch.randn(out.shape, dtype=out.dtype, generator=GEN) (out * out_grad).sum().backward() x_ref = x.detach().clone().requires_grad_() @@ -72,13 +75,15 @@ def test_matmul(precision): fb.output(0, fb.matmul(fb.input(0), fb.input(1), fb.input(2))) module = FunctionModule(fb.function()) - x = torch.randn((N, DIMS_IN), dtype=torch_dtype, requires_grad=True) + x = torch.randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN, requires_grad=True) # single inputs are passed without batch dimension - w = torch.randn((DIMS_OUT, DIMS_IN), dtype=torch_dtype, requires_grad=True) - b = torch.randn((DIMS_OUT,), dtype=torch_dtype, requires_grad=True) + w = torch.randn( + (DIMS_OUT, DIMS_IN), dtype=torch_dtype, generator=GEN, requires_grad=True + ) + b = torch.randn((DIMS_OUT,), dtype=torch_dtype, generator=GEN, requires_grad=True) y = module(x, w, b) assert y.dtype == torch_dtype - y_grad = torch.randn_like(y) + y_grad = torch.randn(y.shape, dtype=y.dtype, generator=GEN) (y * y_grad).sum().backward() refs = [t.detach().clone().requires_grad_() for t in (x, w, b)] @@ -102,7 +107,7 @@ def test_convert_roundtrip(): fb.output(1, fb.convert_float64(fb.relu(single))) module = FunctionModule(fb.function()) - x = torch.randn((N, 3), dtype=torch.float64, requires_grad=True) + x = torch.randn((N, 3), dtype=torch.float64, generator=GEN, requires_grad=True) single_out, out = module(x) assert single_out.dtype == torch.float32 assert out.dtype == torch.float64 @@ -120,5 +125,45 @@ def test_literal_converted_for_float32(): # float64 literals are converted when the instruction dispatches to float32 fb.output(0, fb.matmul(fb.input(0), ms.Value([2.0], [1, 1]), ms.Value([1.0]))) module = FunctionModule(fb.function()) - x = torch.randn((N, 1), dtype=torch.float32) + x = torch.randn((N, 1), dtype=torch.float32, generator=GEN) torch.testing.assert_close(module(x), 2 * x + 1) + + +def test_flow_float32_subnets(tmp_path): + """A flow with float32 subnetworks matches the float64 one with the same + (rounded) weights; the Jacobian and the output stay in float64.""" + ctx = ms.default_context() + flows = { + dtype: ms.Flow( + 4, prefix=f"fp32_test_{name}", subnet_precision=dtype, invert_spline=False + ) + for name, dtype in [("f32", ms.DataType.float32), ("f64", ms.DataType.float64)] + } + for flow in flows.values(): + flow.initialize_globals(ctx, 42) + # randomize all weights (the last layers start at zero) and round to float32 + names_f64 = [n for n in ctx.global_names() if n.startswith("fp32_test_f64")] + for name in names_f64: + weights = ctx.get_global(name).torch() + weights[:] = ( + 0.3 + * torch.randn(weights.shape, generator=GEN, dtype=torch.float32).double() + ) + ctx.get_global(name.replace("f64", "f32")).torch()[:] = weights.float() + + r = torch.rand((N, 4), dtype=torch.float64, generator=GEN) + out32, det32 = flows[ms.DataType.float32].map_forward([r]) + out64, det64 = flows[ms.DataType.float64].map_forward([r]) + assert out32.dtype == det32.dtype == torch.float64 + torch.testing.assert_close(out32, out64, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(det32, det64, rtol=1e-4, atol=1e-5) + + # float32 globals survive saving and loading + ctx.save_globals(str(tmp_path)) + loaded = ms.Context() + loaded.load_globals(str(tmp_path)) + for name in names_f64: + name32 = name.replace("f64", "f32") + global32 = loaded.get_global(name32).torch() + assert global32.dtype == torch.float32 + torch.testing.assert_close(global32, ctx.get_global(name32).torch()) diff --git a/madspace/tests/test_mlp.py b/madspace/tests/test_mlp.py index 56a9f4a6ed..ca28a88844 100644 --- a/madspace/tests/test_mlp.py +++ b/madspace/tests/test_mlp.py @@ -11,9 +11,20 @@ torch.manual_seed(3210) +PRECISIONS = { + "float64": (ms.DataType.float64, torch.float64, 1e-6), + "float32": (ms.DataType.float32, torch.float32, 1e-3), +} + + +@pytest.fixture(params=PRECISIONS.keys()) +def precision(request): + return PRECISIONS[request.param] + + @pytest.fixture -def mlp(): - return ms.MLP(10, 1, 32, 3, ms.MLP.leaky_relu, "") +def mlp(precision): + return ms.MLP(10, 1, 32, 3, ms.MLP.leaky_relu, "", precision[0]) @pytest.fixture @@ -31,9 +42,16 @@ def test_properties(): mlp = ms.MLP(10, 1, 32, 3, ms.MLP.leaky_relu, "") assert mlp.input_dim() == 10 assert mlp.output_dim() == 1 + assert mlp.precision() == ms.DataType.float64 + mlp = ms.MLP(10, 1, precision=ms.DataType.float32) + assert mlp.precision() == ms.DataType.float32 + with pytest.raises(Exception): + ms.MLP(10, 1, precision=ms.DataType.int32) -def test_initialization(mlp_module): +def test_initialization(mlp_module, precision): + for param in mlp_module.global_params.values(): + assert param.dtype == precision[1] assert torch.all(mlp_module.global_params["layer1:weight"] != 0) assert torch.all(mlp_module.global_params["layer1:bias"] != 0) assert torch.all(mlp_module.global_params["layer2:weight"] != 0) @@ -70,14 +88,15 @@ def test_activation(activation): assert grad_me == approx(grad_torch) -def test_training(mlp, mlp_module, module_context): +def test_training(mlp, mlp_module, module_context, precision): + _, dtype, tol = precision mlp_torch = nn.Sequential( nn.Linear(10, 32), nn.LeakyReLU(), nn.Linear(32, 32), nn.LeakyReLU(), nn.Linear(32, 1), - ) + ).to(dtype) with torch.no_grad(): for i in range(3): @@ -105,8 +124,12 @@ def test_training(mlp, mlp_module, module_context): opt_ms = torch.optim.Adam(mlp_module.parameters(), lr=1e-2) opt_torch = torch.optim.Adam(mlp_torch.parameters(), lr=1e-2) + def close(a, b): + return a == approx(b, rel=tol, abs=tol) + for i in range(10): - x = torch.randn((128, 10)) + # the input and the loss are in float64, only the network uses dtype + x = torch.randn((128, 10), dtype=torch.float64) y = x.square().sum(dim=1, keepdim=True) loss_ms = (y - mlp_module(x)).square().mean() @@ -114,33 +137,26 @@ def test_training(mlp, mlp_module, module_context): loss_ms.backward() opt_ms.step() - loss_torch = (y - mlp_torch(x)).square().mean() + loss_torch = (y - mlp_torch(x.to(dtype)).double()).square().mean() opt_torch.zero_grad() loss_torch.backward() opt_torch.step() loss_builtin, mlp_out = builtin_opt.step([x]) - assert loss_ms.item() == approx(loss_torch.item()) - assert loss_builtin.torch().item() == approx(loss_torch.item()) + assert close(loss_ms.item(), loss_torch.item()) + assert close(loss_builtin.torch().item(), loss_torch.item()) with torch.no_grad(): for i in range(3): w_torch = mlp_torch[2 * i].weight.numpy() b_torch = mlp_torch[2 * i].bias.numpy() - assert mlp_module.global_params[f"layer{i+1}:weight"][ - 0 - ].numpy() == approx(w_torch) - assert mlp_module.global_params[f"layer{i+1}:bias"][ - 0 - ].numpy() == approx(b_torch) - assert mlp_module.global_params[f"layer{i+1}:weight"].grad[ - 0 - ].numpy() == approx(mlp_torch[2 * i].weight.grad.numpy()) - assert mlp_module.global_params[f"layer{i+1}:bias"].grad[ - 0 - ].numpy() == approx(mlp_torch[2 * i].bias.grad.numpy()) - assert ctx.get_global(f"layer{i+1}.weight").numpy()[0] == approx( - w_torch - ) - assert ctx.get_global(f"layer{i+1}.bias").numpy()[0] == approx(b_torch) + w_ms = mlp_module.global_params[f"layer{i+1}:weight"] + b_ms = mlp_module.global_params[f"layer{i+1}:bias"] + assert w_ms.dtype == dtype + assert close(w_ms[0].numpy(), w_torch) + assert close(b_ms[0].numpy(), b_torch) + assert close(w_ms.grad[0].numpy(), mlp_torch[2 * i].weight.grad.numpy()) + assert close(b_ms.grad[0].numpy(), mlp_torch[2 * i].bias.grad.numpy()) + assert close(ctx.get_global(f"layer{i+1}.weight").numpy()[0], w_torch) + assert close(ctx.get_global(f"layer{i+1}.bias").numpy()[0], b_torch) From 52b1683edccc0efe3cc3a2ac767b70c15497369d Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 19:42:20 +0200 Subject: [PATCH 06/24] madspace: integer gradient views for gather and select The backward kernels of gather and select declared the (unused) gradient of their integer index input as a float view, while the runtime allocates it as an int32 tensor. The view dtype check rejected this in the backward pass of float32 flows. Also test the gradients of a flow with float32 subnetworks. --- madspace/src/kernels/util.hpp | 4 ++-- madspace/tests/test_float32.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/madspace/src/kernels/util.hpp b/madspace/src/kernels/util.hpp index 5254ce0d7c..5a9cd90c8f 100644 --- a/madspace/src/kernels/util.hpp +++ b/madspace/src/kernels/util.hpp @@ -37,7 +37,7 @@ template KERNELSPEC void backward_kernel_gather( IIn index, FIn output_grad, - FOut index_grad, + IOut index_grad, FOut choices_grad ) { choices_grad.scatter_add(index, output_grad); @@ -78,7 +78,7 @@ KERNELSPEC void backward_kernel_select( ISIn indices, FIn output_grad, FOut input_grad, - FOut indices_grad + IOut indices_grad ) { for (std::size_t i = 0; i < indices.size(); ++i) { input_grad.scatter_add(indices[i], output_grad[i]); diff --git a/madspace/tests/test_float32.py b/madspace/tests/test_float32.py index 7b0e0e42e8..5b2f6ce659 100644 --- a/madspace/tests/test_float32.py +++ b/madspace/tests/test_float32.py @@ -158,6 +158,19 @@ def test_flow_float32_subnets(tmp_path): torch.testing.assert_close(out32, out64, rtol=1e-4, atol=1e-5) torch.testing.assert_close(det32, det64, rtol=1e-4, atol=1e-5) + # gradients with respect to the input and the float32 weights + grads = {} + for dtype, flow in flows.items(): + module = FunctionModule(flow.forward_function(), ctx) + r_grad = r.clone().requires_grad_() + out, det = module(r_grad) + (out.sum() + det.log().sum()).backward() + weight = next(p for n, p in module.global_params.items() if "weight" in n) + assert weight.grad.dtype == weight.dtype + grads[dtype] = (r_grad.grad, weight.grad.double()) + for g32, g64 in zip(grads[ms.DataType.float32], grads[ms.DataType.float64]): + torch.testing.assert_close(g32, g64, rtol=1e-3, atol=1e-4) + # float32 globals survive saving and loading ctx.save_globals(str(tmp_path)) loaded = ms.Context() From 61dc6a445bf85ff4ff9d3910520148a0214db0f6 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 19:42:22 +0200 Subject: [PATCH 07/24] mg7: run card setting madnis.nn_float_precision Selects float32 or float64 for the weights and evaluation of all MadNIS networks (flow, discrete flows, channel weight network). One setting for all of them, as the optimizer trains them together. --- docs/source/run_card.rst | 5 +++++ madgraph/iolibs/template_files/mg7/launch.py | 13 +++++++++++++ madgraph/iolibs/template_files/mg7/run_card.toml | 3 +++ madgraph/various/banner.py | 3 +++ tests/input_files/mg7_run_card.toml | 3 +++ tests/unit_tests/various/test_banner.py | 4 ++++ 6 files changed, 31 insertions(+) diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index 15d3b44b8f..4e34077c38 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -557,3 +557,8 @@ function. The activations are ``relu``, ``leaky_relu``, ``elu``, ``gelu``, ``sig ``cwnet_hidden_dim`` (``"auto"``), ``cwnet_layers`` (``3``), ``cwnet_activation`` (``"leaky_relu"``) Size and activation of the channel weight network. + +``nn_float_precision`` (``"float64"``) + Floating-point precision of the weights and the evaluation of all networks, + ``float32`` or ``float64``. The rest of the computation, including the network + outputs, softmaxes and Jacobians, stays in double precision. diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index e5a961c321..fdfaf17051 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -2923,6 +2923,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: subnet_layers=self.madnis_settings["flow_layers"], subnet_activation=self.activation(madnis_args["flow_activation"]), invert_spline=madnis_args["flow_invert_spline"], + subnet_precision=self.nn_precision(), ) if channel.adaptive_mapping is None: flow.initialize_globals(self.process.contexts[0], seed) @@ -2945,6 +2946,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], subnet_layers=madnis_args["discrete_layers"], subnet_activation=self.activation(madnis_args["discrete_activation"]), + subnet_precision=self.nn_precision(), ) discrete_sym.initialize_globals(self.process.contexts[0], seed) cond_dim += perm_count @@ -2959,6 +2961,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], subnet_layers=madnis_args["discrete_layers"], subnet_activation=self.activation(madnis_args["discrete_activation"]), + subnet_precision=self.nn_precision(), ) discrete_flavor.initialize_globals(self.process.contexts[0], seed) @@ -3039,6 +3042,7 @@ def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork: layers=madnis_args["cwnet_layers"], activation=self.activation(madnis_args["cwnet_activation"]), prefix=f"subproc{self.subproc_id}.cwnet", + precision=self.nn_precision(), ) cwnet.initialize_globals( self.process.contexts[0], self.process.run_seed @@ -3056,6 +3060,15 @@ def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: else: raise ValueError(f"Invalid t-channel mode '{name}'") + def nn_precision(self) -> ms.DataType: + # one precision for all networks: the optimizer trains all of them together + name = self.process.run_card["madnis"]["nn_float_precision"] + precisions = {"float32": ms.DataType.float32, "float64": ms.DataType.float64} + if name in precisions: + return precisions[name] + else: + raise ValueError(f"Invalid network precision '{name}'") + def activation(self, name: str) -> ms.MLP.Activation: activations = { "relu": ms.MLP.relu, diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index c0cacf45ca..5eeb155e7c 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -149,6 +149,9 @@ cwnet_hidden_dim = %(madnis.cwnet_hidden_dim)s cwnet_layers = %(madnis.cwnet_layers)s cwnet_activation = %(madnis.cwnet_activation)s # options: relu, leaky_relu, elu, gelu, sigmoid, softplus +# precision of all networks +nn_float_precision = %(madnis.nn_float_precision)s # options: float32, float64 + # training parameters train_batches = %(madnis.train_batches)s log_interval = %(madnis.log_interval)s diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index e591ba1958..51e4663268 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6913,6 +6913,9 @@ def default_setup(self): self.add_toml_param('madnis', 'cwnet_activation', "leaky_relu", allowed=['relu', 'leaky_relu', 'elu', 'gelu', 'sigmoid', 'softplus'], comment="activation function of the channel weight network") + self.add_toml_param('madnis', 'nn_float_precision', "float64", + allowed=['float32', 'float64'], + comment="floating-point precision of the weights and the evaluation of all networks") self.add_toml_param('madnis', 'train_batches', 1000, comment="number of training batches") self.add_toml_param('madnis', 'log_interval', 100, diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml index 23f1369aab..a36e8ff1a4 100644 --- a/tests/input_files/mg7_run_card.toml +++ b/tests/input_files/mg7_run_card.toml @@ -105,6 +105,9 @@ cwnet_hidden_dim = 64 cwnet_layers = 3 cwnet_activation = "leaky_relu" # options: relu, leaky_relu, elu, gelu, sigmoid, softplus +# precision of all networks +nn_float_precision = "float64" # options: float32, float64 + # training parameters train_batches = 1000 log_interval = 100 diff --git a/tests/unit_tests/various/test_banner.py b/tests/unit_tests/various/test_banner.py index 0e8e6d9f28..818887d8af 100755 --- a/tests/unit_tests/various/test_banner.py +++ b/tests/unit_tests/various/test_banner.py @@ -1862,6 +1862,7 @@ def test_defaults_and_section_access(self): self.assertEqual(rc['beam']['e_cm'], 13000.0) self.assertIs(rc['vegas']['enable'], True) self.assertEqual(rc['madnis']['enable'], 'auto') + self.assertEqual(rc['madnis']['nn_float_precision'], 'float64') self.assertEqual(rc['phasespace']['mode'], 'multichannel') # the "enable" key colliding between [vegas] and [madnis] is stored # independently, so editing one does not affect the other @@ -1887,17 +1888,20 @@ def test_write_is_valid_toml(self): self.assertEqual(data['multiparticles']['photon'], [22]) self.assertEqual(data['cuts']['sqrt_s'], {'min': 0.0}) self.assertEqual(data['madnis']['adam_eps'], 1e-8) + self.assertEqual(data['madnis']['nn_float_precision'], 'float64') def test_round_trip(self): """write() then read() preserves every value""" rc = bannermod.RunCardMG7() rc['generation']['events'] = 250 rc['phasespace']['mode'] = 'flat' + rc['madnis']['nn_float_precision'] = 'float32' out = io.StringIO() rc.write(out, template=self.template) rc2 = bannermod.RunCardMG7(out.getvalue()) self.assertEqual(rc2['generation']['events'], 250) self.assertEqual(rc2['phasespace']['mode'], 'flat') + self.assertEqual(rc2['madnis']['nn_float_precision'], 'float32') self.assertEqual(rc2['multiparticles']['lepton'], rc['multiparticles']['lepton']) self.assertEqual(rc2['cuts']['jet-delta_r'], {'min': 0.4}) From f002872c31ed57853ac76a011404c836fb41ed88 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 19:44:15 +0200 Subject: [PATCH 08/24] madspace: build OpenBLAS with single precision The float32 matmul calls sgemm_ and sgemv_, which the bundled OpenBLAS did not build. --- madspace/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index e334c6c1da..9abf9ce372 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -132,7 +132,7 @@ if(ENABLE_OPENBLAS) -DCMAKE_INSTALL_LIBDIR=lib -DBUILD_SHARED_LIBS=OFF -DBUILD_WITHOUT_LAPACK=ON - -DBUILD_SINGLE=OFF + -DBUILD_SINGLE=ON -DBUILD_DOUBLE=ON -DBUILD_COMPLEX=OFF -DBUILD_COMPLEX16=OFF From 399349e1cc9e849063078f03c998208334e52079 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 20:04:43 +0200 Subject: [PATCH 09/24] madspace: stub location of DataType default arguments The precision arguments default to DataType.float64, which pybind11_stubgen can only render with an enum class location. --- madspace/generate_pyi.py | 1 + 1 file changed, 1 insertion(+) diff --git a/madspace/generate_pyi.py b/madspace/generate_pyi.py index c1581c1d6b..e2067efb7b 100644 --- a/madspace/generate_pyi.py +++ b/madspace/generate_pyi.py @@ -57,6 +57,7 @@ def main(): "--enum-class-locations=Activation:MLP.Activation", "--enum-class-locations=CutMode:Cuts.CutMode", "--enum-class-locations=LRSchedule:AdamOptimizer.LRSchedule", + "--enum-class-locations=DataType:DataType", # Fail the build on unresolvable names/expressions (e.g. an enum # default value needing its own --enum-class-locations entry) # instead of silently emitting a stub with holes in it. From 45b28987e1eec51699dd75521df3e8810be98090 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 09:58:07 +0200 Subject: [PATCH 10/24] madspace: float32 gradients of structural ops, batch sum in GPU broadcast add --- madspace/src/gpu/device.cu | 148 ++++++++++++++++------------ madspace/src/kernels/operations.hpp | 22 +++-- madspace/tests/test_float32.py | 44 +++++++-- 3 files changed, 135 insertions(+), 79 deletions(-) diff --git a/madspace/src/gpu/device.cu b/madspace/src/gpu/device.cu index 0d37e1a8a1..3b19d64a56 100644 --- a/madspace/src/gpu/device.cu +++ b/madspace/src/gpu/device.cu @@ -13,6 +13,50 @@ void wait_for_tensor(const Tensor& tensor) { check_error(gpuStreamSynchronize(reinterpret_cast(*stream))); } } + +// calls func with the kernel policy of a float dtype +template +void dispatch_float(DataType dtype, const char* op, Func&& func) { + switch (dtype) { + case DataType::dt_float64: + func.template operator()(); + break; + case DataType::dt_float32: + func.template operator()(); + break; + default: + throw std::runtime_error(std::string("invalid dtype in ") + op); + } +} + +// sums a contiguous (batch_size, n) source over the batch dimension +template +__global__ void +kernel_batch_sum(std::size_t batch_size, std::size_t n, const F* source, F* output) { + std::size_t i = blockDim.x * blockIdx.x + threadIdx.x; + if (i >= n) { + return; + } + F sum = 0; + for (std::size_t j = 0; j < batch_size; ++j) { + sum += source[j * n + i]; + } + output[i] = sum; +} + +template +void launch_batch_sum(const Tensor& source, Tensor& sum, gpuStream_t stream) { + std::size_t n = sum.shape().product(); + launch_kernel( + kernel_batch_sum, + n, + stream, + source.size(0), + n, + static_cast(source.data()), + static_cast(sum.data()) + ); +} } // namespace std::pair GpuDevice::allocate(std::size_t size, AllocHint hint) const { @@ -107,33 +151,21 @@ void GpuDevice::adam_step( wait_for_tensor(exp_avg); wait_for_tensor(exp_avg_sq); AsyncGpuDevice device(*this, gpuStreamPerThread, 0); - if (parameter.dtype() == DataType::dt_float32) { - tensor_foreach_dynamic, 1, 3>( - {&gradient}, - {¶meter, &exp_avg, &exp_avg_sq}, - gradient.size(0), - device, - static_cast(step_size), - static_cast(beta1), - static_cast(beta2), - static_cast(eps), - static_cast(bias_corr2_sqrt), - static_cast(weight_decay) - ); - } else { - tensor_foreach_dynamic, 1, 3>( + dispatch_float(parameter.dtype(), "adam_step", [&]() { + using F = typename T::FVal; + tensor_foreach_dynamic, 1, 3>( {&gradient}, {¶meter, &exp_avg, &exp_avg_sq}, gradient.size(0), device, - step_size, - beta1, - beta2, - eps, - bias_corr2_sqrt, - weight_decay + static_cast(step_size), + static_cast(beta1), + static_cast(beta2), + static_cast(eps), + static_cast(bias_corr2_sqrt), + static_cast(weight_decay) ); - } + }); check_error(gpuStreamSynchronize(gpuStreamPerThread)); } @@ -326,25 +358,17 @@ void AsyncGpuDevice::tensor_copy(const Tensor& source, Tensor& target) const { if (source.dtype() != target.dtype()) { throw std::runtime_error("invalid dtype in copy"); } - switch (source.dtype()) { - case DataType::dt_float64: - tensor_foreach_dynamic, 1, 1>( - {&source}, {&target}, target.size(0), *this - ); - break; - case DataType::dt_float32: - tensor_foreach_dynamic, 1, 1>( - {&source}, {&target}, target.size(0), *this - ); - break; - case DataType::dt_int32: + if (source.dtype() == DataType::dt_int32) { tensor_foreach_dynamic, 1, 1>( {&source}, {&target}, target.size(0), *this ); - break; - default: - throw std::runtime_error("invalid dtype in copy"); + return; } + dispatch_float(source.dtype(), "copy", [&]() { + tensor_foreach_dynamic, 1, 1>( + {&source}, {&target}, target.size(0), *this + ); + }); } void AsyncGpuDevice::tensor_zero(Tensor& tensor) const { @@ -355,39 +379,41 @@ void AsyncGpuDevice::tensor_zero(Tensor& tensor) const { check_error(gpuMemsetAsync(tensor.data(), 0, tensor.byte_size(), _stream)); return; } - switch (tensor.dtype()) { - case DataType::dt_float64: - tensor_foreach_dynamic, 1, 1>( - {&tensor}, {&tensor}, tensor.size(0), *this - ); - break; - case DataType::dt_float32: - tensor_foreach_dynamic, 1, 1>( - {&tensor}, {&tensor}, tensor.size(0), *this - ); - break; - default: + if (tensor.dtype() == DataType::dt_int32) { tensor_foreach_dynamic, 1, 1>( {&tensor}, {&tensor}, tensor.size(0), *this ); - break; + return; } + dispatch_float(tensor.dtype(), "zero", [&]() { + tensor_foreach_dynamic, 1, 1>( + {&tensor}, {&tensor}, tensor.size(0), *this + ); + }); } void AsyncGpuDevice::tensor_add(const Tensor& source, Tensor& target) const { - if (source.dtype() == DataType::dt_float32 && - target.dtype() == DataType::dt_float32) { - tensor_foreach_dynamic, 1, 1>( - {&source}, {&target}, target.size(0), *this - ); - } else if (source.dtype() == DataType::dt_float64 && - target.dtype() == DataType::dt_float64) { - tensor_foreach_dynamic, 1, 1>( - {&source}, {&target}, target.size(0), *this - ); - } else { + if (source.dtype() != target.dtype()) { throw std::runtime_error("invalid dtype in add"); } + if (target.size(0) == 1 && source.size(0) > 1) { + // a broadcast target (e.g. the gradient of a global) accumulates the + // whole source batch, as on the CPU + auto source_contig = source.contiguous(*this, AllocHint::temporary); + Tensor sum(target.dtype(), target.shape(), *this, AllocHint::temporary); + dispatch_float(source.dtype(), "add", [&]() { + launch_batch_sum(source_contig, sum, _stream); + }); + tensor_add(sum, target); + sum.reset(*this); + source_contig.reset(*this); + return; + } + dispatch_float(source.dtype(), "add", [&]() { + tensor_foreach_dynamic, 1, 1>( + {&source}, {&target}, target.size(0), *this + ); + }); } void AsyncGpuDevice::tensor_cpu(const Tensor& source, Tensor& target) const { diff --git a/madspace/src/kernels/operations.hpp b/madspace/src/kernels/operations.hpp index 1bd937cdbc..1937c8ae50 100644 --- a/madspace/src/kernels/operations.hpp +++ b/madspace/src/kernels/operations.hpp @@ -103,6 +103,11 @@ void backward_batch_foreach( foreach_func(args, input_grads, batch_size, device); } +// gradients have the dtype of float inputs, integer inputs get float64 gradients +inline DataType grad_dtype(const Tensor& input) { + return is_float(input.dtype()) ? input.dtype() : DataType::dt_float64; +} + template void op_stack(const I& instruction, TensorVec& locals, const D& device) { auto& first_shape = locals[instruction.input_indices[0]].shape(); @@ -128,13 +133,12 @@ template void backward_op_stack( const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device ) { - // TODO: differentiate integer and float here (also other backwards) auto& output_grad = local_grads[instruction.output_indices[0]]; for (std::size_t i = 0; std::size_t input_index : instruction.input_indices) { auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float64, + grad_dtype(locals[input_index]), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[i] @@ -185,7 +189,7 @@ void backward_op_unstack( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float64, + grad_dtype(locals[input_index]), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] @@ -212,10 +216,10 @@ void backward_op_pop( const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device ) { auto input_index = instruction.input_indices[0]; - auto input_grad = local_grads[input_index]; + auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float64, + grad_dtype(locals[input_index]), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] @@ -268,7 +272,7 @@ void backward_op_batch_cat( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float64, + grad_dtype(locals[input_index]), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[i] @@ -304,7 +308,7 @@ void backward_op_batch_split( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float64, + grad_dtype(locals[input_index]), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] @@ -356,7 +360,7 @@ void backward_op_cat( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float64, + grad_dtype(locals[input_index]), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[i] @@ -470,7 +474,7 @@ void backward_op_rqs_reshape( auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - DataType::dt_float64, + grad_dtype(locals[input_index]), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] diff --git a/madspace/tests/test_float32.py b/madspace/tests/test_float32.py index 5b2f6ce659..8f089374c8 100644 --- a/madspace/tests/test_float32.py +++ b/madspace/tests/test_float32.py @@ -132,7 +132,8 @@ def test_literal_converted_for_float32(): def test_flow_float32_subnets(tmp_path): """A flow with float32 subnetworks matches the float64 one with the same (rounded) weights; the Jacobian and the output stay in float64.""" - ctx = ms.default_context() + # a private context, so that no globals are left in the default one + ctx = ms.Context() flows = { dtype: ms.Flow( 4, prefix=f"fp32_test_{name}", subnet_precision=dtype, invert_spline=False @@ -151,23 +152,21 @@ def test_flow_float32_subnets(tmp_path): ) ctx.get_global(name.replace("f64", "f32")).torch()[:] = weights.float() + # outputs and gradients with respect to the input and the float32 weights r = torch.rand((N, 4), dtype=torch.float64, generator=GEN) - out32, det32 = flows[ms.DataType.float32].map_forward([r]) - out64, det64 = flows[ms.DataType.float64].map_forward([r]) - assert out32.dtype == det32.dtype == torch.float64 - torch.testing.assert_close(out32, out64, rtol=1e-4, atol=1e-5) - torch.testing.assert_close(det32, det64, rtol=1e-4, atol=1e-5) - - # gradients with respect to the input and the float32 weights - grads = {} + outputs, grads = {}, {} for dtype, flow in flows.items(): module = FunctionModule(flow.forward_function(), ctx) r_grad = r.clone().requires_grad_() out, det = module(r_grad) + assert out.dtype == det.dtype == torch.float64 (out.sum() + det.log().sum()).backward() weight = next(p for n, p in module.global_params.items() if "weight" in n) assert weight.grad.dtype == weight.dtype + outputs[dtype] = (out.detach(), det.detach()) grads[dtype] = (r_grad.grad, weight.grad.double()) + for o32, o64 in zip(outputs[ms.DataType.float32], outputs[ms.DataType.float64]): + torch.testing.assert_close(o32, o64, rtol=1e-4, atol=1e-5) for g32, g64 in zip(grads[ms.DataType.float32], grads[ms.DataType.float64]): torch.testing.assert_close(g32, g64, rtol=1e-3, atol=1e-4) @@ -180,3 +179,30 @@ def test_flow_float32_subnets(tmp_path): global32 = loaded.get_global(name32).torch() assert global32.dtype == torch.float32 torch.testing.assert_close(global32, ctx.get_global(name32).torch()) + + +def test_structural_ops_backward(): + """stack, unstack, cat, batch_cat and batch_split propagate float32 gradients.""" + in_type = batch_type(ms.DataType.float32, [3]) + fb = ms.FunctionBuilder( + ms.NamedTypes([("in", in_type)]), + ms.NamedTypes([("out", batch_type(ms.DataType.float32, [12]))]), + ) + x = fb.input(0) + a, b = fb.unstack(fb.stack([x, fb.relu(x)])) + c = fb.cat([a, b]) + combined, counts = fb.batch_cat([c, c]) + fb.output(0, fb.cat(fb.batch_split(combined, counts))) + module = FunctionModule(fb.function()) + + x = torch.randn((N, 3), dtype=torch.float32, generator=GEN, requires_grad=True) + out = module(x) + out_grad = torch.randn(out.shape, dtype=out.dtype, generator=GEN) + (out * out_grad).sum().backward() + + x_ref = x.detach().clone().requires_grad_() + c_ref = torch.cat([x_ref, torch.relu(x_ref)], dim=1) + out_ref = torch.cat([c_ref, c_ref], dim=1) + (out_ref * out_grad).sum().backward() + torch.testing.assert_close(out, out_ref) + torch.testing.assert_close(x.grad, x_ref.grad) From 1d2311bbfa7459a52e2fd8be40219a6370207118 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 09:58:11 +0200 Subject: [PATCH 11/24] madspace: dispatch helper for float device ops, use dtype_of, review cleanups --- madgraph/iolibs/template_files/mg7/launch.py | 8 +- madspace/include/madspace/compgraphs/type.hpp | 2 +- madspace/src/cpu/device.cpp | 132 ++++++------------ madspace/src/driver/tensor.cpp | 3 +- 4 files changed, 50 insertions(+), 95 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index fdfaf17051..e85b93952f 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -3062,12 +3062,8 @@ def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: def nn_precision(self) -> ms.DataType: # one precision for all networks: the optimizer trains all of them together - name = self.process.run_card["madnis"]["nn_float_precision"] - precisions = {"float32": ms.DataType.float32, "float64": ms.DataType.float64} - if name in precisions: - return precisions[name] - else: - raise ValueError(f"Invalid network precision '{name}'") + # the run card only allows float32 and float64 + return getattr(ms.DataType, self.process.run_card["madnis"]["nn_float_precision"]) def activation(self, name: str) -> ms.MLP.Activation: activations = { diff --git a/madspace/include/madspace/compgraphs/type.hpp b/madspace/include/madspace/compgraphs/type.hpp index dcded6f22a..826b7a47ed 100644 --- a/madspace/include/madspace/compgraphs/type.hpp +++ b/madspace/include/madspace/compgraphs/type.hpp @@ -258,7 +258,7 @@ struct Value { template Value(const std::vector& values, const std::vector& shape = {}) : type{ - std::is_same_v ? DataType::dt_int32 : DataType::dt_float64, + dtype_of(), BatchSize::one, shape.size() == 0 ? std::vector{static_cast(values.size())} : shape diff --git a/madspace/src/cpu/device.cpp b/madspace/src/cpu/device.cpp index 9c4d8f3185..1ee1eb9903 100644 --- a/madspace/src/cpu/device.cpp +++ b/madspace/src/cpu/device.cpp @@ -8,93 +8,73 @@ using namespace madspace::kernels; namespace { +// calls func with the scalar and SIMD kernel policies of a float dtype +template +void dispatch_float(DataType dtype, const char* op, Func&& func) { + switch (dtype) { + case DataType::dt_float64: + func.template operator()(); + break; + case DataType::dt_float32: + func.template operator()(); + break; + default: + throw std::runtime_error(std::string("invalid dtype in ") + op); + } +} + template void tensor_copy_impl(const Tensor& source, Tensor& target, const D& device) { if (source.dtype() != target.dtype()) { throw std::runtime_error("invalid dtype in copy"); } - switch (source.dtype()) { - case DataType::dt_float64: - tensor_foreach_dynamic_single< - kernel_copy, - kernel_copy, - 1, - 1>({&source}, {&target}, target.size(0), device); - break; - case DataType::dt_float32: - tensor_foreach_dynamic_single< - kernel_copy, - kernel_copy, - 1, - 1>({&source}, {&target}, target.size(0), device); - break; - case DataType::dt_int32: + if (source.dtype() == DataType::dt_int32) { tensor_foreach_dynamic_single< kernel_copy_int, kernel_copy_int, 1, 1>({&source}, {&target}, target.size(0), device); - break; - default: - throw std::runtime_error("invalid dtype in copy"); + return; } + dispatch_float(source.dtype(), "copy", [&]() { + tensor_foreach_dynamic_single, kernel_copy, 1, 1>( + {&source}, {&target}, target.size(0), device + ); + }); } template void tensor_zero_impl(Tensor& tensor, const D& device) { - switch (tensor.dtype()) { - case DataType::dt_float64: - tensor_foreach_dynamic_single< - kernel_zero, - kernel_zero, - 1, - 1>({&tensor}, {&tensor}, tensor.size(0), device); - break; - case DataType::dt_float32: - tensor_foreach_dynamic_single< - kernel_zero, - kernel_zero, - 1, - 1>({&tensor}, {&tensor}, tensor.size(0), device); - break; - case DataType::dt_int32: + if (tensor.dtype() == DataType::dt_int32) { tensor_foreach_dynamic_single< kernel_zero_int, kernel_zero_int, 1, 1>({&tensor}, {&tensor}, tensor.size(0), device); - break; - default: - throw std::runtime_error("invalid dtype in zero"); + return; } + dispatch_float(tensor.dtype(), "zero", [&]() { + tensor_foreach_dynamic_single, kernel_zero, 1, 1>( + {&tensor}, {&tensor}, tensor.size(0), device + ); + }); } template void tensor_add_impl(const Tensor& source, Tensor& target, const D& device) { - // a broadcast target (e.g. the gradient of a literal or global) accumulates - // the whole source batch - std::size_t batch_size = std::max(source.size(0), target.size(0)); if (source.dtype() != target.dtype()) { throw std::runtime_error("invalid dtype in add"); } - switch (source.dtype()) { - case DataType::dt_float64: - tensor_foreach_dynamic_single< - kernel_add_inplace, - kernel_add_inplace, - 1, - 1>({&source}, {&target}, batch_size, device); - break; - case DataType::dt_float32: + // a broadcast target (e.g. the gradient of a literal or global) accumulates + // the whole source batch + std::size_t batch_size = std::max(source.size(0), target.size(0)); + dispatch_float(source.dtype(), "add", [&]() { tensor_foreach_dynamic_single< - kernel_add_inplace, - kernel_add_inplace, + kernel_add_inplace, + kernel_add_inplace, 1, 1>({&source}, {&target}, batch_size, device); - break; - default: - throw std::runtime_error("invalid dtype in add"); - } + }); } } // namespace @@ -121,41 +101,21 @@ void CpuDevice::adam_step( double bias_corr2_sqrt, double weight_decay ) const { - if (parameter.dtype() == DataType::dt_float32) { - tensor_foreach_dynamic< - kernel_adam_step, - kernel_adam_step, - 1, - 3>( + dispatch_float(parameter.dtype(), "adam_step", [&]() { + using F = typename C::FVal; + tensor_foreach_dynamic, kernel_adam_step, 1, 3>( {&gradient}, {¶meter, &exp_avg, &exp_avg_sq}, gradient.size(0), *this, - static_cast(step_size), - static_cast(beta1), - static_cast(beta2), - static_cast(eps), - static_cast(bias_corr2_sqrt), - static_cast(weight_decay) + static_cast(step_size), + static_cast(beta1), + static_cast(beta2), + static_cast(eps), + static_cast(bias_corr2_sqrt), + static_cast(weight_decay) ); - } else { - tensor_foreach_dynamic< - kernel_adam_step, - kernel_adam_step, - 1, - 3>( - {&gradient}, - {¶meter, &exp_avg, &exp_avg_sq}, - gradient.size(0), - *this, - step_size, - beta1, - beta2, - eps, - bias_corr2_sqrt, - weight_decay - ); - } + }); } void AsyncCpuDevice::tensor_copy(const Tensor& source, Tensor& target) const { diff --git a/madspace/src/driver/tensor.cpp b/madspace/src/driver/tensor.cpp index e954cab7d6..0c38156485 100644 --- a/madspace/src/driver/tensor.cpp +++ b/madspace/src/driver/tensor.cpp @@ -314,8 +314,7 @@ Tensor Tensor::simd_padded_literal( // element j of lane l is stored at l + lanes * j auto [dtype, buffer] = std::visit( [&](const std::vector& vec) { - auto dtype = - std::is_same_v ? DataType::dt_int32 : DataType::dt_float64; + auto dtype = dtype_of(); std::vector padded(count * lanes); for (std::size_t j = 0; j < count; ++j) { std::fill_n(padded.begin() + j * lanes, lanes, vec[j]); From 04d3c966f9c4fdc5d1557dc01c321da461e3e729 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 10:38:28 +0200 Subject: [PATCH 12/24] madspace: float32 mul, gradient clipper scales in the parameter precision --- madspace/instruction_set.yaml | 1 + madspace/src/compgraphs/function.cpp | 4 ++-- madspace/src/driver/adam_optimizer.cpp | 13 ++++++------- madspace/tests/test_mlp.py | 17 +++++++++++++++-- madspace/tests/test_simd_layout.py | 4 ++-- 5 files changed, 26 insertions(+), 13 deletions(-) diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 8d0378d397..cca6c6b3e6 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -231,6 +231,7 @@ sub: dims: 0 mul: + precision: [float32, float64] inputs: - name: in1 type: [float, ...] diff --git a/madspace/src/compgraphs/function.cpp b/madspace/src/compgraphs/function.cpp index ecaa86994e..f7e07dc382 100644 --- a/madspace/src/compgraphs/function.cpp +++ b/madspace/src/compgraphs/function.cpp @@ -354,7 +354,7 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) { return {arg(0) + arg(1)}; case opcodes::sub: return {arg(0) - arg(1)}; - case opcodes::mul: + case opcodes::mul_float64: return {arg(0) * arg(1)}; case opcodes::sqrt: return {std::sqrt(arg(0))}; @@ -402,7 +402,7 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) { return {var_arg}; } break; - case opcodes::mul: + case opcodes::mul_float64: if (const_val == 1.) { return {var_arg}; } diff --git a/madspace/src/driver/adam_optimizer.cpp b/madspace/src/driver/adam_optimizer.cpp index cf55a256fd..f2a9380fff 100644 --- a/madspace/src/driver/adam_optimizer.cpp +++ b/madspace/src/driver/adam_optimizer.cpp @@ -15,14 +15,13 @@ GradientClipper::GradientClipper(DataType dtype) : NamedVector GradientClipper::build_function_impl( FunctionBuilder& fb, const NamedVector& args ) const { - bool single = args.at(0).type.dtype == DataType::dt_float32; - Value grads_in = single ? fb.convert_float64(args.at(0)) : args.at(0); - Value grad_norm = fb.sqrt(fb.batch_reduce_sum(fb.square(grads_in))); + // only the norm is computed in double precision + Value grads_in = args.at(0); + bool single = grads_in.type.dtype == DataType::dt_float32; + Value grads_sq = fb.square(single ? fb.convert_float64(grads_in) : grads_in); + Value grad_norm = fb.sqrt(fb.batch_reduce_sum(grads_sq)); Value factor = fb.min(fb.div(args.at(1), grad_norm), 1.0); - Value grads_out = fb.mul(grads_in, factor); - if (single) { - grads_out = fb.convert_float32(grads_out); - } + Value grads_out = fb.mul(grads_in, single ? fb.convert_float32(factor) : factor); return { {"gradients_out", grads_out}, {"gradient_norm", grad_norm}, diff --git a/madspace/tests/test_mlp.py b/madspace/tests/test_mlp.py index ca28a88844..3dcf3bf431 100644 --- a/madspace/tests/test_mlp.py +++ b/madspace/tests/test_mlp.py @@ -88,7 +88,8 @@ def test_activation(activation): assert grad_me == approx(grad_torch) -def test_training(mlp, mlp_module, module_context, precision): +@pytest.mark.parametrize("grad_clip", [0.0, 0.01]) +def test_training(mlp, mlp_module, module_context, precision, grad_clip): _, dtype, tol = precision mlp_torch = nn.Sequential( nn.Linear(10, 32), @@ -118,7 +119,10 @@ def test_training(mlp, mlp_module, module_context, precision): fb.output(1, mlp_out) ctx.copy_globals_from(module_context) builtin_opt = ms.AdamOptimizer( - function=fb.function(), context=ctx, learning_rate=1e-2 + function=fb.function(), + context=ctx, + learning_rate=1e-2, + grad_clip_threshold=grad_clip, ) opt_ms = torch.optim.Adam(mlp_module.parameters(), lr=1e-2) @@ -127,6 +131,10 @@ def test_training(mlp, mlp_module, module_context, precision): def close(a, b): return a == approx(b, rel=tol, abs=tol) + # the builtin optimizer clips the RMS gradient, not the norm + n_params = sum(p.numel() for p in mlp_torch.parameters()) + max_norm = grad_clip * n_params**0.5 + for i in range(10): # the input and the loss are in float64, only the network uses dtype x = torch.randn((128, 10), dtype=torch.float64) @@ -135,11 +143,16 @@ def close(a, b): loss_ms = (y - mlp_module(x)).square().mean() opt_ms.zero_grad() loss_ms.backward() + if grad_clip > 0: + clipped = nn.utils.clip_grad_norm_(mlp_module.parameters(), max_norm) + assert clipped > max_norm opt_ms.step() loss_torch = (y - mlp_torch(x.to(dtype)).double()).square().mean() opt_torch.zero_grad() loss_torch.backward() + if grad_clip > 0: + nn.utils.clip_grad_norm_(mlp_torch.parameters(), max_norm) opt_torch.step() loss_builtin, mlp_out = builtin_opt.step([x]) diff --git a/madspace/tests/test_simd_layout.py b/madspace/tests/test_simd_layout.py index e84948a30e..cb95a6e6d1 100644 --- a/madspace/tests/test_simd_layout.py +++ b/madspace/tests/test_simd_layout.py @@ -90,7 +90,7 @@ def body(fb, x): func = build( [("x", ms.batch_float_array(3))], [("y", ms.batch_float_array(3))], body ) - assert instruction_names(func) == ["mul"] + assert instruction_names(func) == ["mul_float64"] def test_single_flavor_integrand_has_no_constant_gathers(tmp_path): @@ -140,7 +140,7 @@ def test_single_flavor_integrand_has_no_constant_gathers(tmp_path): for output in instr.outputs: sources[output.local_index] = instr.instruction.name for instr in func.instructions: - if instr.instruction.name == "mul": + if instr.instruction.name == "mul_float64": assert all( sources.get(arg.local_index) != "full" for arg in instr.inputs ) From c534fdb3a51e4c20e76f3033a682ce599594e945 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 17 Oct 2026 21:08:49 +0200 Subject: [PATCH 13/24] madspace: int32 indices for empty batch in GPU unweight op_unweight returned the empty indices tensor as float64, which the dtype-checked view in batch_gather rejected ("view element type does not match dtype") whenever a generator job had an empty batch. Co-Authored-By: Claude Sonnet 5.5 --- madspace/src/gpu/runtime.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index b2741e55d1..bb9a1b724c 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -1130,7 +1130,7 @@ void op_unweight( if (batch_size == 0) { indices = Tensor( - DataType::dt_float64, {0}, device, instruction.output_alloc_hints[0] + DataType::dt_int32, {0}, device, instruction.output_alloc_hints[0] ); uw_weights = Tensor( DataType::dt_float64, {0}, device, instruction.output_alloc_hints[1] From 3125be0f65fd46eab953ffabcac8117721a68204 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 17 Oct 2026 21:36:49 +0200 Subject: [PATCH 14/24] madspace: single-kernel bias gradient in GPU matmul backward Replace the ones-vector fill + cublas gemv by one block-per-column reduction kernel (double accumulation, no atomics, no temporary storage). cuBLAS split the float32 gemv into dot + reduce kernels, so this removes about 0.3M of the extra float32 kernel launches per 2000 training batches and the ones fill in both precisions. Co-Authored-By: Claude Sonnet 5.5 --- madspace/src/gpu/runtime.cu | 58 ++++++++++++++++--------------------- 1 file changed, 25 insertions(+), 33 deletions(-) diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index bb9a1b724c..9edc320600 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -134,15 +134,6 @@ auto gpublas_gemm(Args... args) { } } -template -auto gpublas_gemv(Args... args) { - if constexpr (std::is_same_v) { - return gpublasSgemv(args...); - } else { - return gpublasDgemv(args...); - } -} - template void op_matmul( const GpuRuntime::Instruction& instruction, @@ -193,11 +184,27 @@ void op_matmul( bias.reset(device); } +// bias_grad[j] += sum_i output_grad[i, j]; the matrix is column-major with +// leading dimension batch_size, so each column is a contiguous segment. One block +// per column: single launch, no atomics (deterministic) and no temporary storage. +// Accumulates in double, as the float32 cuBLAS paths would otherwise split the +// reduction over several kernels. +constexpr int BIAS_GRAD_THREADS = 512; + template -__global__ void kernel_one(std::size_t batch_size, GpuTensorView output) { - me_int_t i = blockDim.x * blockIdx.x + threadIdx.x; - if (i < batch_size) { - output[i] = 1.; +__global__ void kernel_bias_grad( + std::size_t batch_size, const F* __restrict__ output_grad, F* bias_grad +) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + const F* column = output_grad + blockIdx.x * batch_size; + double sum = 0.; + for (std::size_t i = threadIdx.x; i < batch_size; i += BIAS_GRAD_THREADS) { + sum += column[i]; + } + sum = BlockReduce(temp_storage).Sum(sum); + if (threadIdx.x == 0) { + bias_grad[blockIdx.x] += static_cast(sum); } } @@ -290,27 +297,12 @@ void backward_op_matmul( ); // compute bias_grad += sum_i output_grad_ij - Tensor ones(dtype_of(), {batch_size}, device, AllocHint::temporary); - launch_kernel( - kernel_one, batch_size, device.stream(), batch_size, ones.view() - ); - check_error( - gpublas_gemv( - handle, - GPUBLAS_OP_T, - batch_size, - dims_out, - &alpha, - static_cast(output_grad.data()), - batch_size, - static_cast(ones.data()), - 1, - &beta, - static_cast(bias_grad.data()), - 1 - ) + kernel_bias_grad<<>>( + batch_size, + static_cast(output_grad.data()), + static_cast(bias_grad.data()) ); - ones.reset(device); + check_error(); input.reset(device); weight.reset(device); output_grad.reset(device); From 38b420563ecb203c90df2741ef5943aaf9354a02 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 28 Oct 2026 10:35:32 +0100 Subject: [PATCH 15/24] madspace: fuse float32/float64 conversions around network layers Add a GPU-backend lowering pass (fuse_conversions) that rewrites matmul->convert_float64 and the select/rational-quadratic-spline consumers of float32 network outputs into fused ops (matmul_to_float64, select_with_float32, rqs_{forward,inverse}_with_float32), plus a fused backward bias-gradient+convert kernel. The compute graph stays portable; set MADSPACE_NO_FUSION=1 to disable the pass. This removes most of the extra kernel launches of float32 networks. Make float32 the default nn_float_precision in the MG7 run card. Co-Authored-By: Claude Sonnet 5.5 --- docs/source/run_card.rst | 2 +- madgraph/various/banner.py | 2 +- .../include/madspace/compgraphs/function.hpp | 1 + .../include/madspace/compgraphs/optimizer.hpp | 14 ++ madspace/instruction_set.yaml | 83 +++++++++ madspace/src/compgraphs/optimizer.cpp | 86 +++++++++ madspace/src/cpu/runtime.cpp | 141 +++++++++++++++ madspace/src/gpu/runtime.cu | 170 +++++++++++++++++- madspace/src/kernels/nn.hpp | 152 ++++++++++++++-- madspace/src/kernels/util.hpp | 31 ++++ madspace/src/python/madspace.cpp | 10 ++ madspace/tests/test_fused_float32.py | 160 +++++++++++++++++ tests/input_files/mg7_run_card.toml | 2 +- tests/unit_tests/various/test_banner.py | 8 +- 14 files changed, 831 insertions(+), 31 deletions(-) create mode 100644 madspace/tests/test_fused_float32.py diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index 51ec02b486..8b44707844 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -617,7 +617,7 @@ function. The activations are ``relu``, ``leaky_relu``, ``elu``, ``gelu``, ``sig ``cwnet_hidden_dim`` (``"auto"``), ``cwnet_layers`` (``3``), ``cwnet_activation`` (``"leaky_relu"``) Size and activation of the channel weight network. -``nn_float_precision`` (``"float64"``) +``nn_float_precision`` (``"float32"``) Floating-point precision of the weights and the evaluation of all networks, ``float32`` or ``float64``. The rest of the computation, including the network outputs, softmaxes and Jacobians, stays in double precision. diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 593f6c2be8..b18420b36d 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -7027,7 +7027,7 @@ def default_setup(self): self.add_toml_param('madnis', 'cwnet_activation', "leaky_relu", allowed=['relu', 'leaky_relu', 'elu', 'gelu', 'sigmoid', 'softplus'], comment="activation function of the channel weight network") - self.add_toml_param('madnis', 'nn_float_precision', "float64", + self.add_toml_param('madnis', 'nn_float_precision', "float32", allowed=['float32', 'float64'], comment="floating-point precision of the weights and the evaluation of all networks") self.add_toml_param('madnis', 'train_batches', 1000, diff --git a/madspace/include/madspace/compgraphs/function.hpp b/madspace/include/madspace/compgraphs/function.hpp index 9f373f6d81..d2d2ac42c2 100644 --- a/madspace/include/madspace/compgraphs/function.hpp +++ b/madspace/include/madspace/compgraphs/function.hpp @@ -81,6 +81,7 @@ class Function { std::vector _instructions; friend Function sort_breadth_first(const Function& function); + friend Function fuse_conversions(const Function& function); }; std::ostream& operator<<(std::ostream& out, const Value& value); diff --git a/madspace/include/madspace/compgraphs/optimizer.hpp b/madspace/include/madspace/compgraphs/optimizer.hpp index a0dd35a775..9614039625 100644 --- a/madspace/include/madspace/compgraphs/optimizer.hpp +++ b/madspace/include/madspace/compgraphs/optimizer.hpp @@ -33,4 +33,18 @@ class LastUseOfLocals { Function sort_breadth_first(const Function& function); +/** + * Lowering pass for single-precision networks, run by the GPU backend when it + * builds a runtime. Fuses a precision conversion into the instruction that + * produces the converted value: + * - `convert_float32(v)` where `v` is an output of `select`, `rqs_forward` or + * `rqs_inverse` becomes the `*_with_float32` variant of the producer, which + * writes the single-precision copy itself; + * - `convert_float64(matmul_float32(x, w, b))`, with the matmul result used only + * by the conversion, becomes `matmul_to_float64`. + * The graph keeps only the generic instructions and so stays portable; the fused + * instructions never appear in a saved function. + */ +Function fuse_conversions(const Function& function); + } // namespace madspace diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 15044384d5..ec8b0e3ae7 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -2945,6 +2945,23 @@ convert_float64: differentiable: True desc: "Converts a tensor to double precision. Used at the boundaries of single-precision @ref MLP networks." +matmul_to_float64: + inputs: + - name: x + type: [float32, n] + backward_arg: True + - name: weight + type: [float32, single, m, n] + backward_arg: True + - name: bias + type: [float32, single, m] + outputs: + - name: y + type: [float64, m] + custom_op: True + differentiable: True + desc: "Single-precision affine map whose result is converted to double precision. Equivalent to @ref convert_float64 applied to the output of @ref matmul, but the GPU backend fuses the bias addition and the conversion, and the backward pass fuses the gradient conversion with the bias gradient. Used for the last layer of single-precision @ref MLP networks." + matmul: precision: [float32, float64] inputs: @@ -3145,6 +3162,52 @@ rqs_inverse: dims: 2 desc: "Applies a rational-quadratic spline in the inverse direction. Building block of @ref Flow." +rqs_forward_with_float32: + inputs: + - name: input + type: [float, n] + desc: "values to transform" + - name: condition + type: [float, n, 6] + desc: "packed spline parameters" + outputs: + - name: output + type: [float, n] + desc: "transformed values" + - name: det + type: [float, n] + desc: "per-element Jacobian" + - name: output32 + type: [float32, n] + desc: "the transformed values in single precision" + differentiable: True + dims: 2 + vectorized: False + desc: "@ref rqs_forward that also returns its output in single precision, so that the conversion does not need a separate kernel. Produced by the GPU backend's lowering of @ref convert_float32 applied to the output of @ref rqs_forward." + +rqs_inverse_with_float32: + inputs: + - name: input + type: [float, n] + desc: "values to transform" + - name: condition + type: [float, n, 6] + desc: "packed spline parameters" + outputs: + - name: output + type: [float, n] + desc: "transformed values" + - name: det + type: [float, n] + desc: "per-element Jacobian" + - name: output32 + type: [float32, n] + desc: "the transformed values in single precision" + differentiable: True + dims: 2 + vectorized: False + desc: "@ref rqs_inverse that also returns its output in single precision, so that the conversion does not need a separate kernel. Produced by the GPU backend's lowering of @ref convert_float32 applied to the output of @ref rqs_inverse." + softmax: inputs: - name: input @@ -3361,6 +3424,26 @@ select: differentiable: True desc: "Gathers several entries from a float vector by index." +select_with_float32: + inputs: + - name: input + type: [float, n] + desc: "source vector" + - name: indices + type: [int, single, m] + desc: "indices to gather" + backward_arg: True + outputs: + - name: output + type: [float, m] + desc: "gathered entries" + - name: output32 + type: [float32, m] + desc: "the gathered entries in single precision" + differentiable: True + vectorized: False + desc: "@ref select that also returns the gathered entries in single precision, so that the conversion does not need a separate kernel. Produced by the GPU backend's lowering of @ref convert_float32 applied to the output of @ref select." + select_vector: inputs: - name: input diff --git a/madspace/src/compgraphs/optimizer.cpp b/madspace/src/compgraphs/optimizer.cpp index 3435dd3fb7..0269c2043f 100644 --- a/madspace/src/compgraphs/optimizer.cpp +++ b/madspace/src/compgraphs/optimizer.cpp @@ -1,6 +1,7 @@ #include "madspace/compgraphs/optimizer.hpp" #include +#include #include #include #include @@ -80,3 +81,88 @@ Function madspace::sort_breadth_first(const Function& function) { } return func_out; } + +Function madspace::fuse_conversions(const Function& function) { + Function func_out = function; + auto& instructions = func_out._instructions; + std::vector producer(function.locals().size(), -1); + std::vector use_count(function.locals().size(), 0); + for (std::size_t i = 0; i < instructions.size(); ++i) { + for (auto& output : instructions[i].outputs) { + producer.at(output.local_index) = i; + } + for (auto& input : instructions[i].inputs) { + if (input.local_index != -1) { + ++use_count.at(input.local_index); + } + } + } + for (auto& output : function.outputs()) { + ++use_count.at(output.local_index); + } + + auto find_instruction = [](const std::string& name) { + auto find = instruction_set.find(name); + if (find == instruction_set.end()) { + throw std::runtime_error(std::format("Unknown instruction '{}'", name)); + } + return find->second.get(); + }; + + std::vector removed(instructions.size()); + bool any_fused = false; + for (std::size_t i = 0; i < instructions.size(); ++i) { + auto& call = instructions[i]; + const std::string& name = call.instruction->name(); + if (name != "convert_float32" && name != "convert_float64") { + continue; + } + auto& input = call.inputs.at(0); + if (input.local_index == -1) { + continue; + } + int source_index = producer.at(input.local_index); + if (source_index == -1 || removed.at(source_index)) { + continue; + } + auto& source = instructions.at(source_index); + // the fused instruction takes the place of the producer + if (source.stream_index != call.stream_index) { + continue; + } + const std::string& source_name = source.instruction->name(); + if (name == "convert_float32") { + bool fusable = source_name == "select" || source_name == "rqs_forward" || + source_name == "rqs_inverse"; + if (!fusable || source.outputs.at(0).local_index != input.local_index || + input.type.dtype != DataType::dt_float64) { + continue; + } + source.instruction = find_instruction(source_name + "_with_float32"); + source.outputs.push_back(call.outputs.at(0)); + } else { + if (source_name != "matmul_float32" || + use_count.at(input.local_index) != 1) { + continue; + } + source.instruction = find_instruction("matmul_to_float64"); + source.outputs = {call.outputs.at(0)}; + } + producer.at(call.outputs.at(0).local_index) = source_index; + removed.at(i) = true; + any_fused = true; + } + if (!any_fused) { + return func_out; + } + + std::vector kept; + kept.reserve(instructions.size()); + for (std::size_t i = 0; i < instructions.size(); ++i) { + if (!removed.at(i)) { + kept.push_back(std::move(instructions[i])); + } + } + instructions = std::move(kept); + return func_out; +} diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 897f13777e..428c3dea9b 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -377,6 +377,147 @@ void backward_op_matmul( }); } +// matmul with a float32 input and a float64 output. Everything runs in a single job, +// because the CPU runtime offers only one barrier per instruction. +template +void op_matmul_to_float64( + const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device +) { + auto input = locals[instruction.input_indices[0]].contiguous(device); + auto weight = locals[instruction.input_indices[1]].contiguous(device); + auto bias = locals[instruction.input_indices[2]].contiguous(device); + auto& output = locals[instruction.output_indices[0]]; + std::size_t batch_size = input.size(0); + std::size_t dims_in = input.size(1); + std::size_t dims_out = weight.size(1); + output = Tensor(DataType::dt_float64, {batch_size, dims_out}, device); + if (batch_size == 0) { + return; + } + Tensor tmp(DataType::dt_float32, {batch_size, dims_out}, device); + + device.sync_barrier(); + device.submit([=, out = output]() mutable { + auto* tmp_ptr = static_cast(tmp.data()); + auto* bias_ptr = static_cast(bias.data()); + for (std::size_t j = 0; j < dims_out; ++j) { + for (std::size_t i = 0; i < batch_size; ++i) { + tmp_ptr[j * batch_size + i] = bias_ptr[j]; + } + } + gemm( + 'N', + 'T', + batch_size, + dims_out, + dims_in, + 1., + static_cast(input.data()), + batch_size, + static_cast(weight.data()), + dims_out, + 1., + tmp_ptr, + batch_size + ); + auto* out_ptr = static_cast(out.data()); + for (std::size_t i = 0; i < batch_size * dims_out; ++i) { + out_ptr[i] = tmp_ptr[i]; + } + }); +} + +template +void backward_op_matmul_to_float64( + const CpuRuntime::Instruction& instruction, + TensorVec& locals, + TensorVec& local_grads, + const D& device +) { + using F = float; + auto input = locals[instruction.input_indices[0]].contiguous(device); + auto weight = locals[instruction.input_indices[1]].contiguous(device); + auto output_grad = local_grads[instruction.output_indices[0]].contiguous(device); + auto& input_grad = local_grads[instruction.input_indices[0]]; + auto& weight_grad = local_grads[instruction.input_indices[1]]; + auto& bias_grad = local_grads[instruction.input_indices[2]]; + std::size_t batch_size = input.size(0); + std::size_t dims_in = input.size(1); + std::size_t dims_out = weight.size(1); + + if (!input_grad) { + input_grad = + Tensor(dtype_of(), input.shape(), device, AllocHint::local_grad); + } + if (!weight_grad) { + weight_grad = + Tensor(dtype_of(), weight.shape(), device, AllocHint::local_grad); + } + if (!bias_grad) { + bias_grad = Tensor(dtype_of(), {1, dims_out}, device, AllocHint::local_grad); + } + if (batch_size == 0) { + return; + } + Tensor grad_f(dtype_of(), output_grad.shape(), device); + + device.sync_barrier(); + device.submit([=]() mutable { + auto* grad_f_ptr = static_cast(grad_f.data()); + auto* grad_ptr = static_cast(output_grad.data()); + for (std::size_t i = 0; i < batch_size * dims_out; ++i) { + grad_f_ptr[i] = static_cast(grad_ptr[i]); + } + // input_grad += output_grad * weight + gemm( + 'N', + 'N', + batch_size, + dims_in, + dims_out, + 1., + grad_f_ptr, + batch_size, + static_cast(weight.data()), + dims_out, + 1., + static_cast(input_grad.data()), + batch_size + ); + // weight_grad += output_grad.T * input + gemm( + 'T', + 'N', + dims_out, + dims_in, + batch_size, + 1., + grad_f_ptr, + batch_size, + static_cast(input.data()), + batch_size, + 1., + static_cast(weight_grad.data()), + dims_out + ); + // bias_grad += sum_i output_grad_ij + std::vector ones(batch_size, 1.); + gemv( + 'T', + batch_size, + dims_out, + 1., + grad_f_ptr, + batch_size, + ones.data(), + 1, + 1., + static_cast(bias_grad.data()), + 1 + ); + }); +} + template void op_nonzero( const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 9edc320600..f8c02d44c2 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -184,6 +185,89 @@ void op_matmul( bias.reset(device); } +// out[j * batch_size + i] = double(tmp[...] + bias[j]): bias, convert to float64 in +// one pass over the output of a float32 gemm +template +__global__ void kernel_bias_to_double( + std::size_t total, + std::size_t batch_size, + const F* __restrict__ tmp, + const F* __restrict__ bias, + double* __restrict__ out +) { + std::size_t i = blockDim.x * static_cast(blockIdx.x) + threadIdx.x; + if (i < total) { + out[i] = static_cast(tmp[i] + bias[i / batch_size]); + } +} + +// matmul of float32 inputs with the float32 -> float64 conversion of the result +// fused into the bias epilogue (three launches become two) +void op_matmul_to_float64( + const GpuRuntime::Instruction& instruction, + TensorVec& locals, + const AsyncGpuDevice& device +) { + using F = float; + auto input = + locals[instruction.input_indices[0]].contiguous(device, AllocHint::temporary); + auto weight = + locals[instruction.input_indices[1]].contiguous(device, AllocHint::temporary); + auto bias = + locals[instruction.input_indices[2]].contiguous(device, AllocHint::temporary); + auto& output = locals[instruction.output_indices[0]]; + std::size_t batch_size = input.size(0); + std::size_t dims_in = input.size(1); + std::size_t dims_out = weight.size(1); + output = Tensor( + DataType::dt_float64, + {batch_size, dims_out}, + device, + instruction.output_alloc_hints[0] + ); + if (batch_size == 0) { + return; + } + + Tensor tmp(dtype_of(), {batch_size, dims_out}, device, AllocHint::temporary); + gpublasHandle_t handle = instruction.runtime.gpublas_handle(); + check_error(gpublasSetStream(handle, device.stream())); + F alpha = 1., beta = 0.; + check_error( + gpublas_gemm( + handle, + GPUBLAS_OP_N, + GPUBLAS_OP_T, + batch_size, + dims_out, + dims_in, + &alpha, + static_cast(input.data()), + batch_size, + static_cast(weight.data()), + dims_out, + &beta, + static_cast(tmp.data()), + batch_size + ) + ); + std::size_t total = batch_size * dims_out; + launch_kernel( + kernel_bias_to_double, + total, + device.stream(), + total, + batch_size, + static_cast(tmp.data()), + static_cast(bias.data()), + static_cast(output.data()) + ); + tmp.reset(device); + input.reset(device); + weight.reset(device); + bias.reset(device); +} + // bias_grad[j] += sum_i output_grad[i, j]; the matrix is column-major with // leading dimension batch_size, so each column is a contiguous segment. One block // per column: single launch, no atomics (deterministic) and no temporary storage. @@ -208,8 +292,34 @@ __global__ void kernel_bias_grad( } } +// the same reduction for a float64 output gradient, which is also converted to F +// for the gemms: one read of the gradient instead of a separate convert kernel template -void backward_op_matmul( +__global__ void kernel_bias_grad_convert( + std::size_t batch_size, + const double* __restrict__ output_grad, + F* __restrict__ output_grad_f, + F* bias_grad +) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + std::size_t offset = blockIdx.x * batch_size; + double sum = 0.; + for (std::size_t i = threadIdx.x; i < batch_size; i += BIAS_GRAD_THREADS) { + F value = static_cast(output_grad[offset + i]); + output_grad_f[offset + i] = value; + sum += value; + } + sum = BlockReduce(temp_storage).Sum(sum); + if (threadIdx.x == 0) { + bias_grad[blockIdx.x] += static_cast(sum); + } +} + +// OutGradF64: the gradient of the output is float64 (matmul_to_float64) and is +// converted to F together with the bias gradient reduction +template +void backward_matmul_impl( const GpuRuntime::Instruction& instruction, TensorVec& locals, TensorVec& local_grads, @@ -256,6 +366,22 @@ void backward_op_matmul( gpuStream_t stream = device.stream(); check_error(gpublasSetStream(handle, stream)); + if constexpr (OutGradF64) { + // compute bias_grad += sum_i output_grad_ij and convert the gradient + Tensor output_grad_f( + dtype_of(), output_grad.shape(), device, AllocHint::temporary + ); + kernel_bias_grad_convert<<>>( + batch_size, + static_cast(output_grad.data()), + static_cast(output_grad_f.data()), + static_cast(bias_grad.data()) + ); + check_error(); + output_grad.reset(device); + output_grad = output_grad_f; + } + // compute input_grad += output_grad * weight check_error( gpublas_gemm( @@ -296,18 +422,39 @@ void backward_op_matmul( ) ); - // compute bias_grad += sum_i output_grad_ij - kernel_bias_grad<<>>( - batch_size, - static_cast(output_grad.data()), - static_cast(bias_grad.data()) - ); - check_error(); + if constexpr (!OutGradF64) { + // compute bias_grad += sum_i output_grad_ij + kernel_bias_grad<<>>( + batch_size, + static_cast(output_grad.data()), + static_cast(bias_grad.data()) + ); + check_error(); + } input.reset(device); weight.reset(device); output_grad.reset(device); } +template +void backward_op_matmul( + const GpuRuntime::Instruction& instruction, + TensorVec& locals, + TensorVec& local_grads, + const AsyncGpuDevice& device +) { + backward_matmul_impl(instruction, locals, local_grads, device); +} + +void backward_op_matmul_to_float64( + const GpuRuntime::Instruction& instruction, + TensorVec& locals, + TensorVec& local_grads, + const AsyncGpuDevice& device +) { + backward_matmul_impl(instruction, locals, local_grads, device); +} + struct NotMinusOne { __device__ bool operator()(me_int_t val) { return val != -1; } }; @@ -1601,7 +1748,12 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) : // check_error(cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, // &thresh)); - Function function = sort_breadth_first(function_arg); + // MADSPACE_NO_FUSION disables the lowering of precision conversions, to compare + // the fused and the unfused instructions + Function function = sort_breadth_first( + std::getenv("MADSPACE_NO_FUSION") ? function_arg + : fuse_conversions(function_arg) + ); _locals_init.resize(function.locals().size()); _requires_grad_init.resize(function.locals().size()); diff --git a/madspace/src/kernels/nn.hpp b/madspace/src/kernels/nn.hpp index 4cef7f7067..3c928c7bb7 100644 --- a/madspace/src/kernels/nn.hpp +++ b/madspace/src/kernels/nn.hpp @@ -254,9 +254,13 @@ KERNELSPEC void backward_kernel_rqs_find_bin( } template -KERNELSPEC void kernel_rqs_forward( - FIn input, FIn condition, FOut output, FOut det -) { +struct RqsValues { + FVal output; + FVal det; +}; + +template +KERNELSPEC RqsValues rqs_forward_values(FIn input, FIn condition) { FVal width(condition[0]), height(condition[1]); FVal cumwidth(condition[2]), cumheight(condition[3]); FVal derivative_unorm(condition[4]), derivative_plus_one_unorm(condition[5]); @@ -291,20 +295,45 @@ KERNELSPEC void kernel_rqs_forward( auto denominator = delta + (derivative + derivative_plus_one - two_delta) * theta_one_minus_theta; auto out = cumheight + numerator / denominator; - output = where(clamp, input, out); + FVal output_value = where(clamp, input, out); auto derivative_numerator = delta * delta * (derivative_plus_one * theta_theta + two_delta * theta_one_minus_theta + derivative * one_minus_theta * one_minus_theta); - det = where(clamp, 1., derivative_numerator / (denominator * denominator)); + FVal det_value = where(clamp, 1., derivative_numerator / (denominator * denominator)); + return {output_value, det_value}; } template -KERNELSPEC void backward_kernel_rqs_forward( +KERNELSPEC void kernel_rqs_forward( + FIn input, FIn condition, FOut output, FOut det +) { + auto result = rqs_forward_values(input, condition); + output = result.output; + det = result.det; +} + +// the same, with the output also written in single precision +template +KERNELSPEC void kernel_rqs_forward_with_float32( FIn input, FIn condition, - FIn output_grad, - FIn det_grad, + FOut output, + FOut det, + FOut output32 +) { + auto result = rqs_forward_values(input, condition); + output = result.output; + det = result.det; + output32 = FVal(result.output); +} + +template +KERNELSPEC void backward_rqs_forward_impl( + FIn input, + FIn condition, + FVal output_grad, + FVal det_grad, FOut input_grad, FOut condition_grad ) { @@ -437,9 +466,42 @@ KERNELSPEC void backward_kernel_rqs_forward( } template -KERNELSPEC void kernel_rqs_inverse( - FIn input, FIn condition, FOut output, FOut det +KERNELSPEC void backward_kernel_rqs_forward( + FIn input, + FIn condition, + FIn output_grad, + FIn det_grad, + FOut input_grad, + FOut condition_grad +) { + backward_rqs_forward_impl( + input, condition, FVal(output_grad), FVal(det_grad), input_grad, condition_grad + ); +} + +// the gradient of the single-precision copy of the output adds to the output gradient +template +KERNELSPEC void backward_kernel_rqs_forward_with_float32( + FIn input, + FIn condition, + FIn output_grad, + FIn det_grad, + FIn output32_grad, + FOut input_grad, + FOut condition_grad ) { + backward_rqs_forward_impl( + input, + condition, + FVal(output_grad) + FVal(FVal(output32_grad)), + FVal(det_grad), + input_grad, + condition_grad + ); +} + +template +KERNELSPEC RqsValues rqs_inverse_values(FIn input, FIn condition) { FVal height(condition[0]), width(condition[1]); FVal cumheight(condition[2]), cumwidth(condition[3]); FVal derivative_unorm(condition[4]), derivative_plus_one_unorm(condition[5]); @@ -473,7 +535,7 @@ KERNELSPEC void kernel_rqs_inverse( auto discriminant = b * b + 4. * a * c; auto theta = 2. * c / (b + sqrt(discriminant)); auto out = cumwidth + theta * width; - output = where(clamp, input, out); + FVal output_value = where(clamp, input, out); auto one_minus_theta = 1. - theta; auto theta_one_minus_theta = theta * one_minus_theta; @@ -481,15 +543,40 @@ KERNELSPEC void kernel_rqs_inverse( auto derivative_numerator = delta * delta * (derivative_plus_one * theta * theta + two_delta * theta_one_minus_theta + derivative * one_minus_theta * one_minus_theta); - det = where(clamp, 1., denominator * denominator / derivative_numerator); + FVal det_value = where(clamp, 1., denominator * denominator / derivative_numerator); + return {output_value, det_value}; } template -KERNELSPEC void backward_kernel_rqs_inverse( +KERNELSPEC void kernel_rqs_inverse( + FIn input, FIn condition, FOut output, FOut det +) { + auto result = rqs_inverse_values(input, condition); + output = result.output; + det = result.det; +} + +// the same, with the output also written in single precision +template +KERNELSPEC void kernel_rqs_inverse_with_float32( FIn input, FIn condition, - FIn output_grad, - FIn det_grad, + FOut output, + FOut det, + FOut output32 +) { + auto result = rqs_inverse_values(input, condition); + output = result.output; + det = result.det; + output32 = FVal(result.output); +} + +template +KERNELSPEC void backward_rqs_inverse_impl( + FIn input, + FIn condition, + FVal output_grad, + FVal det_grad, FOut input_grad, FOut condition_grad ) { @@ -651,6 +738,41 @@ KERNELSPEC void backward_kernel_rqs_inverse( condition_grad[5] = where(clamp, 0., grad_derivative_plus_one_unorm); } +template +KERNELSPEC void backward_kernel_rqs_inverse( + FIn input, + FIn condition, + FIn output_grad, + FIn det_grad, + FOut input_grad, + FOut condition_grad +) { + backward_rqs_inverse_impl( + input, condition, FVal(output_grad), FVal(det_grad), input_grad, condition_grad + ); +} + +// the gradient of the single-precision copy of the output adds to the output gradient +template +KERNELSPEC void backward_kernel_rqs_inverse_with_float32( + FIn input, + FIn condition, + FIn output_grad, + FIn det_grad, + FIn output32_grad, + FOut input_grad, + FOut condition_grad +) { + backward_rqs_inverse_impl( + input, + condition, + FVal(output_grad) + FVal(FVal(output32_grad)), + FVal(det_grad), + input_grad, + condition_grad + ); +} + template KERNELSPEC void kernel_softmax(FIn input, FOut output) { FVal norm(0.), in_max(0.); diff --git a/madspace/src/kernels/util.hpp b/madspace/src/kernels/util.hpp index 5a9cd90c8f..1cc91d5a3f 100644 --- a/madspace/src/kernels/util.hpp +++ b/madspace/src/kernels/util.hpp @@ -85,6 +85,37 @@ KERNELSPEC void backward_kernel_select( } } +// select with a single-precision copy of the gathered entries +template +KERNELSPEC void kernel_select_with_float32( + FIn input, + ISIn indices, + FOut output, + FOut output32 +) { + for (std::size_t i = 0; i < indices.size(); ++i) { + auto value = input.gather(indices[i]); + output[i] = value; + output32[i] = FVal(value); + } +} + +template +KERNELSPEC void backward_kernel_select_with_float32( + ISIn indices, + FIn output_grad, + FIn output32_grad, + FOut input_grad, + IOut indices_grad +) { + for (std::size_t i = 0; i < indices.size(); ++i) { + input_grad.scatter_add( + indices[i], + FVal(output_grad[i]) + FVal(FVal(output32_grad[i])) + ); + } +} + template KERNELSPEC void kernel_select_vector(FIn input, IIn indices, FOut output) { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index aeec02cdb3..9a78c2be88 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -7,6 +7,7 @@ #include "function_runtime.hpp" #include "instruction_set.hpp" #include "madspace/compgraphs.hpp" +#include "madspace/compgraphs/optimizer.hpp" #include "madspace/driver.hpp" #include "madspace/phasespace.hpp" #include "madspace/util.hpp" @@ -285,6 +286,15 @@ PYBIND11_MODULE(_madspace_py, m) { pydoc::doc("Function::instructions") ); + m.def( + "fuse_conversions", + &fuse_conversions, + py::arg("function"), + "Fuses single-precision conversions into the instructions producing the " + "converted values (the lowering the GPU backend applies when it builds a " + "runtime)." + ); + py::classh device(m, "Device", pydoc::doc("Device")); m.def( "cpu_device", diff --git a/madspace/tests/test_fused_float32.py b/madspace/tests/test_fused_float32.py new file mode 100644 index 0000000000..f4f3ac6249 --- /dev/null +++ b/madspace/tests/test_fused_float32.py @@ -0,0 +1,160 @@ +"""Fused single-precision instructions against their unfused equivalents.""" + +import re + +import pytest +import torch +from madspace_testing import generator, new_context, rand, randn + +import madspace as ms +from madspace.torch import FunctionModule + +N = 211 +DIMS_IN = 7 +DIMS_OUT = 5 + + +def batch_type(dtype, shape): + return ms.Type(dtype, ms.batch_size, shape) + + +def matmul_module(fused): + f32, f64 = ms.DataType.float32, ms.DataType.float64 + fb = ms.FunctionBuilder( + ms.NamedTypes( + [ + ("x", batch_type(f32, [DIMS_IN])), + ("w", ms.Type(f32, ms.BatchSize.one, [DIMS_OUT, DIMS_IN])), + ("b", ms.Type(f32, ms.BatchSize.one, [DIMS_OUT])), + ] + ), + ms.NamedTypes([("y", batch_type(f64, [DIMS_OUT]))]), + ) + x, w, b = fb.input(0), fb.input(1), fb.input(2) + if fused: + y = fb.matmul_to_float64(x, w, b) + else: + y = fb.convert_float64(fb.matmul(x, w, b)) + fb.output(0, y) + return FunctionModule(fb.function(), new_context()) + + +def test_matmul_to_float64(): + g = generator(77) + x = randn((N, DIMS_IN), generator=g, dtype=torch.float32, requires_grad=True) + w = randn((DIMS_OUT, DIMS_IN), generator=g, dtype=torch.float32, requires_grad=True) + b = randn((DIMS_OUT,), generator=g, dtype=torch.float32, requires_grad=True) + y_grad = randn((N, DIMS_OUT), generator=g, dtype=torch.float64) + + results = [] + for fused in (True, False): + xs, ws, bs = (t.detach().clone().requires_grad_() for t in (x, w, b)) + y = matmul_module(fused)(xs, ws, bs) + assert y.dtype == torch.float64 + (y * y_grad).sum().backward() + results.append((y.detach(), xs.grad, ws.grad, bs.grad)) + + y_ref = (x @ w.T + b).double() + torch.testing.assert_close(results[0][0], y_ref, rtol=1e-5, atol=1e-5) + # the gradients through the fused and the unfused graph agree + for fused_t, unfused_t in zip(results[0], results[1]): + torch.testing.assert_close(fused_t, unfused_t, rtol=1e-4, atol=1e-4) + + g_ref = y_grad.float() + torch.testing.assert_close(results[0][1], g_ref @ w.detach(), rtol=1e-4, atol=1e-4) + torch.testing.assert_close( + results[0][2], g_ref.T @ x.detach(), rtol=1e-4, atol=1e-3 + ) + torch.testing.assert_close(results[0][3], g_ref.sum(dim=0), rtol=1e-4, atol=1e-3) + + +FUSED_NAMES = ( + "convert_float32", + "convert_float64", + "matmul_float32", + "matmul_to_float64", + "select_with_float32", + "rqs_forward_with_float32", + "rqs_inverse_with_float32", +) + + +def instruction_counts(function): + names = "|".join(FUSED_NAMES) + counts = {} + for name in re.findall(rf"\b({names})\b", str(function)): + counts[name] = counts.get(name, 0) + 1 + return counts + + +def make_flow(prefix): + return ms.Flow( + 4, + prefix=prefix, + subnet_precision=ms.DataType.float32, + subnet_hidden_dim=16, + subnet_layers=3, + invert_spline=False, + ) + + +@pytest.mark.parametrize("direction", ["forward", "inverse"]) +def test_conversion_lowering_graph(direction): + flow = make_flow("lowering") + function = getattr(flow, f"{direction}_function")() + before = instruction_counts(function) + after = instruction_counts(ms.fuse_conversions(function)) + spline = "rqs_forward" if direction == "forward" else "rqs_inverse" + # every MLP of the flow converts its input and its output + assert before["convert_float32"] == before["convert_float64"] == 4 + assert "convert_float32" not in after and "convert_float64" not in after + assert after["matmul_to_float64"] == 4 + assert after["select_with_float32"] + after[f"{spline}_with_float32"] == 4 + # the unfused matmuls are the hidden layers + assert after["matmul_float32"] == before["matmul_float32"] - 4 + # the original function is not modified, and the pass is idempotent + assert instruction_counts(function) == before + assert instruction_counts(ms.fuse_conversions(ms.fuse_conversions(function))) == after + + +@pytest.mark.parametrize("direction", ["forward", "inverse"]) +def test_conversion_lowering_numerics(direction, monkeypatch): + """The lowered flow agrees with the unlowered one in values and gradients. On a + GPU the reference runs with the lowering disabled.""" + g = generator(2024) + flow = make_flow(f"lowering_{direction}") + ctx = new_context() + flow.initialize_globals(ctx, 42) + # randomize all weights (the last layers start at zero) + for name in ctx.global_names(): + weights = ctx.get_global(name).torch() + weights[:] = 0.3 * randn(weights.shape, generator=g, dtype=torch.float32).to( + weights.dtype + ) + + r = rand((137, 4), generator=g, dtype=torch.float64) + function = getattr(flow, f"{direction}_function")() + results = {} + for lowered in (True, False): + with monkeypatch.context() as mp: + if not lowered: + mp.setenv("MADSPACE_NO_FUSION", "1") + module = FunctionModule( + ms.fuse_conversions(function) if lowered else function, ctx + ) + r_grad = r.clone().requires_grad_() + out, det = module(r_grad) + (out.sum() + det.log().sum()).backward() + results[lowered] = ( + out.detach(), + det.detach(), + r_grad.grad, + {n: p.grad.clone() for n, p in module.global_params.items()}, + ) + for lowered_t, ref_t in zip(results[True][:3], results[False][:3]): + torch.testing.assert_close(lowered_t, ref_t, rtol=1e-12, atol=1e-12) + for name, grad in results[True][3].items(): + torch.testing.assert_close( + grad, results[False][3][name], rtol=1e-4, atol=1e-6 + ) + diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml index b550825e6a..045dd066b3 100644 --- a/tests/input_files/mg7_run_card.toml +++ b/tests/input_files/mg7_run_card.toml @@ -110,7 +110,7 @@ cwnet_layers = 3 cwnet_activation = "leaky_relu" # options: relu, leaky_relu, elu, gelu, sigmoid, softplus # precision of all networks -nn_float_precision = "float64" # options: float32, float64 +nn_float_precision = "float32" # options: float32, float64 # training parameters train_batches = 1000 diff --git a/tests/unit_tests/various/test_banner.py b/tests/unit_tests/various/test_banner.py index c649d8ae45..1dffcd50b1 100755 --- a/tests/unit_tests/various/test_banner.py +++ b/tests/unit_tests/various/test_banner.py @@ -2277,7 +2277,7 @@ def test_defaults_and_section_access(self): self.assertEqual(rc['beam']['e_cm'], 13000.0) self.assertIs(rc['vegas']['enable'], True) self.assertEqual(rc['madnis']['enable'], 'auto') - self.assertEqual(rc['madnis']['nn_float_precision'], 'float64') + self.assertEqual(rc['madnis']['nn_float_precision'], 'float32') self.assertEqual(rc['phasespace']['mode'], 'multichannel') # the "enable" key colliding between [vegas] and [madnis] is stored # independently, so editing one does not affect the other @@ -2306,20 +2306,20 @@ def test_write_is_valid_toml(self): self.assertEqual(data['multiparticles']['photon'], [22]) self.assertEqual(data['cuts']['sqrt_s'], {'min': 0.0}) self.assertEqual(data['madnis']['adam_eps'], 1e-8) - self.assertEqual(data['madnis']['nn_float_precision'], 'float64') + self.assertEqual(data['madnis']['nn_float_precision'], 'float32') def test_round_trip(self): """write() then read() preserves every value""" rc = bannermod.RunCardMG7() rc['generation']['events'] = 250 rc['phasespace']['mode'] = 'flat' - rc['madnis']['nn_float_precision'] = 'float32' + rc['madnis']['nn_float_precision'] = 'float64' out = io.StringIO() rc.write(out, template=self.template) rc2 = bannermod.RunCardMG7(out.getvalue()) self.assertEqual(rc2['generation']['events'], 250) self.assertEqual(rc2['phasespace']['mode'], 'flat') - self.assertEqual(rc2['madnis']['nn_float_precision'], 'float32') + self.assertEqual(rc2['madnis']['nn_float_precision'], 'float64') self.assertEqual(rc2['multiparticles']['lepton'], rc['multiparticles']['lepton']) self.assertEqual(rc2['cuts']['jet-delta_r'], {'min': 0.4}) From a9ceb11a3efb0e41b33a06e65164f1ce4685fa6d Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 28 Oct 2026 13:57:16 +0100 Subject: [PATCH 16/24] madspace: train the MLP test network with a smooth activation With LeakyReLU, a pre-activation within rounding distance of the kink can flip the gradient of an event between madspace and torch in float32, which Adam turns into weight differences of the order of the learning rate. This made test_training[float32-0.0] fail for a few percent of initializations. Co-Authored-By: Claude Opus 5.5 --- madspace/tests/test_mlp.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/madspace/tests/test_mlp.py b/madspace/tests/test_mlp.py index c496c05a95..289fe28222 100644 --- a/madspace/tests/test_mlp.py +++ b/madspace/tests/test_mlp.py @@ -27,7 +27,8 @@ def precision(request): @pytest.fixture def mlp(precision): - return ms.MLP(10, 1, 32, 3, ms.MLP.leaky_relu, "", precision[0]) + # smooth activation: at a kink, float32 rounding can flip the gradient of an event + return ms.MLP(10, 1, 32, 3, ms.MLP.gelu, "", precision[0]) @pytest.fixture @@ -96,9 +97,9 @@ def test_training(mlp, mlp_module, module_context, precision, grad_clip): _, dtype, tol = precision mlp_torch = nn.Sequential( nn.Linear(10, 32), - nn.LeakyReLU(), + nn.GELU(), nn.Linear(32, 32), - nn.LeakyReLU(), + nn.GELU(), nn.Linear(32, 1), ).to(dtype) From 11be642cd903c3750a80e9db6843b658d4f57a8b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 11:57:01 +0100 Subject: [PATCH 17/24] madspace: 64-bit offsets in the x86 SIMD gathers and scatters The offsets index * stride were computed in 32 bits and wrapped for tensors spanning more than 2^31 elements, e.g. a large batch of a wide tensor. They are now 64-bit products of the 32-bit index and stride (_mul_epu32, which needs no AVX512DQ). The float64 gathers keep one instruction per vector; the float32 and int32 ones gather the two halves of a vector separately. Co-Authored-By: Claude Opus 5.5 --- madspace/src/cpu/simd_x86_256.hpp | 73 ++++++++++++-------- madspace/src/cpu/simd_x86_256y.hpp | 58 ++++++++++++---- madspace/src/cpu/simd_x86_512.hpp | 107 ++++++++++++++++++++--------- 3 files changed, 164 insertions(+), 74 deletions(-) diff --git a/madspace/src/cpu/simd_x86_256.hpp b/madspace/src/cpu/simd_x86_256.hpp index 9cc0e3a0b7..30ef451224 100644 --- a/madspace/src/cpu/simd_x86_256.hpp +++ b/madspace/src/cpu/simd_x86_256.hpp @@ -52,15 +52,22 @@ struct BVec { __m256d v; }; -inline __m128i stride_seq(std::size_t stride) { - return _mm_mullo_epi32(_mm_set1_epi32(stride), _mm_set_epi32(3, 2, 1, 0)); -} - -inline __m128i -mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) { - return _mm_add_epi32( - _mm_mullo_epi32(truncate_i64_to_i32(indices), _mm_set1_epi32(index_stride)), - stride_seq(batch_stride) +// 64-bit element offsets index * index_stride + lane * batch_stride of four lanes, +// from the low 32 bits of the 64-bit indices (indices >= 0, strides < 2^32) +inline __m256i mem_indices( + std::size_t batch_stride, + std::size_t index_stride, + __m256i indices, + long long first_lane = 0 +) { + return _mm256_add_epi64( + _mm256_mul_epu32(indices, _mm256_set1_epi64x(index_stride)), + _mm256_mul_epu32( + _mm256_set1_epi64x(batch_stride), + _mm256_add_epi64( + _mm256_set1_epi64x(first_lane), _mm256_set_epi64x(3, 2, 1, 0) + ) + ) ); } @@ -71,7 +78,7 @@ inline FVec vgather( std::size_t index_stride, IVec indices ) { - return _mm256_i32gather_pd( + return _mm256_i64gather_pd( base_ptr, mem_indices(batch_stride, index_stride, indices), 8 ); } @@ -83,7 +90,7 @@ inline IVec vgather( std::size_t index_stride, IVec indices ) { - return _mm256_cvtepi32_epi64(_mm_i32gather_epi32( + return _mm256_cvtepi32_epi64(_mm256_i64gather_epi32( base_ptr, mem_indices(batch_stride, index_stride, indices), 4 )); } @@ -274,13 +281,15 @@ struct BVec32 { __m256 v; }; +// 64-bit offsets of the lower (half = 0) or upper four of the eight lanes +template inline __m256i mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) { - return _mm256_add_epi32( - _mm256_mullo_epi32(indices, _mm256_set1_epi32(index_stride)), - _mm256_mullo_epi32( - _mm256_set1_epi32(batch_stride), _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0) - ) + return mem_indices( + batch_stride, + index_stride, + _mm256_cvtepi32_epi64(_mm256_extracti128_si256(indices, half)), + 4 * half ); } @@ -291,8 +300,13 @@ inline FVec32 vgather( std::size_t index_stride, IVec32 indices ) { - return _mm256_i32gather_ps( - base_ptr, mem_indices32(batch_stride, index_stride, indices), 4 + return _mm256_set_m128( + _mm256_i64gather_ps( + base_ptr, mem_indices32<1>(batch_stride, index_stride, indices), 4 + ), + _mm256_i64gather_ps( + base_ptr, mem_indices32<0>(batch_stride, index_stride, indices), 4 + ) ); } @@ -303,8 +317,13 @@ inline IVec32 vgather( std::size_t index_stride, IVec32 indices ) { - return _mm256_i32gather_epi32( - base_ptr, mem_indices32(batch_stride, index_stride, indices), 4 + return _mm256_set_m128i( + _mm256_i64gather_epi32( + base_ptr, mem_indices32<1>(batch_stride, index_stride, indices), 4 + ), + _mm256_i64gather_epi32( + base_ptr, mem_indices32<0>(batch_stride, index_stride, indices), 4 + ) ); } @@ -334,12 +353,9 @@ inline void vscatter( float values_buf[8]; int indices_buf[8]; _mm256_storeu_ps(values_buf, values); - _mm256_storeu_si256( - reinterpret_cast<__m256i*>(indices_buf), - mem_indices32(batch_stride, index_stride, indices) - ); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(indices_buf), indices); for (int i = 0; i < 8; ++i) { - base_ptr[indices_buf[i]] = values_buf[i]; + base_ptr[index_stride * indices_buf[i] + batch_stride * i] = values_buf[i]; } } @@ -353,12 +369,9 @@ inline void vscatter( int values_buf[8]; int indices_buf[8]; _mm256_storeu_si256(reinterpret_cast<__m256i*>(values_buf), values); - _mm256_storeu_si256( - reinterpret_cast<__m256i*>(indices_buf), - mem_indices32(batch_stride, index_stride, indices) - ); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(indices_buf), indices); for (int i = 0; i < 8; ++i) { - base_ptr[indices_buf[i]] = values_buf[i]; + base_ptr[index_stride * indices_buf[i] + batch_stride * i] = values_buf[i]; } } diff --git a/madspace/src/cpu/simd_x86_256y.hpp b/madspace/src/cpu/simd_x86_256y.hpp index 4d54bc55d0..3789ee7fd9 100644 --- a/madspace/src/cpu/simd_x86_256y.hpp +++ b/madspace/src/cpu/simd_x86_256y.hpp @@ -268,13 +268,17 @@ struct BVec32 { __mmask8 v; }; +// 64-bit offsets of the lower (half = 0) or upper four of the eight lanes +template inline __m256i mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) { - return _mm256_add_epi32( - _mm256_mullo_epi32(indices, _mm256_set1_epi32(index_stride)), - _mm256_mullo_epi32( - _mm256_set1_epi32(batch_stride), _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0) - ) + return _mm256_add_epi64( + mem_indices( + batch_stride, + index_stride, + _mm256_cvtepi32_epi64(_mm256_extracti128_si256(indices, half)) + ), + _mm256_set1_epi64x(4 * half * batch_stride) ); } @@ -285,8 +289,13 @@ inline FVec32 vgather( std::size_t index_stride, IVec32 indices ) { - return _mm256_i32gather_ps( - base_ptr, mem_indices32(batch_stride, index_stride, indices), 4 + return _mm256_set_m128( + _mm256_i64gather_ps( + base_ptr, mem_indices32<1>(batch_stride, index_stride, indices), 4 + ), + _mm256_i64gather_ps( + base_ptr, mem_indices32<0>(batch_stride, index_stride, indices), 4 + ) ); } @@ -297,8 +306,13 @@ inline IVec32 vgather( std::size_t index_stride, IVec32 indices ) { - return _mm256_i32gather_epi32( - base_ptr, mem_indices32(batch_stride, index_stride, indices), 4 + return _mm256_set_m128i( + _mm256_i64gather_epi32( + base_ptr, mem_indices32<1>(batch_stride, index_stride, indices), 4 + ), + _mm256_i64gather_epi32( + base_ptr, mem_indices32<0>(batch_stride, index_stride, indices), 4 + ) ); } @@ -325,8 +339,17 @@ inline void vscatter( IVec32 indices, FVec32 values ) { - _mm256_i32scatter_ps( - base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4 + _mm256_i64scatter_ps( + base_ptr, + mem_indices32<0>(batch_stride, index_stride, indices), + _mm256_castps256_ps128(values), + 4 + ); + _mm256_i64scatter_ps( + base_ptr, + mem_indices32<1>(batch_stride, index_stride, indices), + _mm256_extractf128_ps(values, 1), + 4 ); } @@ -337,8 +360,17 @@ inline void vscatter( IVec32 indices, IVec32 values ) { - _mm256_i32scatter_epi32( - base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4 + _mm256_i64scatter_epi32( + base_ptr, + mem_indices32<0>(batch_stride, index_stride, indices), + _mm256_castsi256_si128(values), + 4 + ); + _mm256_i64scatter_epi32( + base_ptr, + mem_indices32<1>(batch_stride, index_stride, indices), + _mm256_extracti128_si256(values, 1), + 4 ); } diff --git a/madspace/src/cpu/simd_x86_512.hpp b/madspace/src/cpu/simd_x86_512.hpp index 0a80a261dc..1e7021f4c2 100644 --- a/madspace/src/cpu/simd_x86_512.hpp +++ b/madspace/src/cpu/simd_x86_512.hpp @@ -43,19 +43,22 @@ struct BVec { __mmask8 v; }; -inline __m256i stride_seq(std::size_t stride) { - return _mm256_mullo_epi32( - _mm256_set1_epi32(stride), _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0) - ); -} - -inline __m256i -mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) { - return _mm256_add_epi32( - _mm256_mullo_epi32( - _mm512_cvtepi64_epi32(indices), _mm256_set1_epi32(index_stride) - ), - stride_seq(batch_stride) +// 64-bit element offsets index * index_stride + lane * batch_stride of eight lanes, +// from the low 32 bits of the 64-bit indices (indices >= 0, strides < 2^32) +inline __m512i mem_indices( + std::size_t batch_stride, + std::size_t index_stride, + __m512i indices, + long long first_lane = 0 +) { + return _mm512_add_epi64( + _mm512_mul_epu32(indices, _mm512_set1_epi64(index_stride)), + _mm512_mul_epu32( + _mm512_set1_epi64(batch_stride), + _mm512_add_epi64( + _mm512_set1_epi64(first_lane), _mm512_set_epi64(7, 6, 5, 4, 3, 2, 1, 0) + ) + ) ); } @@ -66,7 +69,7 @@ inline FVec vgather( std::size_t index_stride, IVec indices ) { - return _mm512_i32gather_pd( + return _mm512_i64gather_pd( mem_indices(batch_stride, index_stride, indices), base_ptr, 8 ); } @@ -78,8 +81,8 @@ inline IVec vgather( std::size_t index_stride, IVec indices ) { - return _mm512_cvtepi32_epi64(_mm256_i32gather_epi32( - base_ptr, mem_indices(batch_stride, index_stride, indices), 4 + return _mm512_cvtepi32_epi64(_mm512_i64gather_epi32( + mem_indices(batch_stride, index_stride, indices), base_ptr, 4 )); } @@ -108,7 +111,7 @@ inline void vscatter( IVec indices, FVec values ) { - _mm512_i32scatter_pd( + _mm512_i64scatter_pd( base_ptr, mem_indices(batch_stride, index_stride, indices), values, 8 ); } @@ -279,17 +282,30 @@ struct BVec32 { __mmask16 v; }; +// 64-bit offsets of the lower (half = 0) or upper eight of the sixteen lanes +template inline __m512i mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) { - return _mm512_add_epi32( - _mm512_mullo_epi32(indices, _mm512_set1_epi32(index_stride)), - _mm512_mullo_epi32( - _mm512_set1_epi32(batch_stride), - _mm512_set_epi32(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0) - ) + return mem_indices( + batch_stride, + index_stride, + _mm512_cvtepi32_epi64(_mm512_extracti64x4_epi64(indices, half)), + 8 * half ); } +// the halves of a float32 vector (no AVX512DQ for the 32-bit inserts and extracts) +inline __m512 join_halves(__m256 low, __m256 high) { + return _mm512_castpd_ps(_mm512_insertf64x4( + _mm512_castps_pd(_mm512_castps256_ps512(low)), _mm256_castps_pd(high), 1 + )); +} + +template +inline __m256 get_half(__m512 v) { + return _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(v), half)); +} + inline FVec32 vgather( VecTag, float* base_ptr, @@ -297,8 +313,13 @@ inline FVec32 vgather( std::size_t index_stride, IVec32 indices ) { - return _mm512_i32gather_ps( - mem_indices32(batch_stride, index_stride, indices), base_ptr, 4 + return join_halves( + _mm512_i64gather_ps( + mem_indices32<0>(batch_stride, index_stride, indices), base_ptr, 4 + ), + _mm512_i64gather_ps( + mem_indices32<1>(batch_stride, index_stride, indices), base_ptr, 4 + ) ); } @@ -309,8 +330,14 @@ inline IVec32 vgather( std::size_t index_stride, IVec32 indices ) { - return _mm512_i32gather_epi32( - mem_indices32(batch_stride, index_stride, indices), base_ptr, 4 + return _mm512_inserti64x4( + _mm512_castsi256_si512(_mm512_i64gather_epi32( + mem_indices32<0>(batch_stride, index_stride, indices), base_ptr, 4 + )), + _mm512_i64gather_epi32( + mem_indices32<1>(batch_stride, index_stride, indices), base_ptr, 4 + ), + 1 ); } @@ -337,8 +364,17 @@ inline void vscatter( IVec32 indices, FVec32 values ) { - _mm512_i32scatter_ps( - base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4 + _mm512_i64scatter_ps( + base_ptr, + mem_indices32<0>(batch_stride, index_stride, indices), + get_half<0>(values), + 4 + ); + _mm512_i64scatter_ps( + base_ptr, + mem_indices32<1>(batch_stride, index_stride, indices), + get_half<1>(values), + 4 ); } @@ -349,8 +385,17 @@ inline void vscatter( IVec32 indices, IVec32 values ) { - _mm512_i32scatter_epi32( - base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4 + _mm512_i64scatter_epi32( + base_ptr, + mem_indices32<0>(batch_stride, index_stride, indices), + _mm512_extracti64x4_epi64(values, 0), + 4 + ); + _mm512_i64scatter_epi32( + base_ptr, + mem_indices32<1>(batch_stride, index_stride, indices), + _mm512_extracti64x4_epi64(values, 1), + 4 ); } From 88ab19db7518f50609dd29afdc4e241de2554532 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 12:02:47 +0100 Subject: [PATCH 18/24] madspace: add the matmul bias after the gemm matmul prefilled its output with the bias and ran the gemm with beta = 1, while the GPU matmul_to_float64 ran it with beta = 0 and added the bias in its conversion kernel. The two round differently, so the lowered GPU graph was not bitwise equal to the unlowered one (and to MADSPACE_NO_FUSION). Both matmuls now run the gemm with beta = 0 and add the bias in the matmul precision afterwards, on CPU and GPU. The plain matmul stays at two GPU launches, the fused one at two. test_matmul_to_float64 now disables the lowering for its reference, which the GPU runtime would otherwise fuse as well, and requires bitwise equal outputs. Co-Authored-By: Claude Opus 5.5 --- madspace/src/cpu/runtime.cpp | 41 ++++++++++------- madspace/src/gpu/runtime.cu | 67 ++++++++++++++++++---------- madspace/tests/test_fused_float32.py | 27 ++++++----- 3 files changed, 85 insertions(+), 50 deletions(-) diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 428c3dea9b..1ee4e6fd94 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -248,26 +248,40 @@ void op_matrix_element( ); } +// out[j * batch_size + i] = in[...] + bias[j] after a gemm with beta = 0, in the +// precision F of the matmul. matmul and matmul_to_float64 share it to round alike. +template +void add_bias( + const F* in, const F* bias, O* out, std::size_t batch_size, std::size_t dims_out +) { + for (std::size_t j = 0; j < dims_out; ++j) { + for (std::size_t i = 0; i < batch_size; ++i) { + std::size_t k = j * batch_size + i; + out[k] = static_cast(in[k] + bias[j]); + } + } +} + template void op_matmul( const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device ) { auto input = locals[instruction.input_indices[0]].contiguous(device); auto weight = locals[instruction.input_indices[1]].contiguous(device); - auto bias_orig = locals[instruction.input_indices[2]]; - auto bias = bias_orig.contiguous(device); + auto bias = locals[instruction.input_indices[2]].contiguous(device); auto& output = locals[instruction.output_indices[0]]; std::size_t batch_size = input.size(0); std::size_t dims_in = input.size(1); std::size_t dims_out = weight.size(1); output = Tensor(dtype_of(), {batch_size, dims_out}, device); - output.copy_from(bias_orig, device); if (batch_size == 0) { return; } device.sync_barrier(); - device.submit([=]() mutable { + device.submit([=, out = output]() mutable { + auto* out_ptr = static_cast(out.data()); + auto* bias_ptr = static_cast(bias.data()); gemm( 'N', 'T', @@ -279,10 +293,11 @@ void op_matmul( batch_size, static_cast(weight.data()), dims_out, - 1., - static_cast(output.data()), + 0., + out_ptr, batch_size ); + add_bias(out_ptr, bias_ptr, out_ptr, batch_size, dims_out); }); } @@ -400,11 +415,6 @@ void op_matmul_to_float64( device.submit([=, out = output]() mutable { auto* tmp_ptr = static_cast(tmp.data()); auto* bias_ptr = static_cast(bias.data()); - for (std::size_t j = 0; j < dims_out; ++j) { - for (std::size_t i = 0; i < batch_size; ++i) { - tmp_ptr[j * batch_size + i] = bias_ptr[j]; - } - } gemm( 'N', 'T', @@ -416,14 +426,13 @@ void op_matmul_to_float64( batch_size, static_cast(weight.data()), dims_out, - 1., + 0., tmp_ptr, batch_size ); - auto* out_ptr = static_cast(out.data()); - for (std::size_t i = 0; i < batch_size * dims_out; ++i) { - out_ptr[i] = tmp_ptr[i]; - } + add_bias( + tmp_ptr, bias_ptr, static_cast(out.data()), batch_size, dims_out + ); }); } diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index f8c02d44c2..d69e8593f4 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -135,6 +135,34 @@ auto gpublas_gemm(Args... args) { } } +// out[j * batch_size + i] += bias[j], over the output of a gemm +template +__global__ void +kernel_add_bias(std::size_t total, std::size_t batch_size, const F* bias, F* out) { + std::size_t i = blockDim.x * static_cast(blockIdx.x) + threadIdx.x; + if (i < total) { + out[i] += bias[i / batch_size]; + } +} + +// out[j * batch_size + i] = double(tmp[...] + bias[j]): bias, convert to float64 in +// one pass, rounding exactly like kernel_add_bias followed by a conversion +template +__global__ void kernel_bias_to_double( + std::size_t total, + std::size_t batch_size, + const F* __restrict__ tmp, + const F* __restrict__ bias, + double* __restrict__ out +) { + std::size_t i = blockDim.x * static_cast(blockIdx.x) + threadIdx.x; + if (i < total) { + out[i] = static_cast(tmp[i] + bias[i / batch_size]); + } +} + +// the bias is added after the gemm (beta = 0), so that matmul_to_float64 rounds +// exactly like matmul followed by convert_float64 template void op_matmul( const GpuRuntime::Instruction& instruction, @@ -154,14 +182,13 @@ void op_matmul( output = Tensor( dtype_of(), {batch_size, dims_out}, device, instruction.output_alloc_hints[0] ); - output.copy_from(bias, device); if (batch_size == 0) { return; } gpublasHandle_t handle = instruction.runtime.gpublas_handle(); check_error(gpublasSetStream(handle, device.stream())); - F alpha = 1., beta = 1.; + F alpha = 1., beta = 0.; check_error( gpublas_gemm( handle, @@ -180,27 +207,21 @@ void op_matmul( batch_size ) ); + std::size_t total = batch_size * dims_out; + launch_kernel( + kernel_add_bias, + total, + device.stream(), + total, + batch_size, + static_cast(bias.data()), + static_cast(output.data()) + ); input.reset(device); weight.reset(device); bias.reset(device); } -// out[j * batch_size + i] = double(tmp[...] + bias[j]): bias, convert to float64 in -// one pass over the output of a float32 gemm -template -__global__ void kernel_bias_to_double( - std::size_t total, - std::size_t batch_size, - const F* __restrict__ tmp, - const F* __restrict__ bias, - double* __restrict__ out -) { - std::size_t i = blockDim.x * static_cast(blockIdx.x) + threadIdx.x; - if (i < total) { - out[i] = static_cast(tmp[i] + bias[i / batch_size]); - } -} - // matmul of float32 inputs with the float32 -> float64 conversion of the result // fused into the bias epilogue (three launches become two) void op_matmul_to_float64( @@ -1268,9 +1289,8 @@ void op_unweight( gpuStream_t stream = device.stream(); if (batch_size == 0) { - indices = Tensor( - DataType::dt_int32, {0}, device, instruction.output_alloc_hints[0] - ); + indices = + Tensor(DataType::dt_int32, {0}, device, instruction.output_alloc_hints[0]); uw_weights = Tensor( DataType::dt_float64, {0}, device, instruction.output_alloc_hints[1] ); @@ -1751,8 +1771,9 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) : // MADSPACE_NO_FUSION disables the lowering of precision conversions, to compare // the fused and the unfused instructions Function function = sort_breadth_first( - std::getenv("MADSPACE_NO_FUSION") ? function_arg - : fuse_conversions(function_arg) + std::getenv("MADSPACE_NO_FUSION") + ? function_arg + : fuse_conversions(function_arg) ); _locals_init.resize(function.locals().size()); diff --git a/madspace/tests/test_fused_float32.py b/madspace/tests/test_fused_float32.py index f4f3ac6249..3498b9b05f 100644 --- a/madspace/tests/test_fused_float32.py +++ b/madspace/tests/test_fused_float32.py @@ -39,7 +39,7 @@ def matmul_module(fused): return FunctionModule(fb.function(), new_context()) -def test_matmul_to_float64(): +def test_matmul_to_float64(monkeypatch): g = generator(77) x = randn((N, DIMS_IN), generator=g, dtype=torch.float32, requires_grad=True) w = randn((DIMS_OUT, DIMS_IN), generator=g, dtype=torch.float32, requires_grad=True) @@ -48,14 +48,20 @@ def test_matmul_to_float64(): results = [] for fused in (True, False): - xs, ws, bs = (t.detach().clone().requires_grad_() for t in (x, w, b)) - y = matmul_module(fused)(xs, ws, bs) - assert y.dtype == torch.float64 - (y * y_grad).sum().backward() - results.append((y.detach(), xs.grad, ws.grad, bs.grad)) + with monkeypatch.context() as mp: + # the GPU runtime would lower the unfused graph as well + if not fused: + mp.setenv("MADSPACE_NO_FUSION", "1") + xs, ws, bs = (t.detach().clone().requires_grad_() for t in (x, w, b)) + y = matmul_module(fused)(xs, ws, bs) + assert y.dtype == torch.float64 + (y * y_grad).sum().backward() + results.append((y.detach(), xs.grad, ws.grad, bs.grad)) y_ref = (x @ w.T + b).double() torch.testing.assert_close(results[0][0], y_ref, rtol=1e-5, atol=1e-5) + # the fused output rounds exactly like matmul + convert_float64 + torch.testing.assert_close(results[0][0], results[1][0], rtol=0, atol=0) # the gradients through the fused and the unfused graph agree for fused_t, unfused_t in zip(results[0], results[1]): torch.testing.assert_close(fused_t, unfused_t, rtol=1e-4, atol=1e-4) @@ -114,7 +120,9 @@ def test_conversion_lowering_graph(direction): assert after["matmul_float32"] == before["matmul_float32"] - 4 # the original function is not modified, and the pass is idempotent assert instruction_counts(function) == before - assert instruction_counts(ms.fuse_conversions(ms.fuse_conversions(function))) == after + assert ( + instruction_counts(ms.fuse_conversions(ms.fuse_conversions(function))) == after + ) @pytest.mark.parametrize("direction", ["forward", "inverse"]) @@ -154,7 +162,4 @@ def test_conversion_lowering_numerics(direction, monkeypatch): for lowered_t, ref_t in zip(results[True][:3], results[False][:3]): torch.testing.assert_close(lowered_t, ref_t, rtol=1e-12, atol=1e-12) for name, grad in results[True][3].items(): - torch.testing.assert_close( - grad, results[False][3][name], rtol=1e-4, atol=1e-6 - ) - + torch.testing.assert_close(grad, results[False][3][name], rtol=1e-4, atol=1e-6) From 7e899af574f1c69fe4b0a00c24534e034c3225b2 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 12:03:54 +0100 Subject: [PATCH 19/24] madspace: no gradients of integer inputs in structural backward ops The backward passes of stack, unstack, pop, batch_cat, batch_split, cat and rqs_reshape gave integer inputs float64 gradients, while backward_batch_foreach allocates int32 ones for the same inputs. An integer input reaching both kinds of op would then be read through a view of the wrong dtype. Integers carry no gradient, so these ops now skip them. Co-Authored-By: Claude Opus 5.5 --- madspace/src/kernels/operations.hpp | 43 ++++++++++++++++++++++------- 1 file changed, 33 insertions(+), 10 deletions(-) diff --git a/madspace/src/kernels/operations.hpp b/madspace/src/kernels/operations.hpp index 1937c8ae50..73eb2ef158 100644 --- a/madspace/src/kernels/operations.hpp +++ b/madspace/src/kernels/operations.hpp @@ -103,9 +103,11 @@ void backward_batch_foreach( foreach_func(args, input_grads, batch_size, device); } -// gradients have the dtype of float inputs, integer inputs get float64 gradients -inline DataType grad_dtype(const Tensor& input) { - return is_float(input.dtype()) ? input.dtype() : DataType::dt_float64; +// integers carry no gradient, so structural ops skip them (all their inputs share +// one dtype) +template +inline bool skip_int_grad(const I& instruction, const TensorVec& locals) { + return !is_float(locals[instruction.input_indices[0]].dtype()); } template @@ -133,12 +135,15 @@ template void backward_op_stack( const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device ) { + if (skip_int_grad(instruction, locals)) { + return; + } auto& output_grad = local_grads[instruction.output_indices[0]]; for (std::size_t i = 0; std::size_t input_index : instruction.input_indices) { auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - grad_dtype(locals[input_index]), + locals[input_index].dtype(), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[i] @@ -185,11 +190,14 @@ template void backward_op_unstack( const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device ) { + if (skip_int_grad(instruction, locals)) { + return; + } auto input_index = instruction.input_indices[0]; auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - grad_dtype(locals[input_index]), + locals[input_index].dtype(), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] @@ -215,11 +223,14 @@ template void backward_op_pop( const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device ) { + if (skip_int_grad(instruction, locals)) { + return; + } auto input_index = instruction.input_indices[0]; auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - grad_dtype(locals[input_index]), + locals[input_index].dtype(), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] @@ -267,12 +278,15 @@ template void backward_op_batch_cat( const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device ) { + if (skip_int_grad(instruction, locals)) { + return; + } auto output_grad = local_grads[instruction.output_indices[0]]; for (std::size_t i = 0; std::size_t input_index : instruction.input_indices) { auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - grad_dtype(locals[input_index]), + locals[input_index].dtype(), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[i] @@ -303,12 +317,15 @@ template void backward_op_batch_split( const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device ) { + if (skip_int_grad(instruction, locals)) { + return; + } auto& sizes = locals[instruction.input_indices[1]].batch_sizes(); auto input_index = instruction.input_indices[0]; auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - grad_dtype(locals[input_index]), + locals[input_index].dtype(), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] @@ -355,12 +372,15 @@ template void backward_op_cat( const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device ) { + if (skip_int_grad(instruction, locals)) { + return; + } auto& output_grad = local_grads[instruction.output_indices[0]]; for (std::size_t i = 0; std::size_t input_index : instruction.input_indices) { auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - grad_dtype(locals[input_index]), + locals[input_index].dtype(), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[i] @@ -468,13 +488,16 @@ template void backward_op_rqs_reshape( const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device ) { + if (skip_int_grad(instruction, locals)) { + return; + } auto n_dims = instruction.output_shapes[0][0]; auto n_bins = instruction.output_shapes[0][1]; auto input_index = instruction.input_indices[0]; auto& input_grad = local_grads[input_index]; if (!input_grad) { input_grad = Tensor( - grad_dtype(locals[input_index]), + locals[input_index].dtype(), locals[input_index].shape(), device, instruction.input_grad_alloc_hints[0] From e9e86c0852d4882ef5ddbbbb9f6b4f440242d030 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 12:32:39 +0100 Subject: [PATCH 20/24] madspace: fuse network activations into the matmul bias kernel The GPU lowering pass now also rewrites act(matmul(x, w, b)) for the MLP activations (relu, leaky_relu, elu, gelu, sigmoid, softplus) in both precisions into matmul_, which keeps the matmul result as an output for the backward pass. On the GPU it runs the gemm and one kernel that adds the bias and applies the activation (two launches instead of three per hidden layer). Its backward computes the activation gradient inside the bias gradient reduction, which saves a launch there as well. The activations are now structs of value functions, shared by the element-wise kernels and the fused ones, so that the fused instructions round exactly like the unfused ones. The CPU versions, only reached through an explicit fuse_conversions, use the scalar functions and can differ from the SIMD ones of the element-wise kernels in the last bits for the transcendental activations. custom_op in the instruction set can now name an implementation shared by several instructions, with custom_op_args as its leading template arguments. Co-Authored-By: Claude Opus 5.5 --- madspace/generate_headers.py | 21 +- .../include/madspace/compgraphs/optimizer.hpp | 10 +- madspace/instruction_set.yaml | 156 ++++++++++++++ madspace/src/compgraphs/optimizer.cpp | 29 ++- madspace/src/cpu/runtime.cpp | 204 +++++++++++++++--- madspace/src/gpu/runtime.cu | 168 ++++++++++++++- madspace/src/kernels/nn.hpp | 164 ++++++++++---- madspace/src/python/madspace.cpp | 6 +- madspace/tests/test_fused_float32.py | 62 +++++- 9 files changed, 719 insertions(+), 101 deletions(-) diff --git a/madspace/generate_headers.py b/madspace/generate_headers.py index e379291264..26560101d8 100644 --- a/madspace/generate_headers.py +++ b/madspace/generate_headers.py @@ -311,6 +311,19 @@ def instruction_set_mixin(commands, out_dir): f_op.write("\n") +def custom_op_func(prefix, cmd): + """The C++ function of a custom op. custom_op can name an implementation shared + by several instructions instead of True, with the leading template arguments + custom_op_args.""" + custom_op = cmd["custom_op"] + name = custom_op if isinstance(custom_op, str) else cmd["base_name"] + template_args = list(cmd.get("custom_op_args", [])) + if "ctype" in cmd: + template_args.append(cmd["ctype"]) + func = f"{prefix}{name}" + return f"{func}<{', '.join(template_args)}>" if template_args else func + + def runtime_mixin(commands, device, out_dir): with open(out_dir / device / "runtime_mixin.inc", "w") as f: write_autogen(f) @@ -320,9 +333,7 @@ def runtime_mixin(commands, device, out_dir): base_name = cmd["base_name"] suffix = cmd["types_suffix"] if cmd.get("custom_op", False): - func = f"op_{base_name}" - if "ctype" in cmd: - func += f"<{cmd['ctype']}>" + func = custom_op_func("op_", cmd) else: n_inputs = len(cmd["inputs"]) n_outputs = len(cmd["outputs"]) @@ -364,9 +375,7 @@ def runtime_backward_mixin(commands, device, out_dir): base_name = cmd["base_name"] suffix = cmd["types_suffix"] if cmd.get("custom_op", False): - func = f"backward_op_{base_name}" - if "ctype" in cmd: - func += f"<{cmd['ctype']}>" + func = custom_op_func("backward_op_", cmd) f.write( f"case {opcode}:\n" f" {func}(instr, locals, local_grads, device);\n" diff --git a/madspace/include/madspace/compgraphs/optimizer.hpp b/madspace/include/madspace/compgraphs/optimizer.hpp index 9614039625..1c653f85e1 100644 --- a/madspace/include/madspace/compgraphs/optimizer.hpp +++ b/madspace/include/madspace/compgraphs/optimizer.hpp @@ -34,14 +34,16 @@ class LastUseOfLocals { Function sort_breadth_first(const Function& function); /** - * Lowering pass for single-precision networks, run by the GPU backend when it - * builds a runtime. Fuses a precision conversion into the instruction that - * produces the converted value: + * Lowering pass for networks, run by the GPU backend when it builds a runtime. + * Fuses an element-wise instruction into the instruction that produces its input: * - `convert_float32(v)` where `v` is an output of `select`, `rqs_forward` or * `rqs_inverse` becomes the `*_with_float32` variant of the producer, which * writes the single-precision copy itself; * - `convert_float64(matmul_float32(x, w, b))`, with the matmul result used only - * by the conversion, becomes `matmul_to_float64`. + * by the conversion, becomes `matmul_to_float64`; + * - an activation `act(matmul(x, w, b))` (relu, leaky_relu, elu, gelu, sigmoid or + * softplus, in either precision) becomes `matmul_act`, which also keeps the + * matmul result as an output. * The graph keeps only the generic instructions and so stays portable; the fused * instructions never appear in a saved function. */ diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 52c5a050f2..8765b4554a 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -2984,6 +2984,162 @@ matmul: differentiable: True desc: "Affine layer y = weight x + bias. Building block of @ref MLP." +matmul_relu: + precision: [float32, float64] + inputs: + - name: x + type: [float, n] + desc: "input vector" + backward_arg: True + - name: weight + type: [float, single, m, n] + desc: "weight matrix" + backward_arg: True + - name: bias + type: [float, single, m] + desc: "bias vector" + outputs: + - name: linear + type: [float, m] + desc: "the affine transform of the input" + - name: y + type: [float, m] + desc: "the activated affine transform" + custom_op: matmul_activation + custom_op_args: [ReluActivation] + differentiable: True + desc: "@ref matmul followed by @ref relu, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref relu applied to the output of @ref matmul." + +matmul_leaky_relu: + precision: [float32, float64] + inputs: + - name: x + type: [float, n] + desc: "input vector" + backward_arg: True + - name: weight + type: [float, single, m, n] + desc: "weight matrix" + backward_arg: True + - name: bias + type: [float, single, m] + desc: "bias vector" + outputs: + - name: linear + type: [float, m] + desc: "the affine transform of the input" + - name: y + type: [float, m] + desc: "the activated affine transform" + custom_op: matmul_activation + custom_op_args: [LeakyReluActivation] + differentiable: True + desc: "@ref matmul followed by @ref leaky_relu, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref leaky_relu applied to the output of @ref matmul." + +matmul_elu: + precision: [float32, float64] + inputs: + - name: x + type: [float, n] + desc: "input vector" + backward_arg: True + - name: weight + type: [float, single, m, n] + desc: "weight matrix" + backward_arg: True + - name: bias + type: [float, single, m] + desc: "bias vector" + outputs: + - name: linear + type: [float, m] + desc: "the affine transform of the input" + - name: y + type: [float, m] + desc: "the activated affine transform" + custom_op: matmul_activation + custom_op_args: [EluActivation] + differentiable: True + desc: "@ref matmul followed by @ref elu, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref elu applied to the output of @ref matmul." + +matmul_gelu: + precision: [float32, float64] + inputs: + - name: x + type: [float, n] + desc: "input vector" + backward_arg: True + - name: weight + type: [float, single, m, n] + desc: "weight matrix" + backward_arg: True + - name: bias + type: [float, single, m] + desc: "bias vector" + outputs: + - name: linear + type: [float, m] + desc: "the affine transform of the input" + - name: y + type: [float, m] + desc: "the activated affine transform" + custom_op: matmul_activation + custom_op_args: [GeluActivation] + differentiable: True + desc: "@ref matmul followed by @ref gelu, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref gelu applied to the output of @ref matmul." + +matmul_sigmoid: + precision: [float32, float64] + inputs: + - name: x + type: [float, n] + desc: "input vector" + backward_arg: True + - name: weight + type: [float, single, m, n] + desc: "weight matrix" + backward_arg: True + - name: bias + type: [float, single, m] + desc: "bias vector" + outputs: + - name: linear + type: [float, m] + desc: "the affine transform of the input" + - name: y + type: [float, m] + desc: "the activated affine transform" + custom_op: matmul_activation + custom_op_args: [SigmoidActivation] + differentiable: True + desc: "@ref matmul followed by @ref sigmoid, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref sigmoid applied to the output of @ref matmul." + +matmul_softplus: + precision: [float32, float64] + inputs: + - name: x + type: [float, n] + desc: "input vector" + backward_arg: True + - name: weight + type: [float, single, m, n] + desc: "weight matrix" + backward_arg: True + - name: bias + type: [float, single, m] + desc: "bias vector" + outputs: + - name: linear + type: [float, m] + desc: "the affine transform of the input" + - name: y + type: [float, m] + desc: "the activated affine transform" + custom_op: matmul_activation + custom_op_args: [SoftplusActivation] + differentiable: True + desc: "@ref matmul followed by @ref softplus, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref softplus applied to the output of @ref matmul." + relu: precision: [float32, float64] inputs: diff --git a/madspace/src/compgraphs/optimizer.cpp b/madspace/src/compgraphs/optimizer.cpp index 0269c2043f..9576bcc67e 100644 --- a/madspace/src/compgraphs/optimizer.cpp +++ b/madspace/src/compgraphs/optimizer.cpp @@ -1,6 +1,7 @@ #include "madspace/compgraphs/optimizer.hpp" #include +#include #include #include #include @@ -109,12 +110,29 @@ Function madspace::fuse_conversions(const Function& function) { return find->second.get(); }; + // the precision of an activation instruction that fuses into a matmul, or empty + auto activation_precision = [](const std::string& name) -> std::string { + static const std::array activations{ + "relu", "leaky_relu", "elu", "gelu", "sigmoid", "softplus" + }; + for (const auto& activation : activations) { + for (std::string precision : {"float32", "float64"}) { + if (name == activation + "_" + precision) { + return precision; + } + } + } + return ""; + }; + std::vector removed(instructions.size()); bool any_fused = false; for (std::size_t i = 0; i < instructions.size(); ++i) { auto& call = instructions[i]; const std::string& name = call.instruction->name(); - if (name != "convert_float32" && name != "convert_float64") { + std::string precision = activation_precision(name); + if (name != "convert_float32" && name != "convert_float64" && + precision.empty()) { continue; } auto& input = call.inputs.at(0); @@ -131,7 +149,14 @@ Function madspace::fuse_conversions(const Function& function) { continue; } const std::string& source_name = source.instruction->name(); - if (name == "convert_float32") { + if (!precision.empty()) { + if (source_name != "matmul_" + precision) { + continue; + } + // the matmul output stays an output, the activation needs it backward + source.instruction = find_instruction("matmul_" + name); + source.outputs.push_back(call.outputs.at(0)); + } else if (name == "convert_float32") { bool fusable = source_name == "select" || source_name == "rqs_forward" || source_name == "rqs_inverse"; if (!fusable || source.outputs.at(0).local_index != input.local_index || diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 1ee4e6fd94..c196905f65 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -392,6 +392,70 @@ void backward_op_matmul( }); } +// input, weight and bias gradients of a matmul from the gradient of its output, for +// the backward passes that run in a single job +template +void matmul_grads( + const F* output_grad, + const F* input, + const F* weight, + F* input_grad, + F* weight_grad, + F* bias_grad, + std::size_t batch_size, + std::size_t dims_in, + std::size_t dims_out +) { + auto* grad = const_cast(output_grad); + // input_grad += output_grad * weight + gemm( + 'N', + 'N', + batch_size, + dims_in, + dims_out, + 1., + grad, + batch_size, + const_cast(weight), + dims_out, + 1., + input_grad, + batch_size + ); + // weight_grad += output_grad.T * input + gemm( + 'T', + 'N', + dims_out, + dims_in, + batch_size, + 1., + grad, + batch_size, + const_cast(input), + batch_size, + 1., + weight_grad, + dims_out + ); + // bias_grad += sum_i output_grad_ij + std::vector ones(batch_size, 1.); + gemv( + 'T', + batch_size, + dims_out, + 1., + grad, + batch_size, + ones.data(), + 1, + 1., + bias_grad, + 1 + ); +} + // matmul with a float32 input and a float64 output. Everything runs in a single job, // because the CPU runtime offers only one barrier per instruction. template @@ -477,52 +541,126 @@ void backward_op_matmul_to_float64( for (std::size_t i = 0; i < batch_size * dims_out; ++i) { grad_f_ptr[i] = static_cast(grad_ptr[i]); } - // input_grad += output_grad * weight - gemm( - 'N', - 'N', - batch_size, - dims_in, - dims_out, - 1., + matmul_grads( grad_f_ptr, - batch_size, + static_cast(input.data()), static_cast(weight.data()), - dims_out, - 1., static_cast(input_grad.data()), - batch_size + static_cast(weight_grad.data()), + static_cast(bias_grad.data()), + batch_size, + dims_in, + dims_out ); - // weight_grad += output_grad.T * input + }); +} + +// matmul followed by the activation Act, with the outputs linear (before the +// activation) and y. Everything runs in a single job, as in op_matmul_to_float64. +template +void op_matmul_activation( + const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device +) { + using T = CpuTypesT; + auto input = locals[instruction.input_indices[0]].contiguous(device); + auto weight = locals[instruction.input_indices[1]].contiguous(device); + auto bias = locals[instruction.input_indices[2]].contiguous(device); + auto& linear = locals[instruction.output_indices[0]]; + auto& output = locals[instruction.output_indices[1]]; + std::size_t batch_size = input.size(0); + std::size_t dims_in = input.size(1); + std::size_t dims_out = weight.size(1); + linear = Tensor(dtype_of(), {batch_size, dims_out}, device); + output = Tensor(dtype_of(), {batch_size, dims_out}, device); + if (batch_size == 0) { + return; + } + + device.sync_barrier(); + device.submit([=, lin = linear, out = output]() mutable { + auto* lin_ptr = static_cast(lin.data()); + auto* out_ptr = static_cast(out.data()); gemm( - 'T', 'N', + 'T', + batch_size, dims_out, dims_in, - batch_size, 1., - grad_f_ptr, - batch_size, static_cast(input.data()), batch_size, - 1., - static_cast(weight_grad.data()), - dims_out - ); - // bias_grad += sum_i output_grad_ij - std::vector ones(batch_size, 1.); - gemv( - 'T', - batch_size, + static_cast(weight.data()), dims_out, - 1., - grad_f_ptr, - batch_size, - ones.data(), - 1, - 1., + 0., + lin_ptr, + batch_size + ); + add_bias(lin_ptr, static_cast(bias.data()), lin_ptr, batch_size, dims_out); + for (std::size_t k = 0; k < batch_size * dims_out; ++k) { + out_ptr[k] = Act::template forward(lin_ptr[k]); + } + }); +} + +template +void backward_op_matmul_activation( + const CpuRuntime::Instruction& instruction, + TensorVec& locals, + TensorVec& local_grads, + const D& device +) { + using T = CpuTypesT; + auto input = locals[instruction.input_indices[0]].contiguous(device); + auto weight = locals[instruction.input_indices[1]].contiguous(device); + auto linear = locals[instruction.output_indices[0]].contiguous(device); + auto output = locals[instruction.output_indices[1]].contiguous(device); + auto linear_grad = local_grads[instruction.output_indices[0]].contiguous(device); + auto output_grad = local_grads[instruction.output_indices[1]].contiguous(device); + auto& input_grad = local_grads[instruction.input_indices[0]]; + auto& weight_grad = local_grads[instruction.input_indices[1]]; + auto& bias_grad = local_grads[instruction.input_indices[2]]; + std::size_t batch_size = input.size(0); + std::size_t dims_in = input.size(1); + std::size_t dims_out = weight.size(1); + + if (!input_grad) { + input_grad = + Tensor(dtype_of(), input.shape(), device, AllocHint::local_grad); + } + if (!weight_grad) { + weight_grad = + Tensor(dtype_of(), weight.shape(), device, AllocHint::local_grad); + } + if (!bias_grad) { + bias_grad = Tensor(dtype_of(), {1, dims_out}, device, AllocHint::local_grad); + } + if (batch_size == 0) { + return; + } + Tensor grad(dtype_of(), {batch_size, dims_out}, device); + + device.sync_barrier(); + device.submit([=]() mutable { + auto* grad_ptr = static_cast(grad.data()); + auto* lin_ptr = static_cast(linear.data()); + auto* out_ptr = static_cast(output.data()); + auto* lin_grad_ptr = static_cast(linear_grad.data()); + auto* out_grad_ptr = static_cast(output_grad.data()); + // the gradient of linear also gets the contributions of its other uses + for (std::size_t k = 0; k < batch_size * dims_out; ++k) { + grad_ptr[k] = lin_grad_ptr[k] + + Act::template backward(lin_ptr[k], out_ptr[k], out_grad_ptr[k]); + } + matmul_grads( + static_cast(grad_ptr), + static_cast(input.data()), + static_cast(weight.data()), + static_cast(input_grad.data()), + static_cast(weight_grad.data()), static_cast(bias_grad.data()), - 1 + batch_size, + dims_in, + dims_out ); }); } diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index d69e8593f4..f02ad9d056 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -289,6 +290,90 @@ void op_matmul_to_float64( bias.reset(device); } +// linear[j * batch_size + i] += bias[j] and y = Act(linear) in one pass over the +// output of a gemm, rounding exactly like kernel_add_bias and the activation kernel +template +__global__ void kernel_bias_activation( + std::size_t total, + std::size_t batch_size, + const F* __restrict__ bias, + F* __restrict__ linear, + F* __restrict__ y +) { + std::size_t i = blockDim.x * static_cast(blockIdx.x) + threadIdx.x; + if (i < total) { + F value = linear[i] + bias[i / batch_size]; + linear[i] = value; + y[i] = Act::template forward>(value); + } +} + +// matmul followed by the activation Act, with the outputs linear (before the +// activation, for the backward pass) and y: two launches instead of three +template +void op_matmul_activation( + const GpuRuntime::Instruction& instruction, + TensorVec& locals, + const AsyncGpuDevice& device +) { + auto input = + locals[instruction.input_indices[0]].contiguous(device, AllocHint::temporary); + auto weight = + locals[instruction.input_indices[1]].contiguous(device, AllocHint::temporary); + auto bias = + locals[instruction.input_indices[2]].contiguous(device, AllocHint::temporary); + auto& linear = locals[instruction.output_indices[0]]; + auto& output = locals[instruction.output_indices[1]]; + std::size_t batch_size = input.size(0); + std::size_t dims_in = input.size(1); + std::size_t dims_out = weight.size(1); + linear = Tensor( + dtype_of(), {batch_size, dims_out}, device, instruction.output_alloc_hints[0] + ); + output = Tensor( + dtype_of(), {batch_size, dims_out}, device, instruction.output_alloc_hints[1] + ); + if (batch_size == 0) { + return; + } + + gpublasHandle_t handle = instruction.runtime.gpublas_handle(); + check_error(gpublasSetStream(handle, device.stream())); + F alpha = 1., beta = 0.; + check_error( + gpublas_gemm( + handle, + GPUBLAS_OP_N, + GPUBLAS_OP_T, + batch_size, + dims_out, + dims_in, + &alpha, + static_cast(input.data()), + batch_size, + static_cast(weight.data()), + dims_out, + &beta, + static_cast(linear.data()), + batch_size + ) + ); + std::size_t total = batch_size * dims_out; + launch_kernel( + kernel_bias_activation, + total, + device.stream(), + total, + batch_size, + static_cast(bias.data()), + static_cast(linear.data()), + static_cast(output.data()) + ); + input.reset(device); + weight.reset(device); + bias.reset(device); +} + // bias_grad[j] += sum_i output_grad[i, j]; the matrix is column-major with // leading dimension batch_size, so each column is a contiguous segment. One block // per column: single launch, no atomics (deterministic) and no temporary storage. @@ -337,9 +422,42 @@ __global__ void kernel_bias_grad_convert( } } -// OutGradF64: the gradient of the output is float64 (matmul_to_float64) and is -// converted to F together with the bias gradient reduction -template +// the same reduction over the gradient of the input of the activation Act of +// matmul_, linear_grad += Act'(linear, y) * y_grad, which is written +// back for the gemms. It rounds like the activation backward and kernel_bias_grad. +template +__global__ void kernel_bias_grad_activation( + std::size_t batch_size, + const F* __restrict__ linear, + const F* __restrict__ y, + const F* __restrict__ y_grad, + F* __restrict__ linear_grad, + F* bias_grad +) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + std::size_t offset = blockIdx.x * batch_size; + double sum = 0.; + for (std::size_t i = threadIdx.x; i < batch_size; i += BIAS_GRAD_THREADS) { + std::size_t k = offset + i; + F value = linear_grad[k] + + Act::template backward>(linear[k], y[k], y_grad[k]); + linear_grad[k] = value; + sum += value; + } + sum = BlockReduce(temp_storage).Sum(sum); + if (threadIdx.x == 0) { + bias_grad[blockIdx.x] += static_cast(sum); + } +} + +// How the gradient of the matmul output is obtained: as is (matmul), converted +// from float64 (matmul_to_float64) or through an activation struct (matmul_), +// in the last two cases together with the bias gradient reduction +struct OutputGrad {}; +struct OutputGradFromFloat64 {}; + +template void backward_matmul_impl( const GpuRuntime::Instruction& instruction, TensorVec& locals, @@ -387,7 +505,7 @@ void backward_matmul_impl( gpuStream_t stream = device.stream(); check_error(gpublasSetStream(handle, stream)); - if constexpr (OutGradF64) { + if constexpr (std::is_same_v) { // compute bias_grad += sum_i output_grad_ij and convert the gradient Tensor output_grad_f( dtype_of(), output_grad.shape(), device, AllocHint::temporary @@ -401,6 +519,30 @@ void backward_matmul_impl( check_error(); output_grad.reset(device); output_grad = output_grad_f; + } else if constexpr (!std::is_same_v) { + // the gradient of linear (output_grad) and bias_grad through the activation + auto linear = locals[instruction.output_indices[0]].contiguous( + device, AllocHint::temporary + ); + auto y = locals[instruction.output_indices[1]].contiguous( + device, AllocHint::temporary + ); + auto y_grad = local_grads[instruction.output_indices[1]].contiguous( + device, AllocHint::temporary + ); + kernel_bias_grad_activation + <<>>( + batch_size, + static_cast(linear.data()), + static_cast(y.data()), + static_cast(y_grad.data()), + static_cast(output_grad.data()), + static_cast(bias_grad.data()) + ); + check_error(); + linear.reset(device); + y.reset(device); + y_grad.reset(device); } // compute input_grad += output_grad * weight @@ -443,7 +585,7 @@ void backward_matmul_impl( ) ); - if constexpr (!OutGradF64) { + if constexpr (std::is_same_v) { // compute bias_grad += sum_i output_grad_ij kernel_bias_grad<<>>( batch_size, @@ -464,7 +606,7 @@ void backward_op_matmul( TensorVec& local_grads, const AsyncGpuDevice& device ) { - backward_matmul_impl(instruction, locals, local_grads, device); + backward_matmul_impl(instruction, locals, local_grads, device); } void backward_op_matmul_to_float64( @@ -473,7 +615,19 @@ void backward_op_matmul_to_float64( TensorVec& local_grads, const AsyncGpuDevice& device ) { - backward_matmul_impl(instruction, locals, local_grads, device); + backward_matmul_impl( + instruction, locals, local_grads, device + ); +} + +template +void backward_op_matmul_activation( + const GpuRuntime::Instruction& instruction, + TensorVec& locals, + TensorVec& local_grads, + const AsyncGpuDevice& device +) { + backward_matmul_impl(instruction, locals, local_grads, device); } struct NotMinusOne { diff --git a/madspace/src/kernels/nn.hpp b/madspace/src/kernels/nn.hpp index 3c928c7bb7..8dc1bb758d 100644 --- a/madspace/src/kernels/nn.hpp +++ b/madspace/src/kernels/nn.hpp @@ -10,101 +10,167 @@ namespace kernels { // The activations are instantiated in float32 and float64. Literals are cast to // FVal so that float32 code does not get promoted to double. +// +// Each activation is a struct with the value functions forward(x) and +// backward(x, y, g), the latter the input gradient for the input x, the output y +// and the output gradient g. The element-wise kernels and the matmuls with a fused +// activation share them, so that both round alike. + +struct ReluActivation { + template + static KERNELSPEC FVal forward(FVal x) { + using F = FVal; + return where(x < F(0.), F(0.), x); + } + template + static KERNELSPEC FVal backward(FVal x, FVal y, FVal g) { + using F = FVal; + return where(x < F(0.), F(0.), g); + } +}; + +struct LeakyReluActivation { + template + static KERNELSPEC FVal forward(FVal x) { + using F = FVal; + return where(x < F(0.), x * F(0.01), x); + } + template + static KERNELSPEC FVal backward(FVal x, FVal y, FVal g) { + using F = FVal; + return where(x < F(0.), g * F(0.01), g); + } +}; + +struct EluActivation { + template + static KERNELSPEC FVal forward(FVal x) { + using F = FVal; + return where(x < F(0.), expm1(x), x); + } + template + static KERNELSPEC FVal backward(FVal x, FVal y, FVal g) { + using F = FVal; + return where(x < F(0.), g * exp(x), g); + } +}; + +struct GeluActivation { + template + static KERNELSPEC FVal forward(FVal x) { + using F = FVal; + return F(0.5) * x * (F(1.) + erf(F(SQRT_HALF) * x)); + } + template + static KERNELSPEC FVal backward(FVal x, FVal y, FVal g) { + using F = FVal; + F cdf = F(0.5) * (F(1.) + erf(F(SQRT_HALF) * x)); + F pdf = F(0.5 * SQRT_HALF * TWO_DIV_SQRT_PI) * exp(F(-0.5) * x * x); + return g * (cdf + x * pdf); + } +}; + +// the backward pass uses the output +struct SigmoidActivation { + template + static KERNELSPEC FVal forward(FVal x) { + using F = FVal; + return F(1.) / (F(1.) + exp(-x)); + } + template + static KERNELSPEC FVal backward(FVal x, FVal y, FVal g) { + using F = FVal; + return g * y * (F(1.) - y); + } +}; + +struct SoftplusActivation { + template + static KERNELSPEC FVal forward(FVal x) { + using F = FVal; + return where(x < F(20.), log1p(exp(x)), x); + } + template + static KERNELSPEC FVal backward(FVal x, FVal y, FVal g) { + using F = FVal; + F z = exp(x); + return where(x < F(20.), g * z / (F(1.) + z), g); + } +}; template KERNELSPEC void kernel_relu(FIn input, FOut output) { - using F = FVal; - F x(input); - output = where(x < F(0.), F(0.), x); + output = ReluActivation::forward(FVal(input)); } template KERNELSPEC void backward_kernel_relu(FIn input, FIn output_grad, FOut input_grad) { - using F = FVal; - F x(input), g(output_grad); - input_grad += where(x < F(0.), F(0.), g); + FVal x(input); + input_grad += ReluActivation::backward(x, x, FVal(output_grad)); } template KERNELSPEC void kernel_leaky_relu(FIn input, FOut output) { - using F = FVal; - F x(input); - output = where(x < F(0.), x * F(0.01), x); + output = LeakyReluActivation::forward(FVal(input)); } template KERNELSPEC void backward_kernel_leaky_relu( FIn input, FIn output_grad, FOut input_grad ) { - using F = FVal; - F x(input), g(output_grad); - input_grad += where(x < F(0.), g * F(0.01), g); + FVal x(input); + input_grad += LeakyReluActivation::backward(x, x, FVal(output_grad)); } template KERNELSPEC void kernel_elu(FIn input, FOut output) { - using F = FVal; - F x(input); - output = where(x < F(0.), expm1(x), x); + output = EluActivation::forward(FVal(input)); } template KERNELSPEC void backward_kernel_elu(FIn input, FIn output_grad, FOut input_grad) { - using F = FVal; - F x(input), g(output_grad); - input_grad += where(x < F(0.), g * exp(x), g); + FVal x(input); + input_grad += EluActivation::backward(x, x, FVal(output_grad)); } template KERNELSPEC void kernel_gelu(FIn input, FOut output) { - using F = FVal; - F x(input); - output = F(0.5) * x * (F(1.) + erf(F(SQRT_HALF) * x)); + output = GeluActivation::forward(FVal(input)); } template KERNELSPEC void backward_kernel_gelu(FIn input, FIn output_grad, FOut input_grad) { - using F = FVal; - F x(input), g(output_grad); - F cdf = F(0.5) * (F(1.) + erf(F(SQRT_HALF) * x)); - F pdf = F(0.5 * SQRT_HALF * TWO_DIV_SQRT_PI) * exp(F(-0.5) * x * x); - input_grad += g * (cdf + x * pdf); + FVal x(input); + input_grad += GeluActivation::backward(x, x, FVal(output_grad)); } template KERNELSPEC void kernel_sigmoid(FIn input, FOut output) { - using F = FVal; - F x(input); - output = F(1.) / (F(1.) + exp(-x)); + output = SigmoidActivation::forward(FVal(input)); } template KERNELSPEC void backward_kernel_sigmoid( FIn output, FIn output_grad, FOut input_grad ) { - using F = FVal; - F y(output), g(output_grad); - input_grad += g * y * (F(1.) - y); + FVal y(output); + input_grad += SigmoidActivation::backward(y, y, FVal(output_grad)); } template KERNELSPEC void kernel_softplus(FIn input, FOut output) { - using F = FVal; - F x(input); - output = where(x < F(20.), log1p(exp(x)), x); + output = SoftplusActivation::forward(FVal(input)); } template KERNELSPEC void backward_kernel_softplus( FIn input, FIn output_grad, FOut input_grad ) { - using F = FVal; - F x(input), g(output_grad); - F z = exp(x); - input_grad += where(x < F(20.), g * z / (F(1.) + z), g); + FVal x(input); + input_grad += SoftplusActivation::backward(x, x, FVal(output_grad)); } template @@ -300,7 +366,8 @@ KERNELSPEC RqsValues rqs_forward_values(FIn input, FIn condition) auto derivative_numerator = delta * delta * (derivative_plus_one * theta_theta + two_delta * theta_one_minus_theta + derivative * one_minus_theta * one_minus_theta); - FVal det_value = where(clamp, 1., derivative_numerator / (denominator * denominator)); + FVal det_value = + where(clamp, 1., derivative_numerator / (denominator * denominator)); return {output_value, det_value}; } @@ -475,7 +542,12 @@ KERNELSPEC void backward_kernel_rqs_forward( FOut condition_grad ) { backward_rqs_forward_impl( - input, condition, FVal(output_grad), FVal(det_grad), input_grad, condition_grad + input, + condition, + FVal(output_grad), + FVal(det_grad), + input_grad, + condition_grad ); } @@ -543,7 +615,8 @@ KERNELSPEC RqsValues rqs_inverse_values(FIn input, FIn condition) auto derivative_numerator = delta * delta * (derivative_plus_one * theta * theta + two_delta * theta_one_minus_theta + derivative * one_minus_theta * one_minus_theta); - FVal det_value = where(clamp, 1., denominator * denominator / derivative_numerator); + FVal det_value = + where(clamp, 1., denominator * denominator / derivative_numerator); return {output_value, det_value}; } @@ -748,7 +821,12 @@ KERNELSPEC void backward_kernel_rqs_inverse( FOut condition_grad ) { backward_rqs_inverse_impl( - input, condition, FVal(output_grad), FVal(det_grad), input_grad, condition_grad + input, + condition, + FVal(output_grad), + FVal(det_grad), + input_grad, + condition_grad ); } diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 9a78c2be88..4213637897 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -290,9 +290,9 @@ PYBIND11_MODULE(_madspace_py, m) { "fuse_conversions", &fuse_conversions, py::arg("function"), - "Fuses single-precision conversions into the instructions producing the " - "converted values (the lowering the GPU backend applies when it builds a " - "runtime)." + "Fuses single-precision conversions and network activations into the " + "instructions producing their inputs (the lowering the GPU backend applies " + "when it builds a runtime)." ); py::classh device(m, "Device", pydoc::doc("Device")); diff --git a/madspace/tests/test_fused_float32.py b/madspace/tests/test_fused_float32.py index 3498b9b05f..e408387efc 100644 --- a/madspace/tests/test_fused_float32.py +++ b/madspace/tests/test_fused_float32.py @@ -4,7 +4,7 @@ import pytest import torch -from madspace_testing import generator, new_context, rand, randn +from madspace_testing import ON_GPU, generator, new_context, rand, randn import madspace as ms from madspace.torch import FunctionModule @@ -79,6 +79,8 @@ def test_matmul_to_float64(monkeypatch): "convert_float64", "matmul_float32", "matmul_to_float64", + "leaky_relu_float32", + "matmul_leaky_relu_float32", "select_with_float32", "rqs_forward_with_float32", "rqs_inverse_with_float32", @@ -116,8 +118,9 @@ def test_conversion_lowering_graph(direction): assert "convert_float32" not in after and "convert_float64" not in after assert after["matmul_to_float64"] == 4 assert after["select_with_float32"] + after[f"{spline}_with_float32"] == 4 - # the unfused matmuls are the hidden layers - assert after["matmul_float32"] == before["matmul_float32"] - 4 + # the hidden layers fuse their activation + assert before["leaky_relu_float32"] == after["matmul_leaky_relu_float32"] == 8 + assert "matmul_float32" not in after and "leaky_relu_float32" not in after # the original function is not modified, and the pass is idempotent assert instruction_counts(function) == before assert ( @@ -163,3 +166,56 @@ def test_conversion_lowering_numerics(direction, monkeypatch): torch.testing.assert_close(lowered_t, ref_t, rtol=1e-12, atol=1e-12) for name, grad in results[True][3].items(): torch.testing.assert_close(grad, results[False][3][name], rtol=1e-4, atol=1e-6) + + +ACTIVATIONS = ("relu", "leaky_relu", "elu", "gelu", "sigmoid", "softplus") + + +@pytest.mark.parametrize("precision", ["float32", "float64"]) +@pytest.mark.parametrize("activation", ACTIVATIONS) +def test_matmul_activation(activation, precision, monkeypatch): + """An MLP with fused activations agrees with the unfused one. The two round + alike, except that the CPU runs the fused activation with the scalar instead of + the SIMD functions, which differ for the transcendental ones.""" + g = generator(31) + dtype = getattr(ms.DataType, precision) + torch_dtype = getattr(torch, precision) + mlp = ms.MLP( + 5, 3, 8, 3, getattr(ms.MLP, activation), f"act_{activation}_{precision}", dtype + ) + ctx = new_context() + mlp.initialize_globals(ctx, 7) + for name in ctx.global_names(): + weights = ctx.get_global(name).torch() + weights[:] = randn(weights.shape, generator=g, dtype=torch.float64).to( + weights.dtype + ) + function = mlp.function() + lowered = ms.fuse_conversions(function) + assert instruction_counts_any(lowered, f"matmul_{activation}_{precision}") == 2 + assert instruction_counts_any(lowered, f"{activation}_{precision}") == 0 + + x = randn((97, 5), generator=g, dtype=torch.float64) + y_grad = randn((97, 3), generator=g, dtype=torch.float64) + results = {} + for fused in (True, False): + with monkeypatch.context() as mp: + if not fused: + mp.setenv("MADSPACE_NO_FUSION", "1") + module = FunctionModule(lowered if fused else function, ctx) + xs = x.clone().requires_grad_() + y = module(xs) + (y * y_grad).sum().backward() + results[fused] = [ + y.detach(), + xs.grad, + *(p.grad.clone() for p in module.global_params.values()), + ] + exact = ON_GPU or activation in ("relu", "leaky_relu") + tol = 0 if exact else (1e-5 if torch_dtype == torch.float32 else 1e-12) + for fused_t, ref_t in zip(results[True], results[False]): + torch.testing.assert_close(fused_t, ref_t, rtol=tol, atol=tol) + + +def instruction_counts_any(function, name): + return len(re.findall(rf"\b{name}\(", str(function))) From e5a0eea41a1208e480c5028d146899c5001ab4ca Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 12:40:53 +0100 Subject: [PATCH 21/24] madspace: rename fuse_conversions to fuse_instructions The lowering pass fuses network activations as well as conversions now. Co-Authored-By: Claude Opus 5.5 --- madspace/include/madspace/compgraphs/function.hpp | 2 +- madspace/include/madspace/compgraphs/optimizer.hpp | 2 +- madspace/src/compgraphs/optimizer.cpp | 2 +- madspace/src/gpu/runtime.cu | 2 +- madspace/src/python/madspace.cpp | 4 ++-- madspace/tests/test_fused_float32.py | 9 +++++---- 6 files changed, 11 insertions(+), 10 deletions(-) diff --git a/madspace/include/madspace/compgraphs/function.hpp b/madspace/include/madspace/compgraphs/function.hpp index d2d2ac42c2..25ba67fede 100644 --- a/madspace/include/madspace/compgraphs/function.hpp +++ b/madspace/include/madspace/compgraphs/function.hpp @@ -81,7 +81,7 @@ class Function { std::vector _instructions; friend Function sort_breadth_first(const Function& function); - friend Function fuse_conversions(const Function& function); + friend Function fuse_instructions(const Function& function); }; std::ostream& operator<<(std::ostream& out, const Value& value); diff --git a/madspace/include/madspace/compgraphs/optimizer.hpp b/madspace/include/madspace/compgraphs/optimizer.hpp index 1c653f85e1..546902f59a 100644 --- a/madspace/include/madspace/compgraphs/optimizer.hpp +++ b/madspace/include/madspace/compgraphs/optimizer.hpp @@ -47,6 +47,6 @@ Function sort_breadth_first(const Function& function); * The graph keeps only the generic instructions and so stays portable; the fused * instructions never appear in a saved function. */ -Function fuse_conversions(const Function& function); +Function fuse_instructions(const Function& function); } // namespace madspace diff --git a/madspace/src/compgraphs/optimizer.cpp b/madspace/src/compgraphs/optimizer.cpp index 9576bcc67e..b85f1c483c 100644 --- a/madspace/src/compgraphs/optimizer.cpp +++ b/madspace/src/compgraphs/optimizer.cpp @@ -83,7 +83,7 @@ Function madspace::sort_breadth_first(const Function& function) { return func_out; } -Function madspace::fuse_conversions(const Function& function) { +Function madspace::fuse_instructions(const Function& function) { Function func_out = function; auto& instructions = func_out._instructions; std::vector producer(function.locals().size(), -1); diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index f02ad9d056..402fc03d73 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -1927,7 +1927,7 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) : Function function = sort_breadth_first( std::getenv("MADSPACE_NO_FUSION") ? function_arg - : fuse_conversions(function_arg) + : fuse_instructions(function_arg) ); _locals_init.resize(function.locals().size()); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 4213637897..5de606ea41 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -287,8 +287,8 @@ PYBIND11_MODULE(_madspace_py, m) { ); m.def( - "fuse_conversions", - &fuse_conversions, + "fuse_instructions", + &fuse_instructions, py::arg("function"), "Fuses single-precision conversions and network activations into the " "instructions producing their inputs (the lowering the GPU backend applies " diff --git a/madspace/tests/test_fused_float32.py b/madspace/tests/test_fused_float32.py index e408387efc..4480038301 100644 --- a/madspace/tests/test_fused_float32.py +++ b/madspace/tests/test_fused_float32.py @@ -111,7 +111,7 @@ def test_conversion_lowering_graph(direction): flow = make_flow("lowering") function = getattr(flow, f"{direction}_function")() before = instruction_counts(function) - after = instruction_counts(ms.fuse_conversions(function)) + after = instruction_counts(ms.fuse_instructions(function)) spline = "rqs_forward" if direction == "forward" else "rqs_inverse" # every MLP of the flow converts its input and its output assert before["convert_float32"] == before["convert_float64"] == 4 @@ -124,7 +124,8 @@ def test_conversion_lowering_graph(direction): # the original function is not modified, and the pass is idempotent assert instruction_counts(function) == before assert ( - instruction_counts(ms.fuse_conversions(ms.fuse_conversions(function))) == after + instruction_counts(ms.fuse_instructions(ms.fuse_instructions(function))) + == after ) @@ -151,7 +152,7 @@ def test_conversion_lowering_numerics(direction, monkeypatch): if not lowered: mp.setenv("MADSPACE_NO_FUSION", "1") module = FunctionModule( - ms.fuse_conversions(function) if lowered else function, ctx + ms.fuse_instructions(function) if lowered else function, ctx ) r_grad = r.clone().requires_grad_() out, det = module(r_grad) @@ -191,7 +192,7 @@ def test_matmul_activation(activation, precision, monkeypatch): weights.dtype ) function = mlp.function() - lowered = ms.fuse_conversions(function) + lowered = ms.fuse_instructions(function) assert instruction_counts_any(lowered, f"matmul_{activation}_{precision}") == 2 assert instruction_counts_any(lowered, f"{activation}_{precision}") == 0 From a17042819d63227d7445e103c55ba02d4d0fd6a1 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 12:44:59 +0100 Subject: [PATCH 22/24] madspace: atomic accumulation in GPU tensor views, GPU-safe test_float32 GpuTensorView::operator+= was a plain read-modify-write, so the gradient of a broadcast input (batch stride 0) and scatter_add of repeated indices raced between threads. Use atomicAdd for floating-point types. test_float32.py drew from a CPU generator while the default device is the test device; draw on the CPU and move to the device, use the test-device context, and loosen two float32 tolerances that depended on the random draw. Co-Authored-By: Claude Sonnet 5.5 --- madspace/src/gpu/tensor.cuh | 10 +++++++- madspace/tests/test_float32.py | 47 ++++++++++++++++++++++------------ 2 files changed, 40 insertions(+), 17 deletions(-) diff --git a/madspace/src/gpu/tensor.cuh b/madspace/src/gpu/tensor.cuh index 2bbfd73d10..f645280feb 100644 --- a/madspace/src/gpu/tensor.cuh +++ b/madspace/src/gpu/tensor.cuh @@ -5,6 +5,7 @@ #include "madspace/driver/tensor.hpp" #include +#include namespace madspace { namespace gpu { @@ -94,10 +95,17 @@ public: return value; } + // atomic for floating-point types: the gradient of a broadcast input (batch + // stride 0) and the scatter_add of repeated indices accumulate from many threads + // into the same element __device__ T operator+=(T value) requires(_dim == 0) { - *_data += value; + if constexpr (std::is_floating_point_v) { + atomicAdd(_data, value); + } else { + *_data += value; + } return value; } diff --git a/madspace/tests/test_float32.py b/madspace/tests/test_float32.py index db9022e598..ad70dc35bd 100644 --- a/madspace/tests/test_float32.py +++ b/madspace/tests/test_float32.py @@ -5,10 +5,25 @@ import madspace as ms from madspace.torch import FunctionModule +from madspace_testing import DEVICE, new_context # a local generator, so that the global RNG state seen by other tests is unchanged GEN = torch.Generator().manual_seed(1234) + +def _draw(draw, *args, requires_grad=False, **kwargs): + """Random numbers from the CPU generator, moved to the device of the tests.""" + values = draw(*args, device="cpu", **kwargs).to(DEVICE) + return values.requires_grad_(requires_grad) + + +def _randn(*args, **kwargs): + return _draw(torch.randn, *args, **kwargs) + + +def _rand(*args, **kwargs): + return _draw(torch.rand, *args, **kwargs) + N = 103 DIMS_IN = 7 DIMS_OUT = 5 @@ -45,11 +60,11 @@ def test_activation(name, precision): module = FunctionModule(fb.function()) x = ( - 3 * torch.randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN) + 3 * _randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN) ).requires_grad_() out = module(x) assert out.dtype == torch_dtype - out_grad = torch.randn(out.shape, dtype=out.dtype, generator=GEN) + out_grad = _randn(out.shape, dtype=out.dtype, generator=GEN) (out * out_grad).sum().backward() x_ref = x.detach().clone().requires_grad_() @@ -75,15 +90,15 @@ def test_matmul(precision): fb.output(0, fb.matmul(fb.input(0), fb.input(1), fb.input(2))) module = FunctionModule(fb.function()) - x = torch.randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN, requires_grad=True) + x = _randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN, requires_grad=True) # single inputs are passed without batch dimension - w = torch.randn( + w = _randn( (DIMS_OUT, DIMS_IN), dtype=torch_dtype, generator=GEN, requires_grad=True ) - b = torch.randn((DIMS_OUT,), dtype=torch_dtype, generator=GEN, requires_grad=True) + b = _randn((DIMS_OUT,), dtype=torch_dtype, generator=GEN, requires_grad=True) y = module(x, w, b) assert y.dtype == torch_dtype - y_grad = torch.randn(y.shape, dtype=y.dtype, generator=GEN) + y_grad = _randn(y.shape, dtype=y.dtype, generator=GEN) (y * y_grad).sum().backward() refs = [t.detach().clone().requires_grad_() for t in (x, w, b)] @@ -107,7 +122,7 @@ def test_convert_roundtrip(): fb.output(1, fb.convert_float64(fb.relu(single))) module = FunctionModule(fb.function()) - x = torch.randn((N, 3), dtype=torch.float64, generator=GEN, requires_grad=True) + x = _randn((N, 3), dtype=torch.float64, generator=GEN, requires_grad=True) single_out, out = module(x) assert single_out.dtype == torch.float32 assert out.dtype == torch.float64 @@ -125,7 +140,7 @@ def test_literal_converted_for_float32(): # float64 literals are converted when the instruction dispatches to float32 fb.output(0, fb.matmul(fb.input(0), ms.Value([2.0], [1, 1]), ms.Value([1.0]))) module = FunctionModule(fb.function()) - x = torch.randn((N, 1), dtype=torch.float32, generator=GEN) + x = _randn((N, 1), dtype=torch.float32, generator=GEN) torch.testing.assert_close(module(x), 2 * x + 1) @@ -133,7 +148,7 @@ def test_flow_float32_subnets(tmp_path): """A flow with float32 subnetworks matches the float64 one with the same (rounded) weights; the Jacobian and the output stay in float64.""" # a private context, so that no globals are left in the default one - ctx = ms.Context() + ctx = new_context() flows = { dtype: ms.Flow( 4, prefix=f"fp32_test_{name}", subnet_precision=dtype, invert_spline=False @@ -148,12 +163,12 @@ def test_flow_float32_subnets(tmp_path): weights = ctx.get_global(name).torch() weights[:] = ( 0.3 - * torch.randn(weights.shape, generator=GEN, dtype=torch.float32).double() + * _randn(weights.shape, generator=GEN, dtype=torch.float32).double() ) ctx.get_global(name.replace("f64", "f32")).torch()[:] = weights.float() # outputs and gradients with respect to the input and the float32 weights - r = torch.rand((N, 4), dtype=torch.float64, generator=GEN) + r = _rand((N, 4), dtype=torch.float64, generator=GEN) outputs, grads = {}, {} for dtype, flow in flows.items(): module = FunctionModule(flow.forward_function(), ctx) @@ -166,13 +181,13 @@ def test_flow_float32_subnets(tmp_path): outputs[dtype] = (out.detach(), det.detach()) grads[dtype] = (r_grad.grad, weight.grad.double()) for o32, o64 in zip(outputs[ms.DataType.float32], outputs[ms.DataType.float64]): - torch.testing.assert_close(o32, o64, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(o32, o64, rtol=1e-3, atol=1e-4) for g32, g64 in zip(grads[ms.DataType.float32], grads[ms.DataType.float64]): - torch.testing.assert_close(g32, g64, rtol=1e-3, atol=1e-4) + torch.testing.assert_close(g32, g64, rtol=5e-3, atol=1e-4) # float32 globals survive saving and loading ctx.save_globals(str(tmp_path / "globals.tar")) - loaded = ms.Context() + loaded = new_context() loaded.load_globals(str(tmp_path / "globals.tar")) for name in names_f64: name32 = name.replace("f64", "f32") @@ -195,9 +210,9 @@ def test_structural_ops_backward(): fb.output(0, fb.cat(fb.batch_split(combined, counts))) module = FunctionModule(fb.function()) - x = torch.randn((N, 3), dtype=torch.float32, generator=GEN, requires_grad=True) + x = _randn((N, 3), dtype=torch.float32, generator=GEN, requires_grad=True) out = module(x) - out_grad = torch.randn(out.shape, dtype=out.dtype, generator=GEN) + out_grad = _randn(out.shape, dtype=out.dtype, generator=GEN) (out * out_grad).sum().backward() x_ref = x.detach().clone().requires_grad_() From 059fbc6fe3f438ebef84f7d8659562ba4d02a3b7 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 13:22:36 +0100 Subject: [PATCH 23/24] madspace tests: test_float32 uses the default context of the test device Co-Authored-By: Claude Sonnet 5.5 --- madspace/tests/test_float32.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/madspace/tests/test_float32.py b/madspace/tests/test_float32.py index ad70dc35bd..b451a588e8 100644 --- a/madspace/tests/test_float32.py +++ b/madspace/tests/test_float32.py @@ -5,7 +5,7 @@ import madspace as ms from madspace.torch import FunctionModule -from madspace_testing import DEVICE, new_context +from madspace_testing import DEVICE, default_context, new_context # a local generator, so that the global RNG state seen by other tests is unchanged GEN = torch.Generator().manual_seed(1234) @@ -57,7 +57,7 @@ def test_activation(name, precision): ms.NamedTypes([("in", io_type)]), ms.NamedTypes([("out", io_type)]) ) fb.output(0, getattr(fb, name)(fb.input(0))) - module = FunctionModule(fb.function()) + module = FunctionModule(fb.function(), default_context()) x = ( 3 * _randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN) @@ -88,7 +88,7 @@ def test_matmul(precision): ms.NamedTypes([("y", batch_type(ms_dtype, [DIMS_OUT]))]), ) fb.output(0, fb.matmul(fb.input(0), fb.input(1), fb.input(2))) - module = FunctionModule(fb.function()) + module = FunctionModule(fb.function(), default_context()) x = _randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN, requires_grad=True) # single inputs are passed without batch dimension @@ -120,7 +120,7 @@ def test_convert_roundtrip(): single = fb.convert_float32(fb.input(0)) fb.output(0, single) fb.output(1, fb.convert_float64(fb.relu(single))) - module = FunctionModule(fb.function()) + module = FunctionModule(fb.function(), default_context()) x = _randn((N, 3), dtype=torch.float64, generator=GEN, requires_grad=True) single_out, out = module(x) @@ -139,7 +139,7 @@ def test_literal_converted_for_float32(): ) # float64 literals are converted when the instruction dispatches to float32 fb.output(0, fb.matmul(fb.input(0), ms.Value([2.0], [1, 1]), ms.Value([1.0]))) - module = FunctionModule(fb.function()) + module = FunctionModule(fb.function(), default_context()) x = _randn((N, 1), dtype=torch.float32, generator=GEN) torch.testing.assert_close(module(x), 2 * x + 1) @@ -208,7 +208,7 @@ def test_structural_ops_backward(): c = fb.cat([a, b]) combined, counts = fb.batch_cat([c, c]) fb.output(0, fb.cat(fb.batch_split(combined, counts))) - module = FunctionModule(fb.function()) + module = FunctionModule(fb.function(), default_context()) x = _randn((N, 3), dtype=torch.float32, generator=GEN, requires_grad=True) out = module(x) From eb6338c49700f26e941c2709a3d1d2b3ba6a1ca8 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 30 Oct 2026 10:59:29 +0100 Subject: [PATCH 24/24] madspace: review fixes for float32 networks - CPU random_int allocated a float64 tensor for int32 indices, which the new dtype check rejects; replay-buffer sampling aborted on CPU - GPU tensor views decide at launch whether += must be atomic: only views with a zero stride in a threaded dimension (broadcast gradients) use atomicAdd, per-event gradients use plain stores again - fuse_instructions looks up activation names in a static table instead of building strings for every instruction - test_float32: fresh generator per test, so results don't depend on test order; add a random_int + batch_gather test Co-Authored-By: Claude Opus 5.5 --- madspace/src/compgraphs/optimizer.cpp | 30 ++++++------- madspace/src/cpu/runtime.cpp | 2 +- madspace/src/gpu/tensor.cuh | 50 ++++++++++++++------- madspace/tests/test_float32.py | 64 +++++++++++++++++---------- 4 files changed, 87 insertions(+), 59 deletions(-) diff --git a/madspace/src/compgraphs/optimizer.cpp b/madspace/src/compgraphs/optimizer.cpp index b85f1c483c..c43347665e 100644 --- a/madspace/src/compgraphs/optimizer.cpp +++ b/madspace/src/compgraphs/optimizer.cpp @@ -1,7 +1,6 @@ #include "madspace/compgraphs/optimizer.hpp" #include -#include #include #include #include @@ -110,29 +109,26 @@ Function madspace::fuse_instructions(const Function& function) { return find->second.get(); }; - // the precision of an activation instruction that fuses into a matmul, or empty - auto activation_precision = [](const std::string& name) -> std::string { - static const std::array activations{ - "relu", "leaky_relu", "elu", "gelu", "sigmoid", "softplus" - }; - for (const auto& activation : activations) { + // the precision of each activation instruction that fuses into a matmul + static const auto activation_precisions = [] { + std::unordered_map precisions; + for (std::string activation : + {"relu", "leaky_relu", "elu", "gelu", "sigmoid", "softplus"}) { for (std::string precision : {"float32", "float64"}) { - if (name == activation + "_" + precision) { - return precision; - } + precisions.emplace(activation + "_" + precision, precision); } } - return ""; - }; + return precisions; + }(); std::vector removed(instructions.size()); bool any_fused = false; for (std::size_t i = 0; i < instructions.size(); ++i) { auto& call = instructions[i]; const std::string& name = call.instruction->name(); - std::string precision = activation_precision(name); - if (name != "convert_float32" && name != "convert_float64" && - precision.empty()) { + auto find_precision = activation_precisions.find(name); + bool is_activation = find_precision != activation_precisions.end(); + if (name != "convert_float32" && name != "convert_float64" && !is_activation) { continue; } auto& input = call.inputs.at(0); @@ -149,8 +145,8 @@ Function madspace::fuse_instructions(const Function& function) { continue; } const std::string& source_name = source.instruction->name(); - if (!precision.empty()) { - if (source_name != "matmul_" + precision) { + if (is_activation) { + if (source_name != "matmul_" + find_precision->second) { continue; } // the matmul output stays an output, the activation needs it backward diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index c196905f65..8cd1f43da7 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -1146,7 +1146,7 @@ void op_random_int( auto batch_size = locals[instruction.input_indices[0]].batch_sizes()[0]; auto max_val = locals[instruction.input_indices[1]].batch_sizes()[0]; auto& output = locals[instruction.output_indices[0]]; - output = Tensor(DataType::dt_float64, {batch_size}, device); + output = Tensor(DataType::dt_int32, {batch_size}, device); auto flat_view = output.flat_view(0); auto& runtime = instruction.runtime; device.foreach ( diff --git a/madspace/src/gpu/tensor.cuh b/madspace/src/gpu/tensor.cuh index f645280feb..3b52e60948 100644 --- a/madspace/src/gpu/tensor.cuh +++ b/madspace/src/gpu/tensor.cuh @@ -20,9 +20,9 @@ public: using SizeType = std::conditional_t; static const int dim = _dim; - GpuTensorView(const TensorView& view) + GpuTensorView(const TensorView& view, bool atomic = false) requires(packed) - : _data(view.data()) { + : _data(view.data()), _atomic(atomic) { if constexpr (_dim > 0) { std::copy(view.shape(), view.shape() + _dim, std::begin(this->_shape)); std::copy(view.stride(), view.stride() + _dim, std::begin(this->_stride)); @@ -33,15 +33,20 @@ public: __host__ __device__ GpuTensorView(const GpuTensorView& view) requires(!packed) - : _data(view._data), _stride(&view._stride), _shape(&view._shape) {} + : + _data(view._data), + _stride(&view._stride), + _shape(&view._shape), + _atomic(view._atomic) {} GpuTensorView& operator=(GpuTensorView& value) requires(packed) = default; - __device__ GpuTensorView(T* data, std::size_t* stride, std::size_t* shape) + __device__ + GpuTensorView(T* data, std::size_t* stride, std::size_t* shape, bool atomic = false) requires(!packed) - : _data(data), _stride(stride), _shape(shape) {} + : _data(data), _stride(stride), _shape(shape), _atomic(atomic) {} __device__ GpuTensorView(T& value) requires(!packed) @@ -53,13 +58,13 @@ public: __device__ const GpuTensorView operator[](std::size_t index) const requires(_dim != 0) { - return {&_data[index * _stride[0]], &_stride[1], &_shape[1]}; + return {&_data[index * _stride[0]], &_stride[1], &_shape[1], _atomic}; } __device__ GpuTensorView operator[](std::size_t index) requires(_dim != 0) { - return {&_data[index * _stride[0]], &_stride[1], &_shape[1]}; + return {&_data[index * _stride[0]], &_stride[1], &_shape[1], _atomic}; } template @@ -69,7 +74,7 @@ public: T* ptr = _data; int i = 0; ((ptr = &ptr[index * _stride[i++]]), ...); - return {ptr, &_stride[sizeof...(I)], &_shape[sizeof...(I)]}; + return {ptr, &_stride[sizeof...(I)], &_shape[sizeof...(I)], _atomic}; } template @@ -79,7 +84,7 @@ public: T* ptr = _data; int i = 0; ((ptr = &ptr[index * _stride[i++]]), ...); - return {ptr, &_stride[sizeof...(I)], &_shape[sizeof...(I)]}; + return {ptr, &_stride[sizeof...(I)], &_shape[sizeof...(I)], _atomic}; } __device__ operator T() const @@ -95,17 +100,18 @@ public: return value; } - // atomic for floating-point types: the gradient of a broadcast input (batch - // stride 0) and the scatter_add of repeated indices accumulate from many threads - // into the same element + // atomic only if the launch found a zero stride (broadcast gradient), where many + // threads accumulate into the same element __device__ T operator+=(T value) requires(_dim == 0) { if constexpr (std::is_floating_point_v) { - atomicAdd(_data, value); - } else { - *_data += value; + if (_atomic) { + atomicAdd(_data, value); + return value; + } } + *_data += value; return value; } @@ -131,6 +137,7 @@ private: T* _data; SizeType _stride; SizeType _shape; + bool _atomic = false; }; template @@ -159,14 +166,23 @@ template struct get_views; template struct get_views { + std::size_t batch_size; + template auto operator()(TArg&&... args) { return std::make_tuple([&]() { if constexpr (ScalarType) { return GpuScalarView(args); } else { + auto view = + args->template view(); + // a zero stride in a dimension spread over threads aliases elements + bool atomic = false; + for (int i = 0; i < dims; ++i) { + atomic |= view.stride()[i] == 0 && (i > 0 || batch_size > 1); + } return GpuTensorView( - args->template view() + view, atomic ); } }()...); @@ -266,7 +282,7 @@ void tensor_foreach( } // get views to the tensors with the correct types based on the signature of func auto views = std::apply( - get_views(), + get_views{batch_size}, std::tuple_cat(inputs, outputs, std::make_tuple(scalar_args...)) ); auto& first_view = std::get<0>(views); diff --git a/madspace/tests/test_float32.py b/madspace/tests/test_float32.py index b451a588e8..88b9f9aae1 100644 --- a/madspace/tests/test_float32.py +++ b/madspace/tests/test_float32.py @@ -2,13 +2,16 @@ import pytest import torch +from madspace_testing import DEVICE, default_context, new_context import madspace as ms from madspace.torch import FunctionModule -from madspace_testing import DEVICE, default_context, new_context -# a local generator, so that the global RNG state seen by other tests is unchanged -GEN = torch.Generator().manual_seed(1234) + +@pytest.fixture +def gen(): + """A fresh local generator per test, so results don't depend on test order.""" + return torch.Generator().manual_seed(1234) def _draw(draw, *args, requires_grad=False, **kwargs): @@ -24,6 +27,7 @@ def _randn(*args, **kwargs): def _rand(*args, **kwargs): return _draw(torch.rand, *args, **kwargs) + N = 103 DIMS_IN = 7 DIMS_OUT = 5 @@ -50,7 +54,7 @@ def batch_type(dtype, shape): @pytest.mark.parametrize("precision", DTYPES.keys()) @pytest.mark.parametrize("name", ACTIVATIONS.keys()) -def test_activation(name, precision): +def test_activation(name, precision, gen): ms_dtype, torch_dtype, tol = DTYPES[precision] io_type = batch_type(ms_dtype, [DIMS_IN]) fb = ms.FunctionBuilder( @@ -59,12 +63,10 @@ def test_activation(name, precision): fb.output(0, getattr(fb, name)(fb.input(0))) module = FunctionModule(fb.function(), default_context()) - x = ( - 3 * _randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN) - ).requires_grad_() + x = (3 * _randn((N, DIMS_IN), dtype=torch_dtype, generator=gen)).requires_grad_() out = module(x) assert out.dtype == torch_dtype - out_grad = _randn(out.shape, dtype=out.dtype, generator=GEN) + out_grad = _randn(out.shape, dtype=out.dtype, generator=gen) (out * out_grad).sum().backward() x_ref = x.detach().clone().requires_grad_() @@ -75,7 +77,7 @@ def test_activation(name, precision): @pytest.mark.parametrize("precision", DTYPES.keys()) -def test_matmul(precision): +def test_matmul(precision, gen): ms_dtype, torch_dtype, tol = DTYPES[precision] fb = ms.FunctionBuilder( ms.NamedTypes( @@ -90,15 +92,15 @@ def test_matmul(precision): fb.output(0, fb.matmul(fb.input(0), fb.input(1), fb.input(2))) module = FunctionModule(fb.function(), default_context()) - x = _randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN, requires_grad=True) + x = _randn((N, DIMS_IN), dtype=torch_dtype, generator=gen, requires_grad=True) # single inputs are passed without batch dimension w = _randn( - (DIMS_OUT, DIMS_IN), dtype=torch_dtype, generator=GEN, requires_grad=True + (DIMS_OUT, DIMS_IN), dtype=torch_dtype, generator=gen, requires_grad=True ) - b = _randn((DIMS_OUT,), dtype=torch_dtype, generator=GEN, requires_grad=True) + b = _randn((DIMS_OUT,), dtype=torch_dtype, generator=gen, requires_grad=True) y = module(x, w, b) assert y.dtype == torch_dtype - y_grad = _randn(y.shape, dtype=y.dtype, generator=GEN) + y_grad = _randn(y.shape, dtype=y.dtype, generator=gen) (y * y_grad).sum().backward() refs = [t.detach().clone().requires_grad_() for t in (x, w, b)] @@ -109,7 +111,7 @@ def test_matmul(precision): torch.testing.assert_close(t.grad, t_ref.grad, rtol=tol, atol=tol) -def test_convert_roundtrip(): +def test_convert_roundtrip(gen): in_type = batch_type(ms.DataType.float64, [3]) fb = ms.FunctionBuilder( ms.NamedTypes([("in", in_type)]), @@ -122,7 +124,7 @@ def test_convert_roundtrip(): fb.output(1, fb.convert_float64(fb.relu(single))) module = FunctionModule(fb.function(), default_context()) - x = _randn((N, 3), dtype=torch.float64, generator=GEN, requires_grad=True) + x = _randn((N, 3), dtype=torch.float64, generator=gen, requires_grad=True) single_out, out = module(x) assert single_out.dtype == torch.float32 assert out.dtype == torch.float64 @@ -132,7 +134,7 @@ def test_convert_roundtrip(): torch.testing.assert_close(x.grad, (x.detach() > 0).double()) -def test_literal_converted_for_float32(): +def test_literal_converted_for_float32(gen): io_type = batch_type(ms.DataType.float32, [1]) fb = ms.FunctionBuilder( ms.NamedTypes([("in", io_type)]), ms.NamedTypes([("out", io_type)]) @@ -140,11 +142,11 @@ def test_literal_converted_for_float32(): # float64 literals are converted when the instruction dispatches to float32 fb.output(0, fb.matmul(fb.input(0), ms.Value([2.0], [1, 1]), ms.Value([1.0]))) module = FunctionModule(fb.function(), default_context()) - x = _randn((N, 1), dtype=torch.float32, generator=GEN) + x = _randn((N, 1), dtype=torch.float32, generator=gen) torch.testing.assert_close(module(x), 2 * x + 1) -def test_flow_float32_subnets(tmp_path): +def test_flow_float32_subnets(tmp_path, gen): """A flow with float32 subnetworks matches the float64 one with the same (rounded) weights; the Jacobian and the output stay in float64.""" # a private context, so that no globals are left in the default one @@ -162,13 +164,12 @@ def test_flow_float32_subnets(tmp_path): for name in names_f64: weights = ctx.get_global(name).torch() weights[:] = ( - 0.3 - * _randn(weights.shape, generator=GEN, dtype=torch.float32).double() + 0.3 * _randn(weights.shape, generator=gen, dtype=torch.float32).double() ) ctx.get_global(name.replace("f64", "f32")).torch()[:] = weights.float() # outputs and gradients with respect to the input and the float32 weights - r = _rand((N, 4), dtype=torch.float64, generator=GEN) + r = _rand((N, 4), dtype=torch.float64, generator=gen) outputs, grads = {}, {} for dtype, flow in flows.items(): module = FunctionModule(flow.forward_function(), ctx) @@ -196,7 +197,7 @@ def test_flow_float32_subnets(tmp_path): torch.testing.assert_close(global32, ctx.get_global(name32).torch()) -def test_structural_ops_backward(): +def test_structural_ops_backward(gen): """stack, unstack, cat, batch_cat and batch_split propagate float32 gradients.""" in_type = batch_type(ms.DataType.float32, [3]) fb = ms.FunctionBuilder( @@ -210,9 +211,9 @@ def test_structural_ops_backward(): fb.output(0, fb.cat(fb.batch_split(combined, counts))) module = FunctionModule(fb.function(), default_context()) - x = _randn((N, 3), dtype=torch.float32, generator=GEN, requires_grad=True) + x = _randn((N, 3), dtype=torch.float32, generator=gen, requires_grad=True) out = module(x) - out_grad = _randn(out.shape, dtype=out.dtype, generator=GEN) + out_grad = _randn(out.shape, dtype=out.dtype, generator=gen) (out * out_grad).sum().backward() x_ref = x.detach().clone().requires_grad_() @@ -221,3 +222,18 @@ def test_structural_ops_backward(): (out_ref * out_grad).sum().backward() torch.testing.assert_close(out, out_ref) torch.testing.assert_close(x.grad, x_ref.grad) + + +def test_random_int_gather(gen): + """random_int returns int32 indices, as used when sampling the replay buffer.""" + io_type = batch_type(ms.DataType.float64, []) + fb = ms.FunctionBuilder( + ms.NamedTypes([("in", io_type)]), ms.NamedTypes([("out", io_type)]) + ) + size = fb.batch_size([fb.input(0)]) + fb.output(0, fb.batch_gather(fb.random_int(size, size), fb.input(0))) + module = FunctionModule(fb.function(), default_context()) + x = _rand((N,), dtype=torch.float64, generator=gen) + out = module(x) + assert out.shape == x.shape + assert torch.isin(out, x).all()