feat(simulator): support infinite-granularity updates (dw_min=0)

src/aihwkit/simulator/configs/compounds.py

@@ -499,11 +499,38 @@ class TransferCompound(UnitCell):
     define the type of update used for each transfer event.
     """
 
+    def __post_init__(self) -> None:
+        self._assert_nonzero_dw_min()
+
+    def _assert_nonzero_dw_min(self) -> None:
+        """Reject infinite-granularity (``dw_min = 0``) sub-devices.
+
+        Transfer compounds scale the transfer learning rate and the buffer
+        granularity by the sub-device weight granularity (``dw_min``). A zero
+        ``dw_min`` (infinite-granularity mode) collapses these to zero (or
+        divides by zero), which silently disables learning. Infinite
+        granularity is therefore only supported for plain pulsed devices, not
+        for devices placed inside a transfer compound.
+
+        Raises:
+            ConfigError: if any ``unit_cell_device`` has ``dw_min == 0``.
+        """
+        for device in self.unit_cell_devices:
+            if getattr(device, "dw_min", None) == 0:
+                raise ConfigError(
+                    "dw_min=0 (infinite-granularity mode) is not supported for "
+                    "devices inside a transfer compound: the transfer learning "
+                    "rate and buffer granularity scale with the sub-device weight "
+                    "granularity, so dw_min=0 would disable learning."
+                )
+
     def as_bindings(self, data_type: RPUDataType) -> Any:
         """Return a representation of this instance as a simulator bindings object."""
         if not isinstance(self.unit_cell_devices, list):
             raise ConfigError("unit_cell_devices should be a list of devices")
 
+        self._assert_nonzero_dw_min()
+
         n_devices = len(self.unit_cell_devices)
 
         transfer_parameters = parameters_to_bindings(self, data_type)
@@ -775,6 +802,8 @@ def as_bindings(self, data_type: RPUDataType) -> Any:
         if not isinstance(self.unit_cell_devices, list):
             raise ConfigError("unit_cell_devices should be a list of devices")
 
+        self._assert_nonzero_dw_min()
+
         n_devices = len(self.unit_cell_devices)
         if n_devices != 2:
             raise ConfigError("Only 2 devices supported for ChoppedTransferCompound")
@@ -952,6 +981,8 @@ def as_bindings(self, data_type: RPUDataType) -> Any:
         if not isinstance(self.unit_cell_devices, list):
             raise ConfigError("unit_cell_devices should be a list of devices")
 
+        self._assert_nonzero_dw_min()
+
         n_devices = len(self.unit_cell_devices)
         if n_devices != 2:
             raise ConfigError("Only 2 devices supported for ChoppedTransferCompound")

src/aihwkit/simulator/configs/devices.py

@@ -164,7 +164,24 @@ class PulsedDevice(_PrintableMixin):
     """Parameter governing a power-law drift."""
 
     dw_min: float = 0.001
-    """Mean of the minimal update step sizes across devices and directions."""
+    r"""Mean of the minimal update step sizes across devices and directions.
+
+    Note:
+        Setting ``dw_min = 0`` activates **infinite-granularity** mode.
+        In this mode the stochastic pulse-train update is replaced by
+        its exact mean-field limit:
+
+        .. math::
+
+            w \leftarrow w - \mathrm{lr} \cdot (d^\top x) \cdot q(w)
+
+        where :math:`q(w)` is the device-specific response function
+        (the weight-dependent factor normally applied per pulse
+        coincidence), x and d are the forward and backward signals, respetively.
+        Cycle-to-cycle noise and ``dw_min``-related device-to-device variation,
+        like ``dw_min_dtod``, ``dw_min_std``, are dropped, while all other variation parameters,
+        like bounds, gamma, slopes are preserved.
+    """
 
     dw_min_dtod: float = 0.3
     """Device-to-device std deviation of ``dw_min`` (in relative units to

src/rpucuda/cuda/pulsed_weight_updater.cu

@@ -303,6 +303,44 @@ void PulsedWeightUpdater<T>::doDirectUpdate(
   context_->template releaseSharedBuffer<T>(RPU_BUFFER_OUT);
 }
 
+template <typename T>
+template <typename XInputIteratorT, typename DInputIteratorT>
+void PulsedWeightUpdater<T>::doInfiniteGranularityUpdate(
+    XInputIteratorT x_in,
+    DInputIteratorT d_in,
+    PulsedRPUDeviceCuda<T> *pulsed_device,
+    T *dev_weights,
+    const T lr,
+    const int m_batch,
+    const bool x_trans,
+    const bool d_trans) {
+
+  T *fpx_buffer = context_->template getSharedBuffer<T>(RPU_BUFFER_IN, x_size_ * m_batch);
+  T *fpd_buffer = context_->template getSharedBuffer<T>(RPU_BUFFER_OUT, d_size_ * m_batch);
+  T *grad_buffer = context_->template getSharedBuffer<T>(RPU_BUFFER_TEMP_0, x_size_ * d_size_);
+
+  const T *x_out = copyIterator2Buffer(x_in, fpx_buffer, x_size_ * m_batch);
+  const T *d_out = copyIterator2Buffer(d_in, fpd_buffer, d_size_ * m_batch);
+
+  // G = -lr * D^T @ X, written into grad_buffer (beta=0 overwrites).
+  // Uniform gemm path; for m_batch==1 this is a rank-1 update equivalent to ger.
+  RPU::math::gemm<T>(
+      context_, d_trans, !x_trans,
+      d_size_, // M
+      x_size_, // N
+      m_batch, // K
+      -lr, d_out, d_trans ? m_batch : d_size_, x_out, x_trans ? m_batch : x_size_, (T)0.0,
+      grad_buffer, d_size_);
+
+  const int total_size = x_size_ * d_size_;
+  pulsed_device->doInfiniteGranularityUpdate(
+      dev_weights, grad_buffer, context_->getRandomStates(total_size));
+
+  context_->template releaseSharedBuffer<T>(RPU_BUFFER_TEMP_0);
+  context_->template releaseSharedBuffer<T>(RPU_BUFFER_OUT);
+  context_->template releaseSharedBuffer<T>(RPU_BUFFER_IN);
+}
+
 template <typename T>
 bool PulsedWeightUpdater<T>::checkForFPUpdate(
     AbstractRPUDeviceCuda<T> *rpucuda_device_in, const PulsedUpdateMetaParameter<T> &up) {
@@ -358,6 +396,16 @@ void PulsedWeightUpdater<T>::update(
   // safe because of isPulsedDevice
   PulsedRPUDeviceCudaBase<T> *rpucuda_device =
       static_cast<PulsedRPUDeviceCudaBase<T> *>(rpucuda_device_in);
+
+  // Infinite granularity mode (dw_min=0): exact mean-field update
+  if (rpucuda_device->getWeightGranularity() <= (T)0.0) {
+    auto *pulsed_device = static_cast<PulsedRPUDeviceCuda<T> *>(rpucuda_device);
+    T abs_lr = lr < (T)0.0 ? -lr : lr;
+    doInfiniteGranularityUpdate(
+        x_in, d_in, pulsed_device, dev_weights, abs_lr, m_batch, x_trans, d_trans);
+    return;
+  }
+
   bool force_tuning = false;
 
   // check need for init (or re-init)
@@ -428,6 +476,9 @@ void PulsedWeightUpdater<T>::update(
   template void PulsedWeightUpdater<NUM_T>::doDirectUpdate(                                        \
       XITERT, DITERT, AbstractRPUDeviceCuda<NUM_T> *, NUM_T *, const NUM_T,                        \
       const PulsedUpdateMetaParameter<NUM_T> &, const int, const bool, const bool, const NUM_T);   \
+  template void PulsedWeightUpdater<NUM_T>::doInfiniteGranularityUpdate(                           \
+      XITERT, DITERT, PulsedRPUDeviceCuda<NUM_T> *, NUM_T *, const NUM_T, const int, const bool,   \
+      const bool);                                                                                 \
   template void PulsedWeightUpdater<NUM_T>::tuneUpdate(                                            \
       pwukp_t<NUM_T> &, pwukpvec_t<NUM_T> &, XITERT, DITERT, NUM_T *,                              \
       PulsedRPUDeviceCudaBase<NUM_T> *, const PulsedUpdateMetaParameter<NUM_T> &, const NUM_T,     \

src/rpucuda/cuda/pulsed_weight_updater.h

@@ -62,6 +62,17 @@ template <typename T> class PulsedWeightUpdater {
       const bool x_trans,
       const bool d_trans,
       const T beta = (T)1.0);
+
+  template <typename XInputIteratorT, typename DInputIteratorT>
+  void doInfiniteGranularityUpdate(
+      XInputIteratorT x_in,
+      DInputIteratorT d_in,
+      PulsedRPUDeviceCuda<T> *pulsed_device,
+      T *dev_weights,
+      const T lr,
+      const int m_batch,
+      const bool x_trans,
+      const bool d_trans);
   void setVerbosityLevel(int level) { verbose_ = level; };
   void dumpExtra(RPU::state_t &extra, const std::string prefix);
   void loadExtra(const RPU::state_t &extra, const std::string prefix, bool strict);

src/rpucuda/cuda/rpucuda_constantstep_device.cu

@@ -114,6 +114,54 @@ pwukpvec_t<T> ConstantStepRPUDeviceCuda<T>::getUpdateKernels(
 
 #undef ARGS
 
+/*********************************************************************************/
+/* infinite granularity update */
+
+namespace {
+constexpr int IG_THREADS_PER_BLOCK = 256;
+
+inline int ig_num_blocks(int total_size) {
+  return (total_size + IG_THREADS_PER_BLOCK - 1) / IG_THREADS_PER_BLOCK;
+}
+
+template <typename T>
+__global__ void kernelIGConstantStep(
+    T *weights, const T *grad_matrix, const param_t *params, int total_size) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= total_size) {
+    return;
+  }
+
+  T G = grad_matrix[idx];
+  if (G == (T)0.0) {
+    return;
+  }
+
+  param4_t p = reinterpret_cast<const param4_t *>(params)[idx];
+  T wmin = p.x;
+  T scale_down = p.y;
+  T wmax = p.z;
+  T scale_up = p.w;
+
+  T response = (G < (T)0.0) ? scale_down : scale_up;
+  T w = weights[idx];
+  w += G * response;
+  w = (w > wmax) ? wmax : w;
+  w = (w < wmin) ? wmin : w;
+  weights[idx] = w;
+}
+} // namespace
+
+template <typename T>
+void ConstantStepRPUDeviceCuda<T>::doInfiniteGranularityUpdate(
+    T *dev_weights, const T *grad_matrix, curandState_t *dev_states) {
+  UNUSED(dev_states);
+  int total_size = this->x_size_ * this->d_size_;
+  kernelIGConstantStep<T>
+      <<<ig_num_blocks(total_size), IG_THREADS_PER_BLOCK, 0, this->context_->getStream()>>>(
+          dev_weights, grad_matrix, this->get4ParamsData(), total_size);
+}
+
 template class ConstantStepRPUDeviceCuda<float>;
 #ifdef RPU_USE_DOUBLE
 template class ConstantStepRPUDeviceCuda<double>;

src/rpucuda/cuda/rpucuda_constantstep_device.h

@@ -55,6 +55,8 @@ template <typename T> class ConstantStepRPUDeviceCuda : public PulsedRPUDeviceCu
       int use_bo64,
       bool out_trans,
       const PulsedUpdateMetaParameter<T> &up) override;
+  void doInfiniteGranularityUpdate(
+      T *dev_weights, const T *grad_matrix, curandState_t *dev_states) override;
 };
 
 } // namespace RPU

src/rpucuda/cuda/rpucuda_expstep_device.cu

@@ -207,6 +207,103 @@ pwukpvec_t<T> ExpStepRPUDeviceCuda<T>::getUpdateKernels(
   return v;
 }
 
+/*********************************************************************************/
+/* infinite granularity update */
+
+namespace {
+constexpr int IG_THREADS_PER_BLOCK = 256;
+
+inline int ig_num_blocks(int total_size) {
+  return (total_size + IG_THREADS_PER_BLOCK - 1) / IG_THREADS_PER_BLOCK;
+}
+
+template <typename T>
+__device__ __forceinline__ void igStoreWeight(
+    T *weights,
+    T *persistent_weights,
+    const T *write_noise_std,
+    curandState_t *random_states,
+    int idx,
+    T w) {
+  if (persistent_weights == nullptr) {
+    weights[idx] = w;
+    return;
+  }
+
+  persistent_weights[idx] = w;
+  T uw_std = write_noise_std == nullptr ? (T)0.0 : *write_noise_std;
+  if (uw_std > (T)0.0 && random_states != nullptr) {
+    curandState local_state = random_states[idx];
+    weights[idx] = w + uw_std * (T)curand_normal(&local_state);
+    random_states[idx] = local_state;
+  } else {
+    weights[idx] = w;
+  }
+}
+
+template <typename T>
+__global__ void kernelIGExpStep(
+    T *weights,
+    const T *grad_matrix,
+    const param_t *params4,
+    const T *es_par,
+    int total_size,
+    T *persistent_weights,
+    curandState_t *random_states) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= total_size) {
+    return;
+  }
+
+  T G = grad_matrix[idx];
+  if (G == (T)0.0) {
+    return;
+  }
+
+  param4_t p4 = reinterpret_cast<const param4_t *>(params4)[idx];
+  T wmin = p4.x;
+  T wmax = p4.z;
+  T b_diff = wmax - wmin;
+  if (b_diff <= (T)0.0) {
+    return;
+  }
+
+  T A_down = es_par[0];
+  T A_up = es_par[1];
+  T gamma_down = es_par[2];
+  T gamma_up = es_par[3];
+  T a = es_par[4];
+  T b = es_par[5];
+
+  T w = persistent_weights == nullptr ? weights[idx] : persistent_weights[idx];
+  T abs_G = (G > (T)0.0) ? G : -G;
+  T z = (T)2.0 * w / b_diff * a + b;
+  T dw;
+  if (G < (T)0.0) {
+    T y = (T)1.0 - A_down * __expf(-gamma_down * z);
+    dw = (y > (T)0.0) ? -y * (T)p4.y * abs_G : (T)0.0;
+  } else {
+    T y = (T)1.0 - A_up * __expf(gamma_up * z);
+    dw = (y > (T)0.0) ? y * (T)p4.w * abs_G : (T)0.0;
+  }
+
+  w += dw;
+  w = (w > wmax) ? wmax : w;
+  w = (w < wmin) ? wmin : w;
+  igStoreWeight(weights, persistent_weights, es_par + 6, random_states, idx, w);
+}
+} // namespace
+
+template <typename T>
+void ExpStepRPUDeviceCuda<T>::doInfiniteGranularityUpdate(
+    T *dev_weights, const T *grad_matrix, curandState_t *dev_states) {
+  int total_size = this->x_size_ * this->d_size_;
+  kernelIGExpStep<T>
+      <<<ig_num_blocks(total_size), IG_THREADS_PER_BLOCK, 0, this->context_->getStream()>>>(
+          dev_weights, grad_matrix, this->get4ParamsData(), this->getGlobalParamsData(),
+          total_size, this->get1ParamsData(), dev_states);
+}
+
 template class ExpStepRPUDeviceCuda<float>;
 #ifdef RPU_USE_DOUBLE
 template class ExpStepRPUDeviceCuda<double>;

src/rpucuda/cuda/rpucuda_expstep_device.h

@@ -42,7 +42,7 @@ template <typename T> class ExpStepRPUDeviceCuda : public PulsedRPUDeviceCuda<T>
       es_par_arr[3] = par.es_gamma_up;
       es_par_arr[4] = par.es_a;
       es_par_arr[5] = par.es_b;
-      es_par_arr[6] = par.getScaledWriteNoise();
+      es_par_arr[6] = par.write_noise_std * (par.dw_min == (T)0.0 ? (T)1.0 : par.dw_min);
       es_par_arr[7] = par.dw_min_std_add;
       es_par_arr[8] = par.dw_min_std_slope;
       dev_es_par_->assign(es_par_arr);
@@ -65,6 +65,9 @@ template <typename T> class ExpStepRPUDeviceCuda : public PulsedRPUDeviceCuda<T>
                : PulsedRPUDeviceCuda<T>::getWeightGranularityNoise();
   }
 
+  void doInfiniteGranularityUpdate(
+      T *dev_weights, const T *grad_matrix, curandState_t *dev_states) override;
+
 private:
   std::unique_ptr<CudaArray<T>> dev_es_par_ = nullptr;
 };