1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
//! Device Management and Transfer Operations
//!
//! This module handles tensor device placement, transfers between devices
//! (CPU/GPU), and device-specific operations. It provides efficient
//! device-to-device data transfer capabilities.
use super::core::{Tensor, TensorStorage};
use crate::{Device, Result};
// Impl block for methods that need Clone
impl<T: Clone> Tensor<T> {
/// Transfer tensor to specified device
pub fn to(&self, device: Device) -> Result<Self>
where
T: Default + bytemuck::Pod + bytemuck::Zeroable + Send + Sync + 'static,
{
if self.device() == &device {
return Ok(self.clone());
}
match (&self.storage, &device) {
(TensorStorage::Cpu(_array), Device::Cpu) => Ok(self.clone()),
#[cfg(feature = "gpu")]
(TensorStorage::Cpu(array), Device::Gpu(id)) => {
let gpu_buffer = crate::gpu::buffer::GpuBuffer::from_cpu_array(array, *id)?;
Ok(Self {
storage: TensorStorage::Gpu(gpu_buffer),
shape: self.shape().clone(),
device,
requires_grad: self.requires_grad(),
grad: None,
})
}
#[cfg(feature = "gpu")]
(TensorStorage::Gpu(buffer), Device::Cpu) => {
let array = buffer.to_cpu_array()?;
Ok(Self {
storage: TensorStorage::Cpu(array),
shape: self.shape().clone(),
device,
requires_grad: self.requires_grad(),
grad: None,
})
}
// Any remaining combination (e.g. Gpu->Gpu across distinct devices,
// or any ROCm source/target) is a real, constructible input that
// `to()` does not implement. Return an honest error instead of
// panicking via `unreachable!()`, which would crash multi-GPU and
// ROCm builds. Use `transfer_to_device` for the richer transfer path.
#[allow(unreachable_patterns)]
_ => Err(crate::TensorError::unsupported_operation_simple(format!(
"device transfer {:?}->{:?} not supported by Tensor::to; \
use Tensor::to_device for cross-GPU/ROCm transfers",
self.device(),
device
))),
}
}
/// Transfer tensor to a different device
pub fn to_device(&self, target_device: Device) -> Result<Self>
where
T: Clone + Default + Send + Sync + 'static + bytemuck::Pod,
{
if self.device() == &target_device {
return Ok(self.clone());
}
self.transfer_to_device(target_device)
}
/// Move tensor to CPU
pub fn to_cpu(&self) -> Result<Self>
where
T: Clone + Default + Send + Sync + 'static + bytemuck::Pod,
{
self.to_device(Device::Cpu)
}
/// Get GPU context information from this tensor (if it's on GPU)
#[cfg(feature = "gpu")]
pub fn gpu_context_info(&self) -> Option<crate::device::context::GpuContextInfo> {
match &self.storage {
TensorStorage::Gpu(buffer) => Some(crate::device::context::GpuContextInfo {
device: buffer.device.clone(),
queue: buffer.queue.clone(),
}),
_ => None,
}
}
/// Move tensor to GPU with specified ID
#[cfg(feature = "gpu")]
pub fn to_gpu(&self, gpu_id: usize) -> Result<Self>
where
T: Clone + Default + Send + Sync + 'static + bytemuck::Pod,
{
self.to_device(Device::Gpu(gpu_id))
}
/// Internal device transfer implementation
fn transfer_to_device(&self, target_device: Device) -> Result<Self>
where
T: Clone + Default + Send + Sync + 'static + bytemuck::Pod,
{
use crate::device::context::DEVICE_MANAGER;
let _src_ctx = DEVICE_MANAGER.get_context(self.device())?;
let _dst_ctx = DEVICE_MANAGER.get_context(&target_device)?;
match (&self.storage, &target_device) {
// CPU to GPU transfer
#[cfg(feature = "gpu")]
(TensorStorage::Cpu(cpu_array), Device::Gpu(_)) => {
#[cfg(feature = "gpu")]
{
let slice = cpu_array.as_slice().ok_or_else(|| {
crate::TensorError::invalid_argument(
"Cannot convert CPU array to slice".to_string(),
)
})?;
let gpu_buffer =
crate::gpu::buffer::GpuBuffer::from_slice(slice, &target_device)?;
Ok(Self {
storage: TensorStorage::Gpu(gpu_buffer),
shape: self.shape().clone(),
device: target_device,
requires_grad: self.requires_grad(),
grad: None, // Gradients are reset on device transfer
})
}
#[cfg(not(feature = "gpu"))]
{
Err(crate::TensorError::device_error_simple(
"GPU support not compiled",
))
}
}
// GPU to CPU transfer
#[cfg(feature = "gpu")]
(TensorStorage::Gpu(gpu_buffer), Device::Cpu) => {
let cpu_data = gpu_buffer.to_cpu()?;
let array = scirs2_core::ndarray::ArrayD::from_shape_vec(
scirs2_core::ndarray::IxDyn(self.shape().dims()),
cpu_data,
)
.map_err(|e| crate::TensorError::invalid_shape_simple(e.to_string()))?;
Ok(Self {
storage: TensorStorage::Cpu(array),
shape: self.shape().clone(),
device: target_device,
requires_grad: self.requires_grad(),
grad: None, // Gradients are reset on device transfer
})
}
// GPU to GPU transfer (device-to-device)
#[cfg(feature = "gpu")]
(TensorStorage::Gpu(src_buffer), Device::Gpu(_)) => {
let dst_buffer = src_buffer.transfer_to_device(&target_device)?;
Ok(Self {
storage: TensorStorage::Gpu(dst_buffer),
shape: self.shape().clone(),
device: target_device,
requires_grad: self.requires_grad(),
grad: None, // Gradients are reset on device transfer
})
}
// CPU to CPU (should not happen due to early return)
(TensorStorage::Cpu(_), Device::Cpu) => Ok(self.clone()),
// ROCm patterns - Use CPU fallback for now
#[cfg(feature = "rocm")]
(TensorStorage::Cpu(_), Device::Rocm(_)) => {
// ROCm tensor transfer: Fallback to CPU for now
// Future: Implement native ROCm memory transfer kernels
eprintln!("Warning: ROCm tensor transfer using CPU fallback - native implementation pending");
Ok(self.clone()) // Keep on CPU until ROCm support is complete
}
#[cfg(feature = "rocm")]
(TensorStorage::Gpu(_), Device::Rocm(_)) => {
// GPU to ROCm transfer: Go through CPU as intermediate step
// Future: Implement direct GPU<->ROCm memory transfer
eprintln!("Warning: GPU to ROCm transfer using CPU fallback - native implementation pending");
self.to_cpu() // Transfer to CPU for now
}
}
}
/// Copy tensor data from another device (for collective operations)
pub fn copy_from_device(&mut self, src: &Self) -> Result<()>
where
T: Clone + Default + Send + Sync + 'static + bytemuck::Pod,
{
if self.shape() != src.shape() {
return Err(crate::TensorError::ShapeMismatch {
operation: "copy_from_device".to_string(),
expected: self.shape().to_string(),
got: src.shape().to_string(),
context: None,
});
}
let transferred = src.transfer_to_device(*self.device())?;
self.storage = transferred.storage;
Ok(())
}
/// Check if tensor can be transferred to target device
pub fn can_transfer_to(&self, target_device: Device) -> bool
where
T: bytemuck::Pod,
{
match (self.device(), &target_device) {
(Device::Cpu, Device::Cpu) => true,
#[cfg(feature = "gpu")]
(Device::Cpu, Device::Gpu(_)) => true,
#[cfg(feature = "gpu")]
(Device::Gpu(_), Device::Cpu) => true,
#[cfg(feature = "gpu")]
(Device::Gpu(_), Device::Gpu(_)) => true,
#[cfg(feature = "rocm")]
(Device::Rocm(_), _) => true,
#[cfg(feature = "rocm")]
(_, Device::Rocm(_)) => true,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
/// A ROCm transfer target is a real, constructible input that `Tensor::to`
/// does not implement. It must return an honest error instead of panicking
/// through the old `unreachable!()`.
#[cfg(feature = "rocm")]
#[test]
fn cpu_to_rocm_transfer_returns_honest_error_not_panic() {
let tensor = Tensor::<f32>::zeros(&[2, 2]);
let result = tensor.to(Device::Rocm(0));
assert!(
result.is_err(),
"CPU->ROCm transfer via Tensor::to must return an honest error, not succeed"
);
match result {
Err(crate::TensorError::UnsupportedOperation { reason, .. }) => {
assert!(
reason.contains("not supported"),
"error must explain the transfer is unsupported, got: {reason}"
);
}
other => panic!("expected UnsupportedOperation error, got: {other:?}"),
}
}
/// Transferring to the same device is a no-op clone (handled by the early
/// return), and must never reach the error branch.
#[test]
fn cpu_to_cpu_same_device_is_noop_clone() {
let tensor = Tensor::<f32>::zeros(&[3, 3]);
let result = tensor.to(Device::Cpu);
assert!(
result.is_ok(),
"same-device CPU->CPU transfer must succeed as a clone"
);
}
/// `Tensor::to` must not panic for any constructible target device; with
/// only the CPU storage available, a GPU/ROCm target yields a `Result`
/// (Ok for the implemented CPU->GPU path, Err otherwise) — never a panic.
#[cfg(feature = "rocm")]
#[test]
fn rocm_target_does_not_panic() {
let tensor = Tensor::<f32>::zeros(&[1, 4]);
// The call itself must complete without panicking. We only assert it
// returns a Result; the ROCm CPU->ROCm path is unsupported and errs.
let _result: Result<_> = tensor.to(Device::Rocm(1));
}
}