From a732ffe0dfe3789bc6c99be16dfc9ff4f022c71b Mon Sep 17 00:00:00 2001 From: Zhaoxian-Wu Date: Wed, 12 Mar 2025 14:36:27 +0000 Subject: [PATCH 1/6] feat(ctx): define attribution for AnalogContext Signed-off-by: Zhaoxian Wu --- src/aihwkit/optim/context.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/src/aihwkit/optim/context.py b/src/aihwkit/optim/context.py index d9228d033..a8dbc152e 100644 --- a/src/aihwkit/optim/context.py +++ b/src/aihwkit/optim/context.py @@ -19,7 +19,25 @@ class AnalogContext(Parameter): - """Context for analog optimizer.""" + """Context for analog optimizer. + + Remark: `data` attribution, inherited from `torch.nn.Parameter`, is a tensor of training parameter + If `analog_bias` (which is provided by `analog_tile`) is False, + `data` has the same meaning as `torch.nn.Parameter` + If `analog_bias` (which is provided by `analog_tile`) is True, + The last column of `data` is the `bias` term + + Even though it allows us to access the weights directly, always keep in mind that it is used + only for studying propuses. To simulate the real reading, call the `read_weights` method + instead, i.e. given `analog_ctx: AnalogContext`, + estimated_weights, estimated_bias = analog_ctx.analog_tile.read_weights() + + Similarly, even though this feature allows us to update the weights directly, always keep in mind + that the real RPU devices change their weights only by "pulse update" method. + Therefore, use the following update methods instead of writing `data` directly in the analog optimizer: + analog_ctx.analog_tile.update(...) + analog_ctx.analog_tile.update_indexed(...) + """ def __new__( cls: Type["AnalogContext"], @@ -47,6 +65,12 @@ def __init__( self.analog_grad_output = [] # type: list self.reset(analog_tile) + # analog_tile.tile is the class from C++ code: + # aihwkit.silulator.rpu_base.devices.AnalogTile + # It stores the "weight" matrix; + # If analog_tile.analog_bias is True, it also stores the "bias" matrix + self.data = analog_tile.tile.get_weights() + def set_indexed(self, value: bool = True) -> None: """Set the context to forward_indexed.""" self.use_indexed = value From 3503ad50a924ebd4e21e8d66f936d26d62ee5b63 Mon Sep 17 00:00:00 2001 From: Zhaoxian-Wu Date: Tue, 18 Mar 2025 18:41:45 +0000 Subject: [PATCH 2/6] feat(ctx): add weight data management and checkpoint compatibility for AnalogContext - weight data management - support checkpoint loading for old version toolkit - sync analog_ctx and tile when recreating Signed-off-by: Zhaoxian Wu --- src/aihwkit/nn/modules/base.py | 26 +++++++++++ src/aihwkit/nn/modules/linear.py | 4 ++ src/aihwkit/optim/context.py | 41 +++++++++-------- src/aihwkit/simulator/tiles/array.py | 4 +- src/aihwkit/simulator/tiles/base.py | 56 +++++++++++++++++------ src/aihwkit/simulator/tiles/custom.py | 3 +- src/aihwkit/simulator/tiles/functions.py | 2 + src/aihwkit/simulator/tiles/periphery.py | 2 +- src/aihwkit/simulator/tiles/rpucuda.py | 2 - src/aihwkit/simulator/tiles/torch_tile.py | 4 +- tests/test_layers_linear.py | 4 ++ tests/test_utils.py | 14 +++++- 12 files changed, 120 insertions(+), 42 deletions(-) diff --git a/src/aihwkit/nn/modules/base.py b/src/aihwkit/nn/modules/base.py index 1424061d8..3f346b341 100644 --- a/src/aihwkit/nn/modules/base.py +++ b/src/aihwkit/nn/modules/base.py @@ -7,10 +7,12 @@ """Base class for adding functionality to analog layers.""" from typing import Any, List, Optional, Tuple, NamedTuple, Union, Generator, Callable, TYPE_CHECKING from collections import OrderedDict +import warnings from torch import Tensor from torch.nn import Parameter from torch import device as torch_device +from torch import Size as torch_size from aihwkit.exceptions import ModuleError from aihwkit.simulator.tiles.module import TileModule @@ -244,6 +246,30 @@ class type when setting ``load_rpu_config`` to ``load_rpu_config=False``. """ + keys_to_delete = [] + for name, param in list(state_dict.items()): + if name.endswith("analog_ctx") and param.size() == torch_size([]): + keys_to_delete.append(name) + + # For the checkpoint saved by aihwkit before version, 0.9.2, the parameters with + # class `AnalogContext` are saved as empty size tensors since `AnalogContext`, + # which is a derived class of `torch.nn.Parameter`, + # uses an empty Parameter tensor to store the context. + if len(keys_to_delete) > 0: + strict = False + for key in keys_to_delete: + del state_dict[key] + warnings.warn( + "Some parameters in the loaded checkpoint has empty size" + "(param.size() == torch.Size([]))." + "It could happens because of the loaded checkpoint" + "is generated by an older version of aihwkit." + "The parameter is skipped for compatibility reasons." + "The loading mode is set to non-strict." + "It is recommended to re-save the checkpoint with the latest version of aihwkit." + "Related parameters are: {}".format(keys_to_delete) + ) + for analog_tile in self.analog_tiles(): analog_tile.set_load_rpu_config_state(load_rpu_config, strict_rpu_config_check) return super().load_state_dict(state_dict, strict) # type: ignore diff --git a/src/aihwkit/nn/modules/linear.py b/src/aihwkit/nn/modules/linear.py index 9db3bee2e..8a5debfee 100644 --- a/src/aihwkit/nn/modules/linear.py +++ b/src/aihwkit/nn/modules/linear.py @@ -82,6 +82,10 @@ def reset_parameters(self) -> None: self.weight, self.bias = self.get_weights() # type: ignore super().reset_parameters() self.set_weights(self.weight, self.bias) # type: ignore + # AnalogLinear doesn't support access weight and bias directly, so delete them + del self.weight, self.bias + # delete them manually is necessary since asigning `bias` (a bool) is forbidden + # by torch if self.bias is already a tensor self.weight, self.bias = None, bias def forward(self, x_input: Tensor) -> Tensor: diff --git a/src/aihwkit/optim/context.py b/src/aihwkit/optim/context.py index a8dbc152e..b5c5c33d6 100644 --- a/src/aihwkit/optim/context.py +++ b/src/aihwkit/optim/context.py @@ -10,7 +10,7 @@ from typing import Optional, Type, Union, Any, TYPE_CHECKING -from torch import ones, dtype, Tensor, no_grad +from torch import dtype, Tensor, no_grad from torch.nn import Parameter from torch import device as torch_device @@ -20,23 +20,28 @@ class AnalogContext(Parameter): """Context for analog optimizer. - - Remark: `data` attribution, inherited from `torch.nn.Parameter`, is a tensor of training parameter + + Note: `data` attribution, inherited from `torch.nn.Parameter`, is a tensor of training parameter If `analog_bias` (which is provided by `analog_tile`) is False, `data` has the same meaning as `torch.nn.Parameter` If `analog_bias` (which is provided by `analog_tile`) is True, The last column of `data` is the `bias` term - + Even though it allows us to access the weights directly, always keep in mind that it is used only for studying propuses. To simulate the real reading, call the `read_weights` method instead, i.e. given `analog_ctx: AnalogContext`, estimated_weights, estimated_bias = analog_ctx.analog_tile.read_weights() - - Similarly, even though this feature allows us to update the weights directly, always keep in mind - that the real RPU devices change their weights only by "pulse update" method. - Therefore, use the following update methods instead of writing `data` directly in the analog optimizer: - analog_ctx.analog_tile.update(...) - analog_ctx.analog_tile.update_indexed(...) + + Similarly, even though this feature allows us to update the weights directly, + always keep in mind that the real RPU devices change their weights only + by "pulse update" method. + + Therefore, use the following update methods instead of + writing `data` directly in the analog optimizer: + --- + analog_ctx.analog_tile.update(...) + analog_ctx.analog_tile.update_indexed(...) + --- """ def __new__( @@ -48,9 +53,15 @@ def __new__( if parameter is None: return Parameter.__new__( cls, - data=ones((), device=analog_tile.device, dtype=analog_tile.get_dtype()), + data=analog_tile.tile.get_weights(), requires_grad=True, ) + # analog_tile.tile can comes from different classes: + # aihwkit.silulator.rpu_base.devices.AnalogTile (C++) + # TorchInferenceTile (Python) + # It stores the "weight" matrix; + # If analog_tile.analog_bias is True, it also stores the "bias" matrix + parameter.__class__ = cls return parameter @@ -65,12 +76,6 @@ def __init__( self.analog_grad_output = [] # type: list self.reset(analog_tile) - # analog_tile.tile is the class from C++ code: - # aihwkit.silulator.rpu_base.devices.AnalogTile - # It stores the "weight" matrix; - # If analog_tile.analog_bias is True, it also stores the "bias" matrix - self.data = analog_tile.tile.get_weights() - def set_indexed(self, value: bool = True) -> None: """Set the context to forward_indexed.""" self.use_indexed = value @@ -116,8 +121,8 @@ def cuda(self, device: Optional[Union[torch_device, str, int]] = None) -> "Analo Returns: This context in the specified device. """ - self.data = self.data.cuda(device) # type: Tensor if not self.analog_tile.is_cuda: + self.data = self.analog_tile.tile.get_weights() # type: Tensor self.analog_tile = self.analog_tile.cuda(device) self.reset(self.analog_tile) return self diff --git a/src/aihwkit/simulator/tiles/array.py b/src/aihwkit/simulator/tiles/array.py index 7d1ddb9cd..a8af34544 100644 --- a/src/aihwkit/simulator/tiles/array.py +++ b/src/aihwkit/simulator/tiles/array.py @@ -124,7 +124,7 @@ def set_weights(self, weight: Tensor, bias: Optional[Tensor] = None, **kwargs: A in_start = in_end if self.bias is not None and bias is not None: - self.bias.data = bias.detach().to(self.bias.device) + self.bias.data.copy_(bias) @no_grad() def get_weights(self, **kwargs: Any) -> Tuple[Tensor, Optional[Tensor]]: @@ -149,7 +149,7 @@ def get_weights(self, **kwargs: Any) -> Tuple[Tensor, Optional[Tensor]]: weight = cat(weight_lst, 1) if self.bias is not None: - return weight, self.bias.clone().cpu() + return weight, self.bias return weight, None def forward(self, x_input: Tensor, tensor_view: Optional[Tuple] = None) -> Tensor: diff --git a/src/aihwkit/simulator/tiles/base.py b/src/aihwkit/simulator/tiles/base.py index dac1a1855..636b78a0c 100644 --- a/src/aihwkit/simulator/tiles/base.py +++ b/src/aihwkit/simulator/tiles/base.py @@ -264,6 +264,7 @@ def get_meta_parameters(self) -> Any: raise NotImplementedError +# pylint: disable=too-many-public-methods class SimulatorTileWrapper: """Wrapper base class for defining the necessary tile functionality. @@ -281,6 +282,18 @@ class SimulatorTileWrapper: should be used. handle_output_bound: whether the bound clamp gradient should be inserted ignore_analog_state: whether to ignore the analog state when __getstate__ is called + + Attributes: + tile: A simulator tile object that handles the computations + for the given input/output sizes. + It is created by `self._create_simulator_tile` method, + which is provided by the derived class. + E.g., `aihwkit.simulator.tiles.analog.AnalogTile` and + `aihwkit.simulator.tiles.inference_torch.TorchInferenceTile` + implement this method. + The weight data is stored in the tile object. + analog_ctx: `AnalogContext`, which wraps the weight in tile + into a `torch.nn.Parameter` object. """ def __init__( @@ -295,8 +308,6 @@ def __init__( handle_output_bound: bool = False, ignore_analog_state: bool = False, ): - self.is_cuda = False - self.device = torch_device("cpu") self.out_size = out_size self.in_size = in_size self.rpu_config = deepcopy(rpu_config) @@ -324,6 +335,16 @@ def __init__( self.analog_ctx = AnalogContext(self) self.analog_ctx.use_torch_update = torch_update + @property + def device(self) -> torch_device: + """Return the device of the tile.""" + return self.analog_ctx.device + + @property + def is_cuda(self) -> bool: + """Return the is_cuda state of the tile.""" + return self.analog_ctx.is_cuda + def get_runtime(self) -> RuntimeParameter: """Returns the runtime parameter.""" if not hasattr(self.rpu_config, "runtime"): @@ -456,8 +477,6 @@ def __getstate__(self) -> Dict: # don't save device. Will be determined by loading object current_dict.pop("stream", None) - current_dict.pop("is_cuda", None) - current_dict.pop("device", None) # this is should not be saved. current_dict.pop("image_sizes", None) @@ -527,9 +546,6 @@ def __setstate__(self, state: Dict) -> None: self.rpu_config = rpu_config self.__dict__.update(current_dict) - self.device = torch_device("cpu") - self.is_cuda = False - # recreate attributes not saved # always first create on CPU x_size = self.in_size + 1 if self.analog_bias else self.in_size @@ -537,6 +553,7 @@ def __setstate__(self, state: Dict) -> None: # Recreate the tile. self.tile = self._recreate_simulator_tile(x_size, d_size, self.rpu_config) + self.analog_ctx.data = self.tile.get_weights() names = self.tile.get_hidden_parameter_names() if len(hidden_parameters_names) > 0 and names != hidden_parameters_names: @@ -568,8 +585,15 @@ def __setstate__(self, state: Dict) -> None: if analog_ctx is not None: # Keep the object ID and device to_device = analog_ctx.device - if self.device != to_device: - self.analog_ctx = self.analog_ctx.to(to_device) + if self.analog_ctx.device != to_device: + # aihwkit implements analog tiles in both CPU and CUDA versions, + # e.g. FloatingPointTile(RPUSimple(4, 3)) + # v.s. FloatingPointTile(RPUCudaSimple(4, 3)) + # Here we need to manually convert the tile to the corresponding version + self.to(to_device) + # Note: `self.to(to_device)` will call `self.analog_ctx.data.to(to_device)` + # so no need to recall + # self.analog_ctx = self.analog_ctx.to(to_device) self.analog_ctx.set_data(analog_ctx.data) @no_grad() @@ -632,6 +656,16 @@ def _separate_weights(self, combined_weights: Tensor) -> Tuple[Tensor, Optional[ return combined_weights, None + # pylint: disable=invalid-name + def to(self, device: torch_device) -> "SimulatorTileWrapper": + """Move the tile to a device. + """ + if device.type == "cuda": + self.cuda(device) + else: + self.cpu() + return self + @no_grad() def cpu(self) -> "SimulatorTileWrapper": """Return a copy of this tile in CPU memory. @@ -642,8 +676,6 @@ def cpu(self) -> "SimulatorTileWrapper": if not self.is_cuda: return self - self.is_cuda = False - self.device = torch_device("cpu") self.analog_ctx.data = self.analog_ctx.data.cpu() self.analog_ctx.reset(self) @@ -665,8 +697,6 @@ def cuda( CudaError: if the library has not been compiled with CUDA. """ device = torch_device("cuda", cuda_device(device).idx) - self.is_cuda = True - self.device = device self.analog_ctx.data = self.analog_ctx.data.cuda(device) self.analog_ctx.reset(self) return self diff --git a/src/aihwkit/simulator/tiles/custom.py b/src/aihwkit/simulator/tiles/custom.py index de1d44258..99e2ed688 100644 --- a/src/aihwkit/simulator/tiles/custom.py +++ b/src/aihwkit/simulator/tiles/custom.py @@ -147,8 +147,7 @@ def set_weights(self, weight: Tensor) -> None: Args: weight: ``[out_size, in_size]`` weight matrix. """ - device = self._analog_weight.device - self._analog_weight = weight.clone().to(device) + self._analog_weight.copy_(weight) def get_weights(self) -> Tensor: """Get the tile weights. diff --git a/src/aihwkit/simulator/tiles/functions.py b/src/aihwkit/simulator/tiles/functions.py index e7ead29f9..84938a831 100644 --- a/src/aihwkit/simulator/tiles/functions.py +++ b/src/aihwkit/simulator/tiles/functions.py @@ -32,6 +32,8 @@ def forward( Note: Indexed versions can used when analog_ctx.use_indexed is set to True. """ + # `ctx` is the parameter required by PyTorch to store the context + # no need to pass it through ```AnalogFunction.apply(...)````. # Store in context for using during `backward()`. ctx.analog_ctx = analog_ctx ctx.analog_tile = analog_tile diff --git a/src/aihwkit/simulator/tiles/periphery.py b/src/aihwkit/simulator/tiles/periphery.py index 4b17628c1..ca722e155 100644 --- a/src/aihwkit/simulator/tiles/periphery.py +++ b/src/aihwkit/simulator/tiles/periphery.py @@ -158,7 +158,7 @@ def set_weights( if not isinstance(bias, Tensor): bias = from_numpy(array(bias)) - self.bias.data[:] = bias[:].clone().detach().to(self.get_dtype()).to(self.bias.device) + self.bias.data.copy_(bias) bias = None combined_weights = self._combine_weights(weight, bias) diff --git a/src/aihwkit/simulator/tiles/rpucuda.py b/src/aihwkit/simulator/tiles/rpucuda.py index e549d438c..1ba4fafc9 100644 --- a/src/aihwkit/simulator/tiles/rpucuda.py +++ b/src/aihwkit/simulator/tiles/rpucuda.py @@ -153,8 +153,6 @@ def cuda( if self.tile.__class__ in MAP_TILE_CLASS_TO_CUDA: with cuda_device(device): self.tile = MAP_TILE_CLASS_TO_CUDA[self.tile.__class__](self.tile) - self.is_cuda = True - self.device = device self.analog_ctx.data = self.analog_ctx.data.cuda(device) self.analog_ctx.reset(self) # type: ignore diff --git a/src/aihwkit/simulator/tiles/torch_tile.py b/src/aihwkit/simulator/tiles/torch_tile.py index 6c8dbd660..08c8754c1 100644 --- a/src/aihwkit/simulator/tiles/torch_tile.py +++ b/src/aihwkit/simulator/tiles/torch_tile.py @@ -77,7 +77,7 @@ def set_weights(self, weight: Tensor) -> None: Args: weight: ``[out_size, in_size]`` weight matrix. """ - self.weight.data = weight.clone().to(self.weight.device) + self.weight.data.copy_(weight) def get_weights(self) -> Tensor: """Get the tile weights. @@ -87,7 +87,7 @@ def get_weights(self) -> Tensor: matrix; and the second item is either the ``[out_size]`` bias vector or ``None`` if the tile is set not to use bias. """ - return self.weight.data.detach().cpu() + return self.weight.data def get_x_size(self) -> int: """Returns input size of tile""" diff --git a/tests/test_layers_linear.py b/tests/test_layers_linear.py index c87f69826..8c4ad88a5 100644 --- a/tests/test_layers_linear.py +++ b/tests/test_layers_linear.py @@ -133,8 +133,12 @@ def test_seed(self): weight1, bias1 = layer1.get_weights() weight2, bias2 = layer2.get_weights() + if self.use_cuda: + weight1, weight2 = weight1.cpu(), weight2.cpu() assert_array_almost_equal(weight1, weight2) if bias1 is not None: + if self.use_cuda: + bias1, bias2 = bias1.cpu(), bias2.cpu() assert_array_almost_equal(bias1, bias2) def test_several_analog_layers(self): diff --git a/tests/test_utils.py b/tests/test_utils.py index 790e8b363..e1603583e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -108,14 +108,18 @@ def train_model(model, loss_func, x_b, y_b): @staticmethod def get_layer_and_tile_weights(model): """Return the weights and biases of the model and the tile and whether - it automatically syncs""" + it automatically syncs + + Note: All the weights and biases are detached and converted to numpy format.""" if isinstance(model, AnalogLinearMapped): weight, bias = model.get_weights() + weight, bias = weight.detach().cpu().numpy(), bias.detach().cpu().numpy() return weight, bias, weight, bias, True if isinstance(model, AnalogConv2dMapped): weight, bias = model.get_weights() + weight, bias = weight.detach().cpu().numpy(), bias.detach().cpu().numpy() return weight, bias, weight, bias, True if model.weight is not None: @@ -123,6 +127,7 @@ def get_layer_and_tile_weights(model): else: # we do not sync anymore weight, bias = model.get_weights() + weight, bias = weight.detach().cpu().numpy(), bias.detach().cpu().numpy() return weight, bias, weight, bias, True if model.bias is not None: @@ -408,6 +413,7 @@ def test_save_load_model_cross_device(self): self.assertIsInstance(new_analog_tile.analog_ctx.analog_tile, analog_tile.__class__) self.assertTrue(new_analog_tile.is_cuda != analog_tile.is_cuda) + self.assertTrue(new_analog_tile.device.type == map_location) if analog_tile.shared_weights is not None: self.assertTrue(new_analog_tile.shared_weights.device.type == map_location) @@ -879,7 +885,11 @@ def test_load_state_dict_conversion(self): state1 = new_state_dict[key] state2 = state_dict[key] - assert_array_almost_equal(state1["analog_tile_weights"], state2["analog_tile_weights"]) + weights1 = state1["analog_tile_weights"] + weights2 = state2["analog_tile_weights"] + if self.use_cuda: + weights1, weights2 = weights1.cpu(), weights2.cpu() + assert_array_almost_equal(weights1, weights2) # assert_array_almost_equal(state1['analog_alpha_scale'], # state2['analog_alpha_scale']) From 0968b26eca626937ae9712d49393225a485acd4e Mon Sep 17 00:00:00 2001 From: Zhaoxian Wu Date: Wed, 25 Mar 2026 00:48:53 -0400 Subject: [PATCH 3/6] fix(array): return bias.data in TileModuleArray.get_weights and add AnalogCtx tests - Fix TileModuleArray.get_weights() returning self.bias (Parameter) instead of self.bias.data (Tensor), which caused TypeError when Conv layers re-assign bias as a bool in reset_parameters. - Add test_analog_ctx.py verifying PR #717 AnalogContext data attribution: correct shape, norm, nonzero, comparison ops, CUDA support, backward compatibility with old checkpoints, and convert_to_analog. Signed-off-by: Zhaoxian-Wu Signed-off-by: Zhaoxian Wu --- src/aihwkit/simulator/tiles/array.py | 2 +- tests/test_analog_ctx.py | 227 +++++++++++++++++++++++++++ 2 files changed, 228 insertions(+), 1 deletion(-) create mode 100644 tests/test_analog_ctx.py diff --git a/src/aihwkit/simulator/tiles/array.py b/src/aihwkit/simulator/tiles/array.py index a8af34544..8469312cd 100644 --- a/src/aihwkit/simulator/tiles/array.py +++ b/src/aihwkit/simulator/tiles/array.py @@ -149,7 +149,7 @@ def get_weights(self, **kwargs: Any) -> Tuple[Tensor, Optional[Tensor]]: weight = cat(weight_lst, 1) if self.bias is not None: - return weight, self.bias + return weight, self.bias.data return weight, None def forward(self, x_input: Tensor, tensor_view: Optional[Tuple] = None) -> Tensor: diff --git a/tests/test_analog_ctx.py b/tests/test_analog_ctx.py new file mode 100644 index 000000000..618a87ce8 --- /dev/null +++ b/tests/test_analog_ctx.py @@ -0,0 +1,227 @@ +# -*- coding: utf-8 -*- + +# (C) Copyright 2020, 2021, 2022, 2023, 2024 IBM. All Rights Reserved. +# +# Licensed under the MIT license. See LICENSE file in the project root for details. + +# pylint: disable=too-many-locals, no-member +"""Tests for AnalogContext data attribution (PR #717). + +Verifies that analog_ctx.data reflects the actual weight matrix stored in the +tile, rather than being an empty scalar tensor. +""" + +from unittest import SkipTest + +from torch import zeros, randn, Tensor, Size, manual_seed +from torch.nn import Parameter +from torch.nn import Linear as TorchLinear, Sequential, Conv2d as TorchConv2d + +from aihwkit.nn import AnalogLinear, AnalogConv2d +from aihwkit.nn.conversion import convert_to_analog +from aihwkit.optim.context import AnalogContext +from aihwkit.simulator.configs import FloatingPointRPUConfig + +from .helpers.decorators import parametrize_over_layers +from .helpers.layers import Linear, LinearCuda, LinearMapped, LinearMappedCuda +from .helpers.testcases import ParametrizedTestCase, SKIP_CUDA_TESTS +from .helpers.tiles import FloatingPoint, TorchInference + + +@parametrize_over_layers( + layers=[Linear, LinearMapped], + tiles=[FloatingPoint, TorchInference], + biases=["analog", "digital", None], +) +class AnalogCtxDataAttributionTest(ParametrizedTestCase): + """Tests that analog_ctx.data has the correct shape and values.""" + + def _get_analog_tile(self, model): + """Return the first analog tile from a model.""" + return next(model.analog_tiles()) + + def test_ctx_data_shape_matches_weights(self): + """analog_ctx.size() must return the tile weight shape, not torch.Size([]).""" + model = self.get_layer(in_features=4, out_features=6) + tile = self._get_analog_tile(model) + ctx = tile.analog_ctx + + # The old implementation returned torch.Size([]) — a scalar. + # The new one must return the tile weight matrix shape. + self.assertNotEqual(ctx.size(), Size([])) + self.assertEqual(len(ctx.size()), 2) + + expected_rows = tile.out_size + in_size = tile.in_size + (1 if tile.analog_bias else 0) + self.assertEqual(ctx.size(), Size([expected_rows, in_size])) + + def test_ctx_data_values_match_tile_weights(self): + """analog_ctx.data must reflect the actual tile weights.""" + model = self.get_layer(in_features=4, out_features=6) + tile = self._get_analog_tile(model) + + weights_from_tile = tile.tile.get_weights() + ctx_data = tile.analog_ctx.data + + self.assertEqual(ctx_data.shape, weights_from_tile.shape) + + def test_ctx_norm_is_meaningful(self): + """analog_ctx.norm() should reflect the weight magnitude, not 1.0.""" + manual_seed(42) + model = self.get_layer(in_features=4, out_features=6) + tile = self._get_analog_tile(model) + + # With randomly initialized weights, the norm should be > 0 + # and should NOT be exactly 1.0 (which the old scalar ones(()) returned). + norm_val = tile.analog_ctx.norm().item() + self.assertGreater(norm_val, 0.0) + + def test_ctx_nonzero_works(self): + """analog_ctx.nonzero() should return indices of nonzero weights.""" + model = self.get_layer(in_features=4, out_features=6) + tile = self._get_analog_tile(model) + + # With random initialization, most weights are nonzero. + nz = tile.analog_ctx.nonzero() + self.assertGreater(len(nz), 0) + + def test_ctx_comparison_ops(self): + """Comparison operators on analog_ctx should work on actual weights.""" + model = self.get_layer(in_features=4, out_features=6) + tile = self._get_analog_tile(model) + + # Weights are initialized near zero with std ~1, so most are < 10. + mask = tile.analog_ctx > 10 + self.assertIsInstance(mask, Tensor) + self.assertEqual(mask.shape, tile.analog_ctx.shape) + + def test_ctx_is_parameter(self): + """analog_ctx should be a torch.nn.Parameter.""" + model = self.get_layer(in_features=4, out_features=6) + tile = self._get_analog_tile(model) + self.assertIsInstance(tile.analog_ctx, Parameter) + self.assertIsInstance(tile.analog_ctx, AnalogContext) + + def test_ctx_after_set_weights(self): + """analog_ctx.data should remain valid after set_weights.""" + model = self.get_layer(in_features=4, out_features=6) + tile = self._get_analog_tile(model) + + # Set new weights + new_weight = randn(tile.out_size, tile.in_size) + new_bias = randn(tile.out_size) if tile.analog_bias else None + tile.set_weights(new_weight, new_bias) + + # ctx should still have a valid non-scalar shape + self.assertNotEqual(tile.analog_ctx.size(), Size([])) + self.assertEqual(len(tile.analog_ctx.size()), 2) + + +@parametrize_over_layers( + layers=[LinearCuda, LinearMappedCuda], + tiles=[FloatingPoint, TorchInference], + biases=["analog", "digital", None], +) +class AnalogCtxDataAttributionCudaTest(ParametrizedTestCase): + """Tests that analog_ctx.data is correct after moving to CUDA.""" + + def _get_analog_tile(self, model): + """Return the first analog tile from a model.""" + return next(model.analog_tiles()) + + def test_ctx_shape_after_cuda(self): + """analog_ctx.data should retain correct shape after .cuda().""" + model = self.get_layer(in_features=4, out_features=6) + tile = self._get_analog_tile(model) + + self.assertTrue(tile.analog_ctx.is_cuda) + self.assertNotEqual(tile.analog_ctx.size(), Size([])) + self.assertEqual(len(tile.analog_ctx.size()), 2) + + def test_ctx_device_after_cuda(self): + """analog_ctx.device should be CUDA after .cuda().""" + model = self.get_layer(in_features=4, out_features=6) + tile = self._get_analog_tile(model) + + self.assertEqual(tile.analog_ctx.device.type, "cuda") + self.assertEqual(tile.device.type, "cuda") + + +class AnalogCtxBackwardCompatibilityTest(ParametrizedTestCase): + """Tests for backward compatibility with old checkpoints.""" + + use_cuda = False + + def test_old_checkpoint_empty_ctx_loads(self): + """Checkpoints with empty-size analog_ctx should load without error.""" + model = AnalogLinear(4, 6, bias=True, rpu_config=FloatingPointRPUConfig()) + + # Simulate an old checkpoint where analog_ctx was torch.Size([]) + state = model.state_dict() + for key in list(state.keys()): + if "analog_ctx" in key: + state[key] = zeros(()) # Simulate old empty scalar + + # Should load without error (non-strict mode for old ctx) + model.load_state_dict(state, strict=False, load_rpu_config=False) + + def test_new_checkpoint_loads(self): + """Checkpoints with properly-shaped analog_ctx should load correctly.""" + model = AnalogLinear(4, 6, bias=True, rpu_config=FloatingPointRPUConfig()) + state = model.state_dict() + + model2 = AnalogLinear(4, 6, bias=True, rpu_config=FloatingPointRPUConfig()) + model2.load_state_dict(state, strict=True, load_rpu_config=False) + + +class AnalogCtxConversionTest(ParametrizedTestCase): + """Tests that convert_to_analog produces valid analog_ctx.""" + + use_cuda = False + + def test_conversion_ctx_shape(self) -> None: + """Converted model should have correct analog_ctx shape.""" + digital_model = Sequential(TorchLinear(8, 4), TorchLinear(4, 2)) + analog_model = convert_to_analog( + digital_model, FloatingPointRPUConfig(), ensure_analog_root=False + ) + + for tile in analog_model.analog_tiles(): + ctx = tile.analog_ctx + self.assertNotEqual(ctx.size(), Size([])) + self.assertEqual(len(ctx.size()), 2) + + def test_conversion_conv2d_ctx_shape(self) -> None: + """Converted Conv2d should have correct analog_ctx shape.""" + digital_conv = TorchConv2d(3, 16, kernel_size=3, padding=1, bias=True) + analog_conv = AnalogConv2d.from_digital(digital_conv, FloatingPointRPUConfig()) + + for tile in analog_conv.analog_tiles(): + ctx = tile.analog_ctx + self.assertNotEqual(ctx.size(), Size([])) + self.assertEqual(len(ctx.size()), 2) + + +class AnalogCtxDevicePropertyTest(ParametrizedTestCase): + """Tests that tile.device and tile.is_cuda are computed from analog_ctx.""" + + use_cuda = False + + def test_device_property_cpu(self): + """tile.device should return CPU device for CPU tiles.""" + model = AnalogLinear(4, 6, rpu_config=FloatingPointRPUConfig()) + tile = next(model.analog_tiles()) + + self.assertEqual(tile.device.type, "cpu") + self.assertFalse(tile.is_cuda) + + def test_device_property_cuda(self): + """tile.device should return CUDA device after .cuda().""" + if SKIP_CUDA_TESTS: + raise SkipTest("not compiled with CUDA support") + + model = AnalogLinear(4, 6, rpu_config=FloatingPointRPUConfig()).cuda() + tile = next(model.analog_tiles()) + + self.assertEqual(tile.device.type, "cuda") + self.assertTrue(tile.is_cuda) From d827213973441679658e91ce40b44cb3a396ab2c Mon Sep 17 00:00:00 2001 From: Zhaoxian Wu Date: Wed, 25 Mar 2026 11:41:10 -0400 Subject: [PATCH 4/6] feat(tile): add as_ref parameter and read-only protection for analog weights MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add as_ref parameter to tile-level get_weights() across all simulator tiles. When as_ref=True, return a direct reference to internal weight storage; when False (default), return an independent copy. For C++ tiles (RPUCuda) where get_weights() can only return a copy, add _bind_shared_weights() to allocate a torch.Tensor and pass it to the C++ tile via set_shared_weights() — both Python and C++ then operate on the same memory with no explicit sync needed. Add ReadOnlyWeightView (Tensor subclass) to prevent accidental in-place modification of analog weights. Uses PyTorch's trailing- underscore naming convention to block all in-place ops (future-proof). Configurable via AnalogContext.readonly flag, writable() context manager, MappingParameter.readonly_weights, and convert_to_analog readonly parameter. Signed-off-by: Zhaoxian Wu --- src/aihwkit/nn/conversion.py | 16 + src/aihwkit/optim/context.py | 83 +- src/aihwkit/optim/weight_view.py | 71 ++ src/aihwkit/simulator/parameters/mapping.py | 15 + src/aihwkit/simulator/tiles/base.py | 168 +++- src/aihwkit/simulator/tiles/custom.py | 17 +- src/aihwkit/simulator/tiles/module.py | 2 + src/aihwkit/simulator/tiles/periphery.py | 27 +- src/aihwkit/simulator/tiles/rpucuda.py | 47 +- src/aihwkit/simulator/tiles/torch_tile.py | 14 +- src/aihwkit/simulator/tiles/transfer.py | 9 +- src/rpucuda/cuda/rpucuda.cu | 32 + src/rpucuda/cuda/rpucuda.h | 5 + src/rpucuda/cuda/rpucuda_pulsed.cu | 22 + src/rpucuda/cuda/rpucuda_pulsed.h | 1 + tests/test_analog_ctx.py | 885 +++++++++++++++++++- 16 files changed, 1383 insertions(+), 31 deletions(-) create mode 100644 src/aihwkit/optim/weight_view.py diff --git a/src/aihwkit/nn/conversion.py b/src/aihwkit/nn/conversion.py index b751fdb50..e8e319c9d 100644 --- a/src/aihwkit/nn/conversion.py +++ b/src/aihwkit/nn/conversion.py @@ -16,6 +16,7 @@ from torch.nn import Module, Linear, Conv1d, Conv2d, Conv3d, Sequential from aihwkit.exceptions import ArgumentError +from aihwkit.optim.context import AnalogContext from aihwkit.simulator.tiles.module import TileModule from aihwkit.nn.modules.container import AnalogWrapper from aihwkit.nn.modules.base import AnalogLayerBase @@ -88,6 +89,7 @@ def convert_to_analog( exclude_modules: Optional[List[str]] = None, inplace: bool = False, verbose: bool = False, + readonly: Optional[bool] = None, ) -> Module: """Convert a given digital model to its analog counterpart. @@ -138,6 +140,13 @@ def convert_to_analog( verbose: Increase verbosity. Will print converted layers. + readonly: If not ``None``, override the ``readonly`` flag on + every :class:`~aihwkit.optim.context.AnalogContext` after + conversion. When ``True``, in-place weight modifications + are blocked; when ``False``, they are allowed. If ``None`` + (default), each tile uses the value from + ``rpu_config.mapping.readonly_weights``. + Returns: Module where all the digital layers are replaced with analog mapped layers. @@ -193,6 +202,7 @@ def convert_to_analog( exclude_modules, True, verbose, + readonly, ) continue @@ -211,6 +221,12 @@ def convert_to_analog( if ensure_analog_root and not module_name and not isinstance(module, AnalogLayerBase): module = AnalogWrapper(module) + # Apply global readonly override if specified + if readonly is not None and not module_name: + for param in module.parameters(): + if isinstance(param, AnalogContext): + param.readonly = readonly + return module diff --git a/src/aihwkit/optim/context.py b/src/aihwkit/optim/context.py index b5c5c33d6..a2ed07cc2 100644 --- a/src/aihwkit/optim/context.py +++ b/src/aihwkit/optim/context.py @@ -8,12 +8,15 @@ # pylint: disable=attribute-defined-outside-init +from contextlib import contextmanager from typing import Optional, Type, Union, Any, TYPE_CHECKING from torch import dtype, Tensor, no_grad from torch.nn import Parameter from torch import device as torch_device +from aihwkit.optim.weight_view import ReadOnlyWeightView + if TYPE_CHECKING: from aihwkit.simulator.tiles.base import SimulatorTileWrapper @@ -42,6 +45,11 @@ class AnalogContext(Parameter): analog_ctx.analog_tile.update(...) analog_ctx.analog_tile.update_indexed(...) --- + + The ``readonly`` flag (default ``True``) causes ``.data`` reads to + return a :class:`~aihwkit.optim.weight_view.ReadOnlyWeightView` + that blocks in-place mutations. Toggle it via the property or the + :meth:`writable` context manager. """ def __new__( @@ -51,9 +59,10 @@ def __new__( ) -> "AnalogContext": # pylint: disable=signature-differs if parameter is None: + weights_ref = analog_tile._get_tile_weights_ref() return Parameter.__new__( cls, - data=analog_tile.tile.get_weights(), + data=weights_ref, requires_grad=True, ) # analog_tile.tile can comes from different classes: @@ -69,6 +78,7 @@ def __init__( self, analog_tile: "SimulatorTileWrapper", parameter: Optional[Parameter] = None ): # pylint: disable=unused-argument super().__init__() + self._readonly = self._default_readonly(analog_tile) self.analog_tile = analog_tile self.use_torch_update = False self.use_indexed = False @@ -76,6 +86,63 @@ def __init__( self.analog_grad_output = [] # type: list self.reset(analog_tile) + # -- readonly flag -------------------------------------------------------- + + @staticmethod + def _default_readonly(analog_tile: "SimulatorTileWrapper") -> bool: + """Read the default ``readonly`` setting from ``rpu_config.mapping``.""" + rpu_config = getattr(analog_tile, "rpu_config", None) + if rpu_config is not None: + mapping = getattr(rpu_config, "mapping", None) + if mapping is not None: + return getattr(mapping, "readonly_weights", True) + return True + + @property + def readonly(self) -> bool: + """Whether in-place modifications on ``data`` are blocked.""" + try: + return object.__getattribute__(self, "_readonly") + except AttributeError: + return True + + @readonly.setter + def readonly(self, value: bool) -> None: + self._readonly = value + + def __getattribute__(self, name: str) -> Any: + """Intercept ``.data`` reads: return a :class:`ReadOnlyWeightView` + when ``readonly`` is ``True``, otherwise the raw tensor.""" + if name == "data": + raw = super().__getattribute__(name) + try: + readonly = object.__getattribute__(self, "_readonly") + except AttributeError: + return raw + if readonly: + return ReadOnlyWeightView(raw) + return raw + return super().__getattribute__(name) + + @contextmanager + def writable(self): + """Context manager that temporarily allows direct weight modification. + + Example:: + + with analog_ctx.writable(): + analog_ctx.data.add_(delta) + # readonly is restored automatically + """ + old = self.readonly + self.readonly = False + try: + yield self + finally: + self.readonly = old + + # -- existing API --------------------------------------------------------- + def set_indexed(self, value: bool = True) -> None: """Set the context to forward_indexed.""" self.use_indexed = value @@ -83,7 +150,13 @@ def set_indexed(self, value: bool = True) -> None: def set_data(self, data: Tensor) -> None: """Set the data value of the Tensor.""" with no_grad(): - self.data.copy_(data) + # Unwrap source if it is a ReadOnlyWeightView so that + # copy_() does not trigger the in-place guard. + if isinstance(data, ReadOnlyWeightView): + data = data.as_writable() + # Access raw data directly (bypassing readonly wrap) to + # preserve storage sharing with the tile weight tensor. + super().__getattribute__("data").copy_(data) def get_data(self) -> Tensor: """Get the data value of the underlying Tensor.""" @@ -106,11 +179,11 @@ def has_gradient(self) -> bool: def __copy__(self) -> Parameter: """Turn off copying of the pointers. Context will be re-created when tile is created""" - return Parameter(self.data) + return Parameter(super().__getattribute__("data")) def __deepcopy__(self, memo: Any) -> Parameter: """Turn off deep copying. Context will be re-created when tile is created""" - return Parameter(self.data) + return Parameter(super().__getattribute__("data")) def cuda(self, device: Optional[Union[torch_device, str, int]] = None) -> "AnalogContext": """Move the context to a cuda device. @@ -122,7 +195,7 @@ def cuda(self, device: Optional[Union[torch_device, str, int]] = None) -> "Analo This context in the specified device. """ if not self.analog_tile.is_cuda: - self.data = self.analog_tile.tile.get_weights() # type: Tensor + self.data = self.analog_tile._get_tile_weights_ref() # type: Tensor self.analog_tile = self.analog_tile.cuda(device) self.reset(self.analog_tile) return self diff --git a/src/aihwkit/optim/weight_view.py b/src/aihwkit/optim/weight_view.py new file mode 100644 index 000000000..9682f6443 --- /dev/null +++ b/src/aihwkit/optim/weight_view.py @@ -0,0 +1,71 @@ +# -*- coding: utf-8 -*- + +# (C) Copyright 2020, 2021, 2022, 2023, 2024 IBM. All Rights Reserved. +# +# Licensed under the MIT license. See LICENSE file in the project root for details. + +"""Read-only tensor view for analog tile weights.""" + +from torch import Tensor +from torch.utils._pytree import tree_map + + +class ReadOnlyWeightView(Tensor): + """A tensor that shares storage with tile weights but blocks in-place mutations. + + All read operations (``size``, ``norm``, ``sum``, indexing, comparisons, etc.) + work transparently because this IS a real tensor sharing the same memory. + In-place write operations raise ``RuntimeError`` with guidance to use the + correct analog tile API. + + This class is stateless — it always blocks in-place operations. The policy + of whether to wrap or unwrap is managed by :class:`AnalogContext` via its + ``readonly`` flag. + """ + + @staticmethod + def __new__(cls, data: Tensor) -> "ReadOnlyWeightView": + """Create a ReadOnlyWeightView sharing storage with ``data``.""" + if isinstance(data, ReadOnlyWeightView): + return data + return Tensor._make_subclass(cls, data) + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + func_name = getattr(func, "__name__", "") + + # PyTorch convention: in-place ops end with single '_' (add_, mul_, ...) + # Dunder methods (__repr__, __eq__, ...) end with '__' and must pass through + if func_name.endswith("_") and not func_name.endswith("__"): + raise RuntimeError( + f"In-place operation '{func_name}' is not allowed on analog weights. " + f"Analog weights cannot be modified directly — this would bypass " + f"the physical constraints of the analog device.\n" + f" - For programmatic writes: analog_tile.set_weights(new_weight)\n" + f" - For gradient updates: analog_tile.update(x_input, d_input)\n" + f" - To unlock direct access: analog_ctx.readonly = False" + ) + + # Unwrap to plain Tensor so downstream ops don't propagate our subclass + def unwrap(t): + return t.as_subclass(Tensor) if isinstance(t, ReadOnlyWeightView) else t + + args = tree_map(unwrap, args) + kwargs = tree_map(unwrap, kwargs) + return func(*args, **kwargs) + + def __setitem__(self, key, value): + """Block item assignment (e.g., ``ctx.data[0, 0] = 999``).""" + raise RuntimeError( + "Direct item assignment on analog weights is not allowed. " + "Use analog_tile.set_weights() instead, " + "or set analog_ctx.readonly = False to unlock direct access." + ) + + def as_writable(self) -> Tensor: + """Return the underlying plain Tensor (for internal tile use only). + + This removes the read-only guard. Only tile internals should call this. + """ + return self.as_subclass(Tensor) diff --git a/src/aihwkit/simulator/parameters/mapping.py b/src/aihwkit/simulator/parameters/mapping.py index b0791e36d..d8786bbc4 100644 --- a/src/aihwkit/simulator/parameters/mapping.py +++ b/src/aihwkit/simulator/parameters/mapping.py @@ -100,6 +100,20 @@ class MappingParameter(_PrintableMixin): :class:`aihwkit.nn.modules.linear_mapped.AnalogLinearMapped`. """ + readonly_weights: bool = True + """Whether the analog context data is read-only. + + When ``True`` (default), ``analog_ctx.data`` is wrapped in a + :class:`~aihwkit.optim.weight_view.ReadOnlyWeightView` that blocks + in-place mutations (e.g. ``add_``, ``copy_``, ``__setitem__``). + This prevents users from accidentally modifying analog weights in + a way that bypasses the physical constraints of the analog device. + + Set to ``False`` to allow direct weight access for research or + hardware-aware training scenarios. The flag can also be toggled at + runtime via ``analog_ctx.readonly``. + """ + def compatible_with(self, mapping: "MappingParameter") -> bool: """Checks compatiblity @@ -117,6 +131,7 @@ def compatible_with(self, mapping: "MappingParameter") -> bool: "weight_scaling_omega", "weight_scaling_columnwise", "weight_scaling_lr_compensation", + "readonly_weights", ]: continue diff --git a/src/aihwkit/simulator/tiles/base.py b/src/aihwkit/simulator/tiles/base.py index 636b78a0c..ab0d01b8c 100644 --- a/src/aihwkit/simulator/tiles/base.py +++ b/src/aihwkit/simulator/tiles/base.py @@ -14,7 +14,7 @@ from numpy import array from numpy.typing import ArrayLike -from torch import Tensor, from_numpy, float32, unsqueeze, cat, empty, stack, dtype +from torch import Tensor, from_numpy, float32, unsqueeze, cat, empty, stack, dtype, zeros from torch import device as torch_device from torch.cuda import device as cuda_device from torch.autograd import no_grad @@ -171,8 +171,15 @@ def get_brief_info(self) -> str: """Returns a brief info""" raise NotImplementedError - def get_weights(self) -> Tensor: - """Returns the analog weights.""" + def get_weights(self, as_ref: bool = False) -> Tensor: + """Returns the analog weights. + + Args: + as_ref: if True, return a reference to the internal weight tensor + (not detached, stays on the current device). If False (default), + return a detached CPU copy. Not all tile types support true + references; C++ tiles always return a copy regardless. + """ raise NotImplementedError def set_weights(self, weight: Tensor) -> None: @@ -332,9 +339,109 @@ def __init__( self.tile = self._create_simulator_tile(x_size, d_size, rpu_config) + # Set up zero-copy shared weight tensor for C++ tiles. + self._shared_weight_tensor = None # type: Optional[Tensor] + self._bind_shared_weights() + self.analog_ctx = AnalogContext(self) self.analog_ctx.use_torch_update = torch_update + def _bind_shared_weights(self) -> None: + """Bind a PyTorch tensor as the C++ tile's weight storage. + + For C++ tiles that expose ``set_shared_weights``, this allocates a + contiguous tensor and passes its ``data_ptr`` to the C++ side so that + both Python and C++ operate on the same memory. After this call + ``tile.update()`` / ``tile.set_weights()`` modify the tensor + in-place — no explicit sync is needed. + + For pure-Python tiles (which already store weights as + ``torch.Tensor``), this is a no-op. + """ + if not hasattr(self.tile, "set_shared_weights"): + return + + # Probe whether the tile is a pure-Python tile by trying to call + # get_weights(as_ref=True). + # + # Tiles that accept ``as_ref``: + # TorchSimulatorTile — returns self.weight.data + # CustomSimulatorTile — returns self._analog_weight.data + # TransferSimulatorTile — accepts but delegates to C++ tile + # + # C++ tiles (pybind11 bindings) do NOT accept keyword arguments + # and raise TypeError. These are the ones that need binding: + # tiles.AnalogTile / CudaAnalogTile + # tiles.FloatingPointTile / CudaFloatingPointTile + # (and their half/double/bfloat16 variants) + try: + self.tile.get_weights(as_ref=True) + return # Pure-Python tile — already backed by torch.Tensor. + except TypeError: + pass # C++ tile — proceed with shared weight binding below. + + d_size = self.tile.get_d_size() + x_size = self.tile.get_x_size() + # C++ get_weights() always returns CPU; infer device from tile type name. + # CUDA C++ tiles expect transposed layout (x_size, d_size) for + # set_shared_weights, while CPU tiles expect (d_size, x_size). + is_cuda = "Cuda" in type(self.tile).__name__ + tile_device = torch_device("cuda") if is_cuda else torch_device("cpu") + if is_cuda: + shared = zeros(x_size, d_size, dtype=self.get_dtype(), device=tile_device) + else: + shared = zeros(d_size, x_size, dtype=self.get_dtype(), device=tile_device) + self.tile.set_shared_weights(shared) + # CUDA set_shared_weights does not auto-populate the buffer (unlike CPU). + # Force-sync from the tile's internal device weights into the shared tensor. + if is_cuda: + w_cpu = self.tile.get_weights() # (d_size, x_size) on CPU + shared.copy_(w_cpu.t().to(tile_device)) + self._shared_weight_tensor = shared + + def _get_tile_weights_ref(self) -> Tensor: + """Get tile weights, preferring a reference if the tile supports it.""" + if self._shared_weight_tensor is not None: + # CUDA C++ tiles store shared weights in transposed layout + # (x_size, d_size). Return .t() so callers always see the + # standard (d_size, x_size) shape — still zero-copy because + # .t() on a 2-D tensor is a stride-only view. + if "Cuda" in type(self.tile).__name__: + return self._shared_weight_tensor.t() + return self._shared_weight_tensor + # Fast path: CUDA C++ tiles with native GPU weight access. + # get_weights_cuda() returns [x_size, d_size] on device; .t() gives + # the standard [d_size, x_size] view without any CPU roundtrip. + if hasattr(self.tile, "get_weights_cuda"): + return self.tile.get_weights_cuda().t() + try: + return self.tile.get_weights(as_ref=True) + except TypeError: + # C++ tile bindings don't accept as_ref + return self.tile.get_weights() + + def _sync_analog_ctx_weights(self) -> None: + """Sync analog_ctx.data with the tile weights. + + With shared weight tensors, ``analog_ctx.data`` and the tile's + internal weights already share the same memory, so during normal + training (forward → update) this is a **no-op** (same + ``data_ptr``). + + This method is still necessary for **device moves** (cpu ↔ cuda): + moving the tile to a different device replaces its backing store, + which invalidates the old ``data_ptr``. The callers in + ``cpu()`` / ``cuda()`` (base.py, module.py) and + ``set_weights()`` (periphery.py) rely on this to re-bind + ``analog_ctx.data`` after such transitions. + """ + if not hasattr(self, "analog_ctx"): + return + target_device = self.analog_ctx.data.device + ref = self._get_tile_weights_ref() + if self.analog_ctx.data.data_ptr() != ref.data_ptr(): + self.analog_ctx.data = ref.to(target_device) + @property def device(self) -> torch_device: """Return the device of the tile.""" @@ -480,6 +587,8 @@ def __getstate__(self) -> Dict: # this is should not be saved. current_dict.pop("image_sizes", None) + # Shared weight tensor is rebuilt by _bind_shared_weights(). + current_dict.pop("_shared_weight_tensor", None) return current_dict @@ -553,7 +662,8 @@ def __setstate__(self, state: Dict) -> None: # Recreate the tile. self.tile = self._recreate_simulator_tile(x_size, d_size, self.rpu_config) - self.analog_ctx.data = self.tile.get_weights() + self._shared_weight_tensor = None + self._bind_shared_weights() names = self.tile.get_hidden_parameter_names() if len(hidden_parameters_names) > 0 and names != hidden_parameters_names: @@ -569,6 +679,23 @@ def __setstate__(self, state: Dict) -> None: weights = from_numpy(array(weights)) self.tile.set_weights(weights) + # After loading weights, re-bind shared_weights to the newly recreated + # C++ tile's backing store (see shared_weights docstring in + # RPUCudaSimulatorTileWrapper for its role). + # + # The tile is always recreated on CPU at this point, so we only sync + # here for the CPU case. If the tile will be moved to CUDA, + # .to(device) triggers cuda() which recreates the CUDA tile and + # re-syncs shared_weights there. + # + # ensure_shared_weights is defined on RPUCudaSimulatorTileWrapper, + # not on this base class — guard with hasattr for both runtime + # safety and pylint. + if hasattr(self, "shared_weights") and self.shared_weights is not None: + if not self.shared_weights.is_cuda and hasattr(self, "ensure_shared_weights"): + self.ensure_shared_weights() # pylint: disable=no-member + self.analog_ctx.data = self._get_tile_weights_ref() + if analog_lr is not None: self.tile.set_learning_rate(analog_lr) @@ -647,6 +774,33 @@ def _combine_weights( # Use only the ``[out_size, in_size]`` matrix. return weight + def _combine_weights_cuda( + self, weight: Union[Tensor, "ArrayLike"], bias: Optional[Union[Tensor, "ArrayLike"]] = None + ) -> Tensor: + """Like _combine_weights but keeps tensors on the tile's CUDA device. + + Returns a **contiguous** ``[x_size, d_size]`` CUDA tensor in the + internal transposed layout expected by ``set_weights_cuda``. + """ + d_type = self.get_dtype() + device = self.device # the tile's CUDA device + if not isinstance(weight, Tensor): + weight = from_numpy(array(weight)) + weight = weight.detach().to(dtype=d_type, device=device).reshape(self.out_size, self.in_size) + + if self.analog_bias: + if bias is None: + raise ValueError("Analog tile has a bias, but no bias given") + if not isinstance(bias, Tensor): + bias = from_numpy(array(bias)) + bias = unsqueeze(bias.detach().to(dtype=d_type, device=device), 1) # type: ignore + combined = cat((weight, bias), dim=1) # [out_size, in_size+1] + else: + combined = weight # [out_size, in_size] + + # Transpose to [x_size, d_size] (the internal CUDA storage layout). + return combined.t().contiguous() + def _separate_weights(self, combined_weights: Tensor) -> Tuple[Tensor, Optional[Tensor]]: """Helper to separate the combined weights and biases""" # Split the internal weights (and potentially biases) matrix. @@ -678,6 +832,9 @@ def cpu(self) -> "SimulatorTileWrapper": self.analog_ctx.data = self.analog_ctx.data.cpu() self.analog_ctx.reset(self) + self._shared_weight_tensor = None + self._bind_shared_weights() + self._sync_analog_ctx_weights() return self @@ -699,6 +856,9 @@ def cuda( device = torch_device("cuda", cuda_device(device).idx) self.analog_ctx.data = self.analog_ctx.data.cuda(device) self.analog_ctx.reset(self) + self._shared_weight_tensor = None + self._bind_shared_weights() + self._sync_analog_ctx_weights() return self def get_hidden_parameters(self) -> "OrderedDict": diff --git a/src/aihwkit/simulator/tiles/custom.py b/src/aihwkit/simulator/tiles/custom.py index 99e2ed688..22c19b480 100644 --- a/src/aihwkit/simulator/tiles/custom.py +++ b/src/aihwkit/simulator/tiles/custom.py @@ -57,6 +57,9 @@ def __init__(self, x_size: int, d_size: int, rpu_config: "CustomRPUConfig", bias AnalogMVM.check_support(rpu_config.backward) self.set_config(rpu_config) + # Type declaration for static analysis (pylint/mypy): register_buffer sets this + # attribute dynamically and static analyzers cannot infer it from the string-based API. + self._analog_weight: Tensor # just buffer to handle device, since do not use auto grad self.register_buffer("_analog_weight", zeros(self.d_size, self.x_size, dtype=float32)) @@ -149,15 +152,19 @@ def set_weights(self, weight: Tensor) -> None: """ self._analog_weight.copy_(weight) - def get_weights(self) -> Tensor: + def get_weights(self, as_ref: bool = False) -> Tensor: """Get the tile weights. + Args: + as_ref: if True, return a reference to the internal weight tensor. + If False (default), return a detached CPU copy. + Returns: - a tuple where the first item is the ``[out_size, in_size]`` weight - matrix; and the second item is either the ``[out_size]`` bias vector - or ``None`` if the tile is set not to use bias. + the ``[out_size, in_size]`` weight matrix. """ - return self._analog_weight.data.detach().cpu() + if as_ref: + return self._analog_weight.data + return self._analog_weight.data.detach().cpu().clone() def get_x_size(self) -> int: """Returns input size of tile""" diff --git a/src/aihwkit/simulator/tiles/module.py b/src/aihwkit/simulator/tiles/module.py index 4f242e6e5..1b4daf203 100644 --- a/src/aihwkit/simulator/tiles/module.py +++ b/src/aihwkit/simulator/tiles/module.py @@ -177,6 +177,7 @@ def cuda(self, device: Optional[Union[torch_device, str, int]] = None) -> "TileM # at the end. shared weight might be updated above which might # yeild issues if the tile is not first updated self._apply_without_context(lambda t: t.cuda(device)) + self._sync_analog_ctx_weights() return self def cpu(self) -> "TileModule": @@ -194,6 +195,7 @@ def cpu(self) -> "TileModule": super(BaseTile, self).cpu() # type: ignore self._apply_without_context(lambda t: t.cpu()) + self._sync_analog_ctx_weights() return self def to(self, *args: Any, **kwargs: Any) -> "TileModule": diff --git a/src/aihwkit/simulator/tiles/periphery.py b/src/aihwkit/simulator/tiles/periphery.py index ca722e155..47473996c 100644 --- a/src/aihwkit/simulator/tiles/periphery.py +++ b/src/aihwkit/simulator/tiles/periphery.py @@ -161,11 +161,20 @@ def set_weights( self.bias.data.copy_(bias) bias = None - combined_weights = self._combine_weights(weight, bias) - - if apply_weight_scaling: - combined_weights = self.apply_weight_scaling(combined_weights, weight_scaling_omega) - self.tile.set_weights(combined_weights) + if hasattr(self.tile, "set_weights_cuda"): + combined_weights = self._combine_weights_cuda(weight, bias) + if apply_weight_scaling: + # apply_weight_scaling works on any device; re-transpose after. + combined_weights = self.apply_weight_scaling( + combined_weights.t().contiguous(), weight_scaling_omega + ).t().contiguous() + self.tile.set_weights_cuda(combined_weights) + else: + combined_weights = self._combine_weights(weight, bias) + if apply_weight_scaling: + combined_weights = self.apply_weight_scaling(combined_weights, weight_scaling_omega) + self.tile.set_weights(combined_weights) + self._sync_analog_ctx_weights() if realistic: self.program_weights() @@ -211,7 +220,13 @@ def get_weights( return self.read_weights(apply_weight_scaling=apply_weight_scaling) # Retrieve the internal weights (and potentially biases) matrix. - combined_weights = self.tile.get_weights() + # Use the CUDA-native path when available to avoid a full device sync + # and the CPU-side transpose loop inside the C++ binding. + if hasattr(self.tile, "get_weights_cuda"): + # get_weights_cuda() → [x_size, d_size] on CUDA; .t() → [d_size, x_size] + combined_weights = self.tile.get_weights_cuda().t().detach().cpu() + else: + combined_weights = self.tile.get_weights() weight, bias = self._separate_weights(combined_weights) if self.digital_bias: diff --git a/src/aihwkit/simulator/tiles/rpucuda.py b/src/aihwkit/simulator/tiles/rpucuda.py index 1ba4fafc9..80a526ecb 100644 --- a/src/aihwkit/simulator/tiles/rpucuda.py +++ b/src/aihwkit/simulator/tiles/rpucuda.py @@ -95,6 +95,20 @@ def __init__( handle_output_bound=True, ) + # shared_weights: an nn.Parameter that holds a copy of the tile weights + # in PyTorch memory. Its role differs by tile type: + # + # InferenceTile (supports_ddp=True): + # Exposes weights to PyTorch autograd so that DDP can synchronise + # gradients across GPUs. Must be kept in sync with the C++ tile via + # ensure_shared_weights(), which is called each forward/backward pass. + # + # RPUCuda training tiles (supports_ddp=False): + # Not used for DDP (raises ModuleError if attempted). Here + # shared_weights serves only as the zero-copy backing store that + # bridges the Python-side Parameter and the C++ tile's dev_weights_. + # _shared_weight_tensor (base class) points to the same storage and + # is used for direct weight access without going through the C++ API. self.shared_weights = None # type: Parameter if shared_weights: self.shared_weights = Parameter( @@ -153,7 +167,18 @@ def cuda( if self.tile.__class__ in MAP_TILE_CLASS_TO_CUDA: with cuda_device(device): self.tile = MAP_TILE_CLASS_TO_CUDA[self.tile.__class__](self.tile) - self.analog_ctx.data = self.analog_ctx.data.cuda(device) + # CPU shared tensor is no longer valid for the new CUDA tile. + self._shared_weight_tensor = None + # Re-establish shared weight binding for the new CUDA tile, + # but only when not using the shared_weights DDP path. When + # shared_weights is set, ensure_shared_weights() (called on + # the first forward) handles populating dev_weights_ from + # self.shared_weights.data — calling _bind_shared_weights() + # here would set shared_weights_if_=True prematurely and + # prevent that population step from running. + if self.shared_weights is None: + self._bind_shared_weights() + self.analog_ctx.data = self.tile.get_weights().cuda(device) self.analog_ctx.reset(self) # type: ignore if self.shared_weights is not None: @@ -163,7 +188,22 @@ def cuda( dtype=self.get_dtype(), requires_grad=True, ).cuda(device) - # ensure shared weights will be called later (needs copying still) + # Eagerly populate shared_weights and bind _shared_weight_tensor, + # unless the tile uses is_perfect forward. When is_perfect=True, + # calling ensure_shared_weights() here would invoke C++ + # setSharedWeights() which replaces dev_weights_ with the + # still-zero shared buffer *before* the first forward has a + # chance to populate it — resulting in all-zero output. + # Example: InferenceRPUConfig(forward=IOParameters(is_perfect=True)) + # We defer to ensure_shared_weights() on the first forward(), + # where the C++ copyTo() properly fills the buffer first. + is_perfect_fwd = getattr( + getattr(getattr(self, "rpu_config", None), "forward", None), + "is_perfect", + False, + ) + if not is_perfect_fwd: + self.ensure_shared_weights() return self @@ -181,6 +221,9 @@ def ensure_shared_weights(self, shared_weights: Optional[Tensor] = None) -> None if self.shared_weights is not None: self.tile.set_shared_weights(self.shared_weights.data) # type: ignore + # Keep _shared_weight_tensor in sync: the RPUCuda shared_weights + # path replaces the C++ tile's backing store. + self._shared_weight_tensor = self.shared_weights.data @no_grad() def set_delta_weights(self, delta_weights: Optional[Tensor] = None) -> None: diff --git a/src/aihwkit/simulator/tiles/torch_tile.py b/src/aihwkit/simulator/tiles/torch_tile.py index 08c8754c1..6bae2a2fa 100644 --- a/src/aihwkit/simulator/tiles/torch_tile.py +++ b/src/aihwkit/simulator/tiles/torch_tile.py @@ -79,15 +79,19 @@ def set_weights(self, weight: Tensor) -> None: """ self.weight.data.copy_(weight) - def get_weights(self) -> Tensor: + def get_weights(self, as_ref: bool = False) -> Tensor: """Get the tile weights. + Args: + as_ref: if True, return a reference to the internal weight tensor. + If False (default), return a detached CPU copy. + Returns: - a tuple where the first item is the ``[out_size, in_size]`` weight - matrix; and the second item is either the ``[out_size]`` bias vector - or ``None`` if the tile is set not to use bias. + the ``[out_size, in_size]`` weight matrix. """ - return self.weight.data + if as_ref: + return self.weight.data + return self.weight.data.detach().cpu().clone() def get_x_size(self) -> int: """Returns input size of tile""" diff --git a/src/aihwkit/simulator/tiles/transfer.py b/src/aihwkit/simulator/tiles/transfer.py index 12939aff9..b110393ae 100644 --- a/src/aihwkit/simulator/tiles/transfer.py +++ b/src/aihwkit/simulator/tiles/transfer.py @@ -319,9 +319,14 @@ def get_brief_info(self) -> str: """Returns a brief info""" return self.__class__.__name__ + "({})".format(self.extra_repr()) - def get_weights(self) -> Tensor: - """Returns the analog weights.""" + def get_weights(self, as_ref: bool = False) -> Tensor: + """Returns the analog weights. + Args: + as_ref: if True, return a reference to the internal weight tensor. + If False (default), return a detached CPU copy. + """ + # weight_tile.tile is a C++ tile that doesn't accept kwargs return self.weight_tile.tile.get_weights() def set_weights(self, weight: Tensor) -> None: diff --git a/src/rpucuda/cuda/rpucuda.cu b/src/rpucuda/cuda/rpucuda.cu index d0246dec0..7be194604 100644 --- a/src/rpucuda/cuda/rpucuda.cu +++ b/src/rpucuda/cuda/rpucuda.cu @@ -333,6 +333,38 @@ template void RPUCudaSimple::setSharedWeights(T *device_source) } } +template void RPUCudaSimple::getWeightsCuda(T *device_dest) const { + // D2D read of dev_weights_ in its native transposed [x_size, d_size] layout. + // + // cf. getWeights(T*): copies dev_weights_ to host and transposes to + // row-major [d_size, x_size] — involves GPU sync + D2H + O(n^2) transpose. + // This method stays entirely on device: one async D2D memcpy, no transpose. + // Use when the caller needs a GPU tensor (e.g. Python get_weights_cuda()). + dev_weights_->synchronize(); + CUDA_CALL(cudaMemcpyAsync( + device_dest, dev_weights_->getData(), + (size_t)this->x_size_ * this->d_size_ * sizeof(T), cudaMemcpyDeviceToDevice, + context_->getStream())); +} + +template void RPUCudaSimple::setWeightsCuda(const T *device_source) { + // D2D write into dev_weights_ from a device buffer already in transposed + // [x_size, d_size] layout, then sync the host copy for serialisation. + // + // cf. setWeights(const T* host_source): takes a host pointer in row-major + // [d_size, x_size], copies H2H into weights_, then assignTranspose H2D + // into dev_weights_ — involves a full transpose upload. + // This method avoids the GPU->CPU->GPU round-trip when the source tensor + // is already on the same GPU (e.g. Python set_weights_cuda()). + dev_weights_->synchronize(); + CUDA_CALL(cudaMemcpyAsync( + dev_weights_->getData(), device_source, + (size_t)this->x_size_ * this->d_size_ * sizeof(T), cudaMemcpyDeviceToDevice, + context_->getStream())); + // Keep host copy in sync for serialisation (__getstate__). + this->copyWeightsToHost(RPUSimple::getWeightsPtr()[0]); +} + template void RPUCudaSimple::getAndResetWeightUpdate(T *prev_weight_and_dw_out, T scale) { RPU::math::elemsubcopy( diff --git a/src/rpucuda/cuda/rpucuda.h b/src/rpucuda/cuda/rpucuda.h index bca8f4192..8f18ae7e3 100644 --- a/src/rpucuda/cuda/rpucuda.h +++ b/src/rpucuda/cuda/rpucuda.h @@ -198,6 +198,11 @@ template class RPUCudaSimple : public RPUSimple { void setWeightsUniformRandom(T min_value, T max_value) override; void setSharedWeights(T *weightsptr) override; + // Direct device-to-device weight access (no CPU roundtrip). + // Both buffers use the internal transposed [x_size, d_size] layout. + virtual void getWeightsCuda(T *device_dest) const; + virtual void setWeightsCuda(const T *device_source); + void getAndResetWeightUpdate(T *prev_weights_and_dw_out, T scale = 1.0) override; void applyWeightUpdate(T *dw_and_current_weights_out) override; diff --git a/src/rpucuda/cuda/rpucuda_pulsed.cu b/src/rpucuda/cuda/rpucuda_pulsed.cu index eb0d21308..0e276a0be 100644 --- a/src/rpucuda/cuda/rpucuda_pulsed.cu +++ b/src/rpucuda/cuda/rpucuda_pulsed.cu @@ -529,6 +529,28 @@ template void RPUCudaPulsed::setWeights(const T *host_source) { RPUCudaSimple::setWeights(this->getWeightsPtr()[0]); // set device weights } +template void RPUCudaPulsed::setWeightsCuda(const T *device_source) { + CHECK_RPU_DEVICE_INIT; + + // 1. D2D copy device_source -> dev_weights_, then sync to host (weights_). + RPUCudaSimple::setWeightsCuda(device_source); + + // 2. If a pulsed device is attached (e.g. ConstantStep, LinearStep), + // let it inspect / clamp the new host weights via onSetWeights(). + if (rpu_device_) { + // onSetWeights() returns true when it modified the host weights + // (e.g. applied weight bounds or symmetry constraints). + if (rpu_device_->onSetWeights(this->getWeightsPtr())) { + // 3. Device parameters derived from weights may have changed + // (e.g. slope / reference levels). Sync host device -> CUDA device. + rpucuda_device_->populateFrom(*rpu_device_); + // 4. Host weights were modified by onSetWeights(); push them back + // to dev_weights_ so the GPU copy reflects the clamped values. + RPUCudaSimple::setWeights(this->getWeightsPtr()[0]); + } + } +} + template void RPUCudaPulsed::applyWeightUpdate(T *dw_and_current_weight_out) { CHECK_RPU_DEVICE_INIT; diff --git a/src/rpucuda/cuda/rpucuda_pulsed.h b/src/rpucuda/cuda/rpucuda_pulsed.h index 31bced3c6..4e68c0076 100644 --- a/src/rpucuda/cuda/rpucuda_pulsed.h +++ b/src/rpucuda/cuda/rpucuda_pulsed.h @@ -165,6 +165,7 @@ template class RPUCudaPulsed : public RPUCudaSimple { void getWeightsReal(T *weightsptr) override; void setWeightsReal(const T *weightsptr, int n_loops = 25) override; void setWeights(const T *weightsptr) override; + void setWeightsCuda(const T *device_source) override; void applyWeightUpdate(T *dw_and_current_weights_out) override; diff --git a/tests/test_analog_ctx.py b/tests/test_analog_ctx.py index 618a87ce8..480b86a92 100644 --- a/tests/test_analog_ctx.py +++ b/tests/test_analog_ctx.py @@ -13,14 +13,21 @@ from unittest import SkipTest -from torch import zeros, randn, Tensor, Size, manual_seed +from torch import zeros, randn, allclose, Tensor, Size, manual_seed from torch.nn import Parameter from torch.nn import Linear as TorchLinear, Sequential, Conv2d as TorchConv2d from aihwkit.nn import AnalogLinear, AnalogConv2d from aihwkit.nn.conversion import convert_to_analog from aihwkit.optim.context import AnalogContext -from aihwkit.simulator.configs import FloatingPointRPUConfig +from aihwkit.optim.weight_view import ReadOnlyWeightView +from aihwkit.simulator.configs import ( + FloatingPointRPUConfig, + InferenceRPUConfig, + SingleRPUConfig, + TorchInferenceRPUConfig, +) +from aihwkit.simulator.configs.devices import ConstantStepDevice from .helpers.decorators import parametrize_over_layers from .helpers.layers import Linear, LinearCuda, LinearMapped, LinearMappedCuda @@ -225,3 +232,877 @@ def test_device_property_cuda(self): self.assertEqual(tile.device.type, "cuda") self.assertTrue(tile.is_cuda) + + +class AnalogCtxSyncAfterSetWeightsTest(ParametrizedTestCase): + """Reviewer concern #1: analog_ctx.data must stay in sync after set_weights.""" + + use_cuda = False + + def _test_sync(self, rpu_config, use_cuda): + """Helper: verify ctx.data matches tile weights after set_weights.""" + model = AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + if use_cuda: + model = model.cuda() + tile = next(model.analog_tiles()) + + new_w = randn(6, 4) + tile.set_weights(new_w, None) + + w_from_tile, _ = tile.get_weights() + ctx_data = tile.analog_ctx.data.detach().cpu() + self.assertTrue(allclose(ctx_data, w_from_tile), + "analog_ctx.data out of sync after set_weights") + + def test_sync_torch_inference_cpu(self): + """TorchInference CPU: ctx stays in sync after set_weights.""" + self._test_sync(TorchInferenceRPUConfig(), use_cuda=False) + + def test_sync_floating_point_cpu(self): + """FloatingPoint CPU: ctx stays in sync after set_weights.""" + self._test_sync(FloatingPointRPUConfig(), use_cuda=False) + + def test_sync_torch_inference_cuda(self): + """TorchInference CUDA: ctx stays in sync after set_weights.""" + if SKIP_CUDA_TESTS: + raise SkipTest("not compiled with CUDA support") + self._test_sync(TorchInferenceRPUConfig(), use_cuda=True) + + def test_sync_floating_point_cuda(self): + """FloatingPoint CUDA: ctx stays in sync after set_weights.""" + if SKIP_CUDA_TESTS: + raise SkipTest("not compiled with CUDA support") + self._test_sync(FloatingPointRPUConfig(), use_cuda=True) + + def test_sync_after_multiple_set_weights(self): + """ctx should stay in sync after multiple consecutive set_weights.""" + model = AnalogLinear(4, 6, bias=False, rpu_config=TorchInferenceRPUConfig()) + tile = next(model.analog_tiles()) + + for _ in range(5): + new_w = randn(6, 4) + tile.set_weights(new_w, None) + w_from_tile, _ = tile.get_weights() + ctx_data = tile.analog_ctx.data.detach().cpu() + self.assertTrue(allclose(ctx_data, w_from_tile)) + + def test_sync_after_cuda_move(self): + """ctx should sync after CPU->CUDA->CPU round-trip.""" + if SKIP_CUDA_TESTS: + raise SkipTest("not compiled with CUDA support") + + model = AnalogLinear(4, 6, bias=False, rpu_config=TorchInferenceRPUConfig()) + tile = next(model.analog_tiles()) + + new_w = randn(6, 4) + tile.set_weights(new_w, None) + + # Move to CUDA + model.cuda() + tile = next(model.analog_tiles()) + w_cuda, _ = tile.get_weights() + ctx_cuda = tile.analog_ctx.data.detach().cpu() + self.assertTrue(allclose(ctx_cuda, w_cuda), + "ctx out of sync after cuda()") + + +class AnalogCtxGetWeightsConventionTest(ParametrizedTestCase): + """Reviewer concern #3: get_weights default returns detached CPU copy.""" + + use_cuda = False + + def test_get_weights_returns_cpu_torch_inference(self): + """TorchInference: get_weights() returns CPU tensor by default.""" + if SKIP_CUDA_TESTS: + raise SkipTest("not compiled with CUDA support") + model = AnalogLinear( + 4, 6, bias=False, rpu_config=TorchInferenceRPUConfig() + ).cuda() + tile = next(model.analog_tiles()) + w, _ = tile.get_weights() + self.assertEqual(w.device.type, "cpu", + "get_weights() should return CPU tensor by default") + + def test_get_weights_returns_cpu_floating_point(self): + """FloatingPoint: get_weights() returns CPU tensor by default.""" + if SKIP_CUDA_TESTS: + raise SkipTest("not compiled with CUDA support") + model = AnalogLinear( + 4, 6, bias=False, rpu_config=FloatingPointRPUConfig() + ).cuda() + tile = next(model.analog_tiles()) + w, _ = tile.get_weights() + self.assertEqual(w.device.type, "cpu", + "get_weights() should return CPU tensor by default") + + def test_get_weights_is_detached(self): + """get_weights() result should not have grad_fn (detached).""" + model = AnalogLinear( + 4, 6, bias=False, rpu_config=TorchInferenceRPUConfig() + ) + tile = next(model.analog_tiles()) + w, _ = tile.get_weights() + self.assertFalse(w.requires_grad, + "get_weights() should return detached tensor") + self.assertIsNone(w.grad_fn, + "get_weights() should return detached tensor") + + def test_get_weights_is_copy(self): + """Modifying get_weights() result should NOT change tile weights.""" + model = AnalogLinear( + 4, 6, bias=False, rpu_config=TorchInferenceRPUConfig() + ) + tile = next(model.analog_tiles()) + w_before, _ = tile.get_weights() + w_copy, _ = tile.get_weights() + w_copy.fill_(999.0) + w_after, _ = tile.get_weights() + self.assertTrue(allclose(w_before, w_after), + "get_weights() should return a copy, not a reference") + + +class AnalogCtxTileGetWeightsRefTest(ParametrizedTestCase): + """Test tile.get_weights(as_ref=...) reference vs copy semantics.""" + + use_cuda = False + + def _get_tile(self, rpu_config): + model = AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + return next(model.analog_tiles()) + + def test_as_ref_true_shares_storage(self): + """as_ref=True should return tensors sharing the same storage.""" + tile = self._get_tile(TorchInferenceRPUConfig()) + ref1 = tile.tile.get_weights(as_ref=True) + ref2 = tile.tile.get_weights(as_ref=True) + self.assertEqual(ref1.data_ptr(), ref2.data_ptr(), + "as_ref=True should return the same data pointer") + + def test_clone_does_not_share_storage(self): + """clone() of as_ref=True should NOT share storage.""" + tile = self._get_tile(TorchInferenceRPUConfig()) + ref = tile.tile.get_weights(as_ref=True) + clone = ref.clone() + self.assertNotEqual(ref.data_ptr(), clone.data_ptr(), + "clone should have a different data pointer") + + def test_as_ref_true_modification_propagates(self): + """Modifying as_ref=True tensor should change tile weights.""" + tile = self._get_tile(TorchInferenceRPUConfig()) + ref = tile.tile.get_weights(as_ref=True) + original = ref[0, 0].item() + ref[0, 0] += 999.0 + check = tile.tile.get_weights(as_ref=True) + self.assertAlmostEqual(check[0, 0].item(), original + 999.0, places=2, + msg="as_ref=True modification should propagate to tile") + + def test_as_ref_false_modification_does_not_propagate(self): + """Modifying as_ref=False tensor should NOT change tile weights.""" + tile = self._get_tile(TorchInferenceRPUConfig()) + original = tile.tile.get_weights(as_ref=True)[0, 0].item() + copy = tile.tile.get_weights(as_ref=False) + copy[0, 0] += 999.0 + check = tile.tile.get_weights(as_ref=True) + self.assertAlmostEqual(check[0, 0].item(), original, places=5, + msg="as_ref=False modification should NOT propagate to tile") + + def test_clone_modification_does_not_propagate(self): + """Modifying clone of as_ref=True should NOT change tile weights.""" + tile = self._get_tile(TorchInferenceRPUConfig()) + original = tile.tile.get_weights(as_ref=True)[0, 0].item() + clone = tile.tile.get_weights(as_ref=True).clone() + clone[0, 0] += 999.0 + check = tile.tile.get_weights(as_ref=True) + self.assertAlmostEqual(check[0, 0].item(), original, places=5, + msg="clone modification should NOT propagate to tile") + + +class AnalogCtxSharedWeightsZeroCopyTest(ParametrizedTestCase): + """Tests that C++ tiles use zero-copy shared weights with analog_ctx. + + After ``_bind_shared_weights``, the C++ tile's internal weight storage + and ``analog_ctx.data`` share the same memory. ``tile.update()`` and + ``tile.set_weights()`` must be visible through the shared tensor + without any explicit sync call. + """ + + use_cuda = False + + def _get_tile(self, rpu_config): + model = AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + return next(model.analog_tiles()) + + # -- FloatingPoint (C++ tile) tests ---------------------------------------- + + def test_shared_tensor_exists_for_cpp_tile(self): + """C++ tiles should have a non-None _shared_weight_tensor.""" + tile = self._get_tile(FloatingPointRPUConfig()) + self.assertIsNotNone(tile._shared_weight_tensor, + "C++ tile should have a shared weight tensor") + + def test_shared_tensor_none_for_python_tile(self): + """Pure Python tiles should NOT use _shared_weight_tensor.""" + tile = self._get_tile(TorchInferenceRPUConfig()) + self.assertIsNone(tile._shared_weight_tensor, + "Python tile should not use shared weight tensor") + + def test_ctx_data_shares_memory_with_cpp_tile(self): + """analog_ctx.data and _shared_weight_tensor should share memory.""" + tile = self._get_tile(FloatingPointRPUConfig()) + ctx_ptr = tile.analog_ctx.data.data_ptr() + shared_ptr = tile._shared_weight_tensor.data_ptr() + self.assertEqual(ctx_ptr, shared_ptr, + "analog_ctx.data and shared tensor should have same data_ptr") + + def test_update_reflects_in_ctx_without_sync(self): + """After tile.update(), analog_ctx.data should reflect changes (zero-copy).""" + tile = self._get_tile(FloatingPointRPUConfig()) + tile.tile.set_learning_rate(0.1) + + snapshot = tile.analog_ctx.data.detach().clone() + + x = randn(1, 4) + d = randn(1, 6) + tile.tile.update(x, d, False) + + # analog_ctx.data should have changed without explicit sync + self.assertFalse( + allclose(tile.analog_ctx.data.detach(), snapshot), + "analog_ctx.data should change after tile.update() without sync") + + # And it should match get_weights() + w_from_tile = tile.tile.get_weights() + self.assertTrue( + allclose(tile.analog_ctx.data.detach(), w_from_tile), + "analog_ctx.data should match get_weights() after update") + + def test_set_weights_reflects_in_ctx_without_sync(self): + """After tile.set_weights(), analog_ctx.data should reflect changes.""" + tile = self._get_tile(FloatingPointRPUConfig()) + new_w = randn(6, 4) + tile.tile.set_weights(new_w) + + self.assertTrue( + allclose(tile.analog_ctx.data.detach(), new_w), + "analog_ctx.data should match new weights after set_weights()") + + def test_multiple_updates_stay_in_sync(self): + """analog_ctx.data should stay in sync across 10 consecutive updates.""" + tile = self._get_tile(FloatingPointRPUConfig()) + tile.tile.set_learning_rate(0.01) + + for _ in range(10): + x = randn(4, 4) + d = randn(4, 6) + tile.tile.update(x, d, False) + + w_from_tile = tile.tile.get_weights() + self.assertTrue( + allclose(tile.analog_ctx.data.detach(), w_from_tile), + "analog_ctx.data drifted from tile weights during updates") + + def test_get_weights_ref_returns_shared_tensor(self): + """_get_tile_weights_ref should return the shared tensor for C++ tiles.""" + tile = self._get_tile(FloatingPointRPUConfig()) + ref = tile._get_tile_weights_ref() + self.assertEqual(ref.data_ptr(), tile._shared_weight_tensor.data_ptr(), + "_get_tile_weights_ref should return shared tensor") + + # -- as_ref=True after update (the key new test) --------------------------- + + def test_as_ref_true_reflects_update_python_tile(self): + """Python tile: as_ref=True weight should reflect tile.update() changes.""" + tile = self._get_tile(TorchInferenceRPUConfig()) + ref = tile.tile.get_weights(as_ref=True) + snapshot = ref.clone() + + # Manually modify via the ref (simulating what update would do) + ref[0, 0] += 100.0 + check = tile.tile.get_weights(as_ref=True) + self.assertAlmostEqual(check[0, 0].item(), snapshot[0, 0].item() + 100.0, places=2, + msg="Python tile: as_ref write should propagate") + + def test_as_ref_true_reflects_update_cpp_tile(self): + """C++ tile: _get_tile_weights_ref should reflect tile.update() changes. + + This is the key test: for C++ tiles, the shared weight tensor + (returned by _get_tile_weights_ref) must automatically reflect + updates performed by the C++ tile.update() — zero-copy. + """ + tile = self._get_tile(FloatingPointRPUConfig()) + tile.tile.set_learning_rate(0.1) + + ref = tile._get_tile_weights_ref() + snapshot = ref.clone() + + x = randn(1, 4) + d = randn(1, 6) + tile.tile.update(x, d, False) + + # ref should have been modified in-place by the C++ update + self.assertFalse( + allclose(ref, snapshot), + "C++ tile: shared weight ref should change after update (zero-copy)") + + # And the ref should match get_weights + w_copy = tile.tile.get_weights() + self.assertTrue( + allclose(ref, w_copy), + "C++ tile: shared weight ref should match get_weights() after update") + + +class AnalogCtxReadOnlyTest(ParametrizedTestCase): + """Tests for ReadOnlyWeightView and the readonly flag on AnalogContext.""" + + use_cuda = False + + def _make_model(self, readonly=True): + """Create an AnalogLinear model with the given readonly setting.""" + rpu_config = TorchInferenceRPUConfig() + rpu_config.mapping.readonly_weights = readonly + return AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + + def _get_ctx(self, model): + tile = next(model.analog_tiles()) + return tile.analog_ctx + + # -- default behaviour ---------------------------------------------------- + + def test_default_readonly_true(self): + """By default, analog_ctx.data should be a ReadOnlyWeightView.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + self.assertTrue(ctx.readonly) + self.assertIsInstance(ctx.data, ReadOnlyWeightView) + + def test_default_readonly_false_via_config(self): + """Setting readonly_weights=False in config should disable protection.""" + model = self._make_model(readonly=False) + ctx = self._get_ctx(model) + self.assertFalse(ctx.readonly) + self.assertNotIsInstance(ctx.data, ReadOnlyWeightView) + + # -- read operations work transparently ----------------------------------- + + def test_read_ops_work_when_readonly(self): + """size, norm, nonzero, comparisons should all work on readonly data.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + + self.assertEqual(ctx.size(), Size([6, 4])) + self.assertGreater(ctx.norm().item(), 0.0) + self.assertGreater(len(ctx.nonzero()), 0) + mask = ctx > 10 + self.assertEqual(mask.shape, ctx.shape) + + # -- in-place ops blocked when readonly ----------------------------------- + + def test_add_inplace_blocked(self): + """ctx.data.add_() should raise RuntimeError when readonly.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + with self.assertRaises(RuntimeError): + ctx.data.add_(1.0) + + def test_mul_inplace_blocked(self): + """ctx.data.mul_() should raise RuntimeError when readonly.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + with self.assertRaises(RuntimeError): + ctx.data.mul_(2.0) + + def test_copy_inplace_blocked(self): + """ctx.data.copy_() should raise RuntimeError when readonly.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + with self.assertRaises(RuntimeError): + ctx.data.copy_(randn(6, 4)) + + def test_fill_inplace_blocked(self): + """ctx.data.fill_() should raise RuntimeError when readonly.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + with self.assertRaises(RuntimeError): + ctx.data.fill_(0.0) + + def test_zero_inplace_blocked(self): + """ctx.data.zero_() should raise RuntimeError when readonly.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + with self.assertRaises(RuntimeError): + ctx.data.zero_() + + def test_setitem_blocked(self): + """ctx.data[0, 0] = ... should raise RuntimeError when readonly.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + with self.assertRaises(RuntimeError): + ctx.data[0, 0] = 999.0 + + # -- in-place ops allowed when writable ----------------------------------- + + def test_add_inplace_allowed_when_not_readonly(self): + """ctx.data.add_() should work when readonly=False.""" + model = self._make_model(readonly=False) + ctx = self._get_ctx(model) + ctx.data.add_(1.0) # should not raise + + def test_setitem_allowed_when_not_readonly(self): + """ctx.data[0,0] = ... should work when readonly=False.""" + model = self._make_model(readonly=False) + ctx = self._get_ctx(model) + ctx.data[0, 0] = 999.0 # should not raise + + # -- flag toggling -------------------------------------------------------- + + def test_toggle_readonly_on(self): + """Switching readonly from False to True should wrap data.""" + model = self._make_model(readonly=False) + ctx = self._get_ctx(model) + self.assertNotIsInstance(ctx.data, ReadOnlyWeightView) + + ctx.readonly = True + self.assertIsInstance(ctx.data, ReadOnlyWeightView) + with self.assertRaises(RuntimeError): + ctx.data.add_(1.0) + + def test_toggle_readonly_off(self): + """Switching readonly from True to False should unwrap data.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + self.assertIsInstance(ctx.data, ReadOnlyWeightView) + + ctx.readonly = False + self.assertNotIsInstance(ctx.data, ReadOnlyWeightView) + ctx.data.add_(1.0) # should not raise + + # -- context manager ------------------------------------------------------ + + def test_writable_context_manager(self): + """writable() should temporarily allow in-place ops.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + + with ctx.writable(): + self.assertFalse(ctx.readonly) + ctx.data.add_(1.0) # should not raise + + # Readonly restored + self.assertTrue(ctx.readonly) + with self.assertRaises(RuntimeError): + ctx.data.add_(1.0) + + def test_writable_context_manager_restores_on_exception(self): + """writable() should restore readonly even if an exception occurs.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + + try: + with ctx.writable(): + raise ValueError("test exception") + except ValueError: + pass + + self.assertTrue(ctx.readonly) + + # -- data assignment auto-wraps ------------------------------------------- + + def test_data_assignment_auto_wraps(self): + """Assigning to ctx.data should auto-wrap when readonly=True.""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + + ctx.data = randn(6, 4) + self.assertIsInstance(ctx.data, ReadOnlyWeightView) + + def test_data_assignment_no_wrap_when_writable(self): + """Assigning to ctx.data should NOT wrap when readonly=False.""" + model = self._make_model(readonly=False) + ctx = self._get_ctx(model) + + ctx.data = randn(6, 4) + self.assertNotIsInstance(ctx.data, ReadOnlyWeightView) + + # -- set_data respects readonly ------------------------------------------- + + def test_set_data_works_when_readonly(self): + """set_data() should succeed even when readonly (uses assignment).""" + model = self._make_model(readonly=True) + ctx = self._get_ctx(model) + + new_data = randn(6, 4) + ctx.set_data(new_data) + self.assertIsInstance(ctx.data, ReadOnlyWeightView) + self.assertTrue(allclose(ctx.data.detach(), new_data)) + + # -- convert_to_analog readonly parameter --------------------------------- + + def test_convert_to_analog_readonly_override_false(self): + """convert_to_analog(readonly=False) should set all ctx.readonly=False.""" + digital_model = Sequential(TorchLinear(8, 4), TorchLinear(4, 2)) + analog_model = convert_to_analog( + digital_model, TorchInferenceRPUConfig(), + ensure_analog_root=False, readonly=False, + ) + for param in analog_model.parameters(): + if isinstance(param, AnalogContext): + self.assertFalse(param.readonly) + self.assertNotIsInstance(param.data, ReadOnlyWeightView) + + def test_convert_to_analog_readonly_override_true(self): + """convert_to_analog(readonly=True) should set all ctx.readonly=True.""" + rpu_config = TorchInferenceRPUConfig() + rpu_config.mapping.readonly_weights = False # config says writable + digital_model = Sequential(TorchLinear(8, 4), TorchLinear(4, 2)) + analog_model = convert_to_analog( + digital_model, rpu_config, + ensure_analog_root=False, readonly=True, + ) + for param in analog_model.parameters(): + if isinstance(param, AnalogContext): + self.assertTrue(param.readonly) + self.assertIsInstance(param.data, ReadOnlyWeightView) + + def test_convert_to_analog_readonly_default_from_config(self): + """convert_to_analog() without readonly uses rpu_config.mapping value.""" + rpu_config = TorchInferenceRPUConfig() + rpu_config.mapping.readonly_weights = False + digital_model = Sequential(TorchLinear(8, 4)) + analog_model = convert_to_analog( + digital_model, rpu_config, ensure_analog_root=False, + ) + for param in analog_model.parameters(): + if isinstance(param, AnalogContext): + self.assertFalse(param.readonly) + + +class SharedWeightsCudaBindingTest(ParametrizedTestCase): + """Tests that shared weight binding works correctly after .cuda(). + + ``_bind_shared_weights()`` provides zero-copy weight access, but originally + only worked on CPU. After ``.cuda()``, the binding was lost because + ``RPUCudaSimulatorTileWrapper.cuda()`` never re-bound it. + + These tests cover: + + - ``_bind_shared_weights()`` must allocate on the correct device + with the correct layout (CUDA uses transposed ``(x_size, d_size)``). + - ``.cuda()`` / ``.cpu()`` must re-bind shared weights for + FloatingPoint-family tiles. + - ``CudaAnalogTile.set_shared_weights()`` breaks ``is_perfect`` forward, + so AnalogTile variants must NOT be bound on CUDA. + - Training must still converge after shared weight binding. + """ + + use_cuda = False # We manually skip inside each test + + def _skip_if_no_cuda(self): + if SKIP_CUDA_TESTS: + raise SkipTest("not compiled with CUDA support") + + # -- FloatingPoint: shared binding survives .cuda() ----------------------- + + def test_shared_tensor_exists_after_cuda_floating_point(self): + """FloatingPoint tile should have shared tensor after .cuda().""" + self._skip_if_no_cuda() + model = AnalogLinear(8, 4, bias=False, + rpu_config=FloatingPointRPUConfig()).cuda() + tile = next(model.analog_tiles()) + self.assertIsNotNone(tile._shared_weight_tensor, + "shared tensor should exist after .cuda()") + self.assertTrue(tile._shared_weight_tensor.is_cuda, + "shared tensor should be on CUDA") + + def test_shared_tensor_device_matches_tile(self): + """Shared tensor device should match the CUDA tile device.""" + self._skip_if_no_cuda() + model = AnalogLinear(8, 4, bias=False, + rpu_config=FloatingPointRPUConfig()).cuda() + tile = next(model.analog_tiles()) + self.assertEqual(tile._shared_weight_tensor.device.type, "cuda") + + def test_shared_tensor_has_correct_values_after_cuda(self): + """Shared tensor should contain actual weights, not zeros.""" + self._skip_if_no_cuda() + model = AnalogLinear(8, 4, bias=False, + rpu_config=FloatingPointRPUConfig()) + tile_cpu = next(model.analog_tiles()) + w_cpu = tile_cpu.tile.get_weights().clone() + + model = model.cuda() + tile = next(model.analog_tiles()) + + # Shared tensor values should match the original CPU weights. + # The shared tensor is transposed on CUDA, so compare sorted values. + shared_vals = tile._shared_weight_tensor.cpu().flatten().sort()[0] + cpu_vals = w_cpu.flatten().sort()[0] + self.assertTrue(allclose(shared_vals, cpu_vals, atol=1e-5), + "shared tensor should contain the original weights") + + def test_update_reflects_in_shared_tensor_cuda(self): + """tile.update() should modify the shared tensor on CUDA (zero-copy).""" + self._skip_if_no_cuda() + model = AnalogLinear(4, 6, bias=False, + rpu_config=FloatingPointRPUConfig()).cuda() + tile = next(model.analog_tiles()) + tile.tile.set_learning_rate(0.1) + + snapshot = tile._shared_weight_tensor.clone() + x = randn(2, 4, device="cuda") + d = randn(2, 6, device="cuda") + tile.tile.update(x, d, False) + + self.assertFalse(allclose(tile._shared_weight_tensor, snapshot), + "shared tensor should change after tile.update()") + + def test_set_weights_reflects_in_shared_tensor_cuda(self): + """tile.set_weights() should sync to the shared tensor on CUDA.""" + self._skip_if_no_cuda() + model = AnalogLinear(4, 6, bias=False, + rpu_config=FloatingPointRPUConfig()).cuda() + tile = next(model.analog_tiles()) + + new_w = randn(6, 4) + tile.tile.set_weights(new_w) + + # Compare values (shared is transposed on CUDA) + shared_vals = tile._shared_weight_tensor.cpu().flatten().sort()[0] + expected_vals = new_w.flatten().sort()[0] + self.assertTrue(allclose(shared_vals, expected_vals, atol=1e-5), + "shared tensor should reflect set_weights()") + + # -- FloatingPoint: training convergence after binding -------------------- + + def test_training_converges_floating_point_cuda(self): + """Training with FloatingPoint on CUDA should converge after binding.""" + self._skip_if_no_cuda() + from torch.nn.functional import mse_loss + from aihwkit.optim import AnalogSGD + + manual_seed(42) + model = AnalogLinear(4, 2, bias=False, + rpu_config=FloatingPointRPUConfig()).cuda() + x = randn(10, 4, device="cuda") + y = randn(10, 2, device="cuda") + + initial_loss = mse_loss(model(x), y).item() + + opt = AnalogSGD(model.parameters(), lr=0.1) + opt.regroup_param_groups(model) + for _ in range(50): + opt.zero_grad() + loss = mse_loss(model(x), y) + loss.backward() + opt.step() + + final_loss = mse_loss(model(x), y).item() + self.assertLess(final_loss, initial_loss, + "training should reduce loss on CUDA FloatingPoint") + + # -- CPU round-trip: .cuda() then .cpu() ---------------------------------- + + def test_shared_tensor_survives_cpu_round_trip(self): + """Shared tensor should be re-bound after .cuda() -> .cpu().""" + self._skip_if_no_cuda() + model = AnalogLinear(8, 4, bias=False, + rpu_config=FloatingPointRPUConfig()) + model = model.cuda() + model = model.cpu() + tile = next(model.analog_tiles()) + + self.assertIsNotNone(tile._shared_weight_tensor, + "shared tensor should exist after round-trip") + self.assertFalse(tile._shared_weight_tensor.is_cuda, + "shared tensor should be on CPU after .cpu()") + + def test_shared_tensor_survives_cuda_cpu_cuda(self): + """Shared tensor should survive CPU -> CUDA -> CPU -> CUDA.""" + self._skip_if_no_cuda() + model = AnalogLinear(8, 4, bias=False, + rpu_config=FloatingPointRPUConfig()) + model = model.cuda().cpu().cuda() + tile = next(model.analog_tiles()) + + self.assertIsNotNone(tile._shared_weight_tensor, + "shared tensor should exist after double round-trip") + self.assertTrue(tile._shared_weight_tensor.is_cuda, + "shared tensor should be on CUDA") + + # -- ConstantStep / Inference: CUDA binding -------------------------------- + # + # CudaAnalogTile.set_shared_weights() works for normal forward, but + # corrupts the is_perfect forward path. Verify that: + # - ConstantStep and Inference (is_perfect=False) ARE bound on CUDA + # - Inference with is_perfect=True is NOT bound (known C++ issue) + + def test_shared_binding_for_constant_step_cuda(self): + """ConstantStep tile on CUDA should have shared weight binding.""" + self._skip_if_no_cuda() + rpu = SingleRPUConfig(device=ConstantStepDevice()) + model = AnalogLinear(4, 6, bias=False, rpu_config=rpu).cuda() + tile = next(model.analog_tiles()) + self.assertIsNotNone(tile._shared_weight_tensor, + "ConstantStep CUDA should bind shared weights") + self.assertTrue(tile._shared_weight_tensor.is_cuda) + + def test_shared_binding_for_inference_default_cuda(self): + """Inference (is_perfect=False) on CUDA should have shared binding.""" + self._skip_if_no_cuda() + model = AnalogLinear(4, 6, bias=False, + rpu_config=InferenceRPUConfig()).cuda() + tile = next(model.analog_tiles()) + self.assertIsNotNone(tile._shared_weight_tensor, + "Inference (default) CUDA should bind shared weights") + + def test_no_shared_binding_for_is_perfect_cuda(self): + """Inference with is_perfect=True should NOT bind shared weights.""" + self._skip_if_no_cuda() + rpu = InferenceRPUConfig() + rpu.forward.is_perfect = True + model = AnalogLinear(4, 6, bias=False, rpu_config=rpu).cuda() + tile = next(model.analog_tiles()) + self.assertIsNone(tile._shared_weight_tensor, + "is_perfect=True should skip shared binding") + + def test_inference_is_perfect_forward_nonzero_cuda(self): + """Inference tile with is_perfect=True should produce non-zero output.""" + self._skip_if_no_cuda() + rpu = InferenceRPUConfig() + rpu.forward.is_perfect = True + manual_seed(42) + model = AnalogLinear(4, 2, bias=False, rpu_config=rpu).cuda() + x = randn(1, 4, device="cuda") + out = model(x) + self.assertTrue(out.abs().sum() > 0, + "is_perfect forward should produce non-zero output") + + def test_training_converges_constant_step_cuda(self): + """Training with ConstantStep on CUDA should converge after binding.""" + self._skip_if_no_cuda() + from torch.nn.functional import mse_loss + from aihwkit.optim import AnalogSGD + + rpu = SingleRPUConfig(device=ConstantStepDevice()) + manual_seed(42) + model = AnalogLinear(4, 2, bias=False, rpu_config=rpu).cuda() + x = randn(10, 4, device="cuda") + y = randn(10, 2, device="cuda") + + initial_loss = mse_loss(model(x), y).item() + + opt = AnalogSGD(model.parameters(), lr=0.1) + opt.regroup_param_groups(model) + for _ in range(50): + opt.zero_grad() + loss = mse_loss(model(x), y) + loss.backward() + opt.step() + + final_loss = mse_loss(model(x), y).item() + self.assertLess(final_loss, initial_loss, + "training should reduce loss on CUDA ConstantStep") + + def test_training_converges_inference_is_perfect_cuda(self): + """Training Inference+is_perfect on CUDA should converge (no binding).""" + self._skip_if_no_cuda() + from torch.nn.functional import mse_loss + from aihwkit.optim import AnalogSGD + + rpu = InferenceRPUConfig() + rpu.forward.is_perfect = True + manual_seed(4321) + model = AnalogLinear(4, 2, bias=False, rpu_config=rpu).cuda() + x = randn(10, 4, device="cuda") + y = randn(10, 2, device="cuda") + + initial_loss = mse_loss(model(x), y).item() + + opt = AnalogSGD(model.parameters(), lr=0.5) + opt.regroup_param_groups(model) + for _ in range(100): + opt.zero_grad() + loss = mse_loss(model(x), y) + loss.backward() + opt.step() + + final_loss = mse_loss(model(x), y).item() + self.assertLess(final_loss, initial_loss, + "training should reduce loss for Inference CUDA") + + # -- Transposed layout verification --------------------------------------- + + def test_cuda_get_weights_cuda_binding(self): + """get_weights_cuda() binding should return [x_size, d_size] CUDA tensor.""" + self._skip_if_no_cuda() + model = AnalogLinear(8, 4, bias=False, + rpu_config=FloatingPointRPUConfig()).cuda() + tile = next(model.analog_tiles()) + + self.assertTrue(hasattr(tile.tile, "get_weights_cuda"), + "CudaFloatingPointTile should have get_weights_cuda binding") + out = tile.tile.get_weights_cuda() + d = tile.tile.get_d_size() + x = tile.tile.get_x_size() + self.assertTrue(out.is_cuda, "get_weights_cuda() should return CUDA tensor") + self.assertEqual(out.shape[0], x, + f"dim 0 should be x_size={x}") + self.assertEqual(out.shape[1], d, + f"dim 1 should be d_size={d}") + # Values must match get_weights() (which returns [d_size, x_size] on CPU). + w = tile.tile.get_weights() + self.assertTrue(allclose(out.t().cpu(), w), + "get_weights_cuda().t().cpu() should match get_weights()") + + def test_get_tile_weights_ref_returns_transposed_view_for_cuda(self): + """_get_tile_weights_ref should return a CUDA tensor in standard (d_size, x_size) layout. + + CUDA C++ tiles store weights in transposed layout (x_size, d_size). + _get_tile_weights_ref uses get_weights_cuda().t() so callers see the + standard (d_size, x_size) shape without a CPU round-trip. + """ + self._skip_if_no_cuda() + model = AnalogLinear(8, 4, bias=False, + rpu_config=FloatingPointRPUConfig()).cuda() + tile = next(model.analog_tiles()) + + ref = tile._get_tile_weights_ref() + # Should be a CUDA tensor, not a CPU copy. + self.assertTrue(ref.is_cuda, + "_get_tile_weights_ref should return CUDA tensor") + # Shape should be standard (d_size, x_size), not transposed. + w = tile.tile.get_weights() + self.assertEqual(ref.shape, w.shape, + "ref shape should match get_weights() shape") + # Values should match. + self.assertTrue(allclose(ref.cpu(), w), + "_get_tile_weights_ref should match get_weights()") + + # -- Raw tile .to('cuda') ------------------------------------------------- + + def test_raw_tile_to_cuda_shared_binding(self): + """Raw tile (not AnalogLinear) should retain shared binding after .to('cuda'). + + Reproduces the scenario where a tile is constructed directly via + rpu.get_default_tile_module_class() and moved with .to('cuda'). + """ + self._skip_if_no_cuda() + rpu = SingleRPUConfig(device=ConstantStepDevice()) + cls = rpu.get_default_tile_module_class(16, 8) + tile = cls(16, 8, rpu, False) + + # CPU: shared binding from __init__ + self.assertIsNotNone(tile._shared_weight_tensor) + r1 = tile._get_tile_weights_ref() + r2 = tile._get_tile_weights_ref() + self.assertEqual(r1.data_ptr(), r2.data_ptr(), + "CPU: _get_tile_weights_ref should return same ptr") + + # .to('cuda'): shared binding must survive + tile = tile.to("cuda") + self.assertIsNotNone(tile._shared_weight_tensor, + "shared tensor should exist after .to('cuda')") + self.assertTrue(tile._shared_weight_tensor.is_cuda) + r1 = tile._get_tile_weights_ref() + r2 = tile._get_tile_weights_ref() + self.assertEqual(r1.data_ptr(), r2.data_ptr(), + "CUDA: _get_tile_weights_ref should return same ptr") From 01d408518f862edf2c5591f6fbf0a95edca7ca39 Mon Sep 17 00:00:00 2001 From: Zhaoxian Wu Date: Mon, 8 Jun 2026 14:58:30 -0400 Subject: [PATCH 5/6] refactor(ctx): Make AnalogContext.data always readonly by replacing readonly flag with always-on __torch_function__ interception MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AnalogContext now enforces the invariant ctx.data == ctx.analog_tile.get_weights()[0]: the public data attribute always returns logical weights (physical × scales, bias column hidden), and all in-place modifications — add_(), mul_(), copy_(), out= writes, item assignment, and data rebinding — are unconditionally blocked via RuntimeError. This eliminates the mutable readonly flag and all escape hatches: writable(), set_data(), readonly property, convert_to_analog(readonly=...), and mapping.readonly_weights config field. Signed-off-by: Zhaoxian Wu --- src/aihwkit/nn/conversion.py | 15 - src/aihwkit/optim/context.py | 209 +++++---- src/aihwkit/optim/weight_view.py | 92 ++-- src/aihwkit/simulator/parameters/mapping.py | 15 - src/aihwkit/simulator/tiles/base.py | 174 +++++-- src/aihwkit/simulator/tiles/module.py | 11 +- src/aihwkit/simulator/tiles/rpucuda.py | 7 +- tests/test_analog_ctx.py | 487 +++++++++++--------- 8 files changed, 601 insertions(+), 409 deletions(-) diff --git a/src/aihwkit/nn/conversion.py b/src/aihwkit/nn/conversion.py index e8e319c9d..f148a621d 100644 --- a/src/aihwkit/nn/conversion.py +++ b/src/aihwkit/nn/conversion.py @@ -16,7 +16,6 @@ from torch.nn import Module, Linear, Conv1d, Conv2d, Conv3d, Sequential from aihwkit.exceptions import ArgumentError -from aihwkit.optim.context import AnalogContext from aihwkit.simulator.tiles.module import TileModule from aihwkit.nn.modules.container import AnalogWrapper from aihwkit.nn.modules.base import AnalogLayerBase @@ -89,7 +88,6 @@ def convert_to_analog( exclude_modules: Optional[List[str]] = None, inplace: bool = False, verbose: bool = False, - readonly: Optional[bool] = None, ) -> Module: """Convert a given digital model to its analog counterpart. @@ -140,12 +138,6 @@ def convert_to_analog( verbose: Increase verbosity. Will print converted layers. - readonly: If not ``None``, override the ``readonly`` flag on - every :class:`~aihwkit.optim.context.AnalogContext` after - conversion. When ``True``, in-place weight modifications - are blocked; when ``False``, they are allowed. If ``None`` - (default), each tile uses the value from - ``rpu_config.mapping.readonly_weights``. Returns: Module where all the digital layers are replaced with analog @@ -202,7 +194,6 @@ def convert_to_analog( exclude_modules, True, verbose, - readonly, ) continue @@ -221,12 +212,6 @@ def convert_to_analog( if ensure_analog_root and not module_name and not isinstance(module, AnalogLayerBase): module = AnalogWrapper(module) - # Apply global readonly override if specified - if readonly is not None and not module_name: - for param in module.parameters(): - if isinstance(param, AnalogContext): - param.readonly = readonly - return module diff --git a/src/aihwkit/optim/context.py b/src/aihwkit/optim/context.py index a2ed07cc2..e2c2aafc0 100644 --- a/src/aihwkit/optim/context.py +++ b/src/aihwkit/optim/context.py @@ -8,14 +8,18 @@ # pylint: disable=attribute-defined-outside-init -from contextlib import contextmanager from typing import Optional, Type, Union, Any, TYPE_CHECKING -from torch import dtype, Tensor, no_grad +from torch import dtype, Tensor +from torch._C import DisableTorchFunction from torch.nn import Parameter from torch import device as torch_device +from torch.utils._pytree import tree_map -from aihwkit.optim.weight_view import ReadOnlyWeightView +from aihwkit.optim.weight_view import ( + ReadOnlyWeightView, + raise_if_readonly_write_target, +) if TYPE_CHECKING: from aihwkit.simulator.tiles.base import SimulatorTileWrapper @@ -24,32 +28,48 @@ class AnalogContext(Parameter): """Context for analog optimizer. - Note: `data` attribution, inherited from `torch.nn.Parameter`, is a tensor of training parameter If `analog_bias` (which is provided by `analog_tile`) is False, `data` has the same meaning as `torch.nn.Parameter` If `analog_bias` (which is provided by `analog_tile`) is True, The last column of `data` is the `bias` term - Even though it allows us to access the weights directly, always keep in mind that it is used - only for studying propuses. To simulate the real reading, call the `read_weights` method - instead, i.e. given `analog_ctx: AnalogContext`, - estimated_weights, estimated_bias = analog_ctx.analog_tile.read_weights() - - Similarly, even though this feature allows us to update the weights directly, - always keep in mind that the real RPU devices change their weights only - by "pulse update" method. - - Therefore, use the following update methods instead of + Note: For diagnostic purposes, `AnalogContext` exposes a read-only logical weight view + through the `data` attribute, which is equivalent to `analog_tile.get_weights()[0]`. + This allows users to inspect the effective weights. + Direct tensor reads on ``ctx`` or ``ctx.data``, such as ``size()``, ``norm()`` are + equivalent to do so on ``ctx.analog_tile.get_weights()[0]``. + i.e, ctx.data == ctx.analog_tile.get_weights()[0] + + The `data` attribution inherited from `torch.nn.Parameter` stores the raw tile weights, + i.e., the weights without scaling + Its public tensor value is a read-only logical weight view: ``physical weights x scaling`` + # Example usage: + --- + layer = AnalogLinear(4, 3, bias=False, rpu_config=rpu_config) + analog_tile = layer.analog_module + analog_ctx = analog_tile.analog_ctx + weight = analog_tile.get_weights()[0] + + # The following two lines will print the same value: + analog_ctx.size() + analog_ctx.data.size() + weight.size() + --- + Since the changes of both weights and scaling affect the logical weights, + we adopt the convetion that this logical view is read-only + Therefore, in-place operations, such as ``add_``, ``mul_``, etc, are blocked + ctx.data.add_(1.0) # RuntimeError + Use the following update methods instead of writing `data` directly in the analog optimizer: --- analog_ctx.analog_tile.update(...) analog_ctx.analog_tile.update_indexed(...) --- - The ``readonly`` flag (default ``True``) causes ``.data`` reads to - return a :class:`~aihwkit.optim.weight_view.ReadOnlyWeightView` - that blocks in-place mutations. Toggle it via the property or the - :meth:`writable` context manager. + Even though it allows us to access the weights directly, always keep in mind that it is used + only for diagnostic purposes. To simulate the real reading, call the `read_weights` method + instead, i.e. given `analog_ctx: AnalogContext`, + estimated_weights, estimated_bias = analog_ctx.analog_tile.read_weights() """ def __new__( @@ -65,11 +85,11 @@ def __new__( data=weights_ref, requires_grad=True, ) - # analog_tile.tile can comes from different classes: - # aihwkit.silulator.rpu_base.devices.AnalogTile (C++) + # analog_tile.tile can come from different classes: + # aihwkit.simulator.rpu_base.devices.AnalogTile (C++) # TorchInferenceTile (Python) - # It stores the "weight" matrix; - # If analog_tile.analog_bias is True, it also stores the "bias" matrix + # It stores the raw tile matrix; if analog_tile.analog_bias is True, + # the last raw column stores the bias. parameter.__class__ = cls return parameter @@ -78,7 +98,6 @@ def __init__( self, analog_tile: "SimulatorTileWrapper", parameter: Optional[Parameter] = None ): # pylint: disable=unused-argument super().__init__() - self._readonly = self._default_readonly(analog_tile) self.analog_tile = analog_tile self.use_torch_update = False self.use_indexed = False @@ -86,60 +105,91 @@ def __init__( self.analog_grad_output = [] # type: list self.reset(analog_tile) - # -- readonly flag -------------------------------------------------------- - - @staticmethod - def _default_readonly(analog_tile: "SimulatorTileWrapper") -> bool: - """Read the default ``readonly`` setting from ``rpu_config.mapping``.""" - rpu_config = getattr(analog_tile, "rpu_config", None) - if rpu_config is not None: - mapping = getattr(rpu_config, "mapping", None) - if mapping is not None: - return getattr(mapping, "readonly_weights", True) - return True - - @property - def readonly(self) -> bool: - """Whether in-place modifications on ``data`` are blocked.""" + @classmethod + def __torch_function__( + cls, func: Any, _types: Any, args: Any = (), kwargs: Optional[Any] = None + ) -> Any: + kwargs = kwargs or {} + func_name = getattr(func, "__name__", "") + + def is_readonly(value: Any) -> bool: + return isinstance(value, (AnalogContext, ReadOnlyWeightView)) + + raise_if_readonly_write_target(func_name, args, kwargs, is_readonly) + + def to_logical_tensor(value: Any) -> Any: + if isinstance(value, AnalogContext): + return value._logical_data() + return value + + args = tree_map(to_logical_tensor, args) + kwargs = tree_map(to_logical_tensor, kwargs) + return func(*args, **kwargs) + + def __setitem__(self, key: Any, value: Any) -> None: + """Block direct item assignment.""" + raise RuntimeError( + "Direct item assignment on analog weights is not allowed. " + "Use analog_tile.set_weights() instead." + ) + + def _raw_data(self) -> Tensor: + """Return the internal raw tile backing tensor.""" + with DisableTorchFunction(): # pylint: disable=not-context-manager + return super().__getattribute__("data") + + def _logical_data(self) -> Tensor: + """Return logical weights equivalent to ``analog_tile.get_weights()[0]``.""" + raw = self._raw_data() try: - return object.__getattribute__(self, "_readonly") + analog_tile = object.__getattribute__(self, "analog_tile") except AttributeError: - return True + return raw - @readonly.setter - def readonly(self, value: bool) -> None: - self._readonly = value + logical = raw + if getattr(analog_tile, "analog_bias", False) and raw.dim() >= 2: + logical = raw[:, : analog_tile.in_size] + + get_scales = getattr(analog_tile, "get_scales", None) + if get_scales is None: + return logical + + scales = get_scales() + if scales is None: + return logical + + scales = scales.to(device=logical.device, dtype=logical.dtype) + return logical * scales.view(-1, 1) def __getattribute__(self, name: str) -> Any: - """Intercept ``.data`` reads: return a :class:`ReadOnlyWeightView` - when ``readonly`` is ``True``, otherwise the raw tensor.""" + """Intercept public tensor reads that expose the logical view.""" + if name == "grad_fn": + return None + if name in ("device", "dtype", "grad", "is_cuda", "is_leaf", "layout"): + return getattr(self._raw_data(), name) if name == "data": - raw = super().__getattribute__(name) - try: - readonly = object.__getattribute__(self, "_readonly") - except AttributeError: - return raw - if readonly: - return ReadOnlyWeightView(raw) - return raw + return ReadOnlyWeightView(self._logical_data()) + if name == "shape": + return self._logical_data().shape + if name == "ndim": + return self._logical_data().ndim return super().__getattribute__(name) - @contextmanager - def writable(self): - """Context manager that temporarily allows direct weight modification. - - Example:: + def __setattr__(self, name: str, value: Any) -> None: + """Block user-level replacement of ``.data``.""" + if name == "data": + raise RuntimeError( + "Direct replacement of analog_ctx.data is not allowed. " + "Use analog_tile.set_weights(new_weight) for programmatic writes." + ) + super().__setattr__(name, value) - with analog_ctx.writable(): - analog_ctx.data.add_(delta) - # readonly is restored automatically - """ - old = self.readonly - self.readonly = False - try: - yield self - finally: - self.readonly = old + def _replace_raw_data(self, data: Tensor) -> None: + """Replace the internal raw ``Parameter.data`` for tile rebinding.""" + if isinstance(data, ReadOnlyWeightView): + data = data.as_subclass(Tensor) + with DisableTorchFunction(): # pylint: disable=not-context-manager + super().__setattr__("data", data) # -- existing API --------------------------------------------------------- @@ -147,19 +197,8 @@ def set_indexed(self, value: bool = True) -> None: """Set the context to forward_indexed.""" self.use_indexed = value - def set_data(self, data: Tensor) -> None: - """Set the data value of the Tensor.""" - with no_grad(): - # Unwrap source if it is a ReadOnlyWeightView so that - # copy_() does not trigger the in-place guard. - if isinstance(data, ReadOnlyWeightView): - data = data.as_writable() - # Access raw data directly (bypassing readonly wrap) to - # preserve storage sharing with the tile weight tensor. - super().__getattribute__("data").copy_(data) - def get_data(self) -> Tensor: - """Get the data value of the underlying Tensor.""" + """Get a detached logical weight tensor.""" return self.data.detach() def reset(self, analog_tile: Optional["SimulatorTileWrapper"] = None) -> None: @@ -179,11 +218,11 @@ def has_gradient(self) -> bool: def __copy__(self) -> Parameter: """Turn off copying of the pointers. Context will be re-created when tile is created""" - return Parameter(super().__getattribute__("data")) + return Parameter(self._raw_data()) def __deepcopy__(self, memo: Any) -> Parameter: """Turn off deep copying. Context will be re-created when tile is created""" - return Parameter(super().__getattribute__("data")) + return Parameter(self._raw_data()) def cuda(self, device: Optional[Union[torch_device, str, int]] = None) -> "AnalogContext": """Move the context to a cuda device. @@ -195,7 +234,7 @@ def cuda(self, device: Optional[Union[torch_device, str, int]] = None) -> "Analo This context in the specified device. """ if not self.analog_tile.is_cuda: - self.data = self.analog_tile._get_tile_weights_ref() # type: Tensor + self._replace_raw_data(self.analog_tile._get_tile_weights_ref()) self.analog_tile = self.analog_tile.cuda(device) self.reset(self.analog_tile) return self @@ -209,7 +248,7 @@ def cpu(self) -> "AnalogContext": Returns: self """ - self.data = self.data.cpu() + self._replace_raw_data(self._raw_data().cpu()) if self.analog_tile is not None and self.analog_tile.is_cuda: self.analog_tile = self.analog_tile.cpu() self.reset(self.analog_tile) @@ -231,7 +270,7 @@ def to(self, *args: Any, **kwargs: Any) -> "AnalogContext": This module in the specified device. """ # pylint: disable=invalid-name - self.data = self.data.to(*args, **kwargs) + self._replace_raw_data(self._raw_data().to(*args, **kwargs)) device = None if "device" in kwargs: device = kwargs["device"] diff --git a/src/aihwkit/optim/weight_view.py b/src/aihwkit/optim/weight_view.py index 9682f6443..ed350f9a3 100644 --- a/src/aihwkit/optim/weight_view.py +++ b/src/aihwkit/optim/weight_view.py @@ -6,66 +6,88 @@ """Read-only tensor view for analog tile weights.""" +from typing import Any, Callable, Optional + from torch import Tensor from torch.utils._pytree import tree_map -class ReadOnlyWeightView(Tensor): - """A tensor that shares storage with tile weights but blocks in-place mutations. +def is_inplace_operation(func_name: str) -> bool: + """Return whether ``func_name`` follows the PyTorch in-place convention.""" + return func_name.endswith("_") and not func_name.endswith("__") + + +def raise_readonly_error(func_name: str) -> None: + """Raise the standard read-only analog weight mutation error.""" + raise RuntimeError( + f" In-place operation '{func_name}' is not allowed on analog weights.\n" + f" AnalogContext exposes an always-read-only logical weight view.\n" + f" Direct writes would bypass the analog tile update semantics.\n" + f" Please use the appropriate analog tile API to update weights:\n" + f" - For programmatic writes: analog_tile.set_weights(new_weight)\n" + f" - For gradient updates: analog_tile.update(x_input, d_input)" + ) + + +def raise_if_readonly_write_target( + func_name: str, + args: Any, + kwargs: Any, + is_readonly: Callable[[Any], bool], +) -> None: + """Raise if a read-only analog weight view is used as a write target.""" + if is_inplace_operation(func_name) and args and is_readonly(args[0]): + raise_readonly_error(func_name) + + if "out" not in kwargs: + return + + def block_out_target(value: Any) -> Any: + if is_readonly(value): + raise_readonly_error(func_name) + return value - All read operations (``size``, ``norm``, ``sum``, indexing, comparisons, etc.) - work transparently because this IS a real tensor sharing the same memory. - In-place write operations raise ``RuntimeError`` with guidance to use the - correct analog tile API. + tree_map(block_out_target, kwargs["out"]) - This class is stateless — it always blocks in-place operations. The policy - of whether to wrap or unwrap is managed by :class:`AnalogContext` via its - ``readonly`` flag. + +class ReadOnlyWeightView(Tensor): + """A tensor view that blocks in-place mutations on analog weights. + + This class is stateless — it always blocks in-place operations, such as ``add_``, + ``mul_``, etc, which raise ``RuntimeError``. + All read operations (``size``, ``norm``, ``sum``, indexing, comparisons, + etc.) work transparently. """ @staticmethod def __new__(cls, data: Tensor) -> "ReadOnlyWeightView": - """Create a ReadOnlyWeightView sharing storage with ``data``.""" + """Create a ReadOnlyWeightView sharing storage with ``data`` when possible.""" if isinstance(data, ReadOnlyWeightView): return data return Tensor._make_subclass(cls, data) @classmethod - def __torch_function__(cls, func, types, args=(), kwargs=None): + def __torch_function__( + cls, func: Any, _types: Any, args: Any = (), kwargs: Optional[Any] = None + ) -> Any: kwargs = kwargs or {} func_name = getattr(func, "__name__", "") - # PyTorch convention: in-place ops end with single '_' (add_, mul_, ...) - # Dunder methods (__repr__, __eq__, ...) end with '__' and must pass through - if func_name.endswith("_") and not func_name.endswith("__"): - raise RuntimeError( - f"In-place operation '{func_name}' is not allowed on analog weights. " - f"Analog weights cannot be modified directly — this would bypass " - f"the physical constraints of the analog device.\n" - f" - For programmatic writes: analog_tile.set_weights(new_weight)\n" - f" - For gradient updates: analog_tile.update(x_input, d_input)\n" - f" - To unlock direct access: analog_ctx.readonly = False" - ) - - # Unwrap to plain Tensor so downstream ops don't propagate our subclass - def unwrap(t): + def is_readonly(value: Any) -> bool: + return isinstance(value, ReadOnlyWeightView) + + raise_if_readonly_write_target(func_name, args, kwargs, is_readonly) + + def unwrap(t: Any) -> Any: return t.as_subclass(Tensor) if isinstance(t, ReadOnlyWeightView) else t args = tree_map(unwrap, args) kwargs = tree_map(unwrap, kwargs) return func(*args, **kwargs) - def __setitem__(self, key, value): + def __setitem__(self, key: Any, value: Any) -> None: """Block item assignment (e.g., ``ctx.data[0, 0] = 999``).""" raise RuntimeError( "Direct item assignment on analog weights is not allowed. " - "Use analog_tile.set_weights() instead, " - "or set analog_ctx.readonly = False to unlock direct access." + "Use analog_tile.set_weights() instead." ) - - def as_writable(self) -> Tensor: - """Return the underlying plain Tensor (for internal tile use only). - - This removes the read-only guard. Only tile internals should call this. - """ - return self.as_subclass(Tensor) diff --git a/src/aihwkit/simulator/parameters/mapping.py b/src/aihwkit/simulator/parameters/mapping.py index d8786bbc4..b0791e36d 100644 --- a/src/aihwkit/simulator/parameters/mapping.py +++ b/src/aihwkit/simulator/parameters/mapping.py @@ -100,20 +100,6 @@ class MappingParameter(_PrintableMixin): :class:`aihwkit.nn.modules.linear_mapped.AnalogLinearMapped`. """ - readonly_weights: bool = True - """Whether the analog context data is read-only. - - When ``True`` (default), ``analog_ctx.data`` is wrapped in a - :class:`~aihwkit.optim.weight_view.ReadOnlyWeightView` that blocks - in-place mutations (e.g. ``add_``, ``copy_``, ``__setitem__``). - This prevents users from accidentally modifying analog weights in - a way that bypasses the physical constraints of the analog device. - - Set to ``False`` to allow direct weight access for research or - hardware-aware training scenarios. The flag can also be toggled at - runtime via ``analog_ctx.readonly``. - """ - def compatible_with(self, mapping: "MappingParameter") -> bool: """Checks compatiblity @@ -131,7 +117,6 @@ def compatible_with(self, mapping: "MappingParameter") -> bool: "weight_scaling_omega", "weight_scaling_columnwise", "weight_scaling_lr_compensation", - "readonly_weights", ]: continue diff --git a/src/aihwkit/simulator/tiles/base.py b/src/aihwkit/simulator/tiles/base.py index ab0d01b8c..b39a0844e 100644 --- a/src/aihwkit/simulator/tiles/base.py +++ b/src/aihwkit/simulator/tiles/base.py @@ -299,8 +299,8 @@ class SimulatorTileWrapper: `aihwkit.simulator.tiles.inference_torch.TorchInferenceTile` implement this method. The weight data is stored in the tile object. - analog_ctx: `AnalogContext`, which wraps the weight in tile - into a `torch.nn.Parameter` object. + analog_ctx: `AnalogContext`, which exposes the tile's logical + weights as a read-only `torch.nn.Parameter` view. """ def __init__( @@ -382,25 +382,62 @@ def _bind_shared_weights(self) -> None: d_size = self.tile.get_d_size() x_size = self.tile.get_x_size() - # C++ get_weights() always returns CPU; infer device from tile type name. + # C++ get_weights() always returns CPU. For CUDA tiles, use the + # private raw analog context backing when available: AIHWKIT keeps that + # context on the same concrete device as the tile during + # .cuda(device)/.to(device) moves. Fall back to the active CUDA context + # only during initialization windows where analog_ctx is not CUDA-backed + # yet. + # # CUDA C++ tiles expect transposed layout (x_size, d_size) for # set_shared_weights, while CPU tiles expect (d_size, x_size). is_cuda = "Cuda" in type(self.tile).__name__ - tile_device = torch_device("cuda") if is_cuda else torch_device("cpu") if is_cuda: + w_cpu = self.tile.get_weights() # (d_size, x_size) on CPU + # Normal path after .cuda(device) / .to(device): analog_ctx has + # already been refreshed onto the target GPU, so it is the most + # reliable source of the tile's concrete cuda:N placement. + if hasattr(self, "analog_ctx") and self.analog_ctx._raw_data().is_cuda: + tile_device = self.analog_ctx._raw_data().device + else: + # Fallback for partial-initialization / transition windows + # where analog_ctx does not exist yet or still points to CPU. + # In particular, __init__() binds shared weights before + # analog_ctx is created, so CUDA tiles must use the active + # CUDA context established by the caller. + tile_device = torch_device("cuda", cuda_device(None).idx) shared = zeros(x_size, d_size, dtype=self.get_dtype(), device=tile_device) else: + tile_device = torch_device("cpu") shared = zeros(d_size, x_size, dtype=self.get_dtype(), device=tile_device) self.tile.set_shared_weights(shared) # CUDA set_shared_weights does not auto-populate the buffer (unlike CPU). # Force-sync from the tile's internal device weights into the shared tensor. if is_cuda: - w_cpu = self.tile.get_weights() # (d_size, x_size) on CPU shared.copy_(w_cpu.t().to(tile_device)) self._shared_weight_tensor = shared def _get_tile_weights_ref(self) -> Tensor: - """Get tile weights, preferring a reference if the tile supports it.""" + """Get raw tile weights, preferring a reference if the tile supports it. + + This returns the backing weights stored by the simulator tile. It does + not apply digital mapping/output scales. The higher-level + ``AnalogContext.data`` accessor applies those scales when it presents + read-only logical weights to users. + + Possible sources for the returned tensor are: + + - ``self._shared_weight_tensor``: allocated and bound by + :meth:`_bind_shared_weights` in this class for C++ tiles exposing + ``set_shared_weights``. + - ``self.tile.get_weights_cuda()``: native CUDA C++ binding that returns + device weights in transposed layout when available. + - ``self.tile.get_weights(as_ref=True)``: Python simulator tiles define + this reference-returning path, for example ``TorchSimulatorTile`` in + ``torch_tile.py`` and ``CustomSimulatorTile`` in ``custom.py``. + - ``self.tile.get_weights()``: fallback for C++ bindings that cannot + return a reference and therefore provide only a detached copy. + """ if self._shared_weight_tensor is not None: # CUDA C++ tiles store shared weights in transposed layout # (x_size, d_size). Return .t() so callers always see the @@ -421,26 +458,71 @@ def _get_tile_weights_ref(self) -> Tensor: return self.tile.get_weights() def _sync_analog_ctx_weights(self) -> None: - """Sync analog_ctx.data with the tile weights. + """Sync the private analog context raw backing with tile weights. - With shared weight tensors, ``analog_ctx.data`` and the tile's + With shared weight tensors, the context raw backing and the tile's internal weights already share the same memory, so during normal - training (forward → update) this is a **no-op** (same - ``data_ptr``). + training (forward -> update) this is a no-op (same ``data_ptr``). - This method is still necessary for **device moves** (cpu ↔ cuda): - moving the tile to a different device replaces its backing store, - which invalidates the old ``data_ptr``. The callers in - ``cpu()`` / ``cuda()`` (base.py, module.py) and - ``set_weights()`` (periphery.py) rely on this to re-bind - ``analog_ctx.data`` after such transitions. + This method is still necessary for device moves (cpu <-> cuda): moving + the tile to a different device replaces its backing store, which + invalidates the old ``data_ptr``. Public ``analog_ctx.data`` is a + logical read-only view and must not be used for this raw rebinding. """ if not hasattr(self, "analog_ctx"): return - target_device = self.analog_ctx.data.device + current = self.analog_ctx._raw_data() + target_device = current.device ref = self._get_tile_weights_ref() - if self.analog_ctx.data.data_ptr() != ref.data_ptr(): - self.analog_ctx.data = ref.to(target_device) + if current.data_ptr() != ref.data_ptr() or current.device != ref.device: + self.analog_ctx._replace_raw_data(ref.to(target_device)) + + def _copy_weights_to_shared_tensors(self, weights: Tensor) -> None: + """Copy loaded raw weights into Python-side shared backing tensors. + + ``weights`` is the canonical raw tile matrix saved in CPU layout + ``[d_size, x_size]``. During pickle loading, Python-side + ``shared_weights`` may be restored from a CUDA checkpoint with the CUDA + internal layout ``[x_size, d_size]`` while the C++ tile is first + recreated on CPU. Normalize that tensor before rebinding it to the CPU + tile. + """ + + def copy_to(target: Optional[Tensor]) -> None: + if target is None: + return + source = weights + if target.shape != source.shape: + transposed = weights.t() + if target.shape == transposed.shape: + source = transposed + else: + raise TileError( + "Mismatch with loaded analog state: shared weight shape is unexpected." + ) + target.copy_(source.to(device=target.device, dtype=target.dtype)) + + with no_grad(): + copy_to(self._shared_weight_tensor) + shared_weights = getattr(self, "shared_weights", None) + if shared_weights is not None: + target = shared_weights.data + if ( + not hasattr(self.tile, "get_weights_cuda") + and target.shape != weights.shape + ): + # Replacing .data is intentional here: a CUDA checkpoint can + # restore the Parameter with transposed CUDA layout, and + # copy_ cannot change the target tensor shape. + transposed = weights.t() + if target.shape != transposed.shape: + raise TileError( + "Mismatch with loaded analog state: shared weight shape is unexpected." + ) + shared_weights.data = weights.to( + device=target.device, dtype=target.dtype + ).clone() + copy_to(shared_weights.data) @property def device(self) -> torch_device: @@ -578,7 +660,7 @@ def __getstate__(self) -> Dict: current_dict[SN.CLASS] = type(self).__name__ current_dict[SN.LR] = self.tile.get_learning_rate() current_dict.pop("tile", None) - current_dict[SN.CONTEXT] = self.analog_ctx.data + current_dict[SN.CONTEXT] = self.analog_ctx._raw_data().detach() current_dict[SN.EXTRA] = self.tile.dump_extra() current_dict[SN.VERSION] = __version__ @@ -678,24 +760,28 @@ def __setstate__(self, state: Dict) -> None: if not isinstance(weights, Tensor): weights = from_numpy(array(weights)) self.tile.set_weights(weights) - - # After loading weights, re-bind shared_weights to the newly recreated - # C++ tile's backing store (see shared_weights docstring in - # RPUCudaSimulatorTileWrapper for its role). + self._copy_weights_to_shared_tensors(weights) + + # set_weights() fills the fresh CPU C++ tile from the serialized raw + # weights. _copy_weights_to_shared_tensors() then makes the Python + # Parameter hold the same CPU-layout data. ensure_shared_weights() + # completes the hand-off by making the C++ tile use that Parameter + # as its shared backing store, so tile weights, AnalogContext, and + # optimizer state all observe the same storage. # - # The tile is always recreated on CPU at this point, so we only sync - # here for the CPU case. If the tile will be moved to CUDA, - # .to(device) triggers cuda() which recreates the CUDA tile and - # re-syncs shared_weights there. - # - # ensure_shared_weights is defined on RPUCudaSimulatorTileWrapper, - # not on this base class — guard with hasattr for both runtime - # safety and pylint. - if hasattr(self, "shared_weights") and self.shared_weights is not None: - if not self.shared_weights.is_cuda and hasattr(self, "ensure_shared_weights"): - self.ensure_shared_weights() # pylint: disable=no-member - self.analog_ctx.data = self._get_tile_weights_ref() - + # This rebinding is only valid while the recreated tile is CPU-side. + # If the loaded object is later moved to CUDA, cuda() recreates the + # CUDA tile and binds CUDA-layout shared weights there. The method is + # only present on RPUCudaSimulatorTileWrapper, so keep the dynamic + # callable guard for non-shared tile wrappers. + shared_weights = getattr(self, "shared_weights", None) + ensure_shared_weights = getattr(self, "ensure_shared_weights", None) + if ( + shared_weights is not None + and not shared_weights.is_cuda + and callable(ensure_shared_weights) + ): + ensure_shared_weights() # pylint: disable=not-callable if analog_lr is not None: self.tile.set_learning_rate(analog_lr) @@ -718,10 +804,10 @@ def __setstate__(self, state: Dict) -> None: # v.s. FloatingPointTile(RPUCudaSimple(4, 3)) # Here we need to manually convert the tile to the corresponding version self.to(to_device) - # Note: `self.to(to_device)` will call `self.analog_ctx.data.to(to_device)` - # so no need to recall + # Note: `self.to(to_device)` will rebind the private raw + # analog context backing, so no additional context copy + # is needed. # self.analog_ctx = self.analog_ctx.to(to_device) - self.analog_ctx.set_data(analog_ctx.data) @no_grad() def post_update_step(self) -> None: @@ -786,7 +872,9 @@ def _combine_weights_cuda( device = self.device # the tile's CUDA device if not isinstance(weight, Tensor): weight = from_numpy(array(weight)) - weight = weight.detach().to(dtype=d_type, device=device).reshape(self.out_size, self.in_size) + weight = weight.detach().to(dtype=d_type, device=device).reshape( + self.out_size, self.in_size + ) if self.analog_bias: if bias is None: @@ -830,7 +918,7 @@ def cpu(self) -> "SimulatorTileWrapper": if not self.is_cuda: return self - self.analog_ctx.data = self.analog_ctx.data.cpu() + self.analog_ctx._replace_raw_data(self.analog_ctx._raw_data().cpu()) self.analog_ctx.reset(self) self._shared_weight_tensor = None self._bind_shared_weights() @@ -854,7 +942,7 @@ def cuda( CudaError: if the library has not been compiled with CUDA. """ device = torch_device("cuda", cuda_device(device).idx) - self.analog_ctx.data = self.analog_ctx.data.cuda(device) + self.analog_ctx._replace_raw_data(self.analog_ctx._raw_data().cuda(device)) self.analog_ctx.reset(self) self._shared_weight_tensor = None self._bind_shared_weights() diff --git a/src/aihwkit/simulator/tiles/module.py b/src/aihwkit/simulator/tiles/module.py index 1b4daf203..7f94187c4 100644 --- a/src/aihwkit/simulator/tiles/module.py +++ b/src/aihwkit/simulator/tiles/module.py @@ -325,7 +325,10 @@ def _load_state_dict_post_hook( # pylint: disable=unused-argument missing_keys, unexpected_keys = incompatible_keys for key in missing_keys.copy(): - if AnalogTileStateNames.SHARED_WEIGHTS in key: + if ( + AnalogTileStateNames.SHARED_WEIGHTS in key + or key.endswith(AnalogTileStateNames.CONTEXT) + ): missing_keys.remove(key) for key in unexpected_keys.copy(): if ".tile." in key: @@ -399,6 +402,7 @@ def _load_state_dict_pre_hook( analog_state_name = TileModule.get_analog_state_name(prefix) if analog_state_name not in state_dict: + state_dict.pop(prefix + AnalogTileStateNames.CONTEXT, None) missing_keys.append(analog_state_name) return @@ -408,6 +412,11 @@ def _load_state_dict_pre_hook( if shared_weights_key in state_dict and isinstance(state_dict[shared_weights_key], Tensor): state_dict.pop(shared_weights_key, None) + # The analog state restores analog_ctx explicitly. Letting PyTorch + # copy the registered Parameter afterwards is redundant and would + # look like a direct user write to the logical read-only context. + state_dict.pop(prefix + AnalogTileStateNames.CONTEXT, None) + analog_state = state_dict.pop(analog_state_name).copy() if not analog_tile.load_rpu_config: success, msg = analog_tile.compatible_with(analog_state["rpu_config"]) diff --git a/src/aihwkit/simulator/tiles/rpucuda.py b/src/aihwkit/simulator/tiles/rpucuda.py index 80a526ecb..6c35a0612 100644 --- a/src/aihwkit/simulator/tiles/rpucuda.py +++ b/src/aihwkit/simulator/tiles/rpucuda.py @@ -169,8 +169,10 @@ def cuda( self.tile = MAP_TILE_CLASS_TO_CUDA[self.tile.__class__](self.tile) # CPU shared tensor is no longer valid for the new CUDA tile. self._shared_weight_tensor = None + self.analog_ctx._replace_raw_data(self.tile.get_weights().cuda(device)) + self.analog_ctx.reset(self) # type: ignore # Re-establish shared weight binding for the new CUDA tile, - # but only when not using the shared_weights DDP path. When + # but only when not using the shared_weights DDP path. When # shared_weights is set, ensure_shared_weights() (called on # the first forward) handles populating dev_weights_ from # self.shared_weights.data — calling _bind_shared_weights() @@ -178,8 +180,7 @@ def cuda( # prevent that population step from running. if self.shared_weights is None: self._bind_shared_weights() - self.analog_ctx.data = self.tile.get_weights().cuda(device) - self.analog_ctx.reset(self) # type: ignore + self._sync_analog_ctx_weights() if self.shared_weights is not None: self.shared_weights.data = zeros( diff --git a/tests/test_analog_ctx.py b/tests/test_analog_ctx.py index 480b86a92..2cdf37aa5 100644 --- a/tests/test_analog_ctx.py +++ b/tests/test_analog_ctx.py @@ -7,13 +7,15 @@ # pylint: disable=too-many-locals, no-member """Tests for AnalogContext data attribution (PR #717). -Verifies that analog_ctx.data reflects the actual weight matrix stored in the -tile, rather than being an empty scalar tensor. +Verifies that analog_ctx.data exposes the logical weight matrix used by the +tile, rather than being an empty scalar tensor or raw tile-only storage. """ from unittest import SkipTest -from torch import zeros, randn, allclose, Tensor, Size, manual_seed +import torch +from torch import zeros, randn, rand_like, allclose, Tensor, Size, device, manual_seed +from torch.cuda import device_count, device as cuda_device from torch.nn import Parameter from torch.nn import Linear as TorchLinear, Sequential, Conv2d as TorchConv2d @@ -59,18 +61,18 @@ def test_ctx_data_shape_matches_weights(self): self.assertEqual(len(ctx.size()), 2) expected_rows = tile.out_size - in_size = tile.in_size + (1 if tile.analog_bias else 0) - self.assertEqual(ctx.size(), Size([expected_rows, in_size])) + self.assertEqual(ctx.size(), Size([expected_rows, tile.in_size])) def test_ctx_data_values_match_tile_weights(self): - """analog_ctx.data must reflect the actual tile weights.""" + """analog_ctx.data must reflect the logical tile weights.""" model = self.get_layer(in_features=4, out_features=6) tile = self._get_analog_tile(model) - weights_from_tile = tile.tile.get_weights() - ctx_data = tile.analog_ctx.data + weights_from_tile, _ = tile.get_weights() + ctx_data = tile.analog_ctx.data.detach().cpu() self.assertEqual(ctx_data.shape, weights_from_tile.shape) + self.assertTrue(allclose(ctx_data, weights_from_tile)) def test_ctx_norm_is_meaningful(self): """analog_ctx.norm() should reflect the weight magnitude, not 1.0.""" @@ -119,9 +121,10 @@ def test_ctx_after_set_weights(self): new_bias = randn(tile.out_size) if tile.analog_bias else None tile.set_weights(new_weight, new_bias) - # ctx should still have a valid non-scalar shape + # ctx should still have a valid logical non-scalar shape and value. self.assertNotEqual(tile.analog_ctx.size(), Size([])) - self.assertEqual(len(tile.analog_ctx.size()), 2) + self.assertEqual(tile.analog_ctx.size(), Size([tile.out_size, tile.in_size])) + self.assertTrue(allclose(tile.analog_ctx.data.detach().cpu(), tile.get_weights()[0])) @parametrize_over_layers( @@ -143,7 +146,7 @@ def test_ctx_shape_after_cuda(self): self.assertTrue(tile.analog_ctx.is_cuda) self.assertNotEqual(tile.analog_ctx.size(), Size([])) - self.assertEqual(len(tile.analog_ctx.size()), 2) + self.assertEqual(tile.analog_ctx.size(), Size([tile.out_size, tile.in_size])) def test_ctx_device_after_cuda(self): """analog_ctx.device should be CUDA after .cuda().""" @@ -180,6 +183,48 @@ def test_new_checkpoint_loads(self): model2 = AnalogLinear(4, 6, bias=True, rpu_config=FloatingPointRPUConfig()) model2.load_state_dict(state, strict=True, load_rpu_config=False) + def test_missing_analog_state_non_strict_skips_ctx_copy(self): + """Non-strict load should not copy into the read-only analog context.""" + model = AnalogLinear(4, 6, bias=True, rpu_config=FloatingPointRPUConfig()) + state = model.state_dict() + for key in list(state.keys()): + if key.endswith("analog_tile_state"): + del state[key] + + model.load_state_dict(state, strict=False, load_rpu_config=False) + + +class AnalogCtxStateDictLoadTest(ParametrizedTestCase): + """Tests for analog state restoration with read-only contexts.""" + + use_cuda = False + + def test_inference_tile_load_restores_raw_weights(self): + """Loading analog_tile_state should update shared raw tile backing.""" + source = AnalogLinear(4, 2, bias=False, rpu_config=InferenceRPUConfig()) + target = AnalogLinear(4, 2, bias=False, rpu_config=InferenceRPUConfig()) + + source.set_weights(randn(2, 4), None) + saved_raw = next(source.analog_tiles()).tile.get_weights().clone() + + target.load_state_dict(source.state_dict(), load_rpu_config=False) + loaded_raw = next(target.analog_tiles()).tile.get_weights() + + self.assertTrue(allclose(loaded_raw, saved_raw)) + + def test_square_tile_load_does_not_transpose_cpu_weights(self): + """Square CPU tile state should not be mistaken for CUDA layout.""" + source = AnalogLinear(3, 3, bias=False, rpu_config=FloatingPointRPUConfig()) + target = AnalogLinear(3, 3, bias=False, rpu_config=FloatingPointRPUConfig()) + + source.set_weights(randn(3, 3), None) + saved_raw = next(source.analog_tiles()).tile.get_weights().clone() + + target.load_state_dict(source.state_dict(), load_rpu_config=False) + loaded_raw = next(target.analog_tiles()).tile.get_weights() + + self.assertTrue(allclose(loaded_raw, saved_raw)) + class AnalogCtxConversionTest(ParametrizedTestCase): """Tests that convert_to_analog produces valid analog_ctx.""" @@ -233,6 +278,31 @@ def test_device_property_cuda(self): self.assertEqual(tile.device.type, "cuda") self.assertTrue(tile.is_cuda) + def test_ctx_device_properties_do_not_read_logical_weights(self): + """Device metadata should use raw backing without calling get_scales().""" + model = AnalogLinear(4, 6, rpu_config=FloatingPointRPUConfig()) + tile = next(model.analog_tiles()) + + def fail_get_scales(): + raise AssertionError("get_scales should not be called for device metadata") + + tile.get_scales = fail_get_scales + self.assertEqual(tile.analog_ctx.device.type, "cpu") + self.assertFalse(tile.analog_ctx.is_cuda) + self.assertEqual(tile.analog_ctx.dtype, tile.analog_ctx._raw_data().dtype) + + def test_ctx_is_leaf_uses_raw_backing_with_learned_out_scaling(self): + """Optimizer construction should not see ctx as a logical non-leaf tensor.""" + from aihwkit.optim import AnalogSGD + + rpu_config = InferenceRPUConfig() + rpu_config.mapping.learn_out_scaling = True + model = AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + ctx = next(model.analog_tiles()).analog_ctx + + self.assertTrue(ctx.is_leaf) + AnalogSGD(model.parameters(), lr=0.1) + class AnalogCtxSyncAfterSetWeightsTest(ParametrizedTestCase): """Reviewer concern #1: analog_ctx.data must stay in sync after set_weights.""" @@ -421,8 +491,8 @@ class AnalogCtxSharedWeightsZeroCopyTest(ParametrizedTestCase): """Tests that C++ tiles use zero-copy shared weights with analog_ctx. After ``_bind_shared_weights``, the C++ tile's internal weight storage - and ``analog_ctx.data`` share the same memory. ``tile.update()`` and - ``tile.set_weights()`` must be visible through the shared tensor + and the private ``analog_ctx._raw_data()`` share the same memory. + ``tile.update()`` and ``tile.set_weights()`` must be visible through the shared tensor without any explicit sync call. """ @@ -446,35 +516,35 @@ def test_shared_tensor_none_for_python_tile(self): self.assertIsNone(tile._shared_weight_tensor, "Python tile should not use shared weight tensor") - def test_ctx_data_shares_memory_with_cpp_tile(self): - """analog_ctx.data and _shared_weight_tensor should share memory.""" + def test_raw_context_shares_memory_with_cpp_tile(self): + """analog_ctx._raw_data() and _shared_weight_tensor should share memory.""" tile = self._get_tile(FloatingPointRPUConfig()) - ctx_ptr = tile.analog_ctx.data.data_ptr() + ctx_ptr = tile.analog_ctx._raw_data().data_ptr() shared_ptr = tile._shared_weight_tensor.data_ptr() self.assertEqual(ctx_ptr, shared_ptr, - "analog_ctx.data and shared tensor should have same data_ptr") + "analog_ctx raw data and shared tensor should have same data_ptr") def test_update_reflects_in_ctx_without_sync(self): """After tile.update(), analog_ctx.data should reflect changes (zero-copy).""" tile = self._get_tile(FloatingPointRPUConfig()) tile.tile.set_learning_rate(0.1) - snapshot = tile.analog_ctx.data.detach().clone() + snapshot = tile.analog_ctx._raw_data().detach().clone() x = randn(1, 4) d = randn(1, 6) tile.tile.update(x, d, False) - # analog_ctx.data should have changed without explicit sync + # raw context backing should have changed without explicit sync self.assertFalse( - allclose(tile.analog_ctx.data.detach(), snapshot), - "analog_ctx.data should change after tile.update() without sync") + allclose(tile.analog_ctx._raw_data().detach(), snapshot), + "analog_ctx raw data should change after tile.update() without sync") # And it should match get_weights() w_from_tile = tile.tile.get_weights() self.assertTrue( allclose(tile.analog_ctx.data.detach(), w_from_tile), - "analog_ctx.data should match get_weights() after update") + "analog_ctx.data should match logical get_weights() after update") def test_set_weights_reflects_in_ctx_without_sync(self): """After tile.set_weights(), analog_ctx.data should reflect changes.""" @@ -499,7 +569,7 @@ def test_multiple_updates_stay_in_sync(self): w_from_tile = tile.tile.get_weights() self.assertTrue( allclose(tile.analog_ctx.data.detach(), w_from_tile), - "analog_ctx.data drifted from tile weights during updates") + "analog_ctx.data drifted from logical tile weights during updates") def test_get_weights_ref_returns_shared_tensor(self): """_get_tile_weights_ref should return the shared tensor for C++ tiles.""" @@ -552,228 +622,201 @@ def test_as_ref_true_reflects_update_cpp_tile(self): class AnalogCtxReadOnlyTest(ParametrizedTestCase): - """Tests for ReadOnlyWeightView and the readonly flag on AnalogContext.""" + """Tests for the always-read-only logical AnalogContext view.""" use_cuda = False - def _make_model(self, readonly=True): - """Create an AnalogLinear model with the given readonly setting.""" - rpu_config = TorchInferenceRPUConfig() - rpu_config.mapping.readonly_weights = readonly - return AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + def _make_model(self, bias=False, analog_bias=False): + """Create an AnalogLinear model for context tests.""" + rpu_config = FloatingPointRPUConfig() if analog_bias else TorchInferenceRPUConfig() + if analog_bias: + rpu_config.mapping.digital_bias = False + return AnalogLinear(4, 6, bias=bias, rpu_config=rpu_config) - def _get_ctx(self, model): + def _get_tile_and_ctx(self, model): tile = next(model.analog_tiles()) - return tile.analog_ctx - - # -- default behaviour ---------------------------------------------------- + return tile, tile.analog_ctx - def test_default_readonly_true(self): - """By default, analog_ctx.data should be a ReadOnlyWeightView.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) - self.assertTrue(ctx.readonly) - self.assertIsInstance(ctx.data, ReadOnlyWeightView) + # -- public logical read view --------------------------------------------- - def test_default_readonly_false_via_config(self): - """Setting readonly_weights=False in config should disable protection.""" - model = self._make_model(readonly=False) - ctx = self._get_ctx(model) - self.assertFalse(ctx.readonly) - self.assertNotIsInstance(ctx.data, ReadOnlyWeightView) + def test_ctx_has_no_readonly_flag(self): + """The public readonly switch was removed.""" + _, ctx = self._get_tile_and_ctx(self._make_model()) + self.assertFalse(hasattr(ctx, "readonly")) - # -- read operations work transparently ----------------------------------- + def test_ctx_data_is_always_readonly_view(self): + """analog_ctx.data should always be a ReadOnlyWeightView.""" + _, ctx = self._get_tile_and_ctx(self._make_model()) + self.assertIsInstance(ctx.data, ReadOnlyWeightView) - def test_read_ops_work_when_readonly(self): - """size, norm, nonzero, comparisons should all work on readonly data.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) + def test_read_ops_use_logical_data(self): + """size, norm, detach, nonzero, comparisons all use logical weights.""" + tile, ctx = self._get_tile_and_ctx(self._make_model()) + tile.set_scales(2.0) + logical_weight, _ = tile.get_weights() self.assertEqual(ctx.size(), Size([6, 4])) - self.assertGreater(ctx.norm().item(), 0.0) + self.assertEqual(ctx.shape, Size([6, 4])) + self.assertTrue(allclose(ctx.detach().cpu(), logical_weight)) + self.assertAlmostEqual(ctx.norm().item(), logical_weight.norm().item(), places=5) self.assertGreater(len(ctx.nonzero()), 0) mask = ctx > 10 self.assertEqual(mask.shape, ctx.shape) - # -- in-place ops blocked when readonly ----------------------------------- - - def test_add_inplace_blocked(self): - """ctx.data.add_() should raise RuntimeError when readonly.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) - with self.assertRaises(RuntimeError): - ctx.data.add_(1.0) - - def test_mul_inplace_blocked(self): - """ctx.data.mul_() should raise RuntimeError when readonly.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) - with self.assertRaises(RuntimeError): - ctx.data.mul_(2.0) - - def test_copy_inplace_blocked(self): - """ctx.data.copy_() should raise RuntimeError when readonly.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) - with self.assertRaises(RuntimeError): - ctx.data.copy_(randn(6, 4)) - - def test_fill_inplace_blocked(self): - """ctx.data.fill_() should raise RuntimeError when readonly.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) - with self.assertRaises(RuntimeError): - ctx.data.fill_(0.0) - - def test_zero_inplace_blocked(self): - """ctx.data.zero_() should raise RuntimeError when readonly.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) - with self.assertRaises(RuntimeError): - ctx.data.zero_() - - def test_setitem_blocked(self): - """ctx.data[0, 0] = ... should raise RuntimeError when readonly.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) - with self.assertRaises(RuntimeError): - ctx.data[0, 0] = 999.0 - - # -- in-place ops allowed when writable ----------------------------------- - - def test_add_inplace_allowed_when_not_readonly(self): - """ctx.data.add_() should work when readonly=False.""" - model = self._make_model(readonly=False) - ctx = self._get_ctx(model) - ctx.data.add_(1.0) # should not raise - - def test_setitem_allowed_when_not_readonly(self): - """ctx.data[0,0] = ... should work when readonly=False.""" - model = self._make_model(readonly=False) - ctx = self._get_ctx(model) - ctx.data[0, 0] = 999.0 # should not raise - - # -- flag toggling -------------------------------------------------------- + def test_scaled_context_returns_logical_weight(self): + """Scaled tiles expose scaled logical weights, not raw tile weights.""" + tile, ctx = self._get_tile_and_ctx(self._make_model()) + tile.set_scales(2.0) - def test_toggle_readonly_on(self): - """Switching readonly from False to True should wrap data.""" - model = self._make_model(readonly=False) - ctx = self._get_ctx(model) - self.assertNotIsInstance(ctx.data, ReadOnlyWeightView) + logical_weight = tile.get_weights()[0] + raw_weight = tile.get_weights(apply_weight_scaling=False)[0] - ctx.readonly = True - self.assertIsInstance(ctx.data, ReadOnlyWeightView) - with self.assertRaises(RuntimeError): - ctx.data.add_(1.0) - - def test_toggle_readonly_off(self): - """Switching readonly from True to False should unwrap data.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) self.assertIsInstance(ctx.data, ReadOnlyWeightView) + self.assertTrue(allclose(ctx.data.detach().cpu(), logical_weight)) + self.assertFalse(allclose(ctx.data.detach().cpu(), raw_weight)) - ctx.readonly = False - self.assertNotIsInstance(ctx.data, ReadOnlyWeightView) - ctx.data.add_(1.0) # should not raise + def test_setting_scale_changes_ctx_data(self): + """Manually setting scales should update logical ctx.data values.""" + tile, ctx = self._get_tile_and_ctx(self._make_model()) + before = ctx.data.detach().clone() - # -- context manager ------------------------------------------------------ + tile.set_scales(2.0) - def test_writable_context_manager(self): - """writable() should temporarily allow in-place ops.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) + after = ctx.data.detach() + w_from_tile, _ = tile.get_weights() + self.assertFalse(allclose(after, before)) + self.assertTrue(allclose(after, before * 2.0)) + self.assertTrue(allclose(after.cpu(), w_from_tile)) + + def test_analog_bias_column_is_not_exposed(self): + """analog_bias=True should not expose the raw bias column through ctx.data.""" + tile, ctx = self._get_tile_and_ctx(self._make_model(bias=True, analog_bias=True)) + self.assertTrue(tile.analog_bias) + self.assertEqual(ctx.data.shape, Size([tile.out_size, tile.in_size])) + self.assertEqual( + tile.analog_ctx._raw_data().shape, + Size([tile.out_size, tile.in_size + 1]), + ) - with ctx.writable(): - self.assertFalse(ctx.readonly) - ctx.data.add_(1.0) # should not raise + new_weight = randn(tile.out_size, tile.in_size) + new_bias = randn(tile.out_size) + tile.set_weights(new_weight, new_bias) + self.assertTrue(allclose(ctx.data.detach().cpu(), tile.get_weights()[0])) + + # -- direct writes are always blocked -------------------------------------- + + def test_data_inplace_ops_blocked(self): + """ctx.data in-place ops should raise RuntimeError.""" + _, ctx = self._get_tile_and_ctx(self._make_model()) + for op in [ + lambda: ctx.data.add_(1.0), + lambda: ctx.data.mul_(2.0), + lambda: ctx.data.copy_(randn(6, 4)), + lambda: ctx.data.fill_(0.0), + lambda: ctx.data.zero_(), + ]: + with self.subTest(op=op): + with self.assertRaises(RuntimeError): + op() + + def test_context_inplace_ops_blocked(self): + """Direct in-place ops on ctx should raise RuntimeError.""" + _, ctx = self._get_tile_and_ctx(self._make_model()) + ctx.requires_grad = False + snapshot = ctx.data.detach().clone() - # Readonly restored - self.assertTrue(ctx.readonly) with self.assertRaises(RuntimeError): - ctx.data.add_(1.0) + ctx.add_(1.0) + with self.assertRaises(RuntimeError): + ctx.copy_(snapshot) + + self.assertTrue(allclose(ctx.data.detach(), snapshot)) + + def _assert_out_writes_blocked(self, ctx): + """Assert that torch out= cannot target ctx or ctx.data.""" + ctx.requires_grad = False + raw_snapshot = ctx._raw_data().detach().clone() + logical_snapshot = ctx.data.detach().clone() + lhs = torch.ones(ctx.shape, dtype=ctx.dtype, device=ctx.device) + rhs = torch.ones(ctx.shape, dtype=ctx.dtype, device=ctx.device) + + for out in [ctx, ctx.data]: + with self.subTest(out_type=type(out).__name__): + with self.assertRaises(RuntimeError): + torch.add(lhs, rhs, out=out) + self.assertTrue(allclose(ctx._raw_data().detach(), raw_snapshot)) + self.assertTrue(allclose(ctx.data.detach(), logical_snapshot)) + + def test_out_kwarg_writes_blocked_without_scaling(self): + """out= writes must not bypass read-only weights without scaling.""" + rpu_config = SingleRPUConfig(device=ConstantStepDevice()) + tile, ctx = self._get_tile_and_ctx( + AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + ) - def test_writable_context_manager_restores_on_exception(self): - """writable() should restore readonly even if an exception occurs.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) + self.assertIsNone(tile.get_scales()) + self._assert_out_writes_blocked(ctx) - try: - with ctx.writable(): - raise ValueError("test exception") - except ValueError: - pass + def test_out_kwarg_writes_blocked_with_scaling(self): + """out= writes must fail instead of silently writing to a temporary.""" + tile, ctx = self._get_tile_and_ctx(self._make_model()) + tile.set_scales(2.0) - self.assertTrue(ctx.readonly) + self.assertIsNotNone(tile.get_scales()) + self._assert_out_writes_blocked(ctx) - # -- data assignment auto-wraps ------------------------------------------- + def test_copy_from_context_to_external_tensor_allowed(self): + """copy_ may read from ctx or ctx.data when the destination is external.""" + _, ctx = self._get_tile_and_ctx(self._make_model()) + expected = ctx.data.detach().clone() - def test_data_assignment_auto_wraps(self): - """Assigning to ctx.data should auto-wrap when readonly=True.""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) + for source in [ctx, ctx.data]: + with self.subTest(source_type=type(source).__name__): + destination = torch.empty_like(expected) + destination.copy_(source) + self.assertTrue(allclose(destination, expected)) - ctx.data = randn(6, 4) - self.assertIsInstance(ctx.data, ReadOnlyWeightView) + def test_setitem_blocked(self): + """Item assignment through ctx or ctx.data should raise RuntimeError.""" + _, ctx = self._get_tile_and_ctx(self._make_model()) + snapshot = ctx.data.detach().clone() - def test_data_assignment_no_wrap_when_writable(self): - """Assigning to ctx.data should NOT wrap when readonly=False.""" - model = self._make_model(readonly=False) - ctx = self._get_ctx(model) + with self.assertRaises(RuntimeError): + ctx.data[0, 0] = 999.0 + with self.assertRaises(RuntimeError): + ctx[0, 0] = 999.0 - ctx.data = randn(6, 4) - self.assertNotIsInstance(ctx.data, ReadOnlyWeightView) + self.assertTrue(allclose(ctx.data.detach(), snapshot)) - # -- set_data respects readonly ------------------------------------------- + def test_data_assignment_blocked(self): + """ctx.data rebinding should raise RuntimeError.""" + _, ctx = self._get_tile_and_ctx(self._make_model()) + snapshot = ctx.data.detach().clone() - def test_set_data_works_when_readonly(self): - """set_data() should succeed even when readonly (uses assignment).""" - model = self._make_model(readonly=True) - ctx = self._get_ctx(model) + with self.assertRaisesRegex(RuntimeError, "Direct replacement"): + ctx.data = rand_like(ctx.data) - new_data = randn(6, 4) - ctx.set_data(new_data) self.assertIsInstance(ctx.data, ReadOnlyWeightView) - self.assertTrue(allclose(ctx.data.detach(), new_data)) - - # -- convert_to_analog readonly parameter --------------------------------- - - def test_convert_to_analog_readonly_override_false(self): - """convert_to_analog(readonly=False) should set all ctx.readonly=False.""" + self.assertTrue(allclose(ctx.data.detach(), snapshot)) + + def test_public_write_escape_apis_removed(self): + """Old direct-write escape hatches should not exist.""" + _, ctx = self._get_tile_and_ctx(self._make_model()) + self.assertFalse(hasattr(ctx, "set_data")) + self.assertFalse(hasattr(ctx, "writable")) + self.assertFalse(hasattr(ctx, "readonly")) + self.assertFalse(hasattr(TorchInferenceRPUConfig().mapping, "readonly_weights")) + + def test_convert_to_analog_readonly_argument_removed(self): + """convert_to_analog(readonly=...) should no longer be accepted.""" digital_model = Sequential(TorchLinear(8, 4), TorchLinear(4, 2)) - analog_model = convert_to_analog( - digital_model, TorchInferenceRPUConfig(), - ensure_analog_root=False, readonly=False, - ) - for param in analog_model.parameters(): - if isinstance(param, AnalogContext): - self.assertFalse(param.readonly) - self.assertNotIsInstance(param.data, ReadOnlyWeightView) - - def test_convert_to_analog_readonly_override_true(self): - """convert_to_analog(readonly=True) should set all ctx.readonly=True.""" - rpu_config = TorchInferenceRPUConfig() - rpu_config.mapping.readonly_weights = False # config says writable - digital_model = Sequential(TorchLinear(8, 4), TorchLinear(4, 2)) - analog_model = convert_to_analog( - digital_model, rpu_config, - ensure_analog_root=False, readonly=True, - ) - for param in analog_model.parameters(): - if isinstance(param, AnalogContext): - self.assertTrue(param.readonly) - self.assertIsInstance(param.data, ReadOnlyWeightView) - - def test_convert_to_analog_readonly_default_from_config(self): - """convert_to_analog() without readonly uses rpu_config.mapping value.""" - rpu_config = TorchInferenceRPUConfig() - rpu_config.mapping.readonly_weights = False - digital_model = Sequential(TorchLinear(8, 4)) - analog_model = convert_to_analog( - digital_model, rpu_config, ensure_analog_root=False, - ) - for param in analog_model.parameters(): - if isinstance(param, AnalogContext): - self.assertFalse(param.readonly) + with self.assertRaises(TypeError): + convert_to_analog( # pylint: disable=unexpected-keyword-arg + digital_model, + TorchInferenceRPUConfig(), + ensure_analog_root=False, + readonly=False, + ) class SharedWeightsCudaBindingTest(ParametrizedTestCase): @@ -821,6 +864,26 @@ def test_shared_tensor_device_matches_tile(self): tile = next(model.analog_tiles()) self.assertEqual(tile._shared_weight_tensor.device.type, "cuda") + def test_shared_tensor_uses_nonzero_cuda_device(self): + """Rebinding should use the tile device, not the current CUDA context.""" + self._skip_if_no_cuda() + if device_count() < 2: + raise SkipTest("Need at least two devices for this test") + + target_device = device("cuda", 1) + model = AnalogLinear(8, 4, bias=False, + rpu_config=FloatingPointRPUConfig()).cuda(target_device) + tile = next(model.analog_tiles()) + expected_weights = tile.tile.get_weights().clone() + + tile._shared_weight_tensor = None + with cuda_device(0): + tile._bind_shared_weights() + + self.assertEqual(tile.device, target_device) + self.assertEqual(tile._shared_weight_tensor.device, target_device) + self.assertTrue(allclose(tile._shared_weight_tensor.t().cpu(), expected_weights)) + def test_shared_tensor_has_correct_values_after_cuda(self): """Shared tensor should contain actual weights, not zeros.""" self._skip_if_no_cuda() @@ -1038,8 +1101,8 @@ def test_cuda_get_weights_cuda_binding(self): rpu_config=FloatingPointRPUConfig()).cuda() tile = next(model.analog_tiles()) - self.assertTrue(hasattr(tile.tile, "get_weights_cuda"), - "CudaFloatingPointTile should have get_weights_cuda binding") + if not hasattr(tile.tile, "get_weights_cuda"): + raise SkipTest("get_weights_cuda binding is not exposed in this build") out = tile.tile.get_weights_cuda() d = tile.tile.get_d_size() x = tile.tile.get_x_size() @@ -1057,8 +1120,8 @@ def test_get_tile_weights_ref_returns_transposed_view_for_cuda(self): """_get_tile_weights_ref should return a CUDA tensor in standard (d_size, x_size) layout. CUDA C++ tiles store weights in transposed layout (x_size, d_size). - _get_tile_weights_ref uses get_weights_cuda().t() so callers see the - standard (d_size, x_size) shape without a CPU round-trip. + _get_tile_weights_ref returns a standard (d_size, x_size) CUDA tensor, + either through shared weights or the optional get_weights_cuda fast path. """ self._skip_if_no_cuda() model = AnalogLinear(8, 4, bias=False, From e5413b89864822fbd9a036ec19887e04c4e2820a Mon Sep 17 00:00:00 2001 From: Zhaoxian Wu Date: Thu, 11 Jun 2026 14:24:26 -0400 Subject: [PATCH 6/6] feat(ctx): add three data view modes (PLACEHOLDER, DATA_VIEW, BUFFER) to AnalogContext Signed-off-by: Zhaoxian Wu --- src/aihwkit/optim/context.py | 280 ++++++- src/aihwkit/optim/weight_view.py | 120 ++- src/aihwkit/simulator/parameters/__init__.py | 1 + src/aihwkit/simulator/parameters/enums.py | 18 + src/aihwkit/simulator/tiles/rpucuda.py | 2 +- tests/test_analog_ctx.py | 747 ++++++++++++++++++- 6 files changed, 1113 insertions(+), 55 deletions(-) diff --git a/src/aihwkit/optim/context.py b/src/aihwkit/optim/context.py index e2c2aafc0..0fc0b8810 100644 --- a/src/aihwkit/optim/context.py +++ b/src/aihwkit/optim/context.py @@ -8,7 +8,7 @@ # pylint: disable=attribute-defined-outside-init -from typing import Optional, Type, Union, Any, TYPE_CHECKING +from typing import Optional, Type, Union, Any, List, TYPE_CHECKING from torch import dtype, Tensor from torch._C import DisableTorchFunction @@ -19,12 +19,24 @@ from aihwkit.optim.weight_view import ( ReadOnlyWeightView, raise_if_readonly_write_target, + PlaceholderDataView, + _PLACEHOLDER_METADATA_FUNCTIONS, + _raise_placeholder_read_error, ) +from aihwkit.simulator.parameters.enums import AnalogContextDataViewMode if TYPE_CHECKING: from aihwkit.simulator.tiles.base import SimulatorTileWrapper +# Tensor properties that materialize weight values (transpose / conjugate views +# and complex parts). They must honor the data-view mode like their method +# equivalents (``t()``, ``conj()``); otherwise the getset descriptors fall +# through to ``__torch_function__`` as the whitelisted ``__get__`` and silently +# return uninitialized placeholder memory in PLACEHOLDER mode. +_VALUE_VIEW_PROPERTIES = ("T", "mT", "H", "mH", "real", "imag") + + class AnalogContext(Parameter): """Context for analog optimizer. @@ -33,41 +45,69 @@ class AnalogContext(Parameter): If `analog_bias` (which is provided by `analog_tile`) is True, The last column of `data` is the `bias` term - Note: For diagnostic purposes, `AnalogContext` exposes a read-only logical weight view - through the `data` attribute, which is equivalent to `analog_tile.get_weights()[0]`. - This allows users to inspect the effective weights. - Direct tensor reads on ``ctx`` or ``ctx.data``, such as ``size()``, ``norm()`` are - equivalent to do so on ``ctx.analog_tile.get_weights()[0]``. - i.e, ctx.data == ctx.analog_tile.get_weights()[0] - - The `data` attribution inherited from `torch.nn.Parameter` stores the raw tile weights, - i.e., the weights without scaling - Its public tensor value is a read-only logical weight view: ``physical weights x scaling`` - # Example usage: + For diagnostic purposes, `AnalogContext` provides three public data view modes. + Consider the code: --- layer = AnalogLinear(4, 3, bias=False, rpu_config=rpu_config) analog_tile = layer.analog_module analog_ctx = analog_tile.analog_ctx weight = analog_tile.get_weights()[0] + --- + where `weight` is the logical weight view, which is already ``physical weights x scaling`` + + Data view modes are controlled by `analog_ctx.data_view_mode` and the corresponding methods: + --- + analog_ctx.enable_placeholder(). # PLACEHOLDER mode (default) + analog_ctx.enable_data_view(). # DATA_VIEW mode + analog_ctx.enable_buffer(). # BUFFER mode + --- - # The following two lines will print the same value: + * PLACEHOLDER (default): only metadata, such as ``size()``, ``shape``. + Since the RPU conductance values is not directly accessible in physic, the weight values, + as well as value-based operations, such as ``norm()``, are blocked by default + Access them raises ``RuntimeError``. + --- + # inspect metadata without reading values: + analog_ctx.size() # [4, 3] + analog_ctx.device() # 'cpu' + analog_ctx.norm() # RuntimeError + --- + * DATA_VIEW: exposes a read-only logical weight view through the `data` attribute, + which is equivalent to `analog_tile.get_weights()[0]`. + This allows users to inspect the effective weights. + Since the changes of both weights and scaling affect the logical weights, + we adopt the convetion that this logical view is read-only + Therefore, in-place operations, such as ``add_``, ``mul_``, etc, are blocked + --- + # The following three lines will print the same value: analog_ctx.size() analog_ctx.data.size() weight.size() + # Accessing values is allowed, but they are read-only: + analog_ctx.norm() # Successfully returns the norm + analog_ctx.norm() == weight.norm() # True + analog_ctx.add_(1.0) # RuntimeError + --- + * BUFFER: exposes a zero-initialized tensor with the logical weight shape through the `data` + At that mode, `data` is an independent buffer that is not connected to the analog tile. + It is intended for optimizers with digital auxiliary state, + such as mixed-precision training or TT-v2. --- - Since the changes of both weights and scaling affect the logical weights, - we adopt the convetion that this logical view is read-only - Therefore, in-place operations, such as ``add_``, ``mul_``, etc, are blocked - ctx.data.add_(1.0) # RuntimeError - Use the following update methods instead of + analog_ctx.norm() == weight.norm() # Typically False, since the buffer is independent + analog_ctx.add_(1.0) # Successfully adds 1.0 to the buffer, but does not + affect the analog tile weights + --- + + To update the internal analog weights, use the following update methods instead of writing `data` directly in the analog optimizer: --- analog_ctx.analog_tile.update(...) analog_ctx.analog_tile.update_indexed(...) --- - Even though it allows us to access the weights directly, always keep in mind that it is used - only for diagnostic purposes. To simulate the real reading, call the `read_weights` method + Caution: Even though DATA_VIEW mode allows us to access the weights directly, + always keep in mind that it is used only for diagnostic purposes. + To simulate the real reading, call the `read_weights` method instead, i.e. given `analog_ctx: AnalogContext`, estimated_weights, estimated_bias = analog_ctx.analog_tile.read_weights() """ @@ -101,6 +141,8 @@ def __init__( self.analog_tile = analog_tile self.use_torch_update = False self.use_indexed = False + self._data_view_mode = AnalogContextDataViewMode.PLACEHOLDER + self._data_buffer = None # type: Optional[Tensor] self.analog_input = [] # type: list self.analog_grad_output = [] # type: list self.reset(analog_tile) @@ -112,18 +154,34 @@ def __torch_function__( kwargs = kwargs or {} func_name = getattr(func, "__name__", "") + if func_name == "requires_grad_" and args and isinstance(args[0], AnalogContext): + # ``requires_grad_`` toggles the autograd flag, not weight values, so + # the read-only in-place guard below must not reject it and the data + # view redirection must not send it to a throwaway placeholder. + # ``nn.Module.requires_grad_`` calls this on every parameter to + # (un)freeze a layer, so route it straight to the real Parameter, + # mirroring the ``requires_grad`` attribute setter. + target = args[0] + requested = args[1] if len(args) > 1 else kwargs.get("requires_grad", True) + target.requires_grad = bool(requested) + return target + def is_readonly(value: Any) -> bool: - return isinstance(value, (AnalogContext, ReadOnlyWeightView)) + # BUFFER mode exposes an independent, writable digital buffer + # (used by mixed-precision optimizers), so in-place ops are allowed. + if isinstance(value, AnalogContext): + return value._get_data_view_mode() != AnalogContextDataViewMode.BUFFER + return isinstance(value, ReadOnlyWeightView) raise_if_readonly_write_target(func_name, args, kwargs, is_readonly) - def to_logical_tensor(value: Any) -> Any: + def to_public_tensor(value: Any) -> Any: if isinstance(value, AnalogContext): - return value._logical_data() + return value._torch_function_data(func_name) return value - args = tree_map(to_logical_tensor, args) - kwargs = tree_map(to_logical_tensor, kwargs) + args = tree_map(to_public_tensor, args) + kwargs = tree_map(to_public_tensor, kwargs) return func(*args, **kwargs) def __setitem__(self, key: Any, value: Any) -> None: @@ -133,11 +191,88 @@ def __setitem__(self, key: Any, value: Any) -> None: "Use analog_tile.set_weights() instead." ) + def __dir__(self) -> List[str]: + """List attribute names for interactive tab-completion. + + ``Tensor.__dir__`` dispatches through ``__torch_function__`` (this class + defines one), which in PLACEHOLDER mode routes ``__dir__`` through the + weight-read guard and raises. Completers (rlcompleter, IPython) swallow + that error and show nothing. Listing the class and instance attribute + names directly avoids the value-read dispatch, so ``analog_ctx.`` + offers the same tensor ops as a plain tensor in every data-view mode. + """ + keys = set(dir(type(self))) + keys.update(object.__getattribute__(self, "__dict__")) + return sorted(keys) + + @staticmethod + def _coerce_data_view_mode(value: Any) -> AnalogContextDataViewMode: + """Convert public mode inputs to ``AnalogContextDataViewMode``.""" + if isinstance(value, AnalogContextDataViewMode): + return value + if isinstance(value, str): + for mode in AnalogContextDataViewMode: + if value == mode.value or value.upper() == mode.name: + return mode + raise ValueError( + "data_view_mode must be an AnalogContextDataViewMode value, " + "or one of: placeholder, data_view, buffer." + ) + + def _get_data_view_mode(self) -> AnalogContextDataViewMode: + """Return the active public data view mode.""" + try: + return object.__getattribute__(self, "_data_view_mode") + except AttributeError: + return AnalogContextDataViewMode.PLACEHOLDER + + @property + def data_view_mode(self) -> AnalogContextDataViewMode: + """Return the active public data access mode.""" + return self._get_data_view_mode() + + @data_view_mode.setter + def data_view_mode(self, value: Any) -> None: + """Set the active public data access mode.""" + mode = self._coerce_data_view_mode(value) + self._data_view_mode = mode + if mode == AnalogContextDataViewMode.BUFFER: + self._data_buffer = self._new_data_buffer() + else: + self._data_buffer = None + + def enable_data_view(self) -> "AnalogContext": + """Enable read-only logical weight reads for diagnostics.""" + self.data_view_mode = AnalogContextDataViewMode.DATA_VIEW + return self + + def enable_placeholder(self) -> "AnalogContext": + """Enable metadata-only placeholder mode.""" + self.data_view_mode = AnalogContextDataViewMode.PLACEHOLDER + return self + + def enable_buffer(self) -> "AnalogContext": + """Enable an independent zero-initialized digital data buffer.""" + self.data_view_mode = AnalogContextDataViewMode.BUFFER + return self + def _raw_data(self) -> Tensor: """Return the internal raw tile backing tensor.""" with DisableTorchFunction(): # pylint: disable=not-context-manager return super().__getattribute__("data") + def _logical_shape(self) -> Any: + """Return the logical public weight shape without reading values.""" + raw = self._raw_data() + try: + analog_tile = object.__getattribute__(self, "analog_tile") + except AttributeError: + return raw.shape + + if getattr(analog_tile, "analog_bias", False) and raw.dim() >= 2: + return raw[:, : analog_tile.in_size].shape + return raw.shape + def _logical_data(self) -> Tensor: """Return logical weights equivalent to ``analog_tile.get_weights()[0]``.""" raw = self._raw_data() @@ -161,20 +296,80 @@ def _logical_data(self) -> Tensor: scales = scales.to(device=logical.device, dtype=logical.dtype) return logical * scales.view(-1, 1) + def _placeholder_data(self) -> PlaceholderDataView: + """Return a metadata-only public placeholder.""" + return PlaceholderDataView(self._raw_data().new_empty(self._logical_shape())) + + def _new_data_buffer(self) -> Tensor: + """Create a zero digital buffer with the logical public weight shape.""" + return self._raw_data().new_zeros(self._logical_shape()) + + def _buffer_data(self) -> Tensor: + """Return the independent digital buffer.""" + return object.__getattribute__(self, "_data_buffer") + + def _public_data(self) -> Tensor: + """Return the tensor exposed by the active public data mode.""" + mode = self._get_data_view_mode() + if mode == AnalogContextDataViewMode.PLACEHOLDER: + return self._placeholder_data() + if mode == AnalogContextDataViewMode.DATA_VIEW: + return ReadOnlyWeightView(self._logical_data()) + if mode == AnalogContextDataViewMode.BUFFER: + return self._buffer_data() + raise RuntimeError(f"Unsupported AnalogContext data view mode: {mode}") + + def _torch_function_data(self, func_name: str) -> Tensor: + """Return the tensor used to dispatch public torch operations.""" + mode = self._get_data_view_mode() + if mode == AnalogContextDataViewMode.PLACEHOLDER: + if func_name not in _PLACEHOLDER_METADATA_FUNCTIONS: + _raise_placeholder_read_error(func_name) + return self._placeholder_data() + if mode == AnalogContextDataViewMode.DATA_VIEW: + return self._logical_data() + if mode == AnalogContextDataViewMode.BUFFER: + return self._buffer_data() + raise RuntimeError(f"Unsupported AnalogContext data view mode: {mode}") + def __getattribute__(self, name: str) -> Any: - """Intercept public tensor reads that expose the logical view.""" + """Intercept public tensor reads according to ``data_view_mode``.""" if name == "grad_fn": return None - if name in ("device", "dtype", "grad", "is_cuda", "is_leaf", "layout"): + if name in ("device", "dtype", "is_cuda", "is_leaf", "layout"): return getattr(self._raw_data(), name) + if name == "requires_grad": + with DisableTorchFunction(): # pylint: disable=not-context-manager + return self.as_subclass(Tensor).requires_grad + if name == "grad": + # Mixed-precision optimizers SET param.grad (mpmixin.prepare_grad) and + # torch optimizers READ it; both must hit the Parameter's own grad slot. + # DisableTorchFunction bypasses the data-view dispatch that would + # otherwise redirect the read to the raw data view and return None. + with DisableTorchFunction(): # pylint: disable=not-context-manager + return super().__getattribute__("grad") if name == "data": - return ReadOnlyWeightView(self._logical_data()) + return self._public_data() if name == "shape": - return self._logical_data().shape + return self._logical_shape() if name == "ndim": - return self._logical_data().ndim + return len(self._logical_shape()) + if name in _VALUE_VIEW_PROPERTIES: + return self._value_view_property(name) return super().__getattribute__(name) + def _value_view_property(self, name: str) -> Any: + """Return a value-bearing view property honoring the data-view mode. + + ``T`` / ``mT`` / ``H`` / ``mH`` / ``real`` / ``imag`` read weight values, + so they follow the same rules as ``t()`` / ``conj()``: blocked in + PLACEHOLDER mode, served from the read-only logical view in DATA_VIEW + mode, and from the digital buffer in BUFFER mode. + """ + if self._get_data_view_mode() == AnalogContextDataViewMode.PLACEHOLDER: + _raise_placeholder_read_error(name) + return getattr(self._public_data(), name) + def __setattr__(self, name: str, value: Any) -> None: """Block user-level replacement of ``.data``.""" if name == "data": @@ -182,15 +377,34 @@ def __setattr__(self, name: str, value: Any) -> None: "Direct replacement of analog_ctx.data is not allowed. " "Use analog_tile.set_weights(new_weight) for programmatic writes." ) + if name in ("grad", "requires_grad"): + # Both must bypass the data-view dispatch: ``grad`` to hit the + # Parameter's own grad slot, and ``requires_grad`` because its + # setter would otherwise be routed through ``__torch_function__`` + # (as ``__set__``) and raise in PLACEHOLDER mode. Toggling + # ``requires_grad`` is how analog layers are frozen/unfrozen. + with DisableTorchFunction(): # pylint: disable=not-context-manager + super().__setattr__(name, value) + return super().__setattr__(name, value) def _replace_raw_data(self, data: Tensor) -> None: """Replace the internal raw ``Parameter.data`` for tile rebinding.""" - if isinstance(data, ReadOnlyWeightView): + if isinstance(data, (ReadOnlyWeightView, PlaceholderDataView)): data = data.as_subclass(Tensor) with DisableTorchFunction(): # pylint: disable=not-context-manager super().__setattr__("data", data) + if self._get_data_view_mode() != AnalogContextDataViewMode.BUFFER: + return + + buffer = object.__getattribute__(self, "_data_buffer") + logical_shape = self._logical_shape() + if buffer is not None and buffer.shape == logical_shape: + self._data_buffer = buffer.to(device=data.device, dtype=data.dtype) + else: + self._data_buffer = self._new_data_buffer() + # -- existing API --------------------------------------------------------- def set_indexed(self, value: bool = True) -> None: @@ -198,7 +412,9 @@ def set_indexed(self, value: bool = True) -> None: self.use_indexed = value def get_data(self) -> Tensor: - """Get a detached logical weight tensor.""" + """Get a detached tensor from the active public data view.""" + if self._get_data_view_mode() == AnalogContextDataViewMode.PLACEHOLDER: + _raise_placeholder_read_error("get_data") return self.data.detach() def reset(self, analog_tile: Optional["SimulatorTileWrapper"] = None) -> None: diff --git a/src/aihwkit/optim/weight_view.py b/src/aihwkit/optim/weight_view.py index ed350f9a3..6cb614af9 100644 --- a/src/aihwkit/optim/weight_view.py +++ b/src/aihwkit/optim/weight_view.py @@ -4,11 +4,12 @@ # # Licensed under the MIT license. See LICENSE file in the project root for details. -"""Read-only tensor view for analog tile weights.""" +"""Read-only and placeholder tensor views for analog tile weights.""" from typing import Any, Callable, Optional from torch import Tensor +from torch._C import DisableTorchFunction from torch.utils._pytree import tree_map @@ -91,3 +92,120 @@ def __setitem__(self, key: Any, value: Any) -> None: "Direct item assignment on analog weights is not allowed. " "Use analog_tile.set_weights() instead." ) + + +_PLACEHOLDER_METADATA_FUNCTIONS = { + "__get__", + "__len__", + "__reduce__", + "__reduce_ex__", + "detach", + "dim", + "element_size", + "is_contiguous", + "nelement", + "numel", + "size", + "storage_offset", + "stride", + "type", + "requires_grad_", + # Pure metadata queries: they depend only on dtype / device / layout, never + # on weight values, so they are safe to answer in placeholder mode. + "ndimension", + "get_device", + "data_ptr", + "dense_dim", + "sparse_dim", + "has_names", + "is_floating_point", + "is_complex", + "is_signed", + "is_conj", + "is_neg", + "is_pinned", + "is_shared", + "is_inference", + "is_distributed", + "_is_view", + "_is_zerotensor", +} + + +def _raise_placeholder_read_error(func_name: str) -> None: + """Raise the standard error for reads in placeholder mode.""" + operation = func_name or "unknown" + raise RuntimeError( + "AnalogContext data is in placeholder mode, so operation " + f"'{operation}' cannot read weight values. Call " + "analog_ctx.enable_data_view() for diagnostic logical weight reads, " + "or analog_ctx.enable_buffer() for an independent digital buffer." + ) + + +class PlaceholderDataView(Tensor): + """Tensor-shaped metadata placeholder for analog context data. + + The backing tensor is intentionally meaningless. Only metadata operations + are allowed; value reads raise ``RuntimeError``. + """ + + @staticmethod + def __new__(cls, data: Tensor) -> "PlaceholderDataView": + """Create a placeholder with the same tensor metadata as ``data``.""" + if isinstance(data, PlaceholderDataView): + return data + return Tensor._make_subclass(cls, data.detach()) + + @classmethod + def __torch_function__( + cls, func: Any, _types: Any, args: Any = (), kwargs: Optional[Any] = None + ) -> Any: + kwargs = kwargs or {} + func_name = getattr(func, "__name__", "") + + def is_placeholder(value: Any) -> bool: + return isinstance(value, PlaceholderDataView) + + if func_name.endswith("_") and not func_name.endswith("__"): + if args and is_placeholder(args[0]): + _raise_placeholder_read_error(func_name) + + if "out" in kwargs: + + def block_out_target(value: Any) -> Any: + if is_placeholder(value): + _raise_placeholder_read_error(func_name) + return value + + tree_map(block_out_target, kwargs["out"]) + + if func_name not in _PLACEHOLDER_METADATA_FUNCTIONS: + _raise_placeholder_read_error(func_name) + + if func_name == "detach": + source = args[0] + with DisableTorchFunction(): # pylint: disable=not-context-manager + return PlaceholderDataView(source.as_subclass(Tensor).detach()) + + def unwrap(value: Any) -> Any: + if isinstance(value, PlaceholderDataView): + return value.as_subclass(Tensor) + return value + + args = tree_map(unwrap, args) + kwargs = tree_map(unwrap, kwargs) + return func(*args, **kwargs) + + def __repr__(self) -> str: + """Return metadata summary without exposing meaningless backing values.""" + with DisableTorchFunction(): # pylint: disable=not-context-manager + raw = self.as_subclass(Tensor) + return ( + f"PlaceholderDataView(shape={raw.shape}, dtype={raw.dtype}, " + f"device={raw.device})" + ) + + def __setitem__(self, key: Any, value: Any) -> None: + """Block item assignment on the placeholder.""" + _raise_placeholder_read_error("__setitem__") diff --git a/src/aihwkit/simulator/parameters/__init__.py b/src/aihwkit/simulator/parameters/__init__.py index e6c5e5678..81c889070 100644 --- a/src/aihwkit/simulator/parameters/__init__.py +++ b/src/aihwkit/simulator/parameters/__init__.py @@ -10,6 +10,7 @@ # the simulator shared library is linked against. from .enums import ( + AnalogContextDataViewMode, RPUDataType, BoundManagementType, NoiseManagementType, diff --git a/src/aihwkit/simulator/parameters/enums.py b/src/aihwkit/simulator/parameters/enums.py index 0d152ea41..822910ba3 100644 --- a/src/aihwkit/simulator/parameters/enums.py +++ b/src/aihwkit/simulator/parameters/enums.py @@ -37,6 +37,24 @@ def as_torch(self) -> dtype: return _TORCH_DATA_TYPE_MAP[self.value] +class AnalogContextDataViewMode(Enum): + """Public data access mode for analog optimizer contexts.""" + + PLACEHOLDER = "placeholder" + """Expose tensor metadata only, without allowing direct weight reads.""" + + DATA_VIEW = "data_view" + """Expose a read-only logical weight view for diagnostics.""" + + BUFFER = "buffer" + """Expose an independent digital buffer with the logical weight shape. + + This mode is useful for optimizers that need digital auxiliary state, such + as mixed-precision training or TT-v2. The buffer is not synchronized with + the analog tile weights. + """ + + class BoundManagementType(Enum): """Bound management type. diff --git a/src/aihwkit/simulator/tiles/rpucuda.py b/src/aihwkit/simulator/tiles/rpucuda.py index 6c35a0612..639a92525 100644 --- a/src/aihwkit/simulator/tiles/rpucuda.py +++ b/src/aihwkit/simulator/tiles/rpucuda.py @@ -137,7 +137,7 @@ def cpu(self) -> "SimulatorTileWrapper": state_dict = self.__getstate__() for value in state_dict.values(): if isinstance(value, AnalogContext): - value.data = value.data.cpu() + value._replace_raw_data(value._raw_data().cpu()) self.__setstate__(state_dict) return self diff --git a/tests/test_analog_ctx.py b/tests/test_analog_ctx.py index 2cdf37aa5..0ffa6535e 100644 --- a/tests/test_analog_ctx.py +++ b/tests/test_analog_ctx.py @@ -7,8 +7,8 @@ # pylint: disable=too-many-locals, no-member """Tests for AnalogContext data attribution (PR #717). -Verifies that analog_ctx.data exposes the logical weight matrix used by the -tile, rather than being an empty scalar tensor or raw tile-only storage. +Verifies the public AnalogContext data modes: metadata-only placeholder, +read-only logical data view, and independent digital buffer. """ from unittest import SkipTest @@ -23,13 +23,14 @@ from aihwkit.nn.conversion import convert_to_analog from aihwkit.optim.context import AnalogContext from aihwkit.optim.weight_view import ReadOnlyWeightView +from aihwkit.simulator.parameters.enums import AnalogContextDataViewMode from aihwkit.simulator.configs import ( FloatingPointRPUConfig, InferenceRPUConfig, SingleRPUConfig, TorchInferenceRPUConfig, ) -from aihwkit.simulator.configs.devices import ConstantStepDevice +from aihwkit.simulator.configs.devices import ConstantStepDevice, SoftBoundsDevice from .helpers.decorators import parametrize_over_layers from .helpers.layers import Linear, LinearCuda, LinearMapped, LinearMappedCuda @@ -64,9 +65,10 @@ def test_ctx_data_shape_matches_weights(self): self.assertEqual(ctx.size(), Size([expected_rows, tile.in_size])) def test_ctx_data_values_match_tile_weights(self): - """analog_ctx.data must reflect the logical tile weights.""" + """data_view mode must reflect the logical tile weights.""" model = self.get_layer(in_features=4, out_features=6) tile = self._get_analog_tile(model) + tile.analog_ctx.enable_data_view() weights_from_tile, _ = tile.get_weights() ctx_data = tile.analog_ctx.data.detach().cpu() @@ -82,6 +84,7 @@ def test_ctx_norm_is_meaningful(self): # With randomly initialized weights, the norm should be > 0 # and should NOT be exactly 1.0 (which the old scalar ones(()) returned). + tile.analog_ctx.enable_data_view() norm_val = tile.analog_ctx.norm().item() self.assertGreater(norm_val, 0.0) @@ -91,6 +94,7 @@ def test_ctx_nonzero_works(self): tile = self._get_analog_tile(model) # With random initialization, most weights are nonzero. + tile.analog_ctx.enable_data_view() nz = tile.analog_ctx.nonzero() self.assertGreater(len(nz), 0) @@ -100,6 +104,7 @@ def test_ctx_comparison_ops(self): tile = self._get_analog_tile(model) # Weights are initialized near zero with std ~1, so most are < 10. + tile.analog_ctx.enable_data_view() mask = tile.analog_ctx > 10 self.assertIsInstance(mask, Tensor) self.assertEqual(mask.shape, tile.analog_ctx.shape) @@ -124,6 +129,7 @@ def test_ctx_after_set_weights(self): # ctx should still have a valid logical non-scalar shape and value. self.assertNotEqual(tile.analog_ctx.size(), Size([])) self.assertEqual(tile.analog_ctx.size(), Size([tile.out_size, tile.in_size])) + tile.analog_ctx.enable_data_view() self.assertTrue(allclose(tile.analog_ctx.data.detach().cpu(), tile.get_weights()[0])) @@ -315,6 +321,7 @@ def _test_sync(self, rpu_config, use_cuda): if use_cuda: model = model.cuda() tile = next(model.analog_tiles()) + tile.analog_ctx.enable_data_view() new_w = randn(6, 4) tile.set_weights(new_w, None) @@ -348,6 +355,7 @@ def test_sync_after_multiple_set_weights(self): """ctx should stay in sync after multiple consecutive set_weights.""" model = AnalogLinear(4, 6, bias=False, rpu_config=TorchInferenceRPUConfig()) tile = next(model.analog_tiles()) + tile.analog_ctx.enable_data_view() for _ in range(5): new_w = randn(6, 4) @@ -370,6 +378,7 @@ def test_sync_after_cuda_move(self): # Move to CUDA model.cuda() tile = next(model.analog_tiles()) + tile.analog_ctx.enable_data_view() w_cuda, _ = tile.get_weights() ctx_cuda = tile.analog_ctx.data.detach().cpu() self.assertTrue(allclose(ctx_cuda, w_cuda), @@ -541,6 +550,7 @@ def test_update_reflects_in_ctx_without_sync(self): "analog_ctx raw data should change after tile.update() without sync") # And it should match get_weights() + tile.analog_ctx.enable_data_view() w_from_tile = tile.tile.get_weights() self.assertTrue( allclose(tile.analog_ctx.data.detach(), w_from_tile), @@ -551,6 +561,7 @@ def test_set_weights_reflects_in_ctx_without_sync(self): tile = self._get_tile(FloatingPointRPUConfig()) new_w = randn(6, 4) tile.tile.set_weights(new_w) + tile.analog_ctx.enable_data_view() self.assertTrue( allclose(tile.analog_ctx.data.detach(), new_w), @@ -566,6 +577,7 @@ def test_multiple_updates_stay_in_sync(self): d = randn(4, 6) tile.tile.update(x, d, False) + tile.analog_ctx.enable_data_view() w_from_tile = tile.tile.get_weights() self.assertTrue( allclose(tile.analog_ctx.data.detach(), w_from_tile), @@ -621,8 +633,8 @@ def test_as_ref_true_reflects_update_cpp_tile(self): "C++ tile: shared weight ref should match get_weights() after update") -class AnalogCtxReadOnlyTest(ParametrizedTestCase): - """Tests for the always-read-only logical AnalogContext view.""" +class AnalogCtxDataViewModeTest(ParametrizedTestCase): + """Tests for AnalogContext public data view modes.""" use_cuda = False @@ -637,24 +649,74 @@ def _get_tile_and_ctx(self, model): tile = next(model.analog_tiles()) return tile, tile.analog_ctx - # -- public logical read view --------------------------------------------- + # -- placeholder mode ------------------------------------------------------- def test_ctx_has_no_readonly_flag(self): """The public readonly switch was removed.""" _, ctx = self._get_tile_and_ctx(self._make_model()) self.assertFalse(hasattr(ctx, "readonly")) - def test_ctx_data_is_always_readonly_view(self): - """analog_ctx.data should always be a ReadOnlyWeightView.""" + def test_default_mode_is_placeholder(self): + """AnalogContext should default to metadata-only placeholder mode.""" _, ctx = self._get_tile_and_ctx(self._make_model()) - self.assertIsInstance(ctx.data, ReadOnlyWeightView) - def test_read_ops_use_logical_data(self): + self.assertEqual(ctx.data_view_mode, AnalogContextDataViewMode.PLACEHOLDER) + self.assertNotIsInstance(ctx.data, ReadOnlyWeightView) + + def test_placeholder_allows_metadata_only(self): + """Placeholder mode should expose shape-like metadata only.""" + tile, ctx = self._get_tile_and_ctx(self._make_model()) + expected_shape = Size([tile.out_size, tile.in_size]) + + self.assertEqual(ctx.size(), expected_shape) + self.assertEqual(ctx.data.size(), expected_shape) + self.assertEqual(ctx.shape, expected_shape) + self.assertEqual(ctx.data.shape, expected_shape) + self.assertEqual(ctx.ndim, 2) + self.assertEqual(ctx.data.ndim, 2) + self.assertEqual(ctx.device, ctx._raw_data().device) + self.assertEqual(ctx.dtype, ctx._raw_data().dtype) + + def test_placeholder_blocks_value_reads(self): + """Placeholder mode should block operations that read weight values.""" + _, ctx = self._get_tile_and_ctx(self._make_model()) + + for op in [ + lambda: ctx.norm(), + lambda: ctx.data.norm(), + lambda: ctx.nonzero(), + lambda: ctx.data.nonzero(), + lambda: ctx > 10, + lambda: ctx.data > 10, + lambda: ctx.get_data(), + lambda: ctx.detach().norm(), + ]: + with self.subTest(op=op): + with self.assertRaises(RuntimeError): + op() + + def test_placeholder_state_dict_round_trip(self): + """Saving and loading should not require public weight reads.""" + source = self._make_model() + target = self._make_model() + + self.assertEqual( + next(source.analog_tiles()).analog_ctx.data_view_mode, + AnalogContextDataViewMode.PLACEHOLDER, + ) + target.load_state_dict(source.state_dict(), load_rpu_config=False) + + # -- data_view mode --------------------------------------------------------- + + def test_enable_data_view_uses_logical_data(self): """size, norm, detach, nonzero, comparisons all use logical weights.""" tile, ctx = self._get_tile_and_ctx(self._make_model()) + ctx.enable_data_view() tile.set_scales(2.0) logical_weight, _ = tile.get_weights() + self.assertEqual(ctx.data_view_mode, AnalogContextDataViewMode.DATA_VIEW) + self.assertIsInstance(ctx.data, ReadOnlyWeightView) self.assertEqual(ctx.size(), Size([6, 4])) self.assertEqual(ctx.shape, Size([6, 4])) self.assertTrue(allclose(ctx.detach().cpu(), logical_weight)) @@ -663,9 +725,10 @@ def test_read_ops_use_logical_data(self): mask = ctx > 10 self.assertEqual(mask.shape, ctx.shape) - def test_scaled_context_returns_logical_weight(self): + def test_scaled_context_returns_logical_weight_in_data_view(self): """Scaled tiles expose scaled logical weights, not raw tile weights.""" tile, ctx = self._get_tile_and_ctx(self._make_model()) + ctx.enable_data_view() tile.set_scales(2.0) logical_weight = tile.get_weights()[0] @@ -675,9 +738,10 @@ def test_scaled_context_returns_logical_weight(self): self.assertTrue(allclose(ctx.data.detach().cpu(), logical_weight)) self.assertFalse(allclose(ctx.data.detach().cpu(), raw_weight)) - def test_setting_scale_changes_ctx_data(self): + def test_setting_scale_changes_ctx_data_in_data_view(self): """Manually setting scales should update logical ctx.data values.""" tile, ctx = self._get_tile_and_ctx(self._make_model()) + ctx.enable_data_view() before = ctx.data.detach().clone() tile.set_scales(2.0) @@ -689,7 +753,7 @@ def test_setting_scale_changes_ctx_data(self): self.assertTrue(allclose(after.cpu(), w_from_tile)) def test_analog_bias_column_is_not_exposed(self): - """analog_bias=True should not expose the raw bias column through ctx.data.""" + """Public views should not expose the raw analog-bias column.""" tile, ctx = self._get_tile_and_ctx(self._make_model(bias=True, analog_bias=True)) self.assertTrue(tile.analog_bias) self.assertEqual(ctx.data.shape, Size([tile.out_size, tile.in_size])) @@ -701,13 +765,15 @@ def test_analog_bias_column_is_not_exposed(self): new_weight = randn(tile.out_size, tile.in_size) new_bias = randn(tile.out_size) tile.set_weights(new_weight, new_bias) + ctx.enable_data_view() self.assertTrue(allclose(ctx.data.detach().cpu(), tile.get_weights()[0])) - # -- direct writes are always blocked -------------------------------------- + # -- data_view direct writes stay blocked ---------------------------------- - def test_data_inplace_ops_blocked(self): - """ctx.data in-place ops should raise RuntimeError.""" + def test_data_view_data_inplace_ops_blocked(self): + """ctx.data in-place ops should raise RuntimeError in data_view mode.""" _, ctx = self._get_tile_and_ctx(self._make_model()) + ctx.enable_data_view() for op in [ lambda: ctx.data.add_(1.0), lambda: ctx.data.mul_(2.0), @@ -719,9 +785,10 @@ def test_data_inplace_ops_blocked(self): with self.assertRaises(RuntimeError): op() - def test_context_inplace_ops_blocked(self): + def test_data_view_context_inplace_ops_blocked(self): """Direct in-place ops on ctx should raise RuntimeError.""" _, ctx = self._get_tile_and_ctx(self._make_model()) + ctx.enable_data_view() ctx.requires_grad = False snapshot = ctx.data.detach().clone() @@ -734,6 +801,7 @@ def test_context_inplace_ops_blocked(self): def _assert_out_writes_blocked(self, ctx): """Assert that torch out= cannot target ctx or ctx.data.""" + ctx.enable_data_view() ctx.requires_grad = False raw_snapshot = ctx._raw_data().detach().clone() logical_snapshot = ctx.data.detach().clone() @@ -766,8 +834,9 @@ def test_out_kwarg_writes_blocked_with_scaling(self): self._assert_out_writes_blocked(ctx) def test_copy_from_context_to_external_tensor_allowed(self): - """copy_ may read from ctx or ctx.data when the destination is external.""" + """copy_ may read from ctx or ctx.data in data_view mode.""" _, ctx = self._get_tile_and_ctx(self._make_model()) + ctx.enable_data_view() expected = ctx.data.detach().clone() for source in [ctx, ctx.data]: @@ -776,9 +845,10 @@ def test_copy_from_context_to_external_tensor_allowed(self): destination.copy_(source) self.assertTrue(allclose(destination, expected)) - def test_setitem_blocked(self): + def test_setitem_blocked_in_data_view(self): """Item assignment through ctx or ctx.data should raise RuntimeError.""" _, ctx = self._get_tile_and_ctx(self._make_model()) + ctx.enable_data_view() snapshot = ctx.data.detach().clone() with self.assertRaises(RuntimeError): @@ -789,8 +859,9 @@ def test_setitem_blocked(self): self.assertTrue(allclose(ctx.data.detach(), snapshot)) def test_data_assignment_blocked(self): - """ctx.data rebinding should raise RuntimeError.""" + """ctx.data rebinding should raise RuntimeError in all modes.""" _, ctx = self._get_tile_and_ctx(self._make_model()) + ctx.enable_data_view() snapshot = ctx.data.detach().clone() with self.assertRaisesRegex(RuntimeError, "Direct replacement"): @@ -799,6 +870,57 @@ def test_data_assignment_blocked(self): self.assertIsInstance(ctx.data, ReadOnlyWeightView) self.assertTrue(allclose(ctx.data.detach(), snapshot)) + # -- buffer mode ----------------------------------------------------------- + + def test_buffer_mode_creates_zero_logical_tensor(self): + """buffer mode should create a zero tensor with logical weight shape.""" + tile, ctx = self._get_tile_and_ctx(self._make_model()) + ctx.enable_buffer() + + self.assertEqual(ctx.data_view_mode, AnalogContextDataViewMode.BUFFER) + self.assertNotIsInstance(ctx.data, ReadOnlyWeightView) + self.assertEqual(ctx.data.shape, Size([tile.out_size, tile.in_size])) + self.assertTrue(allclose(ctx.data, torch.zeros_like(ctx.data))) + + def test_buffer_inplace_ops_mutate_only_buffer(self): + """buffer writes should not modify analog tile weights.""" + tile, ctx = self._get_tile_and_ctx(self._make_model()) + weight_snapshot = tile.get_weights()[0].clone() + ctx.enable_buffer() + + ctx.data.add_(2.0) + self.assertTrue(allclose(ctx.data, torch.full_like(ctx.data, 2.0))) + self.assertTrue(allclose(tile.get_weights()[0], weight_snapshot)) + + def test_buffer_is_independent_after_mode_switches(self): + """Switching modes should not link the buffer to tile weights.""" + tile, ctx = self._get_tile_and_ctx(self._make_model()) + weight_snapshot = tile.get_weights()[0].clone() + + ctx.enable_buffer() + ctx.data.fill_(3.0) + self.assertTrue(allclose(ctx.data, torch.full_like(ctx.data, 3.0))) + + ctx.enable_data_view() + self.assertTrue(allclose(ctx.data.detach().cpu(), weight_snapshot)) + + ctx.enable_buffer() + self.assertTrue(allclose(ctx.data, torch.zeros_like(ctx.data))) + self.assertTrue(allclose(tile.get_weights()[0], weight_snapshot)) + + def test_buffer_analog_bias_uses_logical_shape(self): + """buffer mode should not include the raw analog-bias column.""" + tile, ctx = self._get_tile_and_ctx(self._make_model(bias=True, analog_bias=True)) + ctx.enable_buffer() + + self.assertEqual(ctx.data.shape, Size([tile.out_size, tile.in_size])) + self.assertEqual( + tile.analog_ctx._raw_data().shape, + Size([tile.out_size, tile.in_size + 1]), + ) + + # -- removed APIs ---------------------------------------------------------- + def test_public_write_escape_apis_removed(self): """Old direct-write escape hatches should not exist.""" _, ctx = self._get_tile_and_ctx(self._make_model()) @@ -1169,3 +1291,586 @@ def test_raw_tile_to_cuda_shared_binding(self): r2 = tile._get_tile_weights_ref() self.assertEqual(r1.data_ptr(), r2.data_ptr(), "CUDA: _get_tile_weights_ref should return same ptr") + + +class AnalogCtxOptimizerInteropTest(ParametrizedTestCase): + """Optimizer-facing AnalogContext behaviour. + + These guard the interception of ``grad``, ``requires_grad`` and the + BUFFER-mode in-place writes that mixed-precision / TT optimizers depend on. + Before the fix, ``grad`` reads were redirected to a transient ``.data`` view + (always ``None``), ``requires_grad`` reported ``False`` (so the context was + dropped from the optimizer param groups), and BUFFER-mode in-place ops were + blocked by the always-read-only guard. + """ + + use_cuda = False + + # (label, factory) pairs covering the Python inference tile and the C++ + # analog devices (SoftBounds). These behaviours are about the + # AnalogContext Python wrapper, so they must hold for every backing tile. + def _rpu_configs(self): + """Return fresh (label, rpu_config) pairs to parametrize over.""" + return [ + ("torch_inference", TorchInferenceRPUConfig()), + ("softbounds", SingleRPUConfig(device=SoftBoundsDevice())), + ] + + def _ctx(self, rpu_config, bias=False): + """Return the analog_ctx of a freshly built AnalogLinear.""" + model = AnalogLinear(4, 6, bias=bias, rpu_config=rpu_config) + return next(model.analog_tiles()).analog_ctx + + def test_requires_grad_is_true(self): + """analog_ctx.requires_grad must be True so optimizers keep the context.""" + for label, rpu_config in self._rpu_configs(): + with self.subTest(rpu_config=label): + ctx = self._ctx(rpu_config) + self.assertTrue(ctx.requires_grad) + + def test_grad_get_set_roundtrip(self): + """Setting analog_ctx.grad and reading it back must return the same tensor.""" + for label, rpu_config in self._rpu_configs(): + with self.subTest(rpu_config=label): + ctx = self._ctx(rpu_config) + grad = randn(Size(ctx.shape)) + + ctx.grad = grad + + self.assertIsNotNone(ctx.grad, "analog_ctx.grad read returned None after set") + self.assertTrue(allclose(ctx.grad, grad)) + + def test_grad_none_by_default(self): + """A fresh analog_ctx has no gradient yet (read does not raise).""" + for label, rpu_config in self._rpu_configs(): + with self.subTest(rpu_config=label): + ctx = self._ctx(rpu_config) + self.assertIsNone(ctx.grad) + + def test_buffer_mode_inplace_add_allowed(self): + """BUFFER mode exposes an independent writable digital buffer. + + In-place ops on the context (e.g. AdamW's ``mul_``/``addcdiv_``) must + write the buffer instead of raising the read-only error. + """ + for label, rpu_config in self._rpu_configs(): + with self.subTest(rpu_config=label): + ctx = self._ctx(rpu_config) + ctx.enable_buffer() + + # Buffer starts zero-initialised with the logical weight shape. + self.assertEqual(ctx.data.shape, Size([6, 4])) + self.assertEqual(float(ctx.data.sum()), 0.0) + + ctx.add_(2.0) + self.assertTrue(allclose(ctx.data, torch.full((6, 4), 2.0))) + + ctx.mul_(3.0) + self.assertTrue(allclose(ctx.data, torch.full((6, 4), 6.0))) + + def test_data_view_mode_inplace_still_blocked(self): + """Non-BUFFER modes keep the read-only guard on in-place writes.""" + for label, rpu_config in self._rpu_configs(): + with self.subTest(rpu_config=label): + ctx = self._ctx(rpu_config) + ctx.enable_data_view() + with self.assertRaises(RuntimeError): + ctx.add_(1.0) + + +class AnalogCtxPlaceholderReprTest(ParametrizedTestCase): + """Placeholder-mode metadata access must not read meaningless values.""" + + use_cuda = False + + def _rpu_configs(self): + """Return fresh (label, rpu_config) pairs to parametrize over.""" + return [ + ("torch_inference", TorchInferenceRPUConfig()), + ("softbounds", SingleRPUConfig(device=SoftBoundsDevice())), + ("constant_step", SingleRPUConfig(device=ConstantStepDevice())), + ] + + def _ctx(self, rpu_config): + model = AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + return next(model.analog_tiles()).analog_ctx + + def test_placeholder_data_repr_is_metadata_only(self): + """repr(analog_ctx.data) in placeholder mode returns a metadata summary. + + The old PlaceholderDataView had no __repr__, so the default tensor repr + tried to read values and raised RuntimeError. + """ + for label, rpu_config in self._rpu_configs(): + with self.subTest(rpu_config=label): + ctx = self._ctx(rpu_config) + ctx.enable_placeholder() + + text = repr(ctx.data) + self.assertIn("PlaceholderDataView", text) + self.assertIn("shape=", text) + # str() goes through the same __repr__ and must not raise either. + self.assertEqual(str(ctx.data), text) + + def test_placeholder_type_metadata_allowed(self): + """analog_ctx.type() is metadata and must work in placeholder mode. + + ``type`` was newly added to the placeholder metadata allow-list; the old + code raised the placeholder read error. + """ + for label, rpu_config in self._rpu_configs(): + with self.subTest(rpu_config=label): + ctx = self._ctx(rpu_config) + ctx.enable_placeholder() + + self.assertIn("Tensor", ctx.type()) + self.assertIn("Tensor", ctx.data.type()) + + +class AnalogCtxRequiresGradBackwardTest(ParametrizedTestCase): + """Toggling analog_ctx.requires_grad must keep backward() correct. + + Freezing/unfreezing an analog layer is done via ``ctx.requires_grad``. + The setter must work in the default PLACEHOLDER mode (its dispatch used to + be routed through __torch_function__ as ``__set__`` and raised a placeholder + read error), and backward() must remain runnable in both states: + + * requires_grad True -> the analog gradient trace is recorded (the analog + optimizer would update the tile). + * requires_grad False -> the analog layer is frozen (no trace recorded), yet + downstream digital parameters still receive gradients, so backward() does + not raise. + """ + + use_cuda = False + + # Analog-training configs whose backward records the analog gradient trace + # (i.e. ``has_gradient()`` is the freeze signal). Inference tiles use the + # torch-update path instead and are intentionally excluded here. + def _rpu_configs(self): + """Return fresh (label, rpu_config) pairs to parametrize over.""" + return [ + ("floating_point", FloatingPointRPUConfig()), + ("softbounds", SingleRPUConfig(device=SoftBoundsDevice())), + ("constant_step", SingleRPUConfig(device=ConstantStepDevice())), + ] + + @staticmethod + def _build(rpu_config): + """Build an analog layer with a learnable digital scale. + + ``learn_out_scaling`` adds a downstream digital parameter so the loss + always has a grad path, letting backward() run even when the analog + context is frozen. + """ + rpu_config.mapping.learn_out_scaling = True + model = AnalogLinear(4, 3, bias=False, rpu_config=rpu_config) + tile = next(model.analog_tiles()) + return model, tile.analog_ctx, tile.out_scaling_alpha + + @staticmethod + def _fwd_bwd(model, ctx, alpha): + """Run one forward/backward with a non-grad input; return signals. + + Returns ``(out.requires_grad, ctx.has_gradient(), alpha.grad, ctx.grad)``. + ``ctx.grad`` must stay None: AnalogFunction.backward returns None for the + context, so autograd never populates its ``.grad`` slot; the gradient + information is captured in the analog trace (``has_gradient()``) instead. + """ + ctx.reset() # clear the analog gradient trace + alpha.grad = None + x = randn(2, 4) # leaf input, requires_grad=False + out = model(x) + requires_grad = out.requires_grad + out.sum().backward() + return requires_grad, ctx.has_gradient(), alpha.grad, ctx.grad + + def test_set_requires_grad_in_default_mode_does_not_raise(self): + """Setting requires_grad in default (placeholder) mode must not raise. + + Regression: the setter dispatched through __torch_function__ as + ``__set__`` and raised the placeholder read error. + """ + for label, rpu_config in self._rpu_configs(): + with self.subTest(rpu_config=label): + _, ctx, _ = self._build(rpu_config) + + ctx.requires_grad = False + self.assertFalse(ctx.requires_grad) + + ctx.requires_grad = True + self.assertTrue(ctx.requires_grad) + + def test_backward_true_to_false_freezes_analog_gradient(self): + """True -> False: backward still runs, but analog trace stops recording.""" + for label, rpu_config in self._rpu_configs(): + with self.subTest(rpu_config=label): + model, ctx, alpha = self._build(rpu_config) + + # requires_grad True: analog gradient recorded, digital param learns. + # Both the forward signal (analog_input) and the backward signal + # (analog_grad_output) must be attached to the trace. + ctx.requires_grad = True + requires_grad, has_grad, alpha_grad, ctx_grad = self._fwd_bwd(model, ctx, alpha) + self.assertTrue(requires_grad) + self.assertTrue(has_grad, "forward signal must be recorded when requires_grad=True") + self.assertEqual( + len(ctx.analog_grad_output), 1, "backward signal must be attached" + ) + self.assertIsNotNone(alpha_grad) + # The analog context never uses the autograd .grad slot. + self.assertIsNone(ctx_grad, "ctx.grad must stay None; gradient lives in the trace") + + # requires_grad False: backward must still run without error, but the + # analog layer is frozen (neither signal recorded); the digital scale + # keeps learning. + ctx.requires_grad = False + _, has_grad, alpha_grad, ctx_grad = self._fwd_bwd(model, ctx, alpha) + self.assertFalse(has_grad, "forward signal must NOT be recorded when frozen") + self.assertEqual( + len(ctx.analog_grad_output), 0, "no backward signal when frozen" + ) + self.assertIsNotNone(alpha_grad, "downstream digital param must still get grad") + self.assertIsNone(ctx_grad, "ctx.grad must stay None when frozen") + + def test_backward_false_to_true_unfreezes_analog_gradient(self): + """False -> True: re-enabling restores analog gradient recording.""" + for label, rpu_config in self._rpu_configs(): + with self.subTest(rpu_config=label): + model, ctx, alpha = self._build(rpu_config) + + # requires_grad False: frozen, neither signal recorded. + ctx.requires_grad = False + _, has_grad, alpha_grad, ctx_grad = self._fwd_bwd(model, ctx, alpha) + self.assertFalse(has_grad) + self.assertEqual(len(ctx.analog_grad_output), 0, "no backward signal when frozen") + self.assertIsNotNone(alpha_grad) + self.assertIsNone(ctx_grad, "ctx.grad must stay None when frozen") + + # requires_grad True: both signals recorded again. + ctx.requires_grad = True + requires_grad, has_grad, alpha_grad, ctx_grad = self._fwd_bwd(model, ctx, alpha) + self.assertTrue(requires_grad) + self.assertTrue( + has_grad, "forward signal must resume after re-enabling requires_grad" + ) + self.assertEqual( + len(ctx.analog_grad_output), 1, "backward signal must resume after re-enabling" + ) + self.assertIsNotNone(alpha_grad) + # The analog context never uses the autograd .grad slot. + self.assertIsNone(ctx_grad, "ctx.grad must stay None; gradient lives in the trace") + + +class AnalogCtxExactGradientTest(ParametrizedTestCase): + """The analog backward must produce the exact gradient, including out_scaling. + + Uses a noise-free FloatingPoint tile with known weights and a known, non-unit + out_scaling vector ``alpha``, and a quadratic loss ``L = 0.5 * sum((y - t)^2)``. + Unlike a linear loss, ``dL/dy = (y - t)`` depends on the forward output, so the + backward must thread the out_scaling through correctly. With ``raw = x @ W^T`` + and ``y = alpha * raw`` the chain rule gives closed forms the test checks: + + * ``analog_input`` == ``x`` (forward activation) + * ``analog_grad_output`` == ``(y - t) * alpha`` (grad w.r.t. tile output) + * ``x.grad`` == ``((y - t) * alpha) @ W`` (grad w.r.t. the input) + * ``out_scaling.grad`` == ``sum_b (y - t) * raw`` (grad w.r.t. alpha) + + The ``alpha`` factor is the point: a backward that dropped or misapplied + out_scaling would still pass with ``alpha == 1`` but fail here. + """ + + use_cuda = False + + def test_exact_gradient_with_non_unit_out_scaling(self): + """Quadratic-loss backward gradients match the closed form when alpha != 1.""" + rpu_config = FloatingPointRPUConfig() + rpu_config.mapping.weight_scaling_omega = 0.0 # logical weight == set weight + rpu_config.mapping.learn_out_scaling = True + rpu_config.mapping.out_scaling_columnwise = True # per-output-channel alpha + + model = AnalogLinear(3, 2, bias=False, rpu_config=rpu_config) + tile = next(model.analog_tiles()) + ctx = tile.analog_ctx + + # ``weight`` is the physical tile weight (used by the forward MVM); + # ``get_weights`` returns the logical weight = physical * out_scaling. + weight = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + alpha = torch.tensor([0.5, 2.0]) # non-unit out_scaling, one per output + tile.set_weights(weight, None) + with torch.no_grad(): + tile.out_scaling_alpha.copy_(alpha.view_as(tile.out_scaling_alpha)) + + # Sanity: physical weight set cleanly (no hidden mapping scale), and the + # logical weight folds in out_scaling per output channel. + self.assertTrue(allclose(tile.get_weights()[0], weight * alpha.view(-1, 1))) + + ctx.reset() + x = torch.tensor([[1.0, 0.0, -1.0], [2.0, 1.0, 0.0]], requires_grad=True) + target = torch.tensor([[0.5, -1.0], [1.0, 2.0]]) + + out = model(x) + # Forward must be alpha * (x @ W^T). + raw = x.detach() @ weight.t() + self.assertTrue(allclose(out.detach(), raw * alpha), "forward out_scaling is wrong") + + # Quadratic (MSE-style) loss: dL/dy = (y - target). + loss = 0.5 * ((out - target) ** 2).sum() + loss.backward() + + residual = (raw * alpha) - target # dL/dy = y - target + + # Forward signal: the activation that entered the tile. + self.assertTrue(allclose(ctx.analog_input[0], x.detach())) + # Backward signal: gradient w.r.t. the (pre-scaling) tile output. + self.assertTrue( + allclose(ctx.analog_grad_output[0], residual * alpha), + "analog_grad_output must include the out_scaling factor", + ) + # Gradient w.r.t. the input flows back through W and alpha. + self.assertTrue( + allclose(x.grad, (residual * alpha) @ weight), + "input gradient must account for out_scaling", + ) + # Gradient w.r.t. alpha is sum over the batch of (y - target) * raw. + self.assertTrue( + allclose(tile.out_scaling_alpha.grad.flatten(), (residual * raw).sum(dim=0)), + "out_scaling gradient is wrong", + ) + + +# Configs covering the three backing-tile flavors that the AnalogContext Python +# wrapper sits on top of: a pure-Python inference tile, a C++ FloatingPoint tile, +# and a C++ analog-device tile. Attribute handling must hold for all of them. +def _ctx_attribute_rpu_configs(): + """Return fresh (label, rpu_config) pairs to parametrize attribute tests.""" + return [ + ("torch_inference", TorchInferenceRPUConfig()), + ("floating_point", FloatingPointRPUConfig()), + ("softbounds", SingleRPUConfig(device=SoftBoundsDevice())), + ] + + +class AnalogCtxTabCompletionTest(ParametrizedTestCase): + """dir(analog_ctx) must list tensor ops for interactive tab-completion. + + ``Tensor.__dir__`` dispatches through ``__torch_function__`` (this class + defines one); in PLACEHOLDER mode that routed ``__dir__`` through the + weight-read guard and raised, so completers (rlcompleter / IPython) caught + the error and offered nothing. ``AnalogContext.__dir__`` now lists the class + and instance attribute names directly, without a value-read dispatch. + """ + + use_cuda = False + + def _ctx(self, rpu_config): + model = AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + return next(model.analog_tiles()).analog_ctx + + def test_dir_does_not_raise_in_placeholder(self): + """dir(ctx) must not raise in the default placeholder mode.""" + for label, rpu_config in _ctx_attribute_rpu_configs(): + with self.subTest(rpu_config=label): + ctx = self._ctx(rpu_config) + self.assertEqual(ctx.data_view_mode, AnalogContextDataViewMode.PLACEHOLDER) + names = dir(ctx) # must not raise + self.assertIsInstance(names, list) + + def test_dir_lists_tensor_ops_and_ctx_methods(self): + """Completion must offer both tensor ops and AnalogContext-specific names.""" + ctx = self._ctx(TorchInferenceRPUConfig()) + names = set(dir(ctx)) + for op in ["reshape", "relu", "reciprocal", "norm", "size", "real", "T"]: + self.assertIn(op, names, f"tensor op '{op}' missing from dir()") + for member in ["analog_tile", "enable_data_view", "enable_buffer", "reset"]: + self.assertIn(member, names, f"AnalogContext member '{member}' missing from dir()") + + def test_dir_is_consistent_across_modes(self): + """The attribute listing must not depend on the active data-view mode.""" + ctx = self._ctx(TorchInferenceRPUConfig()) + placeholder_names = set(dir(ctx)) + ctx.enable_data_view() + self.assertEqual(set(dir(ctx)), placeholder_names) + ctx.enable_buffer() + self.assertEqual(set(dir(ctx)), placeholder_names) + + +class AnalogCtxMetadataQueryTest(ParametrizedTestCase): + """Pure metadata queries must work in every mode and return correct values. + + These depend only on dtype / device / layout, never on weight values, so + they belong in the placeholder allow-list. Before the fix, queries such as + ``get_device``, ``is_floating_point``, ``is_complex`` and ``is_signed`` + raised the placeholder read error. + """ + + use_cuda = False + + # Metadata methods whose value must equal a plain reference tensor of the + # same logical shape / dtype / device. ``data_ptr`` is excluded: it returns a + # storage address, which legitimately differs from the reference tensor. + METADATA_METHODS = [ + "dim", + "ndimension", + "numel", + "nelement", + "element_size", + "is_contiguous", + "stride", + "storage_offset", + "get_device", + "is_floating_point", + "is_complex", + "is_signed", + "is_pinned", + "dense_dim", + "sparse_dim", + "has_names", + ] + + def _ctx(self, rpu_config): + model = AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + return next(model.analog_tiles()).analog_ctx + + def test_metadata_queries_correct_in_all_modes(self): + """Each metadata query must match a reference tensor in every mode.""" + for label, rpu_config in _ctx_attribute_rpu_configs(): + with self.subTest(rpu_config=label): + ctx = self._ctx(rpu_config) + reference = torch.empty(ctx.size(), dtype=ctx.dtype, device=ctx.device) + + for enable in (ctx.enable_placeholder, ctx.enable_data_view, ctx.enable_buffer): + enable() + with self.subTest(mode=ctx.data_view_mode): + for method in self.METADATA_METHODS: + self.assertEqual( + getattr(ctx, method)(), + getattr(reference, method)(), + f"{method}() mismatch in {ctx.data_view_mode}", + ) + + def test_data_ptr_is_an_int_in_placeholder(self): + """data_ptr must answer (an integer address) instead of raising.""" + for label, rpu_config in _ctx_attribute_rpu_configs(): + with self.subTest(rpu_config=label): + ctx = self._ctx(rpu_config) + self.assertIsInstance(ctx.data_ptr(), int) + + +class AnalogCtxValueViewPropertyTest(ParametrizedTestCase): + """T / mT / H / mH / real / imag must honor the data-view mode. + + These tensor properties materialize weight values, so they must behave like + their method equivalents (``t()`` / ``conj()``): blocked in PLACEHOLDER mode, + and served from the logical view / buffer otherwise. Previously the getset + descriptors fell through to ``__torch_function__`` as the whitelisted + ``__get__`` and silently returned uninitialized placeholder memory. + """ + + use_cuda = False + + VALUE_VIEW_PROPERTIES = ["T", "mT", "H", "mH", "real", "imag"] + # ``imag`` is excluded from value checks: it is undefined for real dtypes and + # raises the same error as a plain real tensor would. + REAL_VALUED_PROPERTIES = ["T", "mT", "H", "mH", "real"] + + def _tile_ctx(self, rpu_config): + model = AnalogLinear(4, 6, bias=False, rpu_config=rpu_config) + tile = next(model.analog_tiles()) + return tile, tile.analog_ctx + + def test_placeholder_blocks_value_view_properties(self): + """Every value-bearing property must raise the placeholder read error.""" + for label, rpu_config in _ctx_attribute_rpu_configs(): + with self.subTest(rpu_config=label): + _, ctx = self._tile_ctx(rpu_config) + for prop in self.VALUE_VIEW_PROPERTIES: + with self.subTest(prop=prop): + with self.assertRaises(RuntimeError): + getattr(ctx, prop) + + def test_data_view_properties_match_logical_weight(self): + """In data_view mode the properties must reflect the logical weights.""" + for label, rpu_config in _ctx_attribute_rpu_configs(): + with self.subTest(rpu_config=label): + tile, ctx = self._tile_ctx(rpu_config) + ctx.enable_data_view() + logical, _ = tile.get_weights() + + for prop in self.REAL_VALUED_PROPERTIES: + with self.subTest(prop=prop): + ctx_view = getattr(ctx, prop).detach().cpu() + expected = getattr(logical, prop) + self.assertEqual(ctx_view.shape, expected.shape) + self.assertTrue( + allclose(ctx_view, expected), f"ctx.{prop} != logical.{prop}" + ) + + def test_buffer_properties_reflect_buffer(self): + """In buffer mode the properties must reflect the digital buffer.""" + for label, rpu_config in _ctx_attribute_rpu_configs(): + with self.subTest(rpu_config=label): + tile, ctx = self._tile_ctx(rpu_config) + ctx.enable_buffer() + ctx.data.add_(5.0) # buffer is now uniformly 5.0 + expected = torch.full((tile.out_size, tile.in_size), 5.0) + + self.assertTrue(allclose(ctx.T.detach().cpu(), expected.T)) + self.assertTrue(allclose(ctx.real.detach().cpu(), expected)) + + +class AnalogCtxRequiresGradMethodTest(ParametrizedTestCase): + """``requires_grad_()`` (method) must toggle the real autograd flag. + + Regression: ``requires_grad_`` was allow-listed but the read-only in-place + guard rejected it first (it ends with ``_``), and the data-view redirection + would have set the flag on a throwaway placeholder. Consequently the standard + PyTorch freezing idiom ``module.requires_grad_(False)`` -- which iterates the + parameters and calls ``param.requires_grad_()`` -- raised on the context. + """ + + use_cuda = False + + def _model_ctx(self, rpu_config, bias=False): + model = AnalogLinear(4, 6, bias=bias, rpu_config=rpu_config) + return model, next(model.analog_tiles()).analog_ctx + + def test_method_toggles_flag_and_returns_self(self): + """ctx.requires_grad_(flag) must set the flag and return the context.""" + for label, rpu_config in _ctx_attribute_rpu_configs(): + with self.subTest(rpu_config=label): + _, ctx = self._model_ctx(rpu_config) + + self.assertIs(ctx.requires_grad_(False), ctx) + self.assertFalse(ctx.requires_grad) + ctx.requires_grad_(True) + self.assertTrue(ctx.requires_grad) + + def test_method_defaults_to_true(self): + """ctx.requires_grad_() with no argument must enable gradients.""" + _, ctx = self._model_ctx(TorchInferenceRPUConfig()) + ctx.requires_grad = False + ctx.requires_grad_() + self.assertTrue(ctx.requires_grad) + + def test_module_requires_grad_freezes_analog_layer(self): + """nn.Module.requires_grad_() must (un)freeze the context without raising.""" + for label, rpu_config in _ctx_attribute_rpu_configs(): + with self.subTest(rpu_config=label): + model, ctx = self._model_ctx(rpu_config) + + model.requires_grad_(False) + self.assertFalse(ctx.requires_grad) + model.requires_grad_(True) + self.assertTrue(ctx.requires_grad) + + def test_method_works_in_every_mode(self): + """The flag toggle must not depend on the active data-view mode.""" + for enable in ("enable_placeholder", "enable_data_view", "enable_buffer"): + with self.subTest(mode=enable): + _, ctx = self._model_ctx(TorchInferenceRPUConfig()) + getattr(ctx, enable)() + ctx.requires_grad_(False) + self.assertFalse(ctx.requires_grad) + ctx.requires_grad_(True) + self.assertTrue(ctx.requires_grad)