Tensor math OpenCL lib

2026-04-03 20:30:39 +04:00 · 2025-11-19 18:05:11 +04:00
parent c1874212ae
commit bd8b26c35a
12 changed files with 361 additions and 355 deletions
--- a/src/run.py
+++ b/src/run.py
@@ -1,181 +1,10 @@
 from tensor.tensor import *

+if (MODE == PLATFORM.OPENCL):
+    init("./tensor/")

-def test_matrix_operations():
-    print("=" * 50)
-    print("ТЕСТИРОВАНИЕ БИБЛИОТЕКИ MATRIX")
-    print("=" * 50)
-
-    # Тест создания матриц
-    print("\n1. СОЗДАНИЕ МАТРИЦ:")
-    print("-" * 30)
-
-    # Создание матрицы с заполнением одним значением
-    m1 = Matrix([2, 3], 1.0)
-    print(f"Matrix([2, 3], 1.0) = {m1}")
-
-    # Создание матрицы с разными значениями
-    m2 = Matrix([2, 3], 2.0, 3.0)
-    print(f"Matrix([2, 3], 2.0, 3.0) = {m2}")
-
-    # Создание матрицы для умножения
-    m3 = Matrix([3, 2], 2.0)
-    print(f"Matrix([3, 2], 2.0) = {m3}")
-
-    # Тест получения свойств
-    print("\n2. СВОЙСТВА МАТРИЦ:")
-    print("-" * 30)
-    print(f"m1.get_shape() = {m1.get_shape()}")
-    print(f"m1.get_axes() = {m1.get_axes()}")
-    print(f"m1.get_size() = {m1.get_size()}")
-
-    # Тест доступа к элементам
-    print("\n3. ДОСТУП К ЭЛЕМЕНТАМ:")
-    print("-" * 30)
-    print(f"m1[0] = {m1[0]}")
-    print(f"m1[0, 1] = {m1[0, 1]}")
-
-    # Установка значений
-    m1[0, 1] = 5.0
-    print(f"После m1[0, 1] = 5.0: {m1}")
-
-    # Тест арифметических операций
-    print("\n4. АРИФМЕТИЧЕСКИЕ ОПЕРАЦИИ:")
-    print("-" * 30)
-
-    # Сложение
-    m_add = m1 + m2
-    print(f"m1 + m2 = {m_add}")
-
-    # Вычитание
-    m_sub = m1 - m2
-    print(f"m1 - m2 = {m_sub}")
-
-    # Умножение на скаляр
-    m_mul_scalar = m1 * 2.0
-    print(f"m1 * 2.0 = {m_mul_scalar}")
-
-    # Поэлементное умножение
-    m_mul_element = m1 * m2
-    print(f"m1 * m2 (поэлементно) = {m_mul_element}")
-
-    # Деление на скаляр
-    m_div = m1 / 2.0
-    print(f"m1 / 2.0 = {m_div}")
-
-    # Унарные операторы
-    m_neg = -m1
-    print(f"-m1 = {m_neg}")
-    m_pos = +m1
-    print(f"+m1 = {m_pos}")
-
-    # Тест матричного умножения
-    print("\n5. МАТРИЧНОЕ УМНОЖЕНИЕ:")
-    print("-" * 30)
-    try:
-        m_matmul = m1 @ m3
-        print(f"m1 @ m3 = {m_matmul}")
-    except Exception as e:
-        print(f"Ошибка при матричном умножении: {e}")
-
-    # Тест транспонирования
-    print("\n6. ТРАНСПОНИРОВАНИЕ:")
-    print("-" * 30)
-    m_transposed = m1.t()
-    print(f"m1.t() = {m_transposed}")
-
-    try:
-        m_transpose_method = m1.transpose(0, 1)
-        print(f"m1.transpose(0, 1) = {m_transpose_method}")
-    except Exception as e:
-        print(f"Ошибка при transpose(0, 1): {e}")
-
-    try:
-        m_transpose_list = m1.transpose([0, 1])
-        print(f"m1.transpose([0, 1]) = {m_transpose_list}")
-    except Exception as e:
-        print(f"Ошибка при transpose([0, 1]): {e}")
-
-    # Тест операций на месте
-    print("\n7. ОПЕРАЦИИ НА МЕСТЕ:")
-    print("-" * 30)
-
-    m_test = Matrix([2, 2], 1.0)
-    print(f"Исходная матрица: {m_test}")
-
-    m_test += 2.0
-    print(f"После m_test += 2.0: {m_test}")
-
-    m_test -= 1.0
-    print(f"После m_test -= 1.0: {m_test}")
-
-    m_test *= 3.0
-    print(f"После m_test *= 3.0: {m_test}")
-
-    m_test /= 2.0
-    print(f"После m_test /= 2.0: {m_test}")
-
-    # Тест с вашими матрицами из примера
-    print("\n8. ТЕСТ С ВАШИМИ МАТРИЦАМИ:")
-    print("-" * 30)
-
-    a = Matrix([2, 3], 2)
-    b = Matrix([3, 2], 1)
-
-    print(f"a = {a}")
-    print(f"b = {b}")
-
-    try:
-        result = a @ b
-        print(f"a @ b = {result}")
-    except Exception as e:
-        print(f"Ошибка при a @ b: {e}")
-
-    # Тест обратных операций
-    print("\n9. ОБРАТНЫЕ ОПЕРАЦИИ:")
-    print("-" * 30)
-
-    m_base = Matrix([2, 2], 3.0)
-    print(f"Исходная матрица: {m_base}")
-
-    # Правое сложение
-    m_radd = 2.0 + m_base
-    print(f"2.0 + m_base = {m_radd}")
-
-    # Правое умножение
-    m_rmul = 2.0 * m_base
-    print(f"2.0 * m_base = {m_rmul}")
-
-    # Правое вычитание
-    m_rsub = 10.0 - m_base
-    print(f"10.0 - m_base = {m_rsub}")
-
-    print("\n" + "=" * 50)
-    print("ТЕСТИРОВАНИЕ ЗАВЕРШЕНО")
-    print("=" * 50)
-
-
-def test_edge_cases():
-    print("\n\n10. ТЕСТ ГРАНИЧНЫХ СЛУЧАЕВ:")
-    print("=" * 50)
-
-    try:
-        # Попытка создания с разными параметрами
-        m_empty = Matrix([0, 0])
-        print(f"Matrix([0, 0]) = {m_empty}")
-    except Exception as e:
-        print(f"Ошибка при создании Matrix([0, 0]): {e}")
-
-    try:
-        # Попытка доступа к несуществующему элементу
-        m_test = Matrix([2, 2], 1.0)
-        print(f"Попытка доступа к m_test[5, 5]: ", end="")
-        value = m_test[5, 5]
-        print(value)
-    except Exception as e:
-        print(f"Ошибка: {e}")
-
-
-if __name__ == "__main__":
-    test_matrix_operations()
-    test_edge_cases()
+a = Matrix([1024, 1024], 1)
+a += 1
+b = Matrix([1024, 1024], 1)
+c = a @ b
+print(c)
--- a/src/tensor/Makefile
+++ b/src/tensor/Makefile
@@ -19,27 +19,37 @@ else
 endif

 BUILD_DIR = build
-COMMON_SRC = opencl/opencl.cpp
+COMMON_SRC = 
+OPENCL_SRC = opencl/opencl.cpp

 PYTHON_PATH = $(shell python -c "from sysconfig import get_paths; print(get_paths()['data'])")
-PYTHON_INCLUDE = $(shell python -c "import sysconfig; print(sysconfig.get_config_var('CONFINCLUDEPY'))")
-PYTHON_LIBS = $(PYTHON_PATH)$(SP)libs
+PYTHON_INCLUDE = $(shell python -c "import sysconfig; print(sysconfig.get_config_var([k for k in sysconfig.get_config_vars().keys() if 'INCLUDE' in k][0]))")
+PYTHON_LIB_PATH = $(PYTHON_PATH)$(SP)libs
+PYTHON_LIB = -lpython313 #-lpython3.13
+
 PYBIND_INCLUDE = $(shell python -c "import pybind11; print(pybind11.get_include())")

 OPENCL_INCLUDES = -I"A:/Programs/OpenCL/include" 
-OPENCL_LIB_PATH = -L"A:/Programs/OpenCL/lib" -lOpenCL
-OPENCL_LIB = -L"A:/Programs/OpenCL/lib" -lOpenCL
+OPENCL_LIB_PATH = -L"A:/Programs/OpenCL/lib"
+OPENCL_LIB = -lOpenCL

-.DEFAULT_GOAL := $(TARGET)
+.DEFAULT_GOAL := cpu

 $(BUILD_DIR):
 	$(MKDIR) $(BUILD_DIR)

-$(TARGET): $(COMMON_SRC) main.cpp | $(BUILD_DIR)
-	$(CXX) $(CXXFLAGS) $(OPENCL_INCLUDES) $(OPENCL_LIB_PATH) -o $@ $^ $(OPENCL_LIB)
+cpu: $(COMMON_SRC) main.cpp | $(BUILD_DIR)
+	$(CXX) $(CXXFLAGS) -DUSE_CPU -o $(TARGET) $^ 

-module: $(COMMON_SRC) pybind.cpp | $(BUILD_DIR)
-	$(CXX) $(CXXFLAGS) -shared -fPIC -o tensor.$(SHARED_LIB_EXT) $^ -I"$(PYTHON_INCLUDE)" -L"$(PYTHON_LIBS)" -lpython3.13 -I"$(PYBIND_INCLUDE)"
+opencl: $(COMMON_SRC) $(OPENCL_SRC) main.cpp | $(BUILD_DIR)
+	$(CXX) $(CXXFLAGS) -DUSE_OPENCL $(OPENCL_INCLUDES) $(OPENCL_LIB_PATH) -o $(TARGET) $^ $(OPENCL_LIB)
+
+cpu_module: $(COMMON_SRC) pybind.cpp | $(BUILD_DIR)
+	$(CXX) $(CXXFLAGS) -DUSE_CPU -shared -fPIC -I"$(PYTHON_INCLUDE)" -I"$(PYBIND_INCLUDE)" -L"$(PYTHON_LIB_PATH)" -o tensor.$(SHARED_LIB_EXT) $^ $(PYTHON_LIB)
+	PYTHONPATH=. pybind11-stubgen tensor -o .
+
+opencl_module: $(COMMON_SRC) $(OPENCL_SRC) pybind.cpp | $(BUILD_DIR)
+	$(CXX) $(CXXFLAGS) -DUSE_OPENCL -shared -fPIC -I"$(PYTHON_INCLUDE)" -I"$(PYBIND_INCLUDE)" -L"$(PYTHON_LIB_PATH)" $(OPENCL_INCLUDES) $(OPENCL_LIB_PATH) -o tensor.$(SHARED_LIB_EXT) $^ $(PYTHON_LIB) $(OPENCL_LIB)
 	PYTHONPATH=. pybind11-stubgen tensor -o .

 clean:
--- a/src/tensor/cpu/tensor.hpp
+++ b/src/tensor/cpu/tensor.hpp
@@ -38,8 +38,8 @@ public:
  using ITensor::operator+;
  using ITensor::operator-;

-  Tensor operator+() override;
-  Tensor operator-() override;
+  Tensor operator+() const override;
+  Tensor operator-() const override;

  Tensor &operator+=(const T scalar) override;

--- a/src/tensor/cpu/tensor.tpp
+++ b/src/tensor/cpu/tensor.tpp
@@ -2,6 +2,7 @@

 #include "tensor.hpp"

+#include <iostream>
 #include <random>
 #include <sstream>

@@ -79,13 +80,15 @@ const T &Tensor<T, Dim>::operator()(Indices... indices) const {
 }

 // ===== OPERATORS =====
-template <typename T, int Dim> Tensor<T, Dim> Tensor<T, Dim>::operator+() {
+template <typename T, int Dim>
+Tensor<T, Dim> Tensor<T, Dim>::operator+() const {
  Tensor result = *this;
  for (T &e : result.data_)
    e = +e;
  return result;
 }
-template <typename T, int Dim> Tensor<T, Dim> Tensor<T, Dim>::operator-() {
+template <typename T, int Dim>
+Tensor<T, Dim> Tensor<T, Dim>::operator-() const {
  Tensor result = *this;
  for (T &e : result.data_)
    e = -e;
@@ -156,46 +159,5 @@ Tensor<T, Dim>::operator%(const Tensor &other) const {

 // ===== UTILS =====
 template <typename T, int Dim> std::string Tensor<T, Dim>::toString() const {
-  std::ostringstream oss;
-  if constexpr (Dim == 0) {
-    oss << "Scalar<" << typeid(T).name() << ">: " << data_[0];
-  } else if constexpr (Dim == 1) {
-    oss << "Vector<" << typeid(T).name() << ">(" << shape_[0] << "): [";
-    for (size_t i = 0; i < getSize(); ++i) {
-      oss << data_[i];
-      if (i < getSize() - 1)
-        oss << ", ";
-    }
-    oss << "]";
-  } else if constexpr (Dim == 2) {
-    oss << "Matrix<" << typeid(T).name() << ">(" << shape_[axes_[0]] << "x"
-        << shape_[axes_[1]] << "):";
-    for (size_t i = 0; i < shape_[axes_[0]]; ++i) {
-      oss << "\n  [";
-      for (size_t j = 0; j < shape_[axes_[1]]; ++j) {
-        oss << (*this)(i, j);
-        if (j < shape_[axes_[1]] - 1)
-          oss << ", ";
-      }
-      oss << "]";
-    }
-  } else {
-    oss << "Tensor" << Dim << "D<" << typeid(T).name() << ">" << "[";
-    for (size_t i = 0; i < Dim; ++i) {
-      oss << shape_[axes_[i]];
-      if (i < Dim - 1)
-        oss << "x";
-    }
-    oss << "]: [";
-    size_t show = std::min(getSize(), size_t(10));
-    for (size_t i = 0; i < show; ++i) {
-      oss << data_[i];
-      if (i < show - 1)
-        oss << ", ";
-    }
-    if (getSize() > 10)
-      oss << ", ...";
-    oss << "]";
-  }
-  return oss.str();
+  return ITensor::format(data_);
 }
--- a/src/tensor/main.cpp
+++ b/src/tensor/main.cpp
@@ -1,20 +1,22 @@
-// #include "cpu/tensor.hpp"
+#ifdef USE_OPENCL
 #include "opencl/tensor.hpp"
+OpenCL openCL;
+// TODO: GENERIC KERNELS
+// TODO: Scalar mult
+#elif USE_CPU
+#include "cpu/tensor.hpp"
+#endif

 #include <iostream>

-// TODO: GENERIC KERNELS
-// TODO: Scalar mult
 // TODO: TMult >2

-OpenCL openCL;
-
 int main() {
-  Tensor<float, 2> a = Tensor<float, 2>({8192, 8192}, 1);
-  Tensor<float, 2> b = Tensor<float, 2>({8192, 8192}, 1);
-  auto c = a % b;
-  Tensor<float, 2> d = Tensor<float, 2>(c);
-  d += 1;
-  std::cout << d.toString();
+#ifdef USE_OPENCL
+  openCL.init("./");
+#endif
+
+  Tensor<float, 2> a = Tensor<float, 2>({32, 32}, 2);
+  std::cout << a.toString();
  return 0;
 }
--- a/src/tensor/opencl/opencl.cpp
+++ b/src/tensor/opencl/opencl.cpp
@@ -27,9 +27,9 @@ cl::Program OpenCL::compileProgram(const std::string &file) {
  }
  return program;
 }
-void OpenCL::loadPrograms() {
+void OpenCL::loadPrograms(std::string &programsBasePath) {
  for (const auto &entry : programPaths) {
-    programs[entry.first] = compileProgram(entry.second);
+    programs[entry.first] = compileProgram(programsBasePath + entry.second);
    std::cout << "Loaded program: " << entry.second << std::endl;
  }
 }
@@ -89,10 +89,12 @@ void OpenCL::initializeDevice() {
            << " MB" << std::endl;
 }

-OpenCL::OpenCL() {
+OpenCL::OpenCL() {}
+
+void OpenCL::init(std::string programsBasePath) {
  try {
    initializeDevice();
-    loadPrograms();
+    loadPrograms(programsBasePath);
  } catch (const cl::Error &e) {
    std::cerr << "OpenCL error: " << e.what() << " (" << e.err() << ")"
              << std::endl;
--- a/src/tensor/opencl/opencl.hpp
+++ b/src/tensor/opencl/opencl.hpp
@@ -26,10 +26,10 @@ private:

  std::unordered_map<Program, cl::Program> programs;
  std::unordered_map<Program, std::string> programPaths = {
-      {Program::ATOMIC, "./opencl/kernels/atomic.cl"},
-      {Program::SCALAR, "./opencl/kernels/scalar.cl"},
-      {Program::TENSOR, "./opencl/kernels/tensor.cl"},
-      {Program::FUSION, "./opencl/kernels/fusion.cl"}};
+      {Program::ATOMIC, "opencl/kernels/atomic.cl"},
+      {Program::SCALAR, "opencl/kernels/scalar.cl"},
+      {Program::TENSOR, "opencl/kernels/tensor.cl"},
+      {Program::FUSION, "opencl/kernels/fusion.cl"}};
  std::unordered_map<Method, Program> methodPrograms = {
      {Method::POSITIVE, Program::ATOMIC},
      {Method::NEGATIVE, Program::ATOMIC},
@@ -48,13 +48,15 @@ private:

  std::string readProgram(const std::string &filePath);
  cl::Program compileProgram(const std::string &file);
-  void loadPrograms();
+  void loadPrograms(std::string &programsBasePath);

  void initializeDevice();

 public:
  OpenCL();

+  void init(std::string programsBasePath);
+
  OpenCL(const OpenCL &) = delete;
  OpenCL &operator=(const OpenCL &) = delete;
  OpenCL(OpenCL &&) = delete;
--- a/src/tensor/opencl/tensor.hpp
+++ b/src/tensor/opencl/tensor.hpp
@@ -4,13 +4,12 @@

 #include "../tensor.hpp"

-#include <iostream>
 #include <random>
-#include <sstream>
+
 template <typename T, int Dim> class Tensor : public ITensor<T, Dim> {
 private:
  cl::Buffer *data_ = nullptr;
-  cl::Event event_ = cl::Event();
+  mutable cl::Event event_ = cl::Event();

  class AutoEventList {
  private:
@@ -114,16 +113,10 @@ public:
  const cl::Buffer *getData() const { return data_; }
  const cl::Event &getEvent() const { return event_; }

-  // T &operator[](size_t i);
-  // const T &operator[](size_t i) const;
-  // template <typename... Indices> T &operator()(Indices... indices);
-  // template <typename... Indices> const T &operator()(Indices... indices)
-  // const;
-
  using ITensor::operator+;
  using ITensor::operator-;

-  Tensor operator+() override {
+  Tensor operator+() const override {
    cl::Kernel kernel = openCL.createKernel(OpenCL::Method::POSITIVE);
    kernel.setArg(0, *data_);
    openCL.getQueue().enqueueNDRangeKernel(kernel, cl::NullRange,
@@ -132,7 +125,7 @@ public:
    return *this;
  }

-  Tensor operator-() override {
+  Tensor operator-() const override {
    cl::Kernel kernel = openCL.createKernel(OpenCL::Method::NEGATIVE);
    kernel.setArg(0, *data_);
    openCL.getQueue().enqueueNDRangeKernel(kernel, cl::NullRange,
@@ -191,17 +184,17 @@ public:
      if (shape_[axes_[1]] != other.shape_[other.axes_[0]])
        throw std::invalid_argument(
            "Matrix dimensions must match for multiplication");
-      int m = (int)shape_[axes_[0]];
-      int k = (int)shape_[axes_[1]];
-      int n = (int)other.shape_[other.axes_[1]];
+      size_t m = shape_[axes_[0]];
+      size_t k = shape_[axes_[1]];
+      size_t n = other.shape_[other.axes_[1]];
      Tensor<T, 2> result({m, n});
      cl::Kernel kernel = openCL.createKernel(OpenCL::Method::T_MULT);
      kernel.setArg(0, *data_);
      kernel.setArg(1, *other.getData());
      kernel.setArg(2, *result.getData());
-      kernel.setArg(3, m);
-      kernel.setArg(4, n);
-      kernel.setArg(5, k);
+      kernel.setArg(3, (int)m);
+      kernel.setArg(4, (int)n);
+      kernel.setArg(5, (int)k);
      cl::NDRange global_size(((m + TILE_SIZE - 1) / TILE_SIZE) * TILE_SIZE,
                              ((n + TILE_SIZE - 1) / TILE_SIZE) * TILE_SIZE);
      cl::NDRange local_size(TILE_SIZE, TILE_SIZE);
@@ -214,50 +207,11 @@ public:

  std::string toString() const override {
    std::vector<float> result(getSize());
-    openCL.getQueue().enqueueReadBuffer(
-        *data_, CL_TRUE, 0, getSize() * sizeof(T), result.data(), all(event_));
-    std::ostringstream oss;
-    if constexpr (Dim == 0) {
-      oss << "Scalar<" << typeid(T).name() << ">: " << result[0];
-    } else if constexpr (Dim == 1) {
-      oss << "Vector<" << typeid(T).name() << ">(" << shape_[0] << "): [";
-      for (size_t i = 0; i < getSize(); ++i) {
-        oss << result[i];
-        if (i < getSize() - 1)
-          oss << ", ";
-      }
-      oss << "]";
-    } else if constexpr (Dim == 2) {
-      oss << "Matrix<" << typeid(T).name() << ">(" << shape_[axes_[0]] << "x"
-          << shape_[axes_[1]] << "):";
-      for (size_t i = 0; i < shape_[axes_[0]]; ++i) {
-        oss << "\n  [";
-        for (size_t j = 0; j < shape_[axes_[1]]; ++j) {
-          oss << result[i * shape_[axes_[0]] + j];
-          if (j < shape_[axes_[1]] - 1)
-            oss << ", ";
-        }
-        oss << "]";
-      }
-    } else {
-      oss << "Tensor" << Dim << "D<" << typeid(T).name() << ">" << "[";
-      for (size_t i = 0; i < Dim; ++i) {
-        oss << shape_[axes_[i]];
-        if (i < Dim - 1)
-          oss << "x";
-      }
-      oss << "]: [";
-      size_t show = std::min(getSize(), size_t(10));
-      for (size_t i = 0; i < show; ++i) {
-        oss << result[i];
-        if (i < show - 1)
-          oss << ", ";
-      }
-      if (getSize() > 10)
-        oss << ", ...";
-      oss << "]";
-    }
-    return oss.str();
+    openCL.getQueue().enqueueReadBuffer(*data_, CL_FALSE, 0,
+                                        getSize() * sizeof(T), result.data(),
+                                        all(event_), &event_);
+    event_.wait();
+    return ITensor::format(result);
  }
 };

--- a/src/tensor/pybind.cpp
+++ b/src/tensor/pybind.cpp
@@ -2,10 +2,17 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>

-#include "tensor.hpp"
+#ifdef USE_OPENCL
+#include "opencl/tensor.hpp"
+OpenCL openCL;
+#elif USE_CPU
+#include "cpu/tensor.hpp"
+#endif

 namespace py = pybind11;

+enum class TENSOR_PLATFORM { CPU, OPENCL };
+
 template <typename T, int Dim>
 void register_tensor(py::module &m, const std::string &name) {
  auto tensor = py::class_<Tensor<T, Dim>>(m, name.c_str())
@@ -15,9 +22,9 @@ void register_tensor(py::module &m, const std::string &name) {
                                  const std::vector<T> &>())
                    .def(py::init<const std::array<size_t, Dim> &, T, T>())

-                    .def("get_shape", &TensorInfo<T, Dim>::getShape)
-                    .def("get_axes", &TensorInfo<T, Dim>::getAxes)
-                    .def("get_size", &TensorInfo<T, Dim>::getSize)
+                    .def("get_shape", &Tensor<T, Dim>::getShape)
+                    .def("get_axes", &Tensor<T, Dim>::getAxes)
+                    .def("get_size", &Tensor<T, Dim>::getSize)

                    .def(py::self + py::self)
                    .def(py::self - py::self)
@@ -52,6 +59,7 @@ void register_tensor(py::module &m, const std::string &name) {
        .def("t", &Tensor<T, Dim>::t);
  }

+#ifndef USE_OPENCL
  if constexpr (Dim != 0)
    tensor
        .def(
@@ -91,21 +99,47 @@ void register_tensor(py::module &m, const std::string &name) {
                 t(py::cast<size_t>(indices[I])...) = value;
               }(std::make_index_sequence<Dim>{});
             });
+#endif

-  if constexpr (Dim == 1 || Dim == 2)
+  // if constexpr (Dim == 1 || Dim == 2)
+  if constexpr (Dim == 2)
    tensor.def("__matmul__", &Tensor<T, Dim>::operator%);
 }

 PYBIND11_MODULE(tensor, m) {
  m.doc() = "Tensor math library";

+  py::enum_<TENSOR_PLATFORM>(m, "PLATFORM")
+      .value("CPU", TENSOR_PLATFORM::CPU)
+      .value("OPENCL", TENSOR_PLATFORM::OPENCL)
+      .export_values();
+
+#ifdef USE_OPENCL
+  m.attr("MODE") = TENSOR_PLATFORM::OPENCL;
+#elif USE_CPU
+  m.attr("MODE") = TENSOR_PLATFORM::CPU;
+#endif
+
  register_tensor<float, 0>(m, "Scalar");
  register_tensor<float, 1>(m, "Vector");
  register_tensor<float, 2>(m, "Matrix");
-  // register_tensor<float, 3>(m, "Tensor3");
-  //
-  // register_tensor<int, 0>(m, "iScalar");
-  // register_tensor<int, 1>(m, "iVector");
-  // register_tensor<int, 2>(m, "iMatrix");
-  // register_tensor<int, 3>(m, "iTensor3");
+  register_tensor<float, 3>(m, "Tensor3");
+
+#ifdef USE_OPENCL
+  m.def("init", [](const std::string &programsBasePath) {
+    openCL.init(programsBasePath);
+  });
+#endif
+
+#ifndef USE_OPENCL
+  register_tensor<double, 0>(m, "dScalar");
+  register_tensor<double, 1>(m, "dVector");
+  register_tensor<double, 2>(m, "dMatrix");
+  register_tensor<double, 3>(m, "dTensor3");
+
+  register_tensor<int, 0>(m, "iScalar");
+  register_tensor<int, 1>(m, "iVector");
+  register_tensor<int, 2>(m, "iMatrix");
+  register_tensor<int, 3>(m, "iTensor3");
+#endif
 }
--- a/src/tensor/tensor.hpp
+++ b/src/tensor/tensor.hpp
@@ -3,6 +3,7 @@
 #include <array>
 #include <cstddef>
 #include <string>
+#include <vector>

 template <typename T, int Dim> class Tensor;

@@ -16,6 +17,8 @@ protected:
  void checkItHasSameShape(const ITensor &other) const;
  void checkAxisInDim(int axis) const;

+  std::string format(std::vector<T> data) const;
+
 public:
  typedef class Tensor<T, Dim> Tensor;

@@ -35,8 +38,8 @@ public:
  Tensor &transpose(int axis_a, int axis_b);
  Tensor &t();

-  virtual Tensor operator+() = 0;
-  virtual Tensor operator-() = 0;
+  virtual Tensor operator+() const = 0;
+  virtual Tensor operator-() const = 0;

  virtual Tensor &operator+=(const T scalar) = 0;
  virtual Tensor &operator*=(const T scalar) = 0;
--- a/src/tensor/tensor.pyi
+++ b/src/tensor/tensor.pyi
@@ -4,7 +4,7 @@ Tensor math library
 from __future__ import annotations
 import collections.abc
 import typing
-__all__: list[str] = ['Matrix', 'Scalar', 'Vector']
+__all__: list[str] = ['CPU', 'MODE', 'Matrix', 'OPENCL', 'PLATFORM', 'Scalar', 'Tensor3', 'Vector', 'init']
 class Matrix:
    @typing.overload
    def __add__(self, arg0: Matrix) -> Matrix:
@@ -13,12 +13,6 @@ class Matrix:
    def __add__(self, arg0: typing.SupportsFloat) -> Matrix:
        ...
    @typing.overload
-    def __getitem__(self, arg0: typing.SupportsInt) -> float:
-        ...
-    @typing.overload
-    def __getitem__(self, arg0: tuple) -> float:
-        ...
-    @typing.overload
    def __iadd__(self, arg0: Matrix) -> Matrix:
        ...
    @typing.overload
@@ -71,12 +65,6 @@ class Matrix:
    def __rsub__(self, arg0: typing.SupportsFloat) -> Matrix:
        ...
    @typing.overload
-    def __setitem__(self, arg0: typing.SupportsInt, arg1: typing.SupportsFloat) -> None:
-        ...
-    @typing.overload
-    def __setitem__(self, arg0: tuple, arg1: typing.SupportsFloat) -> None:
-        ...
-    @typing.overload
    def __sub__(self, arg0: Matrix) -> Matrix:
        ...
    @typing.overload
@@ -98,6 +86,43 @@ class Matrix:
    @typing.overload
    def transpose(self, arg0: typing.SupportsInt, arg1: typing.SupportsInt) -> Matrix:
        ...
+class PLATFORM:
+    """
+    Members:
+    
+      CPU
+    
+      OPENCL
+    """
+    CPU: typing.ClassVar[PLATFORM]  # value = <PLATFORM.CPU: 0>
+    OPENCL: typing.ClassVar[PLATFORM]  # value = <PLATFORM.OPENCL: 1>
+    __members__: typing.ClassVar[dict[str, PLATFORM]]  # value = {'CPU': <PLATFORM.CPU: 0>, 'OPENCL': <PLATFORM.OPENCL: 1>}
+    def __eq__(self, other: typing.Any) -> bool:
+        ...
+    def __getstate__(self) -> int:
+        ...
+    def __hash__(self) -> int:
+        ...
+    def __index__(self) -> int:
+        ...
+    def __init__(self, value: typing.SupportsInt) -> None:
+        ...
+    def __int__(self) -> int:
+        ...
+    def __ne__(self, other: typing.Any) -> bool:
+        ...
+    def __repr__(self) -> str:
+        ...
+    def __setstate__(self, state: typing.SupportsInt) -> None:
+        ...
+    def __str__(self) -> str:
+        ...
+    @property
+    def name(self) -> str:
+        ...
+    @property
+    def value(self) -> int:
+        ...
 class Scalar:
    @typing.overload
    def __add__(self, arg0: Scalar) -> Scalar:
@@ -169,6 +194,85 @@ class Scalar:
        ...
    def get_size(self) -> int:
        ...
+class Tensor3:
+    @typing.overload
+    def __add__(self, arg0: Tensor3) -> Tensor3:
+        ...
+    @typing.overload
+    def __add__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    @typing.overload
+    def __iadd__(self, arg0: Tensor3) -> Tensor3:
+        ...
+    @typing.overload
+    def __iadd__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    @typing.overload
+    def __imul__(self, arg0: Tensor3) -> Tensor3:
+        ...
+    @typing.overload
+    def __imul__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    @typing.overload
+    def __init__(self, arg0: typing.Annotated[collections.abc.Sequence[typing.SupportsInt], "FixedSize(3)"]) -> None:
+        ...
+    @typing.overload
+    def __init__(self, arg0: typing.Annotated[collections.abc.Sequence[typing.SupportsInt], "FixedSize(3)"], arg1: typing.SupportsFloat) -> None:
+        ...
+    @typing.overload
+    def __init__(self, arg0: typing.Annotated[collections.abc.Sequence[typing.SupportsInt], "FixedSize(3)"], arg1: collections.abc.Sequence[typing.SupportsFloat]) -> None:
+        ...
+    @typing.overload
+    def __init__(self, arg0: typing.Annotated[collections.abc.Sequence[typing.SupportsInt], "FixedSize(3)"], arg1: typing.SupportsFloat, arg2: typing.SupportsFloat) -> None:
+        ...
+    @typing.overload
+    def __isub__(self, arg0: Tensor3) -> Tensor3:
+        ...
+    @typing.overload
+    def __isub__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    def __itruediv__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    @typing.overload
+    def __mul__(self, arg0: Tensor3) -> Tensor3:
+        ...
+    @typing.overload
+    def __mul__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    def __neg__(self) -> Tensor3:
+        ...
+    def __pos__(self) -> Tensor3:
+        ...
+    def __radd__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    def __repr__(self) -> str:
+        ...
+    def __rmul__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    def __rsub__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    @typing.overload
+    def __sub__(self, arg0: Tensor3) -> Tensor3:
+        ...
+    @typing.overload
+    def __sub__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    def __truediv__(self, arg0: typing.SupportsFloat) -> Tensor3:
+        ...
+    def get_axes(self) -> typing.Annotated[list[int], "FixedSize(3)"]:
+        ...
+    def get_shape(self) -> typing.Annotated[list[int], "FixedSize(3)"]:
+        ...
+    def get_size(self) -> int:
+        ...
+    def t(self) -> Tensor3:
+        ...
+    @typing.overload
+    def transpose(self, arg0: typing.Annotated[collections.abc.Sequence[typing.SupportsInt], "FixedSize(3)"]) -> Tensor3:
+        ...
+    @typing.overload
+    def transpose(self, arg0: typing.SupportsInt, arg1: typing.SupportsInt) -> Tensor3:
+        ...
 class Vector:
    @typing.overload
    def __add__(self, arg0: Vector) -> Vector:
@@ -177,12 +281,6 @@ class Vector:
    def __add__(self, arg0: typing.SupportsFloat) -> Vector:
        ...
    @typing.overload
-    def __getitem__(self, arg0: typing.SupportsInt) -> float:
-        ...
-    @typing.overload
-    def __getitem__(self, arg0: tuple) -> float:
-        ...
-    @typing.overload
    def __iadd__(self, arg0: Vector) -> Vector:
        ...
    @typing.overload
@@ -214,8 +312,6 @@ class Vector:
        ...
    def __itruediv__(self, arg0: typing.SupportsFloat) -> Vector:
        ...
-    def __matmul__(self, arg0: Vector) -> Scalar:
-        ...
    @typing.overload
    def __mul__(self, arg0: Vector) -> Vector:
        ...
@@ -235,12 +331,6 @@ class Vector:
    def __rsub__(self, arg0: typing.SupportsFloat) -> Vector:
        ...
    @typing.overload
-    def __setitem__(self, arg0: typing.SupportsInt, arg1: typing.SupportsFloat) -> None:
-        ...
-    @typing.overload
-    def __setitem__(self, arg0: tuple, arg1: typing.SupportsFloat) -> None:
-        ...
-    @typing.overload
    def __sub__(self, arg0: Vector) -> Vector:
        ...
    @typing.overload
@@ -254,3 +344,8 @@ class Vector:
        ...
    def get_size(self) -> int:
        ...
+def init(arg0: str) -> None:
+    ...
+CPU: PLATFORM  # value = <PLATFORM.CPU: 0>
+MODE: PLATFORM  # value = <PLATFORM.OPENCL: 1>
+OPENCL: PLATFORM  # value = <PLATFORM.OPENCL: 1>
--- a/src/tensor/tensor.tpp
+++ b/src/tensor/tensor.tpp
@@ -2,6 +2,8 @@

 #include "tensor.hpp"

+#include <iomanip>
+#include <sstream>
 #include <stdexcept>

 // ===== UTILS =====
@@ -34,6 +36,117 @@ void ITensor<T, Dim>::checkAxisInDim(int axis) const {
    throw std::invalid_argument("Invalid axis index");
 }

+template <typename T, int Dim>
+std::string ITensor<T, Dim>::format(std::vector<T> data) const {
+  std::ostringstream oss;
+  static auto formatValue = [](T value) -> std::string {
+    std::ostringstream value_oss;
+    if constexpr (std::is_floating_point_v<T>) {
+      value_oss << std::fixed << std::setprecision(3) << value;
+      std::string str = value_oss.str();
+      if (str.find('.') != std::string::npos) {
+        str = str.substr(0, str.find_last_not_of('0') + 1);
+        if (str.back() == '.')
+          str.pop_back();
+      }
+      return str;
+    } else {
+      value_oss << value;
+      return value_oss.str();
+    }
+  };
+
+  if constexpr (Dim == 0) {
+    oss << "Scalar<" << typeid(T).name() << ">: " << formatValue(data[0]);
+  } else if constexpr (Dim == 1) {
+    oss << "Vector<" << typeid(T).name() << ">(" << shape_[0] << "): [";
+    for (size_t i = 0; i < getSize(); ++i) {
+      oss << formatValue(data[i]);
+      if (i < getSize() - 1)
+        oss << ", ";
+    }
+    oss << "]";
+  } else if constexpr (Dim == 2) {
+    const size_t rows = shape_[axes_[0]];
+    const size_t cols = shape_[axes_[1]];
+    oss << "Matrix<" << typeid(T).name() << ">(" << rows << "x" << cols << "):";
+
+    const size_t MAX_FULL_ROWS = 8;
+    const size_t MAX_FULL_COLS = 8;
+    const size_t SHOW_FIRST = 3;
+    const size_t SHOW_LAST = 3;
+    bool show_abbreviated_rows = rows > MAX_FULL_ROWS;
+    bool show_abbreviated_cols = cols > MAX_FULL_COLS;
+    std::vector<std::string> formatted_values;
+    size_t max_width = 0;
+    for (size_t i = 0; i < rows; ++i) {
+      if (show_abbreviated_rows && i >= SHOW_FIRST && i < rows - SHOW_LAST)
+        continue;
+      for (size_t j = 0; j < cols; ++j) {
+        if (show_abbreviated_cols && j >= SHOW_FIRST && j < cols - SHOW_LAST)
+          continue;
+        std::string formatted = formatValue(data[i * cols + j]);
+        formatted_values.push_back(formatted);
+        max_width = std::max(max_width, formatted.length());
+      }
+    }
+    for (size_t i = 0; i < rows; ++i) {
+      if (show_abbreviated_rows && i >= SHOW_FIRST && i < rows - SHOW_LAST) {
+        if (i == SHOW_FIRST) {
+          oss << "\n  ";
+          for (size_t j = 0; j < cols; ++j) {
+            if (show_abbreviated_cols && j >= SHOW_FIRST &&
+                j < cols - SHOW_LAST) {
+              if (j == SHOW_FIRST)
+                oss << std::string(max_width, '.') << ", ";
+              continue;
+            }
+            oss << std::string(max_width, '.');
+            if (!((!show_abbreviated_cols && j == cols - 1) ||
+                  (show_abbreviated_cols &&
+                   ((j < SHOW_FIRST && j == SHOW_FIRST - 1) || j == cols - 1))))
+              oss << ", ";
+          }
+        }
+        continue;
+      }
+      oss << "\n  [";
+      for (size_t j = 0; j < cols; ++j) {
+        if (show_abbreviated_cols && j >= SHOW_FIRST && j < cols - SHOW_LAST) {
+          if (j == SHOW_FIRST)
+            oss << " " << std::setw(max_width) << std::left << "..." << ", ";
+          continue;
+        }
+        std::string formatted = formatValue(data[i * cols + j]);
+        oss << std::setw(max_width) << std::left << formatted;
+        if (!((!show_abbreviated_cols && j == cols - 1) ||
+              (show_abbreviated_cols &&
+               ((j < SHOW_FIRST && j == SHOW_FIRST - 1) || j == cols - 1))))
+          oss << ", ";
+      }
+      oss << "]";
+    }
+  } else {
+    oss << "Tensor" << Dim << "D<" << typeid(T).name() << ">" << "[";
+    for (size_t i = 0; i < Dim; ++i) {
+      oss << shape_[axes_[i]];
+      if (i < Dim - 1)
+        oss << "x";
+    }
+    oss << "]: [";
+    size_t show = std::min(getSize(), size_t(10));
+    for (size_t i = 0; i < show; ++i) {
+      oss << formatValue(data[i]);
+      if (i < show - 1)
+        oss << ", ";
+    }
+    if (getSize() > 10)
+      oss << ", ...";
+    oss << "]";
+  }
+  return oss.str();
+}
+
 // ====== CONSTRUCT =====
 template <typename T, int Dim>
 ITensor<T, Dim>::ITensor(const std::array<size_t, Dim> &shape) {