From 7cc382f0a7d85f22e920c7d007e270d87e12cc0c Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 13 Oct 2026 18:31:30 +0200 Subject: [PATCH 01/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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 )