test_numpy_dtypes.cpp revision 12037
16145Snate@binkert.org/* 26239Snate@binkert.org tests/test_numpy_dtypes.cpp -- Structured and compound NumPy dtypes 36239Snate@binkert.org 46145Snate@binkert.org Copyright (c) 2016 Ivan Smirnov 56239Snate@binkert.org 66239Snate@binkert.org All rights reserved. Use of this source code is governed by a 76239Snate@binkert.org BSD-style license that can be found in the LICENSE file. 86239Snate@binkert.org*/ 96239Snate@binkert.org 106239Snate@binkert.org#include "pybind11_tests.h" 116239Snate@binkert.org#include <pybind11/numpy.h> 126239Snate@binkert.org 136239Snate@binkert.org#ifdef __GNUC__ 146239Snate@binkert.org#define PYBIND11_PACKED(cls) cls __attribute__((__packed__)) 156145Snate@binkert.org#else 166239Snate@binkert.org#define PYBIND11_PACKED(cls) __pragma(pack(push, 1)) cls __pragma(pack(pop)) 176239Snate@binkert.org#endif 186239Snate@binkert.org 196239Snate@binkert.orgnamespace py = pybind11; 206239Snate@binkert.org 216239Snate@binkert.orgstruct SimpleStruct { 226239Snate@binkert.org bool bool_; 236239Snate@binkert.org uint32_t uint_; 246239Snate@binkert.org float float_; 256239Snate@binkert.org long double ldbl_; 266239Snate@binkert.org}; 276145Snate@binkert.org 286145Snate@binkert.orgstd::ostream& operator<<(std::ostream& os, const SimpleStruct& v) { 297039Snate@binkert.org return os << "s:" << v.bool_ << "," << v.uint_ << "," << v.float_ << "," << v.ldbl_; 307039Snate@binkert.org} 316145Snate@binkert.org 328091Snilay@cs.wisc.eduPYBIND11_PACKED(struct PackedStruct { 336145Snate@binkert.org bool bool_; 347002Snate@binkert.org uint32_t uint_; 357002Snate@binkert.org float float_; 368608Snilay@cs.wisc.edu long double ldbl_; 376145Snate@binkert.org}); 386145Snate@binkert.org 396145Snate@binkert.orgstd::ostream& operator<<(std::ostream& os, const PackedStruct& v) { 406145Snate@binkert.org return os << "p:" << v.bool_ << "," << v.uint_ << "," << v.float_ << "," << v.ldbl_; 416145Snate@binkert.org} 426145Snate@binkert.org 436145Snate@binkert.orgPYBIND11_PACKED(struct NestedStruct { 447039Snate@binkert.org SimpleStruct a; 457039Snate@binkert.org PackedStruct b; 467039Snate@binkert.org}); 477039Snate@binkert.org 487039Snate@binkert.orgstd::ostream& operator<<(std::ostream& os, const NestedStruct& v) { 497039Snate@binkert.org return os << "n:a=" << v.a << ";b=" << v.b; 506145Snate@binkert.org} 517039Snate@binkert.org 527039Snate@binkert.orgstruct PartialStruct { 537039Snate@binkert.org bool bool_; 547039Snate@binkert.org uint32_t uint_; 556145Snate@binkert.org float float_; 567039Snate@binkert.org uint64_t dummy2; 577039Snate@binkert.org long double ldbl_; 586145Snate@binkert.org}; 597039Snate@binkert.org 607039Snate@binkert.orgstruct PartialNestedStruct { 617039Snate@binkert.org uint64_t dummy1; 627039Snate@binkert.org PartialStruct a; 637039Snate@binkert.org uint64_t dummy2; 647039Snate@binkert.org}; 657039Snate@binkert.org 667039Snate@binkert.orgstruct UnboundStruct { }; 676145Snate@binkert.org 688091Snilay@cs.wisc.edustruct StringStruct { 698091Snilay@cs.wisc.edu char a[3]; 708091Snilay@cs.wisc.edu std::array<char, 3> b; 718091Snilay@cs.wisc.edu}; 726145Snate@binkert.org 737039Snate@binkert.orgPYBIND11_PACKED(struct StructWithUglyNames { 747039Snate@binkert.org int8_t __x__; 756145Snate@binkert.org uint64_t __y__; 767039Snate@binkert.org}); 776145Snate@binkert.org 787055Snate@binkert.orgenum class E1 : int64_t { A = -1, B = 1 }; 797055Snate@binkert.orgenum E2 : uint8_t { X = 1, Y = 2 }; 807055Snate@binkert.org 816145Snate@binkert.orgPYBIND11_PACKED(struct EnumStruct { 827039Snate@binkert.org E1 e1; 837039Snate@binkert.org E2 e2; 847039Snate@binkert.org}); 857039Snate@binkert.org 867039Snate@binkert.orgstd::ostream& operator<<(std::ostream& os, const StringStruct& v) { 877039Snate@binkert.org os << "a='"; 887039Snate@binkert.org for (size_t i = 0; i < 3 && v.a[i]; i++) os << v.a[i]; 897039Snate@binkert.org os << "',b='"; 907039Snate@binkert.org for (size_t i = 0; i < 3 && v.b[i]; i++) os << v.b[i]; 917039Snate@binkert.org return os << "'"; 926145Snate@binkert.org} 936145Snate@binkert.org 947039Snate@binkert.orgstd::ostream& operator<<(std::ostream& os, const EnumStruct& v) { 957039Snate@binkert.org return os << "e1=" << (v.e1 == E1::A ? "A" : "B") << ",e2=" << (v.e2 == E2::X ? "X" : "Y"); 966145Snate@binkert.org} 977039Snate@binkert.org 987039Snate@binkert.orgtemplate <typename T> 997039Snate@binkert.orgpy::array mkarray_via_buffer(size_t n) { 1006145Snate@binkert.org return py::array(py::buffer_info(nullptr, sizeof(T), 1016145Snate@binkert.org py::format_descriptor<T>::format(), 1027039Snate@binkert.org 1, { n }, { sizeof(T) })); 1037039Snate@binkert.org} 1046145Snate@binkert.org 1057039Snate@binkert.org#define SET_TEST_VALS(s, i) do { \ 1066145Snate@binkert.org s.bool_ = (i) % 2 != 0; \ 1076145Snate@binkert.org s.uint_ = (uint32_t) (i); \ 1087055Snate@binkert.org s.float_ = (float) (i) * 1.5f; \ 1097055Snate@binkert.org s.ldbl_ = (long double) (i) * -2.5L; } while (0) 1106145Snate@binkert.org 1117039Snate@binkert.orgtemplate <typename S> 1127055Snate@binkert.orgpy::array_t<S, 0> create_recarray(size_t n) { 1137039Snate@binkert.org auto arr = mkarray_via_buffer<S>(n); 1146145Snate@binkert.org auto req = arr.request(); 1156145Snate@binkert.org auto ptr = static_cast<S*>(req.ptr); 1167039Snate@binkert.org for (size_t i = 0; i < n; i++) { 1177039Snate@binkert.org SET_TEST_VALS(ptr[i], i); 1186145Snate@binkert.org } 1197039Snate@binkert.org return arr; 1206145Snate@binkert.org} 1216145Snate@binkert.org 1227039Snate@binkert.orgstd::string get_format_unbound() { 1237039Snate@binkert.org return py::format_descriptor<UnboundStruct>::format(); 1246145Snate@binkert.org} 1257039Snate@binkert.org 1267039Snate@binkert.orgpy::array_t<NestedStruct, 0> create_nested(size_t n) { 1276145Snate@binkert.org auto arr = mkarray_via_buffer<NestedStruct>(n); 1287039Snate@binkert.org auto req = arr.request(); 1297039Snate@binkert.org auto ptr = static_cast<NestedStruct*>(req.ptr); 1307039Snate@binkert.org for (size_t i = 0; i < n; i++) { 1317039Snate@binkert.org SET_TEST_VALS(ptr[i].a, i); 1327039Snate@binkert.org SET_TEST_VALS(ptr[i].b, i + 1); 1337039Snate@binkert.org } 1347039Snate@binkert.org return arr; 1357039Snate@binkert.org} 1367039Snate@binkert.org 1377039Snate@binkert.orgpy::array_t<PartialNestedStruct, 0> create_partial_nested(size_t n) { 1387039Snate@binkert.org auto arr = mkarray_via_buffer<PartialNestedStruct>(n); 1397039Snate@binkert.org auto req = arr.request(); 1407039Snate@binkert.org auto ptr = static_cast<PartialNestedStruct*>(req.ptr); 1417039Snate@binkert.org for (size_t i = 0; i < n; i++) { 1427039Snate@binkert.org SET_TEST_VALS(ptr[i].a, i); 1436145Snate@binkert.org } 1446145Snate@binkert.org return arr; 1457027SBrad.Beckmann@amd.com} 1467039Snate@binkert.org 1477039Snate@binkert.orgpy::array_t<StringStruct, 0> create_string_array(bool non_empty) { 1487027SBrad.Beckmann@amd.com auto arr = mkarray_via_buffer<StringStruct>(non_empty ? 4 : 0); 1497027SBrad.Beckmann@amd.com if (non_empty) { 1507027SBrad.Beckmann@amd.com auto req = arr.request(); 1517054Snate@binkert.org auto ptr = static_cast<StringStruct*>(req.ptr); 1527027SBrad.Beckmann@amd.com for (size_t i = 0; i < req.size * req.itemsize; i++) 1537027SBrad.Beckmann@amd.com static_cast<char*>(req.ptr)[i] = 0; 1547027SBrad.Beckmann@amd.com ptr[1].a[0] = 'a'; ptr[1].b[0] = 'a'; 1557027SBrad.Beckmann@amd.com ptr[2].a[0] = 'a'; ptr[2].b[0] = 'a'; 1567027SBrad.Beckmann@amd.com ptr[3].a[0] = 'a'; ptr[3].b[0] = 'a'; 1577027SBrad.Beckmann@amd.com 1587027SBrad.Beckmann@amd.com ptr[2].a[1] = 'b'; ptr[2].b[1] = 'b'; 1597027SBrad.Beckmann@amd.com ptr[3].a[1] = 'b'; ptr[3].b[1] = 'b'; 1607027SBrad.Beckmann@amd.com 1617027SBrad.Beckmann@amd.com ptr[3].a[2] = 'c'; ptr[3].b[2] = 'c'; 1627027SBrad.Beckmann@amd.com } 1637027SBrad.Beckmann@amd.com return arr; 1647027SBrad.Beckmann@amd.com} 1657027SBrad.Beckmann@amd.com 1667027SBrad.Beckmann@amd.compy::array_t<EnumStruct, 0> create_enum_array(size_t n) { 1677563SBrad.Beckmann@amd.com auto arr = mkarray_via_buffer<EnumStruct>(n); 1687027SBrad.Beckmann@amd.com auto ptr = (EnumStruct *) arr.mutable_data(); 1697027SBrad.Beckmann@amd.com for (size_t i = 0; i < n; i++) { 1707027SBrad.Beckmann@amd.com ptr[i].e1 = static_cast<E1>(-1 + ((int) i % 2) * 2); 1717027SBrad.Beckmann@amd.com ptr[i].e2 = static_cast<E2>(1 + (i % 2)); 1727039Snate@binkert.org } 1737039Snate@binkert.org return arr; 1746145Snate@binkert.org} 1756145Snate@binkert.org 1766145Snate@binkert.orgtemplate <typename S> 1776145Snate@binkert.orgpy::list print_recarray(py::array_t<S, 0> arr) { 1787039Snate@binkert.org const auto req = arr.request(); 1796145Snate@binkert.org const auto ptr = static_cast<S*>(req.ptr); 1807039Snate@binkert.org auto l = py::list(); 1816145Snate@binkert.org for (size_t i = 0; i < req.size; i++) { 1826145Snate@binkert.org std::stringstream ss; 1836145Snate@binkert.org ss << ptr[i]; 1846145Snate@binkert.org l.append(py::str(ss.str())); 1857039Snate@binkert.org } 1867039Snate@binkert.org return l; 1876145Snate@binkert.org} 1887039Snate@binkert.org 1896145Snate@binkert.orgpy::list print_format_descriptors() { 1907039Snate@binkert.org const auto fmts = { 1917039Snate@binkert.org py::format_descriptor<SimpleStruct>::format(), 1927039Snate@binkert.org py::format_descriptor<PackedStruct>::format(), 1937039Snate@binkert.org py::format_descriptor<NestedStruct>::format(), 1947039Snate@binkert.org py::format_descriptor<PartialStruct>::format(), 1957039Snate@binkert.org py::format_descriptor<PartialNestedStruct>::format(), 1966145Snate@binkert.org py::format_descriptor<StringStruct>::format(), 1976145Snate@binkert.org py::format_descriptor<EnumStruct>::format() 1987039Snate@binkert.org }; 1997039Snate@binkert.org auto l = py::list(); 2006145Snate@binkert.org for (const auto &fmt : fmts) { 2017039Snate@binkert.org l.append(py::cast(fmt)); 2026145Snate@binkert.org } 2036145Snate@binkert.org return l; 2046145Snate@binkert.org} 2057002Snate@binkert.org 2067039Snate@binkert.orgpy::list print_dtypes() { 2077039Snate@binkert.org const auto dtypes = { 2087039Snate@binkert.org py::str(py::dtype::of<SimpleStruct>()), 2097039Snate@binkert.org py::str(py::dtype::of<PackedStruct>()), 2107039Snate@binkert.org py::str(py::dtype::of<NestedStruct>()), 2117039Snate@binkert.org py::str(py::dtype::of<PartialStruct>()), 2127039Snate@binkert.org py::str(py::dtype::of<PartialNestedStruct>()), 2137039Snate@binkert.org py::str(py::dtype::of<StringStruct>()), 2147811Ssteve.reinhardt@amd.com py::str(py::dtype::of<EnumStruct>()), 2157039Snate@binkert.org py::str(py::dtype::of<StructWithUglyNames>()) 2166145Snate@binkert.org }; 2177039Snate@binkert.org auto l = py::list(); 2187039Snate@binkert.org for (const auto &s : dtypes) { 2197039Snate@binkert.org l.append(s); 2207039Snate@binkert.org } 2217039Snate@binkert.org return l; 2227039Snate@binkert.org} 2237039Snate@binkert.org 2247039Snate@binkert.orgpy::array_t<int32_t, 0> test_array_ctors(int i) { 2257811Ssteve.reinhardt@amd.com using arr_t = py::array_t<int32_t, 0>; 2266145Snate@binkert.org 2277039Snate@binkert.org std::vector<int32_t> data { 1, 2, 3, 4, 5, 6 }; 228 std::vector<size_t> shape { 3, 2 }; 229 std::vector<size_t> strides { 8, 4 }; 230 231 auto ptr = data.data(); 232 auto vptr = (void *) ptr; 233 auto dtype = py::dtype("int32"); 234 235 py::buffer_info buf_ndim1(vptr, 4, "i", 6); 236 py::buffer_info buf_ndim1_null(nullptr, 4, "i", 6); 237 py::buffer_info buf_ndim2(vptr, 4, "i", 2, shape, strides); 238 py::buffer_info buf_ndim2_null(nullptr, 4, "i", 2, shape, strides); 239 240 auto fill = [](py::array arr) { 241 auto req = arr.request(); 242 for (int i = 0; i < 6; i++) ((int32_t *) req.ptr)[i] = i + 1; 243 return arr; 244 }; 245 246 switch (i) { 247 // shape: (3, 2) 248 case 10: return arr_t(shape, strides, ptr); 249 case 11: return py::array(shape, strides, ptr); 250 case 12: return py::array(dtype, shape, strides, vptr); 251 case 13: return arr_t(shape, ptr); 252 case 14: return py::array(shape, ptr); 253 case 15: return py::array(dtype, shape, vptr); 254 case 16: return arr_t(buf_ndim2); 255 case 17: return py::array(buf_ndim2); 256 // shape: (3, 2) - post-fill 257 case 20: return fill(arr_t(shape, strides)); 258 case 21: return py::array(shape, strides, ptr); // can't have nullptr due to templated ctor 259 case 22: return fill(py::array(dtype, shape, strides)); 260 case 23: return fill(arr_t(shape)); 261 case 24: return py::array(shape, ptr); // can't have nullptr due to templated ctor 262 case 25: return fill(py::array(dtype, shape)); 263 case 26: return fill(arr_t(buf_ndim2_null)); 264 case 27: return fill(py::array(buf_ndim2_null)); 265 // shape: (6, ) 266 case 30: return arr_t(6, ptr); 267 case 31: return py::array(6, ptr); 268 case 32: return py::array(dtype, 6, vptr); 269 case 33: return arr_t(buf_ndim1); 270 case 34: return py::array(buf_ndim1); 271 // shape: (6, ) 272 case 40: return fill(arr_t(6)); 273 case 41: return py::array(6, ptr); // can't have nullptr due to templated ctor 274 case 42: return fill(py::array(dtype, 6)); 275 case 43: return fill(arr_t(buf_ndim1_null)); 276 case 44: return fill(py::array(buf_ndim1_null)); 277 } 278 return arr_t(); 279} 280 281py::list test_dtype_ctors() { 282 py::list list; 283 list.append(py::dtype("int32")); 284 list.append(py::dtype(std::string("float64"))); 285 list.append(py::dtype::from_args(py::str("bool"))); 286 py::list names, offsets, formats; 287 py::dict dict; 288 names.append(py::str("a")); names.append(py::str("b")); dict["names"] = names; 289 offsets.append(py::int_(1)); offsets.append(py::int_(10)); dict["offsets"] = offsets; 290 formats.append(py::dtype("int32")); formats.append(py::dtype("float64")); dict["formats"] = formats; 291 dict["itemsize"] = py::int_(20); 292 list.append(py::dtype::from_args(dict)); 293 list.append(py::dtype(names, formats, offsets, 20)); 294 list.append(py::dtype(py::buffer_info((void *) 0, sizeof(unsigned int), "I", 1))); 295 list.append(py::dtype(py::buffer_info((void *) 0, 0, "T{i:a:f:b:}", 1))); 296 return list; 297} 298 299struct TrailingPaddingStruct { 300 int32_t a; 301 char b; 302}; 303 304py::dtype trailing_padding_dtype() { 305 return py::dtype::of<TrailingPaddingStruct>(); 306} 307 308py::dtype buffer_to_dtype(py::buffer& buf) { 309 return py::dtype(buf.request()); 310} 311 312py::list test_dtype_methods() { 313 py::list list; 314 auto dt1 = py::dtype::of<int32_t>(); 315 auto dt2 = py::dtype::of<SimpleStruct>(); 316 list.append(dt1); list.append(dt2); 317 list.append(py::bool_(dt1.has_fields())); list.append(py::bool_(dt2.has_fields())); 318 list.append(py::int_(dt1.itemsize())); list.append(py::int_(dt2.itemsize())); 319 return list; 320} 321 322struct CompareStruct { 323 bool x; 324 uint32_t y; 325 float z; 326}; 327 328py::list test_compare_buffer_info() { 329 py::list list; 330 list.append(py::bool_(py::detail::compare_buffer_info<float>::compare(py::buffer_info(nullptr, sizeof(float), "f", 1)))); 331 list.append(py::bool_(py::detail::compare_buffer_info<unsigned>::compare(py::buffer_info(nullptr, sizeof(int), "I", 1)))); 332 list.append(py::bool_(py::detail::compare_buffer_info<long>::compare(py::buffer_info(nullptr, sizeof(long), "l", 1)))); 333 list.append(py::bool_(py::detail::compare_buffer_info<long>::compare(py::buffer_info(nullptr, sizeof(long), sizeof(long) == sizeof(int) ? "i" : "q", 1)))); 334 list.append(py::bool_(py::detail::compare_buffer_info<CompareStruct>::compare(py::buffer_info(nullptr, sizeof(CompareStruct), "T{?:x:3xI:y:f:z:}", 1)))); 335 return list; 336} 337 338test_initializer numpy_dtypes([](py::module &m) { 339 try { 340 py::module::import("numpy"); 341 } catch (...) { 342 return; 343 } 344 345 // typeinfo may be registered before the dtype descriptor for scalar casts to work... 346 py::class_<SimpleStruct>(m, "SimpleStruct"); 347 348 PYBIND11_NUMPY_DTYPE(SimpleStruct, bool_, uint_, float_, ldbl_); 349 PYBIND11_NUMPY_DTYPE(PackedStruct, bool_, uint_, float_, ldbl_); 350 PYBIND11_NUMPY_DTYPE(NestedStruct, a, b); 351 PYBIND11_NUMPY_DTYPE(PartialStruct, bool_, uint_, float_, ldbl_); 352 PYBIND11_NUMPY_DTYPE(PartialNestedStruct, a); 353 PYBIND11_NUMPY_DTYPE(StringStruct, a, b); 354 PYBIND11_NUMPY_DTYPE(EnumStruct, e1, e2); 355 PYBIND11_NUMPY_DTYPE(TrailingPaddingStruct, a, b); 356 PYBIND11_NUMPY_DTYPE(CompareStruct, x, y, z); 357 358 // ... or after 359 py::class_<PackedStruct>(m, "PackedStruct"); 360 361 PYBIND11_NUMPY_DTYPE_EX(StructWithUglyNames, __x__, "x", __y__, "y"); 362 363 // If uncommented, this should produce a static_assert failure telling the user that the struct 364 // is not a POD type 365// struct NotPOD { std::string v; NotPOD() : v("hi") {}; }; 366// PYBIND11_NUMPY_DTYPE(NotPOD, v); 367 368 m.def("create_rec_simple", &create_recarray<SimpleStruct>); 369 m.def("create_rec_packed", &create_recarray<PackedStruct>); 370 m.def("create_rec_nested", &create_nested); 371 m.def("create_rec_partial", &create_recarray<PartialStruct>); 372 m.def("create_rec_partial_nested", &create_partial_nested); 373 m.def("print_format_descriptors", &print_format_descriptors); 374 m.def("print_rec_simple", &print_recarray<SimpleStruct>); 375 m.def("print_rec_packed", &print_recarray<PackedStruct>); 376 m.def("print_rec_nested", &print_recarray<NestedStruct>); 377 m.def("print_dtypes", &print_dtypes); 378 m.def("get_format_unbound", &get_format_unbound); 379 m.def("create_string_array", &create_string_array); 380 m.def("print_string_array", &print_recarray<StringStruct>); 381 m.def("create_enum_array", &create_enum_array); 382 m.def("print_enum_array", &print_recarray<EnumStruct>); 383 m.def("test_array_ctors", &test_array_ctors); 384 m.def("test_dtype_ctors", &test_dtype_ctors); 385 m.def("test_dtype_methods", &test_dtype_methods); 386 m.def("compare_buffer_info", &test_compare_buffer_info); 387 m.def("trailing_padding_dtype", &trailing_padding_dtype); 388 m.def("buffer_to_dtype", &buffer_to_dtype); 389 m.def("f_simple", [](SimpleStruct s) { return s.uint_ * 10; }); 390 m.def("f_packed", [](PackedStruct s) { return s.uint_ * 10; }); 391 m.def("f_nested", [](NestedStruct s) { return s.a.uint_ * 10; }); 392 m.def("register_dtype", []() { PYBIND11_NUMPY_DTYPE(SimpleStruct, bool_, uint_, float_, ldbl_); }); 393}); 394 395#undef PYBIND11_PACKED 396