test(base): expand source_base unit-test coverage

source/source_base/module_container/ATen/core/tensor_buffer.cpp

@@ -1,26 +1,33 @@
 #include <ATen/core/tensor_buffer.h>
-
 #include <base/core/cpu_allocator.h>
 #include <base/macros/macros.h>
 
 #if defined(__CUDA) || defined(__ROCM)
 #include <base/core/gpu_allocator.h>
 #endif
 
-namespace container {
+namespace container
+{
 
 // Construct a new TensorBuffer object.
-TensorBuffer::TensorBuffer(base::core::Allocator* alloc, void* data_ptr) : alloc_(alloc), data_(data_ptr), owns_memory_(true) {}
+TensorBuffer::TensorBuffer(base::core::Allocator* alloc, void* data_ptr)
+    : alloc_(alloc), data_(data_ptr), owns_memory_(true)
+{
+}
 
 // Construct a new TensorBuffer object.
 // Note, this is a reference TensorBuffer, does not own memory itself.
-TensorBuffer::TensorBuffer(void* data_ptr) : alloc_(), data_(data_ptr), owns_memory_(false) {}
+TensorBuffer::TensorBuffer(void* data_ptr) : alloc_(), data_(data_ptr), owns_memory_(false)
+{
+}
 
-// Class members are initialized in the order of their declaration, 
+// Class members are initialized in the order of their declaration,
 // rather than the order they appear in the initialization list!
-TensorBuffer::TensorBuffer(base::core::Allocator* alloc, size_t size) {
-    alloc_ = alloc; 
-    if (size > 0) {
+TensorBuffer::TensorBuffer(base::core::Allocator* alloc, size_t size)
+{
+    alloc_ = alloc;
+    if (size > 0)
+    {
         data_ = alloc_->allocate(size);
         owns_memory_ = true;
         allocated_bytes_ = size;
@@ -29,106 +36,154 @@ TensorBuffer::TensorBuffer(base::core::Allocator* alloc, size_t size) {
 
 // Move constructor.
 TensorBuffer::TensorBuffer(TensorBuffer&& other) noexcept
-        : alloc_(other.alloc_),
-          data_(other.data_), 
-          owns_memory_(other.owns_memory_),
-          allocated_bytes_(other.allocated_bytes_)
+    : alloc_(other.alloc_), data_(other.data_), owns_memory_(other.owns_memory_),
+      allocated_bytes_(other.allocated_bytes_)
 {
     // Reset the other TensorBuffer.
+    other.alloc_ = nullptr;
     other.data_ = nullptr;
     other.owns_memory_ = false;
     other.allocated_bytes_ = 0;
 }
 
 // Destroy the TensorBuffer object.
-TensorBuffer::~TensorBuffer() {
-    if (this->OwnsMemory() && data_ != nullptr) {
+TensorBuffer::~TensorBuffer()
+{
+    if (this->OwnsMemory() && data_ != nullptr)
+    {
         alloc_->free(data_);
     }
-    if (alloc_ != nullptr) {
+    if (alloc_ != nullptr)
+    {
         delete alloc_;
     }
 }
 
 // Get the raw data pointer.
-void* TensorBuffer::data() const { return data_; }
+void* TensorBuffer::data() const
+{
+    return data_;
+}
 
 // Get the total number of bytes allocated for the buffer.
 // This method returns the total number of bytes allocated for the buffer by the allocator
 // associated with the TensorBuffer. If the buffer is not yet allocated, the function returns 0.
-size_t TensorBuffer::GetAllocatedBytes() const {
+size_t TensorBuffer::GetAllocatedBytes() const
+{
     return allocated_bytes_;
 }
 
 // Get the root TensorBuffer object.
 // If this TensorBuffer is a sub-buffer of another TensorBuffer, returns that
 // TensorBuffer. Otherwise, returns this.
-TensorBuffer* TensorBuffer::root_buffer() { return this; } // Implementation goes here.
+TensorBuffer* TensorBuffer::root_buffer()
+{
+    return this;
+} // Implementation goes here.
 
 // Get the Allocator object used in this class.
-base::core::Allocator* TensorBuffer::allocator() const {
+base::core::Allocator* TensorBuffer::allocator() const
+{
     return alloc_;
 }
 
 // Check whether this TensorBuffer owns the underlying memory.
-bool TensorBuffer::OwnsMemory() const { return this->owns_memory_; }
+bool TensorBuffer::OwnsMemory() const
+{
+    return this->owns_memory_;
+}
 
 // Get the type of device used by the TensorBuffer.
-DeviceType TensorBuffer::GetDeviceType() const {
-    if (alloc_ != nullptr) {
+DeviceType TensorBuffer::GetDeviceType() const
+{
+    if (alloc_ != nullptr)
+    {
         return alloc_->GetDeviceType();
     }
     return DeviceType::UnKnown;
 }
 
-void TensorBuffer::resize(size_t size) {
+void TensorBuffer::resize(size_t size)
+{
     // Allocate a new buffer.
     void* new_data = this->alloc_->allocate(size);
 
     // Free the old buffer.
-    if (this->OwnsMemory()) {
+    if (this->OwnsMemory())
+    {
         this->alloc_->free(data_);
     }
 
     // Update the internal state.
     this->data_ = new_data;
     this->owns_memory_ = true;
+    this->allocated_bytes_ = size;
 }
+TensorBuffer& TensorBuffer::operator=(const TensorBuffer& other)
+{
+    if (this == &other)
+    {
+        return *this;
+    }
 
-TensorBuffer& TensorBuffer::operator=(const TensorBuffer& other) {
-    if (this->OwnsMemory()) {
+    if (this->OwnsMemory())
+    {
         this->alloc_->free(data_);
     }
 
     delete this->alloc_;
-    if (other.GetDeviceType() == DeviceType::CpuDevice) {
+    this->alloc_ = nullptr;
+    this->data_ = nullptr;
+    this->owns_memory_ = false;
+    this->allocated_bytes_ = 0;
+
+    if (other.GetDeviceType() == DeviceType::CpuDevice)
+    {
         this->alloc_ = new base::core::CPUAllocator();
     }
-    #if defined(__CUDA) || defined(__ROCM)
-    else if (other.GetDeviceType() == DeviceType::GpuDevice) {
+#if defined(__CUDA) || defined(__ROCM)
+    else if (other.GetDeviceType() == DeviceType::GpuDevice)
+    {
         this->alloc_ = new base::core::GPUAllocator();
     }
-    #endif // __CUDA || __ROCM
-
+#endif // __CUDA || __ROCM
+    else
+    {
+        // `other` has no allocator: it is either a moved-from buffer or a non-owning
+        // reference buffer built from a raw pointer. There is nothing to allocate from,
+        // so leave this buffer empty instead of dereferencing the freed allocator.
+        return *this;
+    }
 
     this->data_ = this->alloc_->allocate(other.GetAllocatedBytes());
     this->owns_memory_ = true;
+    this->allocated_bytes_ = other.GetAllocatedBytes();
     return *this;
 }
 
-TensorBuffer& TensorBuffer::operator=(TensorBuffer&& other) noexcept {
-    if (this->OwnsMemory()) {
+TensorBuffer& TensorBuffer::operator=(TensorBuffer&& other) noexcept
+{
+    if (this == &other)
+    {
+        return *this;
+    }
+
+    if (this->OwnsMemory())
+    {
         this->alloc_->free(data_);
     }
     delete this->alloc_;
     this->alloc_ = other.alloc_;
     this->data_ = other.data_;
     this->owns_memory_ = other.owns_memory_;
+    this->allocated_bytes_ = other.allocated_bytes_;
 
     // Reset the other TensorBuffer.
+    other.alloc_ = nullptr;
     other.data_ = nullptr;
     other.owns_memory_ = false;
+    other.allocated_bytes_ = 0;
     return *this;
 }
 
-}  // namespace container
+} // namespace container

source/source_base/module_container/ATen/kernels/memory_impl.cpp

@@ -102,7 +102,7 @@ struct resize_memory<T, DEVICE_GPU> {
 
 template <typename T>
 struct set_memory<T, DEVICE_GPU> {
-    void operator()(T* arr, const int var, const size_t& size) {}
+    void operator()(T* arr, const T& var, const size_t& size) {}
 };
 
 template <typename T>
@@ -207,4 +207,4 @@ template struct delete_memory<std::complex<double>, DEVICE_GPU>;
 #endif
 
 } // namespace kernels
-} // namespace container
+} // namespace container

source/source_base/module_container/ATen/kernels/test/lapack_test.cpp

@@ -56,12 +56,9 @@ TYPED_TEST(LapackTest, Trtri) {
 }
 
 TYPED_TEST(LapackTest, Potrf) {
-
-    return;
     using Type = typename std::tuple_element<0, decltype(TypeParam())>::type;
     using Device = typename std::tuple_element<1, decltype(TypeParam())>::type;
 
-    blas_gemm<Type, Device> gemmCalculator;
     lapack_potrf<Type, Device> potrfCalculator;
     set_matrix<Type, Device> setMatrixCalculator;
 
@@ -71,25 +68,51 @@ TYPED_TEST(LapackTest, Potrf) {
                                  static_cast<Type>(2.0), static_cast<Type>(3.0), static_cast<Type>(6.0)}).to_device<Device>());
 
     Tensor B = A;
-    Tensor C = B;
-    C.zero();
-
-    const char transa = 'N';
-    const char transb = 'C';
-    const int m = 3;
-    const int n = 3;
-    const int k = 3;
-    const Type alpha = static_cast<Type>(1.0);
-    const Type beta  = static_cast<Type>(0.0);
     // Note all blas and lapack operators within container are column major!
     // For this reason, we should employ 'L' instead of 'U' in the subsequent line.
     potrfCalculator('L', dim, B.data<Type>(), dim);
-    // Keep the upper triangle of B
-    setMatrixCalculator('U', B.data<Type>(), dim);
-    // A = U**T * U
-    gemmCalculator(transa, transb, m, n, k, &alpha, B.to_device<DEVICE_CPU>().data<Type>(), k, B.to_device<DEVICE_CPU>().data<Type>(), n, &beta, C.to_device<DEVICE_CPU>().data<Type>(), n);
+    // B may live on an accelerator, so pull it back before inspecting elements on the host.
+    const Tensor factorized = B.to_device<DEVICE_CPU>();
+    EXPECT_GT(std::abs(factorized.data<Type>()[0]), 0.0);
+    EXPECT_GT(std::abs(factorized.data<Type>()[4]), 0.0);
+    EXPECT_GT(std::abs(factorized.data<Type>()[8]), 0.0);
+
+    setMatrixCalculator('L', B.data<Type>(), dim);
+    const Tensor masked = B.to_device<DEVICE_CPU>();
+    EXPECT_EQ(masked.data<Type>()[1], static_cast<Type>(0.0));
+    EXPECT_EQ(masked.data<Type>()[2], static_cast<Type>(0.0));
+    EXPECT_EQ(masked.data<Type>()[5], static_cast<Type>(0.0));
+}
+
+TYPED_TEST(LapackTest, GetrfGetriGetrs) {
+    using Type = typename std::tuple_element<0, decltype(TypeParam())>::type;
+    // This test drives the wrappers with host stack buffers, so it is pinned to the CPU
+    // backend; handing these pointers to the cuSolver path would be an invalid device pointer.
+    using Device = DEVICE_CPU;
 
-    EXPECT_EQ(A, C);
+    const int dim = 2;
+    const int rhs_count = 1;
+    const int workspace_size = 8;
+    Type matrix[4] = {static_cast<Type>(4.0),
+                      static_cast<Type>(2.0),
+                      static_cast<Type>(1.0),
+                      static_cast<Type>(3.0)};
+    int pivots[dim] = {0};
+
+    lapack_getrf<Type, Device>()(dim, dim, matrix, dim, pivots);
+
+    Type factorized[4] = {matrix[0], matrix[1], matrix[2], matrix[3]};
+    Type rhs[2] = {static_cast<Type>(1.0), static_cast<Type>(1.0)};
+    lapack_getrs<Type, Device>()('N', dim, rhs_count, factorized, dim, pivots, rhs, dim);
+    EXPECT_NEAR(std::abs(rhs[0] - static_cast<Type>(0.2)), 0.0, 1.0e-6);
+    EXPECT_NEAR(std::abs(rhs[1] - static_cast<Type>(0.2)), 0.0, 1.0e-6);
+
+    Type workspace[workspace_size];
+    lapack_getri<Type, Device>()(dim, matrix, dim, pivots, workspace, workspace_size);
+    EXPECT_NEAR(std::abs(matrix[0] - static_cast<Type>(0.3)), 0.0, 1.0e-6);
+    EXPECT_NEAR(std::abs(matrix[1] - static_cast<Type>(-0.2)), 0.0, 1.0e-6);
+    EXPECT_NEAR(std::abs(matrix[2] - static_cast<Type>(-0.1)), 0.0, 1.0e-6);
+    EXPECT_NEAR(std::abs(matrix[3] - static_cast<Type>(0.4)), 0.0, 1.0e-6);
 }
 
 // lapack_geqrf_inplace,

source/source_base/module_container/ATen/kernels/test/linalg_test.cpp

@@ -203,5 +203,60 @@ TYPED_TEST(LinalgTest, Reduce) {
     EXPECT_EQ(A_reduce, expected);
 }
 
+template <typename T>
+void test_integer_linalg_kernels()
+{
+    const int count = 4;
+    const T alpha = static_cast<T>(2);
+    const T beta = static_cast<T>(3);
+    const T x[count] = {static_cast<T>(1), static_cast<T>(2), static_cast<T>(3), static_cast<T>(4)};
+    const T y[count] = {static_cast<T>(4), static_cast<T>(3), static_cast<T>(2), static_cast<T>(1)};
+    T output[count] = {};
+
+    kernels::add<T, DEVICE_CPU>()(count, alpha, x, beta, y, output);
+    EXPECT_EQ(output[0], static_cast<T>(14));
+    EXPECT_EQ(output[3], static_cast<T>(11));
+
+    kernels::mul<T, DEVICE_CPU>()(count, alpha, x, output);
+    EXPECT_EQ(output[0], static_cast<T>(2));
+    EXPECT_EQ(output[3], static_cast<T>(8));
+
+    kernels::mul<T, DEVICE_CPU>()(count, alpha, x, y, output);
+    EXPECT_EQ(output[0], static_cast<T>(8));
+    EXPECT_EQ(output[3], static_cast<T>(8));
+
+    kernels::div<T, DEVICE_CPU>()(count, alpha, x, y, output);
+    EXPECT_EQ(output[0], static_cast<T>(0));
+    EXPECT_EQ(output[3], static_cast<T>(8));
+
+    kernels::fma<T, DEVICE_CPU>()(count, alpha, x, y, beta, x, output);
+    EXPECT_EQ(output[0], static_cast<T>(11));
+    EXPECT_EQ(output[3], static_cast<T>(20));
+
+    const std::vector<int> permutation = {0};
+    const std::vector<int64_t> shape = {count};
+    kernels::transpose<T, DEVICE_CPU>()(permutation, shape, shape, x, output);
+    EXPECT_EQ(output[2], x[2]);
+
+    const std::vector<int64_t> unit_stride = {1};
+    kernels::stride<T, DEVICE_CPU>()(unit_stride, shape, shape, x, output);
+    EXPECT_EQ(output[2], x[2]);
+
+    kernels::inflate<T, DEVICE_CPU>()(unit_stride, shape, shape, x, output);
+    EXPECT_EQ(output[2], x[2]);
+
+    const int64_t output_count = 2;
+    const int64_t inner_dimension = 2;
+    kernels::reduce<T, DEVICE_CPU>()(output_count, inner_dimension, x, output);
+    EXPECT_EQ(output[0], static_cast<T>(3));
+    EXPECT_EQ(output[1], static_cast<T>(7));
+}
+
+TEST(LinalgIntegerTest, CoversIntegerKernelInstantiations)
+{
+    test_integer_linalg_kernels<int>();
+    test_integer_linalg_kernels<int64_t>();
+}
+
 } // namespace kernels
-} // namespace container
+} // namespace container