Skip to content

Commit 9defc8b

Browse files
tpadioleautpadioleau
committed
Add save_npy function
1 parent f1fef3e commit 9defc8b

9 files changed

Lines changed: 529 additions & 0 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ target_sources(
134134
src/ddc/discrete_element.cpp
135135
src/ddc/discrete_space.cpp
136136
src/ddc/discrete_vector.cpp
137+
src/ddc/save_npy.cpp
137138
src/ddc/print.cpp
138139
PUBLIC
139140
FILE_SET core_public
@@ -176,6 +177,7 @@ target_sources(
176177
src/ddc/print.hpp
177178
src/ddc/real_type.hpp
178179
src/ddc/reducer.hpp
180+
src/ddc/save_npy.hpp
179181
src/ddc/scope_guard.hpp
180182
src/ddc/sparse_discrete_domain.hpp
181183
src/ddc/strided_discrete_domain.hpp

‎src/ddc/ddc.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,3 +80,4 @@ namespace ddc {
8080

8181
// Output
8282
#include "print.hpp"
83+
#include "save_npy.hpp"

‎src/ddc/save_npy.cpp‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
// Copyright (C) The DDC development team, see COPYRIGHT.md file
2+
//
3+
// SPDX-License-Identifier: MIT
4+
5+
#include <bit>
6+
#include <cstddef>
7+
#include <cstdint>
8+
#include <filesystem>
9+
#include <fstream>
10+
#include <numeric>
11+
#include <stdexcept>
12+
#include <string>
13+
#include <vector>
14+
15+
#include "save_npy.hpp"
16+
17+
namespace ddc::detail {
18+
19+
NpyByteOrder get_byte_order(std::size_t itemsize) noexcept
20+
{
21+
if (itemsize == 1) {
22+
return NpyByteOrder::not_applicable;
23+
}
24+
25+
if (std::endian::native == std::endian::little) {
26+
return NpyByteOrder::little_endian;
27+
}
28+
29+
return NpyByteOrder::big_endian;
30+
}
31+
32+
std::string NpyDtype::str() const
33+
{
34+
return std::string(1, static_cast<char>(byte_order)) + static_cast<char>(kind)
35+
+ std::to_string(itemsize);
36+
}
37+
38+
// See specification at https://numpy.org/neps/nep-0001-npy-format.html#format-specification-version-1-0
39+
void save_npy(std::ostream& os, NpyArrayView const& view)
40+
{
41+
// Build shape string: (d0, d1, ..., dN,)
42+
std::string shape_str = "(";
43+
for (std::size_t ext : view.shape) {
44+
shape_str += std::to_string(ext);
45+
shape_str += ", ";
46+
}
47+
shape_str += ")";
48+
49+
std::string const header_dict
50+
= std::string("{'descr': '") + view.dtype.str() + "', 'fortran_order': "
51+
+ (view.fortran_order ? "True" : "False") + ", 'shape': " + shape_str + ", }";
52+
53+
// Pad header to a multiple of 16
54+
std::size_t const non_padded_header_len = header_dict.size() + 1;
55+
// magic(6) + major(1) + minor(1) + header_len(2) + header
56+
std::size_t const padding = 16 - (6 + 1 + 1 + 2 + non_padded_header_len) % 16;
57+
if (non_padded_header_len + padding > std::numeric_limits<std::uint16_t>::max()) {
58+
throw std::runtime_error("save_npy: header too large for npy v1.0.");
59+
}
60+
auto const header_len = static_cast<std::uint16_t>(non_padded_header_len + padding);
61+
62+
// magic string
63+
os.write("\x93NUMPY", 6);
64+
// major version
65+
os.put(1);
66+
// minor version
67+
os.put(0);
68+
69+
// header length
70+
os.write(reinterpret_cast<char const*>(&header_len), sizeof(header_len));
71+
// header + padding + newline
72+
os.write(header_dict.data(), header_dict.size());
73+
os.write(" ", padding);
74+
os.put('\n');
75+
76+
// Raw data
77+
std::size_t const n_elems
78+
= std::accumulate(view.shape.begin(), view.shape.end(), 1ULL, std::multiplies<> {});
79+
os.write(reinterpret_cast<char const*>(view.data), n_elems * view.dtype.itemsize);
80+
}
81+
82+
void save_npy(std::filesystem::path const& filename, NpyArrayView const& view)
83+
{
84+
std::ofstream file(filename, std::ios::binary);
85+
file.exceptions(std::ios::failbit | std::ios::badbit);
86+
87+
save_npy(file, view);
88+
}
89+
90+
} // namespace ddc::detail

‎src/ddc/save_npy.hpp‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
// Copyright (C) The DDC development team, see COPYRIGHT.md file
2+
//
3+
// SPDX-License-Identifier: MIT
4+
5+
#pragma once
6+
7+
#include <complex>
8+
#include <cstddef>
9+
#include <filesystem>
10+
#include <iosfwd>
11+
#include <string>
12+
#include <type_traits>
13+
#include <utility>
14+
#include <vector>
15+
16+
#include <Kokkos_Core.hpp>
17+
18+
namespace ddc::detail {
19+
20+
enum class NpyByteOrder : char { little_endian = '<', big_endian = '>', not_applicable = '|' };
21+
22+
NpyByteOrder get_byte_order(std::size_t itemsize) noexcept;
23+
24+
enum class NpyKind : char {
25+
boolean = 'b',
26+
signed_int = 'i',
27+
unsigned_int = 'u',
28+
floating_point = 'f',
29+
complex = 'c',
30+
other = 'V',
31+
};
32+
33+
struct NpyDtype
34+
{
35+
NpyByteOrder byte_order;
36+
NpyKind kind;
37+
std::size_t itemsize; // in bytes
38+
39+
std::string str() const;
40+
};
41+
42+
template <typename T>
43+
NpyDtype convert_to_npy_dtype()
44+
{
45+
std::size_t const itemsize = sizeof(T);
46+
NpyByteOrder const byte_order = get_byte_order(itemsize);
47+
NpyKind kind;
48+
49+
if constexpr (std::is_same_v<T, bool>) {
50+
// ── Single-byte / untyped ─────────────────────────────────────────
51+
kind = NpyKind::boolean;
52+
} else if constexpr (std::is_same_v<T, std::byte>) {
53+
// std::byte → raw byte buffer, no arithmetic meaning
54+
kind = NpyKind::other;
55+
} else if constexpr (std::is_same_v<T, char>) {
56+
// char is a distinct type; its signedness is implementation-defined
57+
kind = std::is_signed_v<char> ? NpyKind::signed_int : NpyKind::unsigned_int;
58+
} else if constexpr (
59+
std::is_same_v<T, std::complex<float>> || std::is_same_v<T, std::complex<double>>) {
60+
// ── Complex ───────────────────────────────────────────────────────
61+
// NumPy 'c' dtype stores interleaved real+imag, same layout as std::complex
62+
kind = NpyKind::complex;
63+
} else if constexpr (std::is_floating_point_v<T>) {
64+
// ── Floating-point ────────────────────────────────────────────────
65+
static_assert(
66+
!std::is_same_v<T, long double>,
67+
"long double is platform-specific (80/96/128-bit); cast to double first.");
68+
kind = NpyKind::floating_point;
69+
} else if constexpr (std::is_signed_v<T>) {
70+
// ── Integers ──────────────────────────────────────────────────────
71+
kind = NpyKind::signed_int;
72+
} else if constexpr (std::is_unsigned_v<T>) {
73+
kind = NpyKind::unsigned_int;
74+
} else {
75+
static_assert(sizeof(T) == 0, "Unsupported type for NpyDtype::of<T>()");
76+
}
77+
78+
return {.byte_order = byte_order, .kind = kind, .itemsize = itemsize};
79+
}
80+
81+
struct NpyArrayView
82+
{
83+
void const* data;
84+
NpyDtype dtype;
85+
std::vector<std::size_t> shape;
86+
bool fortran_order;
87+
};
88+
89+
void save_npy(std::ostream& os, NpyArrayView const& view);
90+
91+
void save_npy(std::filesystem::path const& filename, NpyArrayView const& view);
92+
93+
} // namespace ddc::detail
94+
95+
namespace ddc::experimental {
96+
97+
template <typename T, typename Extents, typename Layout, typename Accessor>
98+
void save_npy(std::ostream& os, Kokkos::mdspan<T, Extents, Layout, Accessor> const& mds)
99+
{
100+
static_assert(
101+
std::is_same_v<Layout, Kokkos::layout_left>
102+
|| std::is_same_v<Layout, Kokkos::layout_right>,
103+
"save_npy: only contiguous layouts supported.");
104+
static_assert(
105+
std::is_same_v<Accessor, Kokkos::default_accessor<T>>
106+
|| std::is_same_v<Accessor, Kokkos::default_accessor<T const>>,
107+
"save_npy: non-host-accessible accessor. Use create_mirror_view + deep_copy first.");
108+
109+
std::vector<std::size_t> shape(Extents::rank());
110+
for (std::size_t i = 0; i < Extents::rank(); ++i) {
111+
shape[i] = mds.extent(i);
112+
}
113+
114+
ddc::detail::save_npy(
115+
os,
116+
ddc::detail::NpyArrayView {
117+
.data = mds.data_handle(),
118+
.dtype = ddc::detail::convert_to_npy_dtype<std::remove_const_t<T>>(),
119+
.shape = std::move(shape),
120+
.fortran_order = std::is_same_v<Layout, Kokkos::layout_left>,
121+
});
122+
}
123+
124+
template <typename T, typename Extents, typename Layout, typename Accessor>
125+
void save_npy(
126+
std::filesystem::path const& filename,
127+
Kokkos::mdspan<T, Extents, Layout, Accessor> const& mds)
128+
{
129+
static_assert(
130+
std::is_same_v<Layout, Kokkos::layout_left>
131+
|| std::is_same_v<Layout, Kokkos::layout_right>,
132+
"save_npy: only contiguous layouts supported.");
133+
static_assert(
134+
std::is_same_v<Accessor, Kokkos::default_accessor<T>>
135+
|| std::is_same_v<Accessor, Kokkos::default_accessor<T const>>,
136+
"save_npy: non-host-accessible accessor. Use create_mirror_view + deep_copy first.");
137+
138+
std::vector<std::size_t> shape(Extents::rank());
139+
for (std::size_t i = 0; i < Extents::rank(); ++i) {
140+
shape[i] = mds.extent(i);
141+
}
142+
143+
ddc::detail::save_npy(
144+
filename,
145+
ddc::detail::NpyArrayView {
146+
.data = mds.data_handle(),
147+
.dtype = ddc::detail::convert_to_npy_dtype<std::remove_const_t<T>>(),
148+
.shape = std::move(shape),
149+
.fortran_order = std::is_same_v<Layout, Kokkos::layout_left>,
150+
});
151+
}
152+
153+
} // namespace ddc::experimental

‎tests/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ include(GoogleTest)
66

77
add_subdirectory(discrete_space)
88

9+
add_subdirectory(save_npy)
10+
911
add_library(ddc_gtest_main)
1012
target_compile_features(ddc_gtest_main PRIVATE cxx_std_20)
1113
target_link_libraries(ddc_gtest_main PRIVATE DDC::core GTest::gtest)
@@ -41,6 +43,7 @@ add_executable(
4143
reducer.cpp
4244
relocatable_device_code.cpp
4345
relocatable_device_code_initialization.cpp
46+
save_npy.cpp
4447
sparse_discrete_domain.cpp
4548
strided_discrete_domain.cpp
4649
tagged_vector.cpp

‎tests/save_npy.cpp‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
// Copyright (C) The DDC development team, see COPYRIGHT.md file
2+
//
3+
// SPDX-License-Identifier: MIT
4+
5+
// #include <array>
6+
#include <cmath>
7+
// #include <complex>
8+
#include <stdexcept>
9+
10+
#include <ddc/ddc.hpp>
11+
12+
#include <gtest/gtest.h>
13+
14+
#include <Kokkos_Core.hpp>
15+
16+
// namespace {
17+
18+
// std::array constexpr ns {2, 3, 4};
19+
// int constexpr n = ns[0] * ns[1] * ns[2];
20+
21+
// template <typename T>
22+
// constexpr T make_value()
23+
// {
24+
// std::complex<double> const base_value(2.3, 0.4);
25+
// if constexpr (
26+
// std::is_same_v<T, std::complex<float>> || std::is_same_v<T, std::complex<double>>) {
27+
// return T(base_value);
28+
// } else if constexpr (std::is_floating_point_v<T>) {
29+
// return T(std::real(base_value));
30+
// } else {
31+
// return T(std::llround(std::real(base_value)));
32+
// }
33+
// }
34+
35+
// } // namespace
36+
37+
// template <typename T>
38+
// struct SaveNpyTest : public ::testing::Test
39+
// {
40+
// using data_type = T;
41+
// };
42+
43+
// using SaveNpyTypes = ::testing::Types<
44+
// float,
45+
// double,
46+
// std::complex<float>,
47+
// std::complex<double>,
48+
// char,
49+
// signed char,
50+
// signed short,
51+
// signed int,
52+
// signed long,
53+
// signed long long,
54+
// unsigned char,
55+
// unsigned short,
56+
// unsigned int,
57+
// unsigned long,
58+
// unsigned long long>;
59+
60+
// TYPED_TEST_SUITE(SaveNpyTest, SaveNpyTypes);
61+
62+
// TYPED_TEST(SaveNpyTest, SaveNpy0d)
63+
// {
64+
// using data_type = typename TestFixture::data_type;
65+
// std::string const label(typeid(data_type).name());
66+
// Kokkos::View<data_type, Kokkos::HostSpace> const alloc(label);
67+
// Kokkos::deep_copy(alloc, make_value<data_type>());
68+
// Kokkos::mdspan const view(alloc.data());
69+
70+
// ddc::experimental::save_npy("test.npy", view);
71+
// }
72+
73+
// TYPED_TEST(SaveNpyTest, SaveNpy1d)
74+
// {
75+
// using data_type = typename TestFixture::data_type;
76+
// std::string const label(typeid(data_type).name());
77+
// Kokkos::View<data_type*, Kokkos::HostSpace> const alloc(label, n);
78+
// Kokkos::deep_copy(alloc, make_value<data_type>());
79+
// Kokkos::mdspan const view(alloc.data(), n);
80+
81+
// ddc::experimental::save_npy("test_" + label + "_1d.npy", view);
82+
// }
83+
84+
// TYPED_TEST(SaveNpyTest, SaveNpy3d)
85+
// {
86+
// using data_type = typename TestFixture::data_type;
87+
// std::string const label(typeid(data_type).name());
88+
// Kokkos::View<data_type*, Kokkos::HostSpace> const alloc(label, n);
89+
// Kokkos::deep_copy(alloc, make_value<data_type>());
90+
// Kokkos::mdspan const view(alloc.data(), ns);
91+
92+
// ddc::experimental::save_npy("test_" + label + "_3d.npy", view);
93+
// }
94+
95+
TEST(SaveNpy, LargeHeader)
96+
{
97+
ddc::detail::NpyArrayView const view {
98+
.data = nullptr,
99+
.dtype = ddc::detail::convert_to_npy_dtype<char>(),
100+
.shape = std::vector<std::size_t>(25'000, 0),
101+
.fortran_order = true,
102+
};
103+
EXPECT_THROW(ddc::detail::save_npy("test.npy", view), std::runtime_error);
104+
}

0 commit comments

Comments
 (0)