Allow querying the allocator for the buffer size (#1404)

mlx/backend/metal/allocator.cpp

@@ -241,6 +241,10 @@ void MetalAllocator::free(Buffer buffer) {
   }
 }
 
+size_t MetalAllocator::size(Buffer buffer) const {
+  return static_cast<MTL::Buffer*>(buffer.ptr())->length();
+}
+
 MetalAllocator& allocator() {
   // By creating the |allocator_| on heap, the destructor of MetalAllocator will
   // not be called on exit and all the buffers will be leaked. This is necessary

mlx/backend/metal/allocator.h

@@ -56,6 +56,7 @@ class MetalAllocator : public allocator::Allocator {
  public:
   virtual Buffer malloc(size_t size, bool allow_swap = false) override;
   virtual void free(Buffer buffer) override;
+  virtual size_t size(Buffer buffer) const override;
   size_t get_active_memory() {
     return active_memory_;
   };