apply_triu_tril_single Class — pytorch Architecture

Source Code

aten/src/ATen/native/TriangularOps.cpp lines 42–85

template <typename scalar_t>
void apply_triu_tril_single(
    scalar_t* result,
    const scalar_t* self,
    bool inplace,
    int64_t k,
    int64_t n,
    int64_t m,
    int64_t res_row_stride,
    int64_t res_col_stride,
    int64_t self_row_stride,
    int64_t self_col_stride,
    bool upper) {
  constexpr int64_t zero = 0;
  k = std::clamp(k, -n, m); // Clamp k to [-n, m] to prevent i + k arithmetic overflow, especially if k approaches INT64_MAX/INT64_MIN.

  if (upper) {
    parallel_for(0, n, 0, [&](int64_t start, int64_t end) {
      for (int64_t i : c10::irange(start, end)) {
        for (int64_t j = 0; j < std::min(m, i + k); j++) {
          result[i * res_row_stride + j * res_col_stride] = static_cast<scalar_t>(0);
        }
        if (!inplace) {  // copy the rest of the self if not inplace
          for (int64_t j = std::max(zero, i + k); j < m; j++) {
            result[i * res_row_stride + j * res_col_stride] = c10::load(&self[i * self_row_stride + j * self_col_stride]);
          }
        }
      }
    });
  } else {
    parallel_for(0, n, 0, [&](int64_t start, int64_t end) {
      for (int64_t i : c10::irange(start, end)) {
        for (int64_t j = std::max(zero, i + k + 1); j < m; j++) {
          result[i * res_row_stride + j * res_col_stride] = static_cast<scalar_t>(0);
        }
        if (!inplace) {  // copy the rest of the self if not inplace
          for (int64_t j = zero; j < std::min(m, i + k + 1); j++) {
            result[i * res_row_stride + j * res_col_stride] = c10::load(&self[i * self_row_stride + j * self_col_stride]);
          }
        }
      }
    });
  }
}