Efficient implementation of Tensor::ones() for metal
#2512
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
candle-metal-kernels/src/lib.rs
Outdated
|
|
||
| encoder.use_resource(output, metal::MTLResourceUsage::Write); | ||
|
|
||
| let grid_size = MTLSize { width: length as u64, height: 1, depth: 1 }; |
There was a problem hiding this comment.
Is length here correct? Shouldn't it divided by the number of threads? Anyway better to use utils::linear_split as it should do the proper thing here.
There was a problem hiding this comment.
this was stupid of me. Resolved.
LaurentMazare
approved these changes
Oct 1, 2024
|
@LaurentMazare |
|
Yes please. |
LaurentMazare
approved these changes
Oct 1, 2024
|
Thanks! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
While working with
Tensor::ones()I realized the implementation is significantly slower thanpytorchcounterpart inmetal. The original implementation would create aCpuStorageand move it to Gpu.This implementation introduces a metal kernel function for
Tensor::ones()impl.Benchmarks (original vs new implementation)
Benchmark Code ran on Mac M1 Pro 16GB
u8u8u8u32u32u32f32f32f32bf16bf16bf16f16f16f16i64i64i64Looks like we are getting some significant performance boost.
Implementation Details:
MetalDevice::ones_implnow introduces a code-path similar to therandomimplementation - falls back toCpuStoragebased implementation forDType::f64fill- following the pattern implemented fornormalmetal kernelscandle-coretest cases to test forbf16andf16typescandle-metal-kernelsfor theconstant_fillrelated code introducedPlease let me know if I'm missing something!