| 148 | } |
| 149 | |
| 150 | inline void build_kernel( |
| 151 | std::ostream& os, |
| 152 | const std::string& kernel_name, |
| 153 | const std::vector<array>& inputs, |
| 154 | const std::vector<array>& outputs, |
| 155 | const std::vector<array>& tape, |
| 156 | const std::function<bool(size_t)>& is_constant, |
| 157 | bool contiguous, |
| 158 | int ndim) { |
| 159 | NodeNamer namer; |
| 160 | |
| 161 | #ifdef _MSC_VER |
| 162 | // Export the symbol |
| 163 | os << "__declspec(dllexport) "; |
| 164 | #endif |
| 165 | |
| 166 | // Start the kernel |
| 167 | os << "void " << kernel_name |
| 168 | << "(int* shape, int64_t** strides, void** args) {" << std::endl; |
| 169 | |
| 170 | // Add the input arguments |
| 171 | int cnt = 0; |
| 172 | int strides_index = 1; |
| 173 | for (size_t i = 0; i < inputs.size(); ++i) { |
| 174 | // Skip constants from the input list |
| 175 | if (is_constant(i)) { |
| 176 | continue; |
| 177 | } |
| 178 | |
| 179 | const auto& x = inputs[i]; |
| 180 | auto& xname = namer.get_name(x); |
| 181 | |
| 182 | auto tstr = get_type_string(x.dtype()); |
| 183 | os << " " << tstr << "* " << xname << " = (" << tstr << "*)args[" << cnt++ |
| 184 | << "];" << std::endl; |
| 185 | // Scalars and contiguous need no strides |
| 186 | if (!is_scalar(x) && !contiguous) { |
| 187 | os << " const int64_t* " << xname << "_strides = strides[" |
| 188 | << strides_index++ << "];" << std::endl; |
| 189 | } |
| 190 | } |
| 191 | |
| 192 | // Add the output arguments |
| 193 | for (auto& x : outputs) { |
| 194 | auto tstr = get_type_string(x.dtype()); |
| 195 | os << " " << tstr << "* " << namer.get_name(x) << " = (" << tstr |
| 196 | << "*)args[" << cnt++ << "];" << std::endl; |
| 197 | } |
| 198 | // Add output size |
| 199 | if (contiguous) { |
| 200 | os << " const size_t size = (size_t)args[" << cnt++ << "];" << std::endl; |
| 201 | } |
| 202 | |
| 203 | if (contiguous) { |
| 204 | os << " for (size_t i = 0; i < size; ++i) {" << std::endl; |
| 205 | } else { |
| 206 | for (int d = 0; d < ndim; ++d) { |
| 207 | os << " for (int i" << d << " = 0; i" << d << " < shape[" << d |
no test coverage detected