| 142 | } // namespace |
| 143 | |
| 144 | CustomKernelFunction cuda_kernel( |
| 145 | const std::string& name, |
| 146 | const std::vector<std::string>& input_names, |
| 147 | const std::vector<std::string>& output_names, |
| 148 | const std::string& source, |
| 149 | const std::string& header, |
| 150 | bool ensure_row_contiguous, |
| 151 | int shared_memory) { |
| 152 | if (output_names.empty()) { |
| 153 | throw std::invalid_argument( |
| 154 | "[custom_kernel] Must specify at least one output."); |
| 155 | } |
| 156 | |
| 157 | std::vector<std::tuple<bool, bool, bool>> shape_infos; |
| 158 | for (auto& n : input_names) { |
| 159 | std::tuple<bool, bool, bool> shape_info; |
| 160 | std::get<0>(shape_info) = source.find(n + "_shape") != std::string::npos; |
| 161 | std::get<1>(shape_info) = source.find(n + "_strides") != std::string::npos; |
| 162 | std::get<2>(shape_info) = source.find(n + "_ndim") != std::string::npos; |
| 163 | shape_infos.push_back(shape_info); |
| 164 | } |
| 165 | |
| 166 | return [=, shape_infos = std::move(shape_infos)]( |
| 167 | const std::vector<array>& inputs, |
| 168 | const std::vector<Shape>& output_shapes, |
| 169 | const std::vector<Dtype>& output_dtypes, |
| 170 | std::tuple<int, int, int> grid, |
| 171 | std::tuple<int, int, int> threadgroup, |
| 172 | const std::vector<std::pair<std::string, TemplateArg>>& |
| 173 | template_args = {}, |
| 174 | std::optional<float> init_value = std::nullopt, |
| 175 | bool verbose = false, |
| 176 | StreamOrDevice s_ = {}) { |
| 177 | if (inputs.size() != input_names.size()) { |
| 178 | std::ostringstream msg; |
| 179 | msg << "[custom_kernel] Expected `inputs` to have size " |
| 180 | << input_names.size() << " but got size " << inputs.size() << "." |
| 181 | << std::endl; |
| 182 | throw std::invalid_argument(msg.str()); |
| 183 | } |
| 184 | if (output_shapes.size() != output_names.size()) { |
| 185 | std::ostringstream msg; |
| 186 | msg << "[custom_kernel] Expected `output_shapes` to have size " |
| 187 | << output_names.size() << " but got size " << output_shapes.size() |
| 188 | << "." << std::endl; |
| 189 | throw std::invalid_argument(msg.str()); |
| 190 | } |
| 191 | if (output_dtypes.size() != output_names.size()) { |
| 192 | std::ostringstream msg; |
| 193 | msg << "[custom_kernel] Expected `output_dtypes` to have size " |
| 194 | << output_names.size() << " but got size " << output_dtypes.size() |
| 195 | << "." << std::endl; |
| 196 | throw std::invalid_argument(msg.str()); |
| 197 | } |
| 198 | |
| 199 | auto s = to_stream(s_); |
| 200 | if (s.device != Device::gpu) { |
| 201 | throw std::invalid_argument("[custom_kernel] Only supports the GPU."); |
nothing calls this directly
no test coverage detected