Skip to content

Commit f8ebd0a

Browse files
tpadioleautpadioleau
committed
Add save_npy function
1 parent f5b37e7 commit f8ebd0a

6 files changed

Lines changed: 379 additions & 1 deletion

File tree

‎CMakeLists.txt‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,14 +126,15 @@ endif()
126126
add_library(ddc_core)
127127
add_library(DDC::core ALIAS ddc_core)
128128
configure_file(cmake/config.hpp.in generated/ddc/config.hpp NO_SOURCE_PERMISSIONS @ONLY)
129-
target_compile_features(ddc_core PUBLIC cxx_std_17)
129+
target_compile_features(ddc_core INTERFACE cxx_std_17 PRIVATE cxx_std_20)
130130
target_link_libraries(ddc_core PUBLIC Kokkos::kokkos)
131131
target_sources(
132132
ddc_core
133133
PRIVATE
134134
src/ddc/discrete_space.cpp
135135
src/ddc/discrete_element.cpp
136136
src/ddc/discrete_vector.cpp
137+
src/ddc/save_npy.cpp
137138
src/ddc/print.cpp
138139
INTERFACE
139140
FILE_SET HEADERS
@@ -172,6 +173,7 @@ target_sources(
172173
src/ddc/print.hpp
173174
src/ddc/real_type.hpp
174175
src/ddc/reducer.hpp
176+
src/ddc/save_npy.hpp
175177
src/ddc/scope_guard.hpp
176178
src/ddc/sparse_discrete_domain.hpp
177179
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: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
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 <Kokkos_Core.hpp>
16+
17+
#include "save_npy.hpp"
18+
19+
namespace ddc::detail {
20+
21+
NpyByteOrder get_byte_order() noexcept
22+
{
23+
return (std::endian::native == std::endian::little) ? NpyByteOrder::little_endian
24+
: NpyByteOrder::big_endian;
25+
}
26+
27+
std::string NpyDtype::str() const
28+
{
29+
return std::string(1, static_cast<char>(byte_order)) + static_cast<char>(kind)
30+
+ std::to_string(itemsize);
31+
}
32+
33+
void save_npy(std::ostream& os, NpyArrayView const& view)
34+
{
35+
// Build shape string: (d0, d1, ..., dN,)
36+
std::string shape_str = "(";
37+
for (std::size_t ext : view.shape) {
38+
shape_str += std::to_string(ext);
39+
shape_str += ", ";
40+
}
41+
shape_str += ")";
42+
43+
std::string header_dict = std::string("{'descr': '") + view.dtype.str()
44+
+ "', 'fortran_order': " + (view.fortran_order ? "True" : "False")
45+
+ ", 'shape': " + shape_str + ", }";
46+
47+
// Pad header to a multiple of 64 bytes
48+
constexpr std::size_t prefix_size = 6 + 1 + 1 + 2; // magic + major + minor + hlen
49+
std::size_t const total_header = prefix_size + header_dict.size() + 1; // +1 for '\n'
50+
std::size_t const padded = ((total_header + 63) / 64) * 64;
51+
header_dict += std::string(padded - total_header, ' ');
52+
header_dict += '\n';
53+
54+
if (header_dict.size() > std::numeric_limits<std::uint16_t>::max()) {
55+
throw std::runtime_error("save_npy: header too large for npy v1.0.");
56+
}
57+
auto const hlen = static_cast<std::uint16_t>(header_dict.size());
58+
59+
// Magic + version
60+
os.write("\x93NUMPY", 6);
61+
std::uint8_t const major = 1;
62+
std::uint8_t const minor = 0;
63+
os.write(reinterpret_cast<char const*>(&major), 1);
64+
os.write(reinterpret_cast<char const*>(&minor), 1);
65+
66+
// Header length + content
67+
os.write(reinterpret_cast<char const*>(&hlen), sizeof(hlen));
68+
os.write(header_dict.data(), header_dict.size());
69+
70+
// Raw data
71+
std::size_t const n_elems
72+
= std::accumulate(view.shape.begin(), view.shape.end(), 1ULL, std::multiplies<> {});
73+
os.write(reinterpret_cast<char const*>(view.data), n_elems * view.dtype.itemsize);
74+
}
75+
76+
void save_npy(std::filesystem::path const& filename, NpyArrayView const& view)
77+
{
78+
std::ofstream file(filename, std::ios::binary);
79+
file.exceptions(std::ios::failbit | std::ios::badbit);
80+
81+
save_npy(file, view);
82+
}
83+
84+
} // namespace ddc::detail

‎src/ddc/save_npy.hpp‎

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

‎tests/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ add_executable(
4141
reducer.cpp
4242
relocatable_device_code.cpp
4343
relocatable_device_code_initialization.cpp
44+
save_npy.cpp
4445
sparse_discrete_domain.cpp
4546
strided_discrete_domain.cpp
4647
tagged_vector.cpp

0 commit comments

Comments
 (0)