Added implementation to calculate Gram and Row Gram matrices

cpp/open3d/core/CMakeLists.txt

@@ -48,6 +48,7 @@ target_sources(core PRIVATE
     linalg/LeastSquares.cpp
     linalg/LU.cpp
     linalg/Matmul.cpp
+    linalg/Gramian.cpp
     linalg/Solve.cpp
     linalg/SVD.cpp
     linalg/Tri.cpp
@@ -83,6 +84,7 @@ target_sources(core_impl PRIVATE
     linalg/LeastSquaresCPU.cpp
     linalg/LUCPU.cpp
     linalg/MatmulCPU.cpp
+    linalg/GramianCPU.cpp
     linalg/SolveCPU.cpp
     linalg/SVDCPU.cpp
     linalg/TriCPU.cpp
@@ -102,6 +104,7 @@ open3d_sycl_target_sources(core_impl PRIVATE
     linalg/LeastSquaresSYCL.cpp
     linalg/LUSYCL.cpp
     linalg/MatmulSYCL.cpp
+    linalg/GramianSYCL.cpp
     linalg/SolveSYCL.cpp
     linalg/SVDSYCL.cpp
     linalg/TriSYCL.cpp
@@ -129,6 +132,7 @@ if (BUILD_CUDA_MODULE)
         linalg/LinalgUtils.cpp
         linalg/LUCUDA.cpp
         linalg/MatmulCUDA.cpp
+        linalg/GramianCUDA.cpp
         linalg/SolveCUDA.cpp
         linalg/SVDCUDA.cpp
         linalg/TriCUDA.cu

cpp/open3d/core/Tensor.cpp

@@ -25,6 +25,7 @@
 #include "open3d/core/kernel/IndexReduction.h"
 #include "open3d/core/kernel/Kernel.h"
 #include "open3d/core/linalg/Det.h"
+#include "open3d/core/linalg/Gramian.h"
 #include "open3d/core/linalg/Inverse.h"
 #include "open3d/core/linalg/LU.h"
 #include "open3d/core/linalg/LeastSquares.h"
@@ -1057,6 +1058,40 @@ Tensor Tensor::T() const {
     }
 }
 
+Tensor Tensor::Gram() const {
+    int64_t n_dims = NumDims();
+    if (n_dims <= 0) {
+        return *this;
+    } else if (n_dims == 1) {
+        return Reshape({-1, 1}).Gram();
+    } else if (n_dims == 2) {
+        Tensor output;
+        core::Gram(*this, output);
+        return output;
+    } else {
+        utility::LogError(
+                "Tensor::Gram() expects a Tensor with <= 2 dimensions, but the "
+                "Tensor has {} dimensions.");
+    }
+}
+
+Tensor Tensor::RowGram() const {
+    int64_t n_dims = NumDims();
+    if (n_dims <= 0) {
+        return *this;
+    } else if (n_dims == 1) {
+        return Reshape({-1, 1}).RowGram();
+    } else if (n_dims == 2) {
+        Tensor output;
+        core::RowGram(*this, output);
+        return output;
+    } else {
+        utility::LogError(
+                "Tensor::RowGram() expects a Tensor with <= 2 dimensions, but "
+                "the Tensor has {} dimensions.");
+    }
+}
+
 double Tensor::Det() const {
     AssertTensorDtypes(*this, {Float32, Float64});
     return core::Det(*this);

cpp/open3d/core/Tensor.h

@@ -608,6 +608,24 @@ class Tensor : public IsDevice {
     /// 0-D and 1-D Tensor remains the same.
     Tensor T() const;
 
+    /// \brief Computes the Gram matrix (Gramian) of this tensor. Expects input
+    /// to be <= 2-D Tensor.
+    ///
+    /// \return The Gramian of the given tensor. For a 2D tensor of shape {m,
+    /// n}, this will be a square matrix of shape {n, n}. 1D tensors of shape
+    /// {n} are reshaped to {1, n} before the operation. For 0D, the returned
+    /// tensor is the same as the input.
+    Tensor Gram() const;
+
+    /// \brief Computes the row Gram matrix of this tensor. Expects input to be
+    /// <= 2-D Tensor.
+    ///
+    /// \return The Row Gramian of the given tensor. For a 2D tensor of shape
+    /// {m, n}, this will be a square matrix of shape {m, m}. 1D tensors of
+    /// shape {n} are reshaped to {1, n} before the operation. For 0D, the
+    /// returned tensor is the same as the input.
+    Tensor RowGram() const;
+
     /// \brief Compute the determinant of a 2D square tensor.
     /// \return returns the determinant of the matrix (double).
     double Det() const;

cpp/open3d/core/linalg/Gramian.cpp

@@ -0,0 +1,128 @@
+// ----------------------------------------------------------------------------
+// -                        Open3D: www.open3d.org                            -
+// ----------------------------------------------------------------------------
+// Copyright (c) 2018-2024 www.open3d.org
+// SPDX-License-Identifier: MIT
+// ----------------------------------------------------------------------------
+
+#include "open3d/core/linalg/Gramian.h"
+
+#include <unordered_map>
+
+#include "open3d/core/CUDAUtils.h"
+
+namespace open3d {
+namespace core {
+
+void Gram(const Tensor& A, Tensor& output) {
+    const Device device = A.GetDevice();
+    const Dtype dtype_original = A.GetDtype();
+    Dtype dtype;
+
+    if (dtype_original != core::Float32 && dtype_original != core::Float64) {
+        utility::LogDebug("Converting to Float32 dtype to from {}.",
+                          dtype_original.ToString());
+        dtype = core::Float32;
+    } else {
+        dtype = dtype_original;
+    }
+
+    // Check shapes
+    SizeVector A_shape = A.GetShape();
+
+    if (A_shape.size() != 2) {
+        utility::LogError("Tensor A must be 2D, but got {}D.", A_shape.size());
+    }
+
+    // Dispatch to backends
+    int64_t m = A_shape[0];
+    int64_t n = A_shape[1];
+
+    if (m == 0 || n == 0) {
+        utility::LogError(
+                "Tensor shapes should not contain dimensions with zero.");
+    }
+
+    Tensor A_contiguous = A.Contiguous().To(dtype);
+    void* A_data = A_contiguous.GetDataPtr();
+
+    output = Tensor::Empty({n, n}, dtype, device);
+    void* B_data = output.GetDataPtr();
+
+    if (device.IsSYCL()) {
+#ifdef BUILD_SYCL_MODULE
+        GramSYCL(A_data, B_data, m, n, dtype, device);
+#else
+        utility::LogError("Unimplemented device.");
+#endif
+    } else if (device.IsCUDA()) {
+#ifdef BUILD_CUDA_MODULE
+        CUDAScopedDevice scoped_device(device);
+        GramCUDA(A_data, B_data, m, n, dtype, device);
+#else
+        utility::LogError("Unimplemented device.");
+#endif
+    } else {
+        GramCPU(A_data, B_data, m, n, dtype);
+    }
+
+    output = output.To(dtype_original);
+}
+
+void RowGram(const Tensor& A, Tensor& output) {
+    const Device device = A.GetDevice();
+    const Dtype dtype_original = A.GetDtype();
+    Dtype dtype;
+
+    if (dtype_original != core::Float32 && dtype_original != core::Float64) {
+        utility::LogDebug("Converting to Float32 dtype to from {}.",
+                          dtype_original.ToString());
+        dtype = core::Float32;
+    } else {
+        dtype = dtype_original;
+    }
+
+    // Check shapes
+    SizeVector A_shape = A.GetShape();
+
+    if (A_shape.size() != 2) {
+        utility::LogError("Tensor A must be 2D, but got {}D.", A_shape.size());
+    }
+
+    // Dispatch to backends
+    int64_t m = A_shape[0];
+    int64_t n = A_shape[1];
+
+    if (m == 0 || n == 0) {
+        utility::LogError(
+                "Tensor shapes should not contain dimensions with zero.");
+    }
+
+    Tensor A_contiguous = A.Contiguous().To(dtype);
+    void* A_data = A_contiguous.GetDataPtr();
+
+    output = Tensor::Empty({m, m}, dtype, device);
+    void* B_data = output.GetDataPtr();
+
+    if (device.IsSYCL()) {
+#ifdef BUILD_SYCL_MODULE
+        RowGramSYCL(A_data, B_data, m, n, dtype, device);
+#else
+        utility::LogError("Unimplemented device.");
+#endif
+    } else if (device.IsCUDA()) {
+#ifdef BUILD_CUDA_MODULE
+        CUDAScopedDevice scoped_device(device);
+        RowGramCUDA(A_data, B_data, m, n, dtype, device);
+#else
+        utility::LogError("Unimplemented device.");
+#endif
+    } else {
+        RowGramCPU(A_data, B_data, m, n, dtype);
+    }
+
+    output = output.To(dtype_original);
+}
+
+}  // namespace core
+}  // namespace open3d

cpp/open3d/core/linalg/Gramian.h

@@ -0,0 +1,52 @@
+// ----------------------------------------------------------------------------
+// -                        Open3D: www.open3d.org                            -
+// ----------------------------------------------------------------------------
+// Copyright (c) 2018-2024 www.open3d.org
+// SPDX-License-Identifier: MIT
+// ----------------------------------------------------------------------------
+
+#pragma once
+
+#include "open3d/core/Tensor.h"
+
+namespace open3d {
+namespace core {
+
+/// Computes the gram matrix of a Tensor (B = A.T @ A)
+void Gram(const Tensor& A, Tensor& B);
+
+/// Computes the row gram matrix of a Tensor (B = A @ A.T)
+void RowGram(const Tensor& A, Tensor& B);
+
+#ifdef BUILD_SYCL_MODULE
+void GramSYCL(void* A_data,
+              void* B_data,
+              int64_t m,
+              int64_t n,
+              Dtype dtype,
+              const Device& device);
+void RowGramSYCL(void* A_data,
+                 void* B_data,
+                 int64_t m,
+                 int64_t n,
+                 Dtype dtype,
+                 const Device& device);
+#endif
+#ifdef BUILD_CUDA_MODULE
+void GramCUDA(void* A_data,
+              void* B_data,
+              int64_t m,
+              int64_t n,
+              Dtype dtype,
+              const Device& device);
+void RowGramCUDA(void* A_data,
+                 void* B_data,
+                 int64_t m,
+                 int64_t n,
+                 Dtype dtype,
+                 const Device& device);
+#endif
+void GramCPU(void* A_data, void* B_data, int64_t m, int64_t n, Dtype dtype);
+void RowGramCPU(void* A_data, void* B_data, int64_t m, int64_t n, Dtype dtype);
+}  // namespace core
+}  // namespace open3d

cpp/open3d/core/linalg/GramianCPU.cpp

@@ -0,0 +1,35 @@
+// ----------------------------------------------------------------------------
+// -                        Open3D: www.open3d.org                            -
+// ----------------------------------------------------------------------------
+// Copyright (c) 2018-2024 www.open3d.org
+// SPDX-License-Identifier: MIT
+// ----------------------------------------------------------------------------
+
+#include "open3d/core/linalg/BlasWrapper.h"
+#include "open3d/core/linalg/Gramian.h"
+#include "open3d/core/linalg/LinalgUtils.h"
+
+namespace open3d {
+namespace core {
+void GramCPU(void* A_data, void* B_data, int64_t m, int64_t n, Dtype dtype) {
+    DISPATCH_LINALG_DTYPE_TO_TEMPLATE(dtype, [&]() {
+        scalar_t alpha = 1, beta = 0;
+        gemm_cpu<scalar_t>(CblasColMajor, CblasNoTrans, CblasTrans, n, n, m,
+                           alpha, static_cast<const scalar_t*>(A_data), n,
+                           static_cast<const scalar_t*>(A_data), n, beta,
+                           static_cast<scalar_t*>(B_data), n);
+    });
+}
+
+void RowGramCPU(void* A_data, void* B_data, int64_t m, int64_t n, Dtype dtype) {
+    DISPATCH_LINALG_DTYPE_TO_TEMPLATE(dtype, [&]() {
+        scalar_t alpha = 1, beta = 0;
+        gemm_cpu<scalar_t>(CblasColMajor, CblasTrans, CblasNoTrans, m, m, n,
+                           alpha, static_cast<const scalar_t*>(A_data), n,
+                           static_cast<const scalar_t*>(A_data), n, beta,
+                           static_cast<scalar_t*>(B_data), m);
+    });
+}
+
+}  // namespace core
+}  // namespace open3d

cpp/open3d/core/linalg/GramianCUDA.cpp

@@ -0,0 +1,55 @@
+// ----------------------------------------------------------------------------
+// -                        Open3D: www.open3d.org                            -
+// ----------------------------------------------------------------------------
+// Copyright (c) 2018-2024 www.open3d.org
+// SPDX-License-Identifier: MIT
+// ----------------------------------------------------------------------------
+
+#include "open3d/core/linalg/BlasWrapper.h"
+#include "open3d/core/linalg/LinalgUtils.h"
+#include "open3d/core/linalg/Matmul.h"
+#include "open3d/utility/Logging.h"
+
+namespace open3d {
+namespace core {
+
+void GramCUDA(void* A_data,
+              void* B_data,
+              int64_t m,
+              int64_t n,
+              Dtype dtype,
+              const Device& device) {
+    cublasHandle_t handle = CuBLASContext::GetInstance().GetHandle(device);
+    DISPATCH_LINALG_DTYPE_TO_TEMPLATE(dtype, [&]() {
+        scalar_t alpha = 1, beta = 0;
+        OPEN3D_CUBLAS_CHECK(
+                gemm_cuda<scalar_t>(handle, CUBLAS_OP_N, CUBLAS_OP_T, n, n, m,
+                                    &alpha,
+                                    static_cast<const scalar_t*>(A_data), n,
+                                    static_cast<const scalar_t*>(A_data), n,
+                                    &beta, static_cast<scalar_t*>(B_data), n),
+                "cuda gemm failed");
+    });
+}
+
+void RowGramCUDA(void* A_data,
+                 void* B_data,
+                 int64_t m,
+                 int64_t n,
+                 Dtype dtype,
+                 const Device& device) {
+    cublasHandle_t handle = CuBLASContext::GetInstance().GetHandle(device);
+    DISPATCH_LINALG_DTYPE_TO_TEMPLATE(dtype, [&]() {
+        scalar_t alpha = 1, beta = 0;
+        OPEN3D_CUBLAS_CHECK(
+                gemm_cuda<scalar_t>(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, m, n,
+                                    &alpha,
+                                    static_cast<const scalar_t*>(A_data), n,
+                                    static_cast<const scalar_t*>(A_data), n,
+                                    &beta, static_cast<scalar_t*>(B_data), m),
+                "cuda gemm failed");
+    });
+}
+
+}  // namespace core
+}  // namespace open3d