def cuda(
self,
device: Optional[Union[torch_device, str,
int]] = None) -> 'BaseTile':
"""Return a copy of this tile in CUDA memory.
Args:
device: CUDA device
Returns:
Self with the underlying C++ tile moved to CUDA memory.
Raises:
CudaError: if the library has not been compiled with CUDA.
"""
if not cuda.is_compiled():
raise CudaError('aihwkit has not been compiled with CUDA support')
device = torch_device('cuda', cuda_device(device).idx)
if self.is_cuda and device != self.device:
raise CudaError(
'Cannot switch CUDA devices of existing Cuda tiles')
if isinstance(self.tile, tiles.AnalogTile):
with cuda_device(device):
self.tile = tiles.CudaAnalogTile(self.tile)
self.is_cuda = True
self.device = device
self.analog_ctx.cuda(device)
return self
def __init__(self, source_tile: AnalogTile):
if not cuda.is_compiled():
raise CudaError('aihwkit has not been compiled with CUDA support')
# Create a new instance of the rpu config.
new_rpu_config = deepcopy(source_tile.rpu_config)
# Create the tile, replacing the simulator tile.
super().__init__(source_tile.out_size, source_tile.in_size, new_rpu_config,
source_tile.bias, source_tile.in_trans, source_tile.out_trans)
self.tile = tiles.CudaAnalogTile(source_tile.tile)
# Set the cuda properties
self.stream = current_stream()
self.device = torch_device(current_device())
def __init__(self,
out_size: int,
in_size: int,
resistive_device: Optional[BaseResistiveDevice] = None,
bias: bool = False,
in_trans: bool = False,
out_trans: bool = False):
if not cuda.is_compiled():
raise RuntimeError(
'aihwkit has not been compiled with CUDA support')
super().__init__(out_size, in_size, resistive_device, bias, in_trans,
out_trans)
self.tile = tiles.CudaAnalogTile(self.tile)
self.stream = current_stream()
self.device = torch_device(current_device())
