backward_down_column_cdist Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/DistanceOpsKernel.cpp lines 392–414

  template <typename F>
  static void backward_down_column_cdist(const scalar_t * t1, const scalar_t * t2, scalar_t * res, const scalar_t * grad_k, const scalar_t * dist_k, const Vec& pvec, int64_t r1, int64_t r2, int64_t m, int64_t d, int64_t gs, int64_t l1_size, int64_t l2_size, int64_t count = Vec::size()) {
    const scalar_t * t1_end = t1 + l1_size;
    const scalar_t * t2_end = t2 + l2_size;

    for ([[maybe_unused]] const auto l : c10::irange(d)) {
      for (; t1 != t1_end; t1 += m, res += m) {
        const Vec vec_t1 = Vec::loadu(t1, count);
        Vec res_vec = Vec::loadu(res, count);

        for (const scalar_t * t2_curr = t2; t2_curr != t2_end; t2_curr += m, grad_k += gs, dist_k += 1) {
          const Vec vec_t2 = Vec::loadu(t2_curr, count);
          Vec res = F::backward(vec_t1 - vec_t2, *grad_k, *dist_k, pvec);
          res_vec = res_vec + res;
        }

        res_vec.store(res, count);
      }
      t1_end += l1_size;
      t2_end += l2_size;
      t2 += l2_size;
    }
  }