| 144 | MLIRCUDAExecutable::~MLIRCUDAExecutable() {} |
| 145 | |
| 146 | void MLIRCUDAExecutable::FuncCache::exec( |
| 147 | const JITExecutor* fusion_opr, const MLIRCUDAExecutable* cuda_exe) { |
| 148 | Func* func; |
| 149 | { |
| 150 | MGB_LOCK_GUARD(mtx); |
| 151 | auto ins = cn2func.insert({fusion_opr->comp_node(), {}}); |
| 152 | func = &ins.first->second; |
| 153 | if (ins.second) { |
| 154 | MGB_CUDA_CU_CHECK(cuModuleLoadData(&func->module, kernel_data.data())); |
| 155 | MGB_CUDA_CU_CHECK(cuModuleGetFunction( |
| 156 | &func->func, func->module, cuda_exe->m_kernel_name.c_str())); |
| 157 | int min_grid_size = 0; |
| 158 | MGB_CUDA_CU_CHECK(cuOccupancyMaxPotentialBlockSize( |
| 159 | &min_grid_size, &func->block_size, func->func, nullptr, 0, 0)); |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | mgb_assert( |
| 164 | fusion_opr->args().outputs.size() == 1, |
| 165 | "Currently only support 1 outputs, got %zu", |
| 166 | fusion_opr->args().outputs.size()); |
| 167 | int out_dim = fusion_opr->args().outputs[0].from->layout().ndim; |
| 168 | DType dtype = fusion_opr->args().outputs[0].from->layout().dtype; |
| 169 | |
| 170 | switch (out_dim) { |
| 171 | #define cb(_ndim) \ |
| 172 | case _ndim: \ |
| 173 | setup_and_launch_dim<_ndim>(dtype, fusion_opr, func->func, func->block_size); \ |
| 174 | break; |
| 175 | cb(1); |
| 176 | cb(2); |
| 177 | cb(3); |
| 178 | cb(4); |
| 179 | #undef cb |
| 180 | } |
| 181 | } |
| 182 | |
| 183 | #endif // MGB_CUDA |
| 184 | #endif // MGB_JIT && MGB_JIT_MLIR |