Skip to content

Commit 630033a

Browse files
EricLBuehlerEricLBuehler
committed
tmp
1 parent 2d197d5 commit 630033a

6 files changed

Lines changed: 189 additions & 161 deletions

File tree

‎candle-core/src/metal_backend/device.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use std::collections::HashMap;
55
use std::path::Path;
66
use std::sync::{Arc, Mutex, RwLock};
77

8-
use super::{current_pool, register_pool_allocation, MetalError};
8+
use super::{buffer_offset_for_output, current_pool, register_pool_allocation, MetalError};
99

1010
// iOS and macOS have different storage modes for shared buffers.
1111
// due to the GPU/CPU management differences.
@@ -295,6 +295,10 @@ impl MetalDevice {
295295
Ok(new_buffer)
296296
}
297297

298+
pub fn buffer_offset<'a>(&self, buffer: &'a Arc<Buffer>) -> super::BufferOffset<'a> {
299+
buffer_offset_for_output(buffer)
300+
}
301+
298302
pub fn allocate_zeros(&self, size_in_bytes: usize) -> Result<Arc<Buffer>> {
299303
let buffer = self.allocate_buffer(
300304
size_in_bytes as NSUInteger,

‎candle-core/src/metal_backend/mod.rs‎

Lines changed: 55 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,24 @@ fn take_pool_allocation(buffer: &Arc<Buffer>) -> Option<Arc<MetalPoolAllocation>
6464
})
6565
}
6666

67+
fn peek_pool_allocation(buffer: &Arc<Buffer>) -> Option<Arc<MetalPoolAllocation>> {
68+
let ptr = Arc::as_ptr(buffer) as usize;
69+
pool_registry()
70+
.lock()
71+
.ok()
72+
.and_then(|registry| registry.get(&ptr).and_then(|allocs| allocs.last().cloned()))
73+
}
74+
75+
fn buffer_offset_for_output(buffer: &Arc<Buffer>) -> BufferOffset<'_> {
76+
let offset = peek_pool_allocation(buffer)
77+
.map(|alloc| alloc.offset())
78+
.unwrap_or(0);
79+
BufferOffset {
80+
buffer,
81+
offset_in_bytes: offset,
82+
}
83+
}
84+
6785
pub fn buffer_o<'a>(storage: &'a MetalStorage, l: &Layout) -> BufferOffset<'a> {
6886
storage.buffer_offset(l)
6987
}
@@ -189,7 +207,7 @@ impl BackendStorage for MetalStorage {
189207
name,
190208
el,
191209
src,
192-
&buffer,
210+
buffer_offset_for_output(&buffer),
193211
mul as f32,
194212
add as f32,
195213
)
@@ -209,7 +227,7 @@ impl BackendStorage for MetalStorage {
209227
layout.dims(),
210228
src,
211229
layout.stride(),
212-
&buffer,
230+
buffer_offset_for_output(&buffer),
213231
mul as f32,
214232
add as f32,
215233
)
@@ -243,7 +261,7 @@ impl BackendStorage for MetalStorage {
243261
name,
244262
el,
245263
src,
246-
&buffer,
264+
buffer_offset_for_output(&buffer),
247265
pow as f32,
248266
)
249267
.map_err(MetalError::from)?;
@@ -262,7 +280,7 @@ impl BackendStorage for MetalStorage {
262280
layout.dims(),
263281
src,
264282
layout.stride(),
265-
&buffer,
283+
buffer_offset_for_output(&buffer),
266284
pow as f32,
267285
)
268286
.map_err(MetalError::from)?;
@@ -295,7 +313,7 @@ impl BackendStorage for MetalStorage {
295313
name,
296314
el,
297315
src,
298-
&buffer,
316+
buffer_offset_for_output(&buffer),
299317
alpha as f32,
300318
)
301319
.map_err(MetalError::from)?;
@@ -314,7 +332,7 @@ impl BackendStorage for MetalStorage {
314332
layout.dims(),
315333
src,
316334
layout.stride(),
317-
&buffer,
335+
buffer_offset_for_output(&buffer),
318336
alpha as f32,
319337
)
320338
.map_err(MetalError::from)?;
@@ -398,7 +416,7 @@ impl BackendStorage for MetalStorage {
398416
src_dims,
399417
dst_el,
400418
src,
401-
&buffer,
419+
buffer_offset_for_output(&buffer),
402420
)
403421
.map_err(MetalError::from)?;
404422

@@ -454,7 +472,7 @@ impl BackendStorage for MetalStorage {
454472
&stride,
455473
dst_el,
456474
src,
457-
&buffer,
475+
buffer_offset_for_output(&buffer),
458476
)
459477
.map_err(MetalError::from)?;
460478

@@ -638,7 +656,7 @@ impl BackendStorage for MetalStorage {
638656
kernel_name,
639657
el_count,
640658
src,
641-
&buffer,
659+
buffer_offset_for_output(&buffer),
642660
)
643661
.map_err(MetalError::from)?;
644662
} else {
@@ -691,7 +709,7 @@ impl BackendStorage for MetalStorage {
691709
layout.dims(),
692710
src,
693711
layout.stride(),
694-
&buffer,
712+
buffer_offset_for_output(&buffer),
695713
)
696714
.map_err(MetalError::from)?;
697715
}
@@ -785,7 +803,7 @@ impl BackendStorage for MetalStorage {
785803
kernel_name,
786804
el_count,
787805
src,
788-
&buffer,
806+
buffer_offset_for_output(&buffer),
789807
)
790808
.map_err(MetalError::from)?;
791809
}
@@ -861,7 +879,7 @@ impl BackendStorage for MetalStorage {
861879
kernel_name,
862880
el_count,
863881
src,
864-
&buffer,
882+
buffer_offset_for_output(&buffer),
865883
)
866884
.map_err(MetalError::from)?;
867885
}
@@ -926,7 +944,7 @@ impl BackendStorage for MetalStorage {
926944
crate::bail!("Metal strided unary {name} {dtype:?} not implemented")
927945
}
928946
};
929-
let dst = BufferOffset::zero_offset(&buffer);
947+
let dst = buffer_offset_for_output(&buffer);
930948
candle_metal_kernels::call_unary_strided(
931949
&device.device,
932950
&command_buffer,
@@ -1001,7 +1019,7 @@ impl BackendStorage for MetalStorage {
10011019
t_l.stride(),
10021020
f,
10031021
f_l.stride(),
1004-
&buffer,
1022+
buffer_offset_for_output(&buffer),
10051023
)
10061024
.map_err(MetalError::from)?;
10071025
Ok(Self::new(buffer, device, el, dtype))
@@ -1044,7 +1062,7 @@ impl BackendStorage for MetalStorage {
10441062
strides,
10451063
(k_size, stride, padding, dilation),
10461064
src,
1047-
&dst,
1065+
buffer_offset_for_output(&dst),
10481066
)
10491067
.map_err(MetalError::from)?;
10501068
let col = Self::new(dst, device, dst_el, self.dtype);
@@ -1139,8 +1157,8 @@ impl BackendStorage for MetalStorage {
11391157
&[b_size, l_in, c_out, k_size],
11401158
params.k_size,
11411159
params.stride,
1142-
BufferOffset::zero_offset(&col.buffer),
1143-
&buffer,
1160+
col.base_offset(),
1161+
buffer_offset_for_output(&buffer),
11441162
)
11451163
.map_err(MetalError::from)?;
11461164
buffer
@@ -1178,7 +1196,7 @@ impl BackendStorage for MetalStorage {
11781196
layout.start_offset() * self.dtype.size_in_bytes(),
11791197
&k.buffer,
11801198
k_layout.start_offset() * k.dtype.size_in_bytes(),
1181-
&buffer,
1199+
buffer_offset_for_output(&buffer),
11821200
)
11831201
.map_err(MetalError::from)?;
11841202
buffer
@@ -1231,7 +1249,7 @@ impl BackendStorage for MetalStorage {
12311249
layout.stride(),
12321250
(h_k, w_k, stride, padding, dilation),
12331251
src,
1234-
&dst,
1252+
buffer_offset_for_output(&dst),
12351253
)
12361254
.map_err(MetalError::from)?;
12371255
let col = Self::new(dst, device, dst_el, self.dtype);
@@ -1323,7 +1341,7 @@ impl BackendStorage for MetalStorage {
13231341
},
13241342
&self.buffer,
13251343
&kernel.buffer,
1326-
&buffer,
1344+
buffer_offset_for_output(&buffer),
13271345
)
13281346
.map_err(MetalError::from)?;
13291347
Ok(Self::new(buffer, self.device.clone(), dst_el, self.dtype))
@@ -1366,7 +1384,7 @@ impl BackendStorage for MetalStorage {
13661384
w_stride,
13671385
h_stride,
13681386
&self.buffer,
1369-
&buffer,
1387+
buffer_offset_for_output(&buffer),
13701388
)
13711389
.map_err(MetalError::from)?;
13721390
Ok(Self::new(buffer, self.device.clone(), dst_el, self.dtype))
@@ -1409,7 +1427,7 @@ impl BackendStorage for MetalStorage {
14091427
w_stride,
14101428
h_stride,
14111429
&self.buffer,
1412-
&buffer,
1430+
buffer_offset_for_output(&buffer),
14131431
)
14141432
.map_err(MetalError::from)?;
14151433
Ok(Self::new(buffer, self.device.clone(), dst_el, self.dtype))
@@ -1453,7 +1471,7 @@ impl BackendStorage for MetalStorage {
14531471
out_w,
14541472
out_h,
14551473
src,
1456-
&buffer,
1474+
buffer_offset_for_output(&buffer),
14571475
)
14581476
.map_err(MetalError::from)?;
14591477
Ok(Self::new(buffer, self.device.clone(), dst_el, self.dtype))
@@ -1495,7 +1513,7 @@ impl BackendStorage for MetalStorage {
14951513
dim,
14961514
src,
14971515
ids,
1498-
&buffer,
1516+
buffer_offset_for_output(&buffer),
14991517
)
15001518
.map_err(MetalError::from)?;
15011519
Ok(Self::new(buffer, device.clone(), dst_el, dtype))
@@ -1653,7 +1671,7 @@ impl BackendStorage for MetalStorage {
16531671
src_l.stride(),
16541672
src,
16551673
ids,
1656-
&buffer,
1674+
buffer_offset_for_output(&buffer),
16571675
)
16581676
.map_err(MetalError::from)?;
16591677
Ok(Self::new(buffer, device.clone(), dst_el, dtype))
@@ -1715,7 +1733,7 @@ impl BackendStorage for MetalStorage {
17151733
dim,
17161734
src,
17171735
ids,
1718-
&acc.buffer,
1736+
buffer_o(&acc, l),
17191737
)
17201738
.map_err(MetalError::from)?;
17211739
Ok(acc)
@@ -1754,7 +1772,7 @@ impl BackendStorage for MetalStorage {
17541772
rhs_l.stride(),
17551773
rhs_l.start_offset() * rhs.dtype.size_in_bytes(),
17561774
&rhs.buffer,
1757-
&buffer,
1775+
buffer_offset_for_output(&buffer),
17581776
)
17591777
.map_err(MetalError::from)?;
17601778

@@ -1900,13 +1918,19 @@ impl MetalStorage {
19001918
.map(|allocation| allocation.pool())
19011919
}
19021920

1903-
fn buffer_offset<'a>(&'a self, layout: &Layout) -> BufferOffset<'a> {
1921+
fn base_offset(&self) -> BufferOffset<'_> {
19041922
BufferOffset {
19051923
buffer: &self.buffer,
1906-
offset_in_bytes: self.offset_bytes + layout.start_offset() * self.dtype.size_in_bytes(),
1924+
offset_in_bytes: self.offset_bytes,
19071925
}
19081926
}
19091927

1928+
fn buffer_offset<'a>(&'a self, layout: &Layout) -> BufferOffset<'a> {
1929+
let mut base = self.base_offset();
1930+
base.offset_in_bytes += layout.start_offset() * self.dtype.size_in_bytes();
1931+
base
1932+
}
1933+
19101934
fn determine_pool(storages: &[&Self]) -> Result<Option<Arc<MetalTensorPool>>> {
19111935
let mut pool: Option<Arc<MetalTensorPool>> = None;
19121936
for storage in storages {
@@ -2029,7 +2053,7 @@ impl MetalStorage {
20292053
el_count,
20302054
lhs,
20312055
rhs,
2032-
&buffer,
2056+
buffer_offset_for_output(&buffer),
20332057
)
20342058
.map_err(MetalError::from)?;
20352059
(buffer, dtype)
@@ -2130,7 +2154,7 @@ impl MetalStorage {
21302154
lhs_l.stride(),
21312155
rhs,
21322156
rhs_l.stride(),
2133-
&buffer,
2157+
buffer_offset_for_output(&buffer),
21342158
)
21352159
.map_err(MetalError::from)?;
21362160
(buffer, dtype)

0 commit comments

Comments
 (0)