Skip to content

Commit

Permalink
Fix matrix data offset for large matrices (#4823)
Browse files Browse the repository at this point in the history
* Fix matrix data offset for large matrices

* Fix overflow in cudamatrix too
  • Loading branch information
nshmyrev authored Apr 26, 2023
1 parent 9a8588a commit 1918508
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 13 deletions.
14 changes: 7 additions & 7 deletions src/cudamatrix/cu-matrix.h
Original file line number Diff line number Diff line change
Expand Up @@ -231,7 +231,7 @@ class CuMatrixBase {
bool ApproxEqual(const CuMatrixBase<Real> &other, float tol = 0.01) const;

/// Get size of matrix in bytes
MatrixIndexT SizeInBytes() const { return num_rows_*stride_*sizeof(Real); }
size_t SizeInBytes() const { return static_cast<size_t>(num_rows_)*static_cast<size_t>(stride_)*sizeof(Real); }

// Copy functions. These do not resize.
template<typename OtherReal>
Expand Down Expand Up @@ -670,29 +670,29 @@ class CuMatrixBase {
inline const CuSubVector<Real> Row(MatrixIndexT i) const {
KALDI_ASSERT(static_cast<UnsignedMatrixIndexT>(i) <
static_cast<UnsignedMatrixIndexT>(num_rows_));
return CuSubVector<Real>(data_ + (i * stride_), NumCols());
return CuSubVector<Real>(data_ + (static_cast<size_t>(i) * static_cast<size_t>(stride_)), NumCols());
}

inline CuSubVector<Real> Row(MatrixIndexT i) {
KALDI_ASSERT(static_cast<UnsignedMatrixIndexT>(i) <
static_cast<UnsignedMatrixIndexT>(num_rows_));
return CuSubVector<Real>(data_ + (i * stride_), NumCols());
return CuSubVector<Real>(data_ + (static_cast<size_t>(i) * static_cast<size_t>(stride_)), NumCols());
}

inline CuValue<Real> operator() (MatrixIndexT r, MatrixIndexT c) {
KALDI_PARANOID_ASSERT(static_cast<UnsignedMatrixIndexT>(r) <
static_cast<UnsignedMatrixIndexT>(num_rows_) &&
static_cast<UnsignedMatrixIndexT>(c) <
static_cast<UnsignedMatrixIndexT>(num_cols_));
return CuValue<Real>(data_ + r * stride_ + c);
return CuValue<Real>(data_ + static_cast<size_t>(r) * static_cast<size_t>(stride_) + c);
}

inline Real operator() (MatrixIndexT r, MatrixIndexT c) const {
KALDI_PARANOID_ASSERT(static_cast<UnsignedMatrixIndexT>(r) <
static_cast<UnsignedMatrixIndexT>(num_rows_) &&
static_cast<UnsignedMatrixIndexT>(c) <
static_cast<UnsignedMatrixIndexT>(num_cols_));
return CuValue<Real>(data_ + r * stride_ + c); // will be casted to Real.
return CuValue<Real>(data_ + static_cast<size_t>(r) * static_cast<size_t>(stride_) + c); // will be casted to Real.
}

Real Sum() const;
Expand Down Expand Up @@ -737,10 +737,10 @@ class CuMatrixBase {

/// Get raw row pointer (const). Warning: may return a pointer to GPU memory. Use at
/// your own risk.
inline const Real* RowData(MatrixIndexT r) const { return data_ + r * stride_; }
inline const Real* RowData(MatrixIndexT r) const { return data_ + static_cast<size_t>(r) * static_cast<size_t>(stride_); }
/// Get raw row pointer. Warning: may return a pointer to GPU memory. Use at
/// your own risk.
inline Real* RowData(MatrixIndexT r) { return data_ + r * stride_; }
inline Real* RowData(MatrixIndexT r) { return data_ + static_cast<size_t>(r) * static_cast<size_t>(stride_); }
/// Return data pointer (const). Warning: may return a pointer to GPU memory.
/// Use at your own risk.
inline const Real *Data() const { return data_; }
Expand Down
12 changes: 6 additions & 6 deletions src/matrix/kaldi-matrix.h
Original file line number Diff line number Diff line change
Expand Up @@ -87,14 +87,14 @@ class MatrixBase {
inline Real* RowData(MatrixIndexT i) {
KALDI_ASSERT(static_cast<UnsignedMatrixIndexT>(i) <
static_cast<UnsignedMatrixIndexT>(num_rows_));
return data_ + i * stride_;
return data_ + static_cast<size_t>(i) * static_cast<size_t>(stride_);
}

/// Returns pointer to data for one row (const)
inline const Real* RowData(MatrixIndexT i) const {
KALDI_ASSERT(static_cast<UnsignedMatrixIndexT>(i) <
static_cast<UnsignedMatrixIndexT>(num_rows_));
return data_ + i * stride_;
return data_ + static_cast<size_t>(i) * static_cast<size_t>(stride_);
}

/// Indexing operator, non-const
Expand All @@ -104,7 +104,7 @@ class MatrixBase {
static_cast<UnsignedMatrixIndexT>(num_rows_) &&
static_cast<UnsignedMatrixIndexT>(c) <
static_cast<UnsignedMatrixIndexT>(num_cols_));
return *(data_ + r * stride_ + c);
return *(data_ + static_cast<size_t>(r) * static_cast<size_t>(stride_) + c);
}
/// Indexing operator, provided for ease of debugging (gdb doesn't work
/// with parenthesis operator).
Expand All @@ -117,7 +117,7 @@ class MatrixBase {
static_cast<UnsignedMatrixIndexT>(num_rows_) &&
static_cast<UnsignedMatrixIndexT>(c) <
static_cast<UnsignedMatrixIndexT>(num_cols_));
return *(data_ + r * stride_ + c);
return *(data_ + static_cast<size_t>(r) * static_cast<size_t>(stride_) + c);
}

/* Basic setting-to-special values functions. */
Expand Down Expand Up @@ -188,14 +188,14 @@ class MatrixBase {
inline const SubVector<Real> Row(MatrixIndexT i) const {
KALDI_ASSERT(static_cast<UnsignedMatrixIndexT>(i) <
static_cast<UnsignedMatrixIndexT>(num_rows_));
return SubVector<Real>(data_ + (i * stride_), NumCols());
return SubVector<Real>(data_ + (static_cast<size_t>(i) * static_cast<size_t>(stride_)), NumCols());
}

/// Return specific row of matrix.
inline SubVector<Real> Row(MatrixIndexT i) {
KALDI_ASSERT(static_cast<UnsignedMatrixIndexT>(i) <
static_cast<UnsignedMatrixIndexT>(num_rows_));
return SubVector<Real>(data_ + (i * stride_), NumCols());
return SubVector<Real>(data_ + (static_cast<size_t>(i) * static_cast<size_t>(stride_)), NumCols());
}

/// Return a sub-part of matrix.
Expand Down

0 comments on commit 1918508

Please sign in to comment.