# -*- coding: utf-8 -*-
# (C) Copyright 2020, 2021, 2022, 2023, 2024 IBM. All Rights Reserved.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
# pylint: disable=too-many-instance-attributes
"""Mapping parameters for resistive processing units."""
from typing import Type, Optional
from dataclasses import dataclass, fields, field
from aihwkit.exceptions import ConfigError
from .base import RPUConfigBase
from .helpers import _PrintableMixin
[docs]@dataclass
class MappingParameter(_PrintableMixin):
"""Parameter related to hardware design and the mapping of logical
weight matrices to physical tiles.
Caution:
Some of these parameters have only an effect for modules that
support tile mappings.
"""
digital_bias: bool = True
"""Whether the bias term is handled by the analog tile or kept in
digital.
Note:
Default is having a *digital* bias so that bias values are
*not* stored onto the analog crossbar. This needs to be
supported by the chip design. Set to False if the analog bias
is instead situated on the the crossbar itself (as an extra
column)
Note:
``digital_bias`` is supported by *all* analog modules.
"""
weight_scaling_omega: float = 0.0
"""omega_scale is a user defined parameter used to scale the weights
while remapping these to cover the full range of values allowed.
By default, no remapping is performed. If values > 0.0 are supplied
the abs-max of the weight is scaled to that value.
"""
weight_scaling_columnwise: bool = False
"""Whether the weight matrix will be remapped column-wise over
the maximum device allowed value."""
weight_scaling_lr_compensation: bool = False
"""Whether to adjust the LR to compensate for the mapping factors
that are not learned.
The learning rate will be divided
for a tile individually by the mean of the mapping scales that are
determined by the ``weight_scaling_omega`` setting.
Otherwise the gradient information will be divided isntead before
the update.
"""
learn_out_scaling: bool = False
"""Define (additional) out scales that are learnable parameter
used to scale the output."""
out_scaling_columnwise: bool = False
"""Whether the learnable out scaling parameter enabled by
``learn_out_scaling`` is a scalar (``False``) or learned for
each output (``True``).
"""
max_input_size: int = 512
"""Maximal input size (number of columns) of the weight matrix
that is handled on a single analog tile.
If the logical weight matrix size exceeds this size it will be
split and mapped onto multiple analog tiles.
Caution:
Only relevant for ``Mapped`` modules such as
:class:`aihwkit.nn.modules.linear_mapped.AnalogLinearMapped`.
"""
max_output_size: int = 512
"""Maximal output size (number of rows) of the weight matrix
that is handled on a single analog tile.
If the logical weight matrix size exceeds this size it will be
split and mapped onto multiple analog tiles.
Caution:
Only relevant for ``Mapped`` modules such as
:class:`aihwkit.nn.modules.linear_mapped.AnalogLinearMapped`.
"""
[docs] def compatible_with(self, mapping: "MappingParameter") -> bool:
"""Checks compatiblity
Args:
mapping: param to check
Returns:
success: if compatible
"""
if mapping == self:
return True
for key in fields(mapping):
if key.name in [
"weight_scaling_omega",
"weight_scaling_columnwise",
"weight_scaling_lr_compensation",
]:
continue
if mapping.__dict__[key.name] != self.__dict__[key.name]:
return False
return True
[docs]@dataclass
class MappableRPU(RPUConfigBase, _PrintableMixin):
"""Defines the mapping parameters and utility factories"""
tile_array_class: Optional[Type] = None
"""Tile array class that correspond to the RPUConfig.
This is used to build logical arrays of tiles. Needs to be defined
in the derived class.
"""
mapping: MappingParameter = field(default_factory=MappingParameter)
"""Parameter related to mapping weights to tiles for supporting modules."""
[docs] def get_default_tile_module_class(self, out_size: int = 0, in_size: int = 0) -> Type:
"""Returns the default TileModule class.
Args:
out_size: overall output size
in_size: overall output size
Raises:
ConfigError: in case tile array is not defined.
"""
if self.tile_array_class is None:
ConfigError("RPUConfig does not support any tile array class")
if self.tile_array_class is None or (
self.mapping.max_input_size == 0 and self.mapping.max_output_size == 0
):
return self.tile_class
if self.mapping.max_input_size < in_size or self.mapping.max_output_size < out_size:
return self.tile_array_class
return self.tile_class