1 /**
2 * Copyright 2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "debug/data_dump/npy_header.h"
18
19 #include <utility>
20 #include <sstream>
21
22 #include "utils/hash_map.h"
23 #include "mindspore/core/ir/dtype.h"
24 #include "mindspore/core/utils/log_adapter.h"
25 #include "mindspore/core/utils/convert_utils_base.h"
26
27 namespace mindspore {
28 namespace {
29 // npy file header start information
30 const char kMagicPrefix[] = "\x93NUMPY";
31 // magical length include kMagicPrefix length and version length
32 const size_t kMagicLen = 6;
33 const size_t kArrayAlign = 64;
34
35 // first: header_length_type, second: encoding_type
36 // header_length_type: 1 represents 2 bytes; 2 and 3 represents 4 bytes
37 // encoding_type: 1 and 2 represents 'latin1'; 3 represents 'utf8'
38 using version_type = std::pair<int, int>;
39
40 // data type description
41 // byteorder char: '<' is little endian; '>' is big endian; '|' is ignore(no change to byte order)
42 // type char: 'b' represents bool; 'u' represents uint; 'i' represents int; 'f' represents float
43 struct DtypeDescr {
44 char byteorder;
45 char type;
46 size_t length;
47
48 std::string str() const;
49 };
50
51 // npy file header description, includes data type description, fortran_order and array shape
52 // fortran_order: true represents the array data Fortran-contiguous; false represents the array data C-contiguity
53 struct NpyHeader {
54 public:
55 DtypeDescr dtype_descr;
56 bool fortran_order;
57 ShapeVector shape;
58
59 std::string str() const;
60
61 private:
62 std::string fortran_order_to_str() const;
63 std::string shape_to_str() const;
64 };
65
str() const66 std::string DtypeDescr::str() const {
67 std::ostringstream buffer;
68 buffer << "\'" << byteorder << type << length << "\'";
69 return buffer.str();
70 }
71
str() const72 std::string NpyHeader::str() const {
73 const std::string first_field = "'descr': ";
74 const std::string second_field = "'fortran_order': ";
75 const std::string third_field = "'shape': ";
76 std::ostringstream buffer;
77 buffer << "{" << first_field << dtype_descr.str() << ", " << second_field << fortran_order_to_str() << ", "
78 << third_field << shape_to_str() << ", }";
79 return buffer.str();
80 }
81
fortran_order_to_str() const82 std::string NpyHeader::fortran_order_to_str() const { return fortran_order ? "True" : "False"; }
83
shape_to_str() const84 std::string NpyHeader::shape_to_str() const {
85 std::ostringstream buffer;
86 buffer << "(";
87 for (const auto i : shape) {
88 buffer << std::to_string(i) << ",";
89 }
90 buffer << ")";
91 return buffer.str();
92 }
93
94 // dtype description corresponding to tensor type
95 const mindspore::HashMap<TypeId, DtypeDescr> type_desc_map = {
96 {kNumberTypeBool, DtypeDescr{'|', 'b', 1}}, {kNumberTypeInt8, DtypeDescr{'|', 'i', 1}},
97 {kNumberTypeInt16, DtypeDescr{'<', 'i', 2}}, {kNumberTypeInt32, DtypeDescr{'<', 'i', 4}},
98 {kNumberTypeInt64, DtypeDescr{'<', 'i', 8}}, {kNumberTypeUInt8, DtypeDescr{'|', 'u', 1}},
99 {kNumberTypeUInt16, DtypeDescr{'<', 'u', 2}}, {kNumberTypeUInt32, DtypeDescr{'<', 'u', 4}},
100 {kNumberTypeUInt64, DtypeDescr{'<', 'u', 8}}, {kNumberTypeFloat16, DtypeDescr{'<', 'f', 2}},
101 {kNumberTypeFloat32, DtypeDescr{'<', 'f', 4}}, {kNumberTypeFloat64, DtypeDescr{'<', 'f', 8}},
102 {kNumberTypeBFloat16, DtypeDescr{'<', 'T', 2}}, {kNumberTypeComplex128, DtypeDescr{'<', 'c', 16}},
103 {kNumberTypeComplex64, DtypeDescr{'<', 'c', 8}},
104 };
105 } // namespace
106
int_to_byte(size_t number,char * byte,size_t length)107 void int_to_byte(size_t number, char *byte, size_t length) {
108 const size_t byte_len = 8;
109 const size_t mask = 0xff;
110 for (size_t i = 0; i < length; i++) {
111 byte[i] = (number >> (i * byte_len)) & mask;
112 }
113 }
114
GenerateNpyHeader(const ShapeVector & shape,TypeId type_id,bool fortran_order)115 std::string GenerateNpyHeader(const ShapeVector &shape, TypeId type_id, bool fortran_order) {
116 auto type_desc = type_desc_map.find(type_id);
117 if (type_desc == type_desc_map.end()) {
118 MS_LOG(WARNING) << "Not support dump the " << TypeIdToType(type_id)->ToString() << " data to npy file.";
119 return std::string();
120 }
121
122 NpyHeader npy_header{type_desc->second, fortran_order, shape};
123 std::string header_str = npy_header.str();
124 version_type version{1, 0};
125 const size_t header_len = header_str.length();
126 const size_t version_len = 2;
127 const size_t max_len = 65535;
128 size_t length_len = 2;
129 size_t total_len = kMagicLen + version_len + length_len + header_len + 1;
130 if (total_len > max_len) {
131 version = {2, 0};
132 length_len = 4;
133 total_len = kMagicLen + version_len + length_len + header_len + 1;
134 }
135
136 const size_t pad_len = kArrayAlign - total_len % kArrayAlign;
137 const size_t padding_header_len = header_len + pad_len + 1;
138 const std::string padding(pad_len, ' ');
139 const std::string end_line = "\n";
140 char *length_byte = new char[length_len];
141 int_to_byte(padding_header_len, length_byte, length_len);
142
143 std::ostringstream out;
144 (void)out.write(kMagicPrefix, SizeToLong(kMagicLen));
145 (void)out.put(version.first);
146 (void)out.put(version.second);
147 (void)out.write(length_byte, SizeToLong(length_len));
148 out << header_str << padding << end_line;
149 delete[] length_byte;
150 return out.str();
151 }
152 } // namespace mindspore
153