normal_fill_VSX Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/DistributionTemplates.h lines 164–203

template <typename scalar_t, typename RNG>
void normal_fill_VSX(const TensorBase &self, const scalar_t mean, const scalar_t std, RNG generator) {
  float *data = self.data_ptr<float>();
  auto size = self.numel();
  std::lock_guard<std::mutex> lock(generator->mutex_);
  for (const auto i : c10::irange(size)) {
    at::uniform_real_distribution<scalar_t> uniform(0, 1);
    data[i] = uniform(generator);
  }

  using Vec = Vectorized<float>;
  const Vec two_pi = Vec(2.0f * c10::pi<double>);
  const Vec one = Vec(1.0f);
  const Vec minus_two = Vec(-2.0f);
  const Vec var_vec  = Vec(std);
  const Vec mean_vec = Vec(mean);

  for (int64_t i = 0; i < size - 15; i += 16) {
    if(Vec::size()==8) {
      normal_fill_16_VSX(data + i, two_pi, one, minus_two, mean_vec, var_vec);
    }
    else{
      normal_fill_16<scalar_t>(data + i, mean, std);
    }
  }
  if (size % 16 != 0) {
    // Recompute the last 16 values.
    data = data + size - 16;
    for (const auto i : c10::irange(16)) {
      at::uniform_real_distribution<scalar_t> uniform(0, 1);
      data[i] = uniform(generator);
    }
    if(Vec::size()==8){
      normal_fill_16_VSX(data, two_pi, one, minus_two, mean_vec, var_vec);
    }
    else{
      normal_fill_16<scalar_t>(data, mean, std);
    }
  }
}