| 30 | |
| 31 | template <int out_dim, typename ctype> |
| 32 | void setup_and_launch(const JITExecutor* fusion_opr, CUfunction func, int block_size) { |
| 33 | auto&& args = fusion_opr->args(); |
| 34 | size_t num_memrefs = args.inputs.size() + args.outputs.size(); |
| 35 | std::vector<StridedMemRefType<ctype, out_dim>> param_holders(num_memrefs); |
| 36 | std::vector<void*> params; |
| 37 | |
| 38 | auto set_params = [¶m_holders, ¶ms]( |
| 39 | size_t idx, void* ptr, |
| 40 | const megdnn::TensorLayout& layout) { |
| 41 | auto& desc = param_holders[idx]; |
| 42 | desc.basePtr = static_cast<ctype*>(ptr); |
| 43 | params.push_back(&(desc.basePtr)); |
| 44 | desc.data = static_cast<ctype*>(ptr); |
| 45 | params.push_back(&(desc.data)); |
| 46 | desc.offset = 0; |
| 47 | params.push_back(&(desc.offset)); |
| 48 | for (size_t i = 0; i < layout.ndim; i++) { |
| 49 | desc.sizes[i] = layout.shape[i]; |
| 50 | params.push_back(&(desc.sizes[i])); |
| 51 | desc.strides[i] = layout.stride[i]; |
| 52 | params.push_back(&(desc.strides[i])); |
| 53 | } |
| 54 | }; |
| 55 | |
| 56 | size_t idx = 0; |
| 57 | for (const auto& arg : args.inputs) { |
| 58 | set_params(idx++, arg.from->dev_tensor().raw_ptr(), arg.from->layout()); |
| 59 | } |
| 60 | |
| 61 | int64_t nr_elements = 0; |
| 62 | for (const auto& arg : args.outputs) { |
| 63 | if (nr_elements == 0) { |
| 64 | nr_elements = arg.from->layout().total_nr_elems(); |
| 65 | } else { |
| 66 | mgb_assert( |
| 67 | static_cast<size_t>(nr_elements) == arg.layout.total_nr_elems(), |
| 68 | "The number of elements of outputs mismatch, expected: " |
| 69 | "%zu got: %zu(%s)", |
| 70 | static_cast<size_t>(nr_elements), |
| 71 | arg.from->layout().total_nr_elems(), |
| 72 | arg.from->layout().to_string().c_str()); |
| 73 | } |
| 74 | |
| 75 | set_params(idx++, arg.from->dev_tensor().raw_ptr(), arg.from->layout()); |
| 76 | } |
| 77 | |
| 78 | mgb_assert( |
| 79 | param_holders.size() == num_memrefs, |
| 80 | "calling push_back method of param_holders is unsafe as it " |
| 81 | "might cause reallocation of std::vector"); |
| 82 | |
| 83 | const CompNodeEnv& env = CompNodeEnv::from_comp_node(fusion_opr->comp_node()); |
| 84 | |
| 85 | int64_t grid_size; |
| 86 | if (nr_elements <= block_size) { |
| 87 | block_size = nr_elements; |
| 88 | grid_size = 1; |
| 89 | } else { |
nothing calls this directly
no test coverage detected