MCPcopy Create free account
hub / github.com/UbiquitousLearning/mllm / buildOpAndSubmitTask

Method buildOpAndSubmitTask

mllm/engine/Context.cpp:49–97  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

47}
48
49std::vector<Tensor> Context::buildOpAndSubmitTask(OpTypes op_type, const BaseOpOptionsBase& base_options,
50 const std::vector<Tensor>& inputs, DeviceTypes special_device) {
51 MLLM_TRACY_ZONE_SCOPED;
52 auto device = special_device != kDeviceTypes_End ? special_device : inputs[0].device();
53
54 // If input device and special device are different, prefer non-CPU device
55 if (special_device != kDeviceTypes_End && special_device != inputs[0].device()) {
56 auto input_device = inputs[0].device();
57 if (input_device == kCPU && special_device != kCPU) {
58 // Use special device (non-CPU) over input device (CPU)
59 device = special_device;
60 } else if (special_device == kCPU && input_device != kCPU) {
61 // Use input device (non-CPU) over special device (CPU)
62 device = input_device;
63 } else {
64 // Both are non-CPU or both are CPU, use special device as originally intended
65 device = special_device;
66 }
67 }
68
69 auto op = getBackend(device)->createOp(op_type, base_options);
70 auto task = Task::createExecuteOpTask(op, inputs, {});
71
72 auto this_thread = thisThread();
73 if (this_thread->trace_mode) {
74 // Submit!
75 // At this moment, heart pounding like thunder
76 // Tasks racing through kernels, swift as lightning
77 // Threads await, fate hanging by a thread
78 // Success or failure in this one moment
79 task->custom_context_ptr = this_thread->ir_context.get();
80 dispatcherManager()->submit(Dispatcher::trace_dispatcher_id, task);
81
82 // Everything is Ok. Bravo! You did it.
83 // Return what we need.
84 return task->outputs;
85 } else {
86 // Submit!
87 // At this moment, heart pounding like thunder
88 // Tasks racing through kernels, swift as lightning
89 // Threads await, fate hanging by a thread
90 // Success or failure in this one moment
91 dispatcherManager()->submit(static_cast<int32_t>(device), task);
92
93 // Everything is Ok. Bravo! You did it.
94 // Return what we need.
95 return task->outputs;
96 }
97}
98
99uint32_t Context::getUUID() {
100 uint32_t ret = custom_uuid_giver_;

Callers 15

radixAttnSWAwSinkFunction · 0.80
flashAttnSWAwSinkFunction · 0.80
operator[]Method · 0.80
sliceMethod · 0.80
zerosMethod · 0.80
onesMethod · 0.80
arangeMethod · 0.80
randomMethod · 0.80
operator+Method · 0.80
operator-Method · 0.80
operator*Method · 0.80
operator/Method · 0.80

Calls 5

thisThreadFunction · 0.85
createOpMethod · 0.80
submitMethod · 0.80
deviceMethod · 0.45
getMethod · 0.45

Tested by 2

radixAttnSWAwSinkFunction · 0.64
flashAttnSWAwSinkFunction · 0.64