| 242 | } |
| 243 | |
| 244 | std::vector<array> precompiled_cuda_kernel( |
| 245 | const std::string& name, |
| 246 | const std::string& compiled_source, |
| 247 | const std::vector<array>& inputs, |
| 248 | const std::vector<Shape>& output_shapes, |
| 249 | const std::vector<Dtype>& output_dtypes, |
| 250 | const std::vector<ScalarArg>& scalars, |
| 251 | std::tuple<int, int, int> grid, |
| 252 | std::tuple<int, int, int> threadgroup, |
| 253 | int shared_memory, |
| 254 | std::optional<float> init_value, |
| 255 | bool ensure_row_contiguous, |
| 256 | StreamOrDevice s) { |
| 257 | std::vector<std::tuple<bool, bool, bool>> shape_infos( |
| 258 | inputs.size(), {false, false, false}); |
| 259 | return array::make_arrays( |
| 260 | output_shapes, |
| 261 | output_dtypes, |
| 262 | std::make_shared<CustomKernel>( |
| 263 | to_stream(s), |
| 264 | name, |
| 265 | compiled_source, |
| 266 | grid, |
| 267 | threadgroup, |
| 268 | shape_infos, |
| 269 | ensure_row_contiguous, |
| 270 | init_value, |
| 271 | scalars, |
| 272 | true, |
| 273 | shared_memory), |
| 274 | inputs); |
| 275 | } |
| 276 | |
| 277 | void CustomKernel::eval_gpu( |
| 278 | const std::vector<array>& inputs, |