| 3563 | */ |
| 3564 | template <typename IndexType, class Func> |
| 3565 | CUresult parallel_for(ExecutionPolicy policy, IndexType begin, IndexType end, |
| 3566 | Lambda<Func> const& lambda) { |
| 3567 | using namespace jitify; |
| 3568 | |
| 3569 | if (policy.location == HOST) { |
| 3570 | #ifdef _OPENMP |
| 3571 | #pragma omp parallel for |
| 3572 | #endif |
| 3573 | for (IndexType i = begin; i < end; i++) { |
| 3574 | lambda._func(i); |
| 3575 | } |
| 3576 | return CUDA_SUCCESS; // FIXME - replace with non-CUDA enum type? |
| 3577 | } |
| 3578 | |
| 3579 | thread_local static JitCache kernel_cache(policy.cache_size); |
| 3580 | |
| 3581 | std::vector<std::string> arg_decls; |
| 3582 | arg_decls.push_back("I begin, I end"); |
| 3583 | arg_decls.insert(arg_decls.end(), lambda._capture._arg_decls.begin(), |
| 3584 | lambda._capture._arg_decls.end()); |
| 3585 | |
| 3586 | std::stringstream source_ss; |
| 3587 | source_ss << "parallel_for_program\n"; |
| 3588 | for (auto const& header : policy.headers) { |
| 3589 | std::string header_name = header.substr(0, header.find("\n")); |
| 3590 | source_ss << "#include <" << header_name << ">\n"; |
| 3591 | } |
| 3592 | source_ss << "template<typename I>\n" |
| 3593 | "__global__\n" |
| 3594 | "void parallel_for_kernel(" |
| 3595 | << reflection::reflect_list(arg_decls) |
| 3596 | << ") {\n" |
| 3597 | " I i0 = threadIdx.x + blockDim.x*blockIdx.x;\n" |
| 3598 | " for( I i=i0+begin; i<end; i+=blockDim.x*gridDim.x ) {\n" |
| 3599 | " " |
| 3600 | << "\t" << lambda._func_string << ";\n" |
| 3601 | << " }\n" |
| 3602 | "}\n"; |
| 3603 | |
| 3604 | Program program = kernel_cache.program(source_ss.str(), policy.headers, |
| 3605 | policy.options, policy.file_callback); |
| 3606 | |
| 3607 | std::vector<void*> arg_ptrs; |
| 3608 | arg_ptrs.push_back(&begin); |
| 3609 | arg_ptrs.push_back(&end); |
| 3610 | arg_ptrs.insert(arg_ptrs.end(), lambda._capture._arg_ptrs.begin(), |
| 3611 | lambda._capture._arg_ptrs.end()); |
| 3612 | |
| 3613 | size_t n = end - begin; |
| 3614 | dim3 block(policy.block_size); |
| 3615 | dim3 grid((unsigned int)std::min((n - 1) / block.x + 1, size_t(65535))); |
| 3616 | cudaSetDevice(policy.device); |
| 3617 | return program.kernel("parallel_for_kernel") |
| 3618 | .instantiate<IndexType>() |
| 3619 | .configure(grid, block, 0, policy.stream) |
| 3620 | .launch(arg_ptrs); |
| 3621 | } |
| 3622 | |