How does Candle's metal implementation work?

[nihalpasham]

A couple of quick questions with regards to candle-metal-kernels. I'd like to know if my understanding as it relates to launching candle-metal-kernels is correct?

• Candle uses an abstraction named Storage, which represents the underlying data store for every Tensor. Storage can be one of three different types: CPUStorage, CudaStorage, or MetalStorage.
• Invoking an operation such as matmul will result in the corresponding Storage implementation's matmul being called. For Apple GPUs, this is MetalStorage, which sets up the necessary GPU resources (like GPU buffers and command buffers) to call the corresponding candle-metal-kernel. Again, in the case of matmul, this is a call to call_gemm(...).

Additionally,

• It appears that call_gemm uses Metal Flash Attention kernels to perform the matmul. Is this correct?
• However, the implementation of call_gemm seems only to encode (i.e., gemm) commands into the command buffer. I don't see a call to commit the command buffer (i.e., enqueue the command buffer). Where does this happen?

[nihalpasham]

@LaurentMazare, @MilkFather, @tomsanbear - was not sure if this page is actively monitored. So, I thought I'd tag contributors to the candle-metal-kernels crate on this thread. Hope this is Ok. Please let me know if you'd like me to remove you from the thread.

[tomsanbear]

Hey @nihalpasham

The logic to commit the buffer is in the metal backend source code. The backend currently will encode up to ~50 encoders into the command buffer then will commit the buffer after that. Commits will also occur any time that a shared memory access occurs (to_cpu or copy operation atm).

[nihalpasham]

Thank you, this helps!

Just to confirm - the methods below are the ones that manage commits > 50 logic.

fn new(ordinal: usize) -> Result<Self> {
        let device = metal::Device::all().swap_remove(ordinal);
        let command_queue = device.new_command_queue();
        let command_buffer = command_queue.new_command_buffer().to_owned();
        command_buffer.enqueue();
        let command_buffer = Arc::new(RwLock::new(command_buffer));
        let command_buffer_index = Arc::new(RwLock::new(0));
        let kernels = Arc::new(Kernels::new());
        let buffers = Arc::new(RwLock::new(HashMap::new()));
        let compute_per_buffer = match std::env::var("CANDLE_METAL_COMPUTE_PER_BUFFER") {
            Ok(val) => val.parse()?,
            _ => 50,
        };
        let seed = Arc::new(Mutex::new(device.new_buffer_with_data(
            [299792458].as_ptr() as *const c_void,
            4,
            MTLResourceOptions::StorageModeManaged,
        )));
        Ok(Self {
            id: DeviceId::new(),
            device,
            command_queue,
            command_buffer,
            command_buffer_index,
            compute_per_buffer,
            buffers,
            kernels,
            seed,
        })
    }

---*---*---

pub fn command_buffer(&self) -> Result<CommandBuffer> {
        let mut command_buffer_lock = self.command_buffer.try_write().map_err(MetalError::from)?;
        let mut command_buffer = command_buffer_lock.to_owned();
        let mut index = self
            .command_buffer_index
            .try_write()
            .map_err(MetalError::from)?;
        if *index > self.compute_per_buffer {
            command_buffer.commit();
            command_buffer = self.command_queue.new_command_buffer().to_owned();
            *command_buffer_lock = command_buffer.clone();
            *index = 0;

            self.drop_unused_buffers()?;
        }
        *index += 1;
        Ok(command_buffer)
    }

Commits that occur on shared memory access: this is the part I'm trying to wrap my head around. I see a to_cpu impl which I presume needs to explicitly be called when we want move data to the cpu.

However, I couldn't find the logic for triggering a commit on a copy operation, for example if we were to test a simple matmul and check its output like so. I guess an implicit copy occurs when assigning the result of the matmul (although Tensor's only derive Clone not Copy). How is this wired to the commit logic?

fn test_metal_kernel_launch() {
        let x = 2048usize;
        let y = 512usize;
        let device = Device::new_metal(0).unwrap();
        let a = Tensor::randn(0f32, 1., (x, y), &device).unwrap();
        let b = Tensor::randn(0f32, 1., (y, x), &device).unwrap();
        let result = a.matmul(&b).unwrap();
        println!("result: {}", result);   
    }

[tomsanbear]

Something that's worth noting in the implementation, is that the backend is "lazy" in that unless you are moving the data to the CPU, the operations are not committed. There exists a method on the device trait called "synchronize" which will force a commit if one needs to trigger it manually.

Under the hood this "synchronize" functionality is called for various operations, but the majority of those are ones that are either moving data from or to the cpu, so under the hood they are making calls to commit the buffer and then a call to "wait until completed" (wait_until_completed in the rust case on metal::Device).

[nihalpasham]

I figured out what I was missing.

Something that's worth noting in the implementation, is that the backend is "lazy" in that unless you are moving the data to the CPU, the operations are not committed.

Yep, that makes sense. Commands are encoded in a command buffer but only committed when one of two conditions is met:

• The commit threshold is exceeded (more than 50 commands) or
• An operation indirectly triggers a commit, such as printing the actual data of a tensor.

In my previous example with a simple matmul operation, I kept adding breakpoints to find the indirect manual trigger that commits the corresponding command buffer. However, I overlooked inspecting the print statement itself. Tensor's implement the Display trait, whose implementation ultimately ends up calling the to_cpu method for the underlying Storage type. In other words, the implementation is indeed lazy.

Curious question: I came across the synchronize method and tried to use VSCode to find all references in the project, but found none. Is this method intended for internal use only?

[tomsanbear]

Not sure why it didn't come up in references, but we use for benchmarks since we don't perform a to_cpu access on that data. Overall though it's just an escape hatch in the case that something isn't triggering cpu memory transfer, I can't see a real use for it for consumers though which is why you won't see it in examples 🤷