#pragma once
|
#include <ATen/Utils.h>
|
#include <c10/util/ArrayRef.h>
|
|
#include <vector>
|
|
namespace at {
|
/// MatrixRef - Like an ArrayRef, but with an extra recorded strides so that
|
/// we can easily view it as a multidimensional array.
|
///
|
/// Like ArrayRef, this class does not own the underlying data, it is expected
|
/// to be used in situations where the data resides in some other buffer.
|
///
|
/// This is intended to be trivially copyable, so it should be passed by
|
/// value.
|
///
|
/// For now, 2D only (so the copies are actually cheap, without having
|
/// to write a SmallVector class) and contiguous only (so we can
|
/// return non-strided ArrayRef on index).
|
///
|
/// P.S. dimension 0 indexes rows, dimension 1 indexes columns
|
template<typename T>
|
class MatrixRef {
|
public:
|
typedef size_t size_type;
|
|
private:
|
/// Underlying ArrayRef
|
ArrayRef<T> arr;
|
|
/// Stride of dim 0 (outer dimension)
|
size_type stride0;
|
|
// Stride of dim 1 is assumed to be 1
|
|
public:
|
/// Construct an empty Matrixref.
|
/*implicit*/ MatrixRef() : arr(nullptr), stride0(0) {}
|
|
/// Construct an MatrixRef from an ArrayRef and outer stride.
|
/*implicit*/ MatrixRef(ArrayRef<T> arr, size_type stride0)
|
: arr(arr), stride0(stride0) {
|
TORCH_CHECK(arr.size() % stride0 == 0, "MatrixRef: ArrayRef size ", arr.size(), " not divisible by stride ", stride0)
|
}
|
|
/// @}
|
/// @name Simple Operations
|
/// @{
|
|
/// empty - Check if the matrix is empty.
|
bool empty() const { return arr.empty(); }
|
|
const T *data() const { return arr.data(); }
|
|
/// size - Get size a dimension
|
size_t size(size_t dim) const {
|
if (dim == 0) {
|
return arr.size() / stride0;
|
} else if (dim == 1) {
|
return stride0;
|
} else {
|
TORCH_CHECK(0, "MatrixRef: out of bounds dimension ", dim, "; expected 0 or 1");
|
}
|
}
|
|
size_t numel() const {
|
return arr.size();
|
}
|
|
/// equals - Check for element-wise equality.
|
bool equals(MatrixRef RHS) const {
|
return stride0 == RHS.stride0 && arr.equals(RHS.arr);
|
}
|
|
/// @}
|
/// @name Operator Overloads
|
/// @{
|
ArrayRef<T> operator[](size_t Index) const {
|
return arr.slice(Index*stride0, stride0);
|
}
|
|
/// Disallow accidental assignment from a temporary.
|
///
|
/// The declaration here is extra complicated so that "arrayRef = {}"
|
/// continues to select the move assignment operator.
|
template <typename U>
|
typename std::enable_if<std::is_same<U, T>::value, MatrixRef<T>>::type &
|
operator=(U &&Temporary) = delete;
|
|
/// Disallow accidental assignment from a temporary.
|
///
|
/// The declaration here is extra complicated so that "arrayRef = {}"
|
/// continues to select the move assignment operator.
|
template <typename U>
|
typename std::enable_if<std::is_same<U, T>::value, MatrixRef<T>>::type &
|
operator=(std::initializer_list<U>) = delete;
|
|
};
|
|
} // end namespace at
|
