| 57 | } |
| 58 | |
| 59 | void* get_cuda_stream(mgbComputeNode_t device) { |
| 60 | void* rst = nullptr; |
| 61 | #if MGB_CUDA |
| 62 | auto* cn = reinterpret_cast<mgb::CompNode*>(device); |
| 63 | MGB_TRY { rst = CompNodeEnv::from_comp_node(*cn).cuda_env().stream; } |
| 64 | MGB_CATCH(MegBrainError & exc, { |
| 65 | mgb_log_error("failed to get stream: %s", exc.what()); |
| 66 | }) |
| 67 | #else |
| 68 | mgb_log_error("megbrain compiled without cuda support!"); |
| 69 | #endif |
| 70 | return rst; |
| 71 | } |
| 72 | |
| 73 | MGB_API DeviceLocator get_physical_location(mgbComputeNode_t device) { |
| 74 | auto location = reinterpret_cast<CompNode*>(device)->locator().to_physical(); |