MCPcopy Create free account
hub / github.com/MegEngine/MegCC / write_helper

Function write_helper

compiler/test/kernel/common/src/rng.cpp:76–146  ·  view source on GitHub ↗

namespace

Source from the content-addressed store, hash-verified

74};
75} // namespace
76static void write_helper(const TensorND& tensor, std::function<float()> gen) {
77 auto t_dtype = tensor.layout.dtype;
78
79 if (t_dtype.enumv() == DTypeEnum::Float32) {
80 FastIter<float> iter(tensor, tensor.layout.is_contiguous());
81 for (size_t i = 0; i < tensor.layout.total_nr_elems(); ++i) {
82 *iter = gen();
83 ++iter;
84 }
85 } else if (
86 t_dtype.enumv() == DTypeEnum::Uint8 ||
87 t_dtype.enumv() == DTypeEnum::Quantized8Asymm) {
88 auto iter_tensor = tensor;
89 iter_tensor.layout.dtype = dtype::Uint8();
90 FastIter<uint8_t> iter(iter_tensor, tensor.layout.is_contiguous());
91 for (size_t i = 0; i < tensor.layout.total_nr_elems(); ++i) {
92 int val = (int)gen();
93 val = val > 255 ? 255 : val;
94 val = val < 0 ? 0 : val;
95 *iter = (uint8_t)val;
96 ++iter;
97 }
98 } else if (
99 t_dtype.enumv() == DTypeEnum::Int8 ||
100 t_dtype.enumv() == DTypeEnum::QuantizedS8) {
101 auto iter_tensor = tensor;
102 iter_tensor.layout.dtype = dtype::Int8();
103 FastIter<int8_t> iter(iter_tensor, tensor.layout.is_contiguous());
104 for (size_t i = 0; i < tensor.layout.total_nr_elems(); ++i) {
105 int val = (int)gen();
106 val = val > 127 ? 127 : val;
107 val = val < -128 ? -128 : val;
108 *iter = (int8_t)val;
109 ++iter;
110 }
111 } else if (
112 t_dtype.enumv() == DTypeEnum::Int32 ||
113 t_dtype.enumv() == DTypeEnum::QuantizedS32) {
114 auto iter_tensor = tensor;
115 iter_tensor.layout.dtype = dtype::Int32();
116 FastIter<int> iter(iter_tensor, tensor.layout.is_contiguous());
117 for (size_t i = 0; i < tensor.layout.total_nr_elems(); ++i) {
118 *iter = gen();
119 ++iter;
120 }
121 } else if (t_dtype.enumv() == DTypeEnum::Int16) {
122 auto iter_tensor = tensor;
123 iter_tensor.layout.dtype = dtype::Int16();
124 FastIter<int16_t> iter(iter_tensor, tensor.layout.is_contiguous());
125 for (size_t i = 0; i < tensor.layout.total_nr_elems(); ++i) {
126 *iter = gen();
127 ++iter;
128 }
129 } else if (t_dtype.enumv() == DTypeEnum::Float16) {
130 FastIter<megdnn::dt_float16> iter(tensor, tensor.layout.is_contiguous());
131 for (size_t i = 0; i < tensor.layout.total_nr_elems(); ++i) {
132 auto val = megdnn::dt_float16(gen());
133

Callers 1

genMethod · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected