Implement mta_execute_model - #306
Conversation
This function handles consistency checking of inputs and outputs, and unit conversion of inputs and outputs.
| { | ||
| let mut cache = ATOMIC_TYPES_CACHE.lock().expect("ATOMIC_TYPES_CACHE lock poisoned"); | ||
| let atomic_types = &capabilities.atomic_types; | ||
| let valid_types = cache.get_or_insert_ref(atomic_types.as_slice(), || { |
There was a problem hiding this comment.
Is atomic_types guaranteed to be sorted? (From the comment, you want the key to be sorted.)
There was a problem hiding this comment.
no, I forgot to sort them in there. I've made this part of the type contract so we don't get it wrong next time.
|
|
||
| /// Validate that the inputs to a model are consistent with its capabilities. | ||
| /// | ||
| /// This checks that: |
There was a problem hiding this comment.
Maybe worth checking also that the device is in the model's supported devices?
| for (parameter, gradient) in block.gradients() { | ||
| let grad_device = gradient.values().device()?; | ||
| let grad_copy = gradient.values().copy(grad_device)?; | ||
| let mut dlpack = grad_copy.as_dlpack(device, None, dlpack_version)?; |
There was a problem hiding this comment.
| let mut dlpack = grad_copy.as_dlpack(device, None, dlpack_version)?; | |
| let mut dlpack = grad_copy.as_dlpack(grad_device, None, dlpack_version)?; |
| let src_ptr = unsafe { dlpack_to_device_ptr(*tensor) }; | ||
| let n_elements: i64 = tensor.shape().iter().product(); | ||
| let elem_size = tensor.dtype().bits as usize / 8; | ||
| let num_bytes = n_elements as usize * elem_size; |
There was a problem hiding this comment.
Worth checking that the tensor is contiguous?
| outputs: *mut *mut mts_tensormap_t, | ||
| outputs_count: usize, | ||
| ) -> Result<(), Error> { | ||
| let capabilities = model.capabilities()?; |
There was a problem hiding this comment.
Can these be cached? Doing string allocation, JSON conversion, etc at every call might be wasteful.
There was a problem hiding this comment.
I think we can try to get them at every call for now, and I'll leave a TODO note to measure the overhead of this
1f4f620 to
1637cfb
Compare
2 participants
This should be the last building block for the new C API!
Contributor (creator of pull-request) checklist
Issue referenced (for PRs that solve an issue)?Reviewer checklist
CHANGELOG updated with public API or any other important changes?