Enable tensor view operation

README.md

@@ -5,7 +5,6 @@ Here are the supported features:
 #### 1. Basic infrastructure
 
 - [x] GPU Launcher
-- [x] Device Allocator
 - [x] Caching Allocator
 - [x] Tensor Implementation
 - [x] Tensor Iterator
@@ -14,14 +13,15 @@ Here are the supported features:
 
 #### 2. GPU Operator
 
-> Structured operator:
+> Basic operator:
 
 - [x] from_numpy/to_numpy
 - [x] add/sub/mul/div
 - [x] permute/contiguous/copy
 - [x] sum/mean
 - [x] sort/topk
-- [ ] slice/concat/split/stack
+- [x] slice/view
+- [ ] concat/split/stack
 
 > Neural network operator:
 

src/core/include/tensor.h

@@ -223,11 +223,12 @@ class Tensor {
     Tensor &fill_(const any_t &value);
     int64_t offset(const std::vector<int64_t> &indices) const;
     Tensor contiguous() const;
-    Tensor as_strided(const std::vector<int64_t> sizes, const std::vector<int64_t> strides, int64_t storage_offset = 0) const;
+    Tensor as_strided(std::vector<int64_t> sizes, std::vector<int64_t> strides, int64_t storage_offset = 0) const;
     Tensor permute(const std::vector<int64_t> dims) const;
     Tensor slice(int64_t dim, std::optional<int64_t> start, std::optional<int64_t> end, int64_t step) const;
     Tensor select(int64_t dim, int64_t index) const;
     Tensor narrow(int64_t dim, int64_t start, int64_t length) const;
+    Tensor view(std::vector<int64_t> sizes) const;
 
     Tensor _half() const;
     Tensor _float() const;

src/core/tensor.cpp

@@ -116,12 +116,21 @@ Tensor Tensor::contiguous() const {
     return gpu::clone(*this);
 }
 
-Tensor Tensor::as_strided(const std::vector<int64_t> sizes,
-                          const std::vector<int64_t> strides,
-                          int64_t storage_offset) const {
-    // in-bounds check
+Tensor Tensor::as_strided(std::vector<int64_t> sizes,
+                          std::vector<int64_t> strides, int64_t storage_offset) const {
     auto ndim = sizes.size();
-    CHECK_FAIL(ndim == strides.size());
+    bool has_strides = strides.size() > 0;
+    if (has_strides) {
+        CHECK_FAIL(ndim == strides.size());
+    } else {
+        strides.reserve(ndim);
+        int64_t cumprod = 1;
+        for (int i = ndim - 1; i >= 0; i--) {
+            strides[i] = cumprod;
+            cumprod *= sizes[i];
+        }
+    }
+    // in-bounds check
     auto [min_offset, max_offset] = compute_offset_range(sizes, strides, ndim);
     min_offset += storage_offset;
     max_offset += storage_offset;
@@ -235,6 +244,29 @@ Tensor Tensor::narrow(int64_t dim, int64_t start, int64_t length) const {
     return this->slice(dim, start, start + length, 1);
 }
 
+Tensor Tensor::view(std::vector<int64_t> sizes) const {
+    CHECK_FAIL(this->is_contiguous_);
+    int64_t cumprod = 1;
+    bool has_neg_dim = false;
+    int64_t neg_dim = -1;
+    for (auto i = 0; i < sizes.size(); ++i) {
+        auto size = sizes[i];
+        if (size < 0) {
+            CHECK_FAIL(has_neg_dim == false);
+            has_neg_dim = true;
+            neg_dim = i;
+        } else {
+            cumprod *= size;
+        }
+    }
+    if (has_neg_dim) {
+        sizes[neg_dim] = this->numel_ / cumprod;
+        cumprod *= sizes[neg_dim];
+    }
+    CHECK_FAIL(cumprod == this->numel_);
+    return this->as_strided(sizes, {});
+}
+
 Tensor Tensor::_half() const {
     return gpu::convert(*this, ScalarType::Half);
 }

src/register.cpp

@@ -115,14 +115,30 @@ PYBIND11_MODULE(kfunca, m) {
         })
         .def("storage_ref_count", &Tensor::storage_ref_count)
         .def("contiguous", &Tensor::contiguous)
-        .def("permute", &Tensor::permute)
+        .def("permute", [](Tensor &self, py::args args) {
+            CHECK_FAIL(args.size() == self.dim());
+            std::vector<int64_t> dims;
+            for (auto arg : args) {
+                int64_t idx = arg.cast<int64_t>();
+                dims.emplace_back(idx);
+            }
+            return self.permute(dims);
+        })
+        .def("view", [](Tensor &self, py::args args) {
+            std::vector<int64_t> dims;
+            for (auto arg : args) {
+                int64_t idx = arg.cast<int64_t>();
+                dims.emplace_back(idx);
+            }
+            return self.view(dims);
+        })
         .def("sort", &Tensor::sort)
         .def("topk", &Tensor::topk)
         .def("__getitem__", [](Tensor &self, py::object key) {
             Tensor output = self;
             if (py::isinstance<py::tuple>(key)) {
                 auto t = key.cast<py::tuple>();
-                CHECK_FAIL(t.size() == self.dim());
+                CHECK_FAIL(t.size() <= self.dim());
                 int dim = 0;
                 for (auto item : t) {
                     if (py::isinstance<py::slice>(item)) {
@@ -132,7 +148,7 @@ PYBIND11_MODULE(kfunca, m) {
                         output = output.slice(dim, start, end, step);
                         dim++;
                     } else if (py::isinstance<py::int_>(item)) {
-                        size_t idx = item.cast<size_t>();
+                        int64_t idx = item.cast<int64_t>();
                         output = output.select(dim, idx);
                     }
                 }
@@ -142,7 +158,7 @@ PYBIND11_MODULE(kfunca, m) {
                 s.compute(self.shape(0), &start, &end, &step, &len);
                 output = output.slice(0, start, end, step);
             } else {
-                size_t idx = key.cast<size_t>();
+                int64_t idx = key.cast<int64_t>();
                 output = output.select(0, idx);
             }
             return output;

test/common.py

@@ -0,0 +1,11 @@
+import kfunca
+import torch
+import numpy as np
+
+
+def assert_allclose(tensor_a, tensor_b, atol=1e-3, rtol=1e-3):
+    if not isinstance(tensor_a, np.ndarray):
+        tensor_a = tensor_a.contiguous().numpy()
+    if not isinstance(tensor_b, np.ndarray):
+        tensor_b = tensor_b.contiguous().numpy()
+    assert(np.allclose(tensor_a, tensor_b, rtol=atol, atol=rtol) == True)

test/test_gemm.py

@@ -1,5 +1,6 @@
 import kfunca
 import numpy as np
+from common import assert_allclose
 
 print(kfunca.__file__)
 
@@ -13,7 +14,7 @@ def test_gemm_base(self):
         print(a_gpu.sizes(), b_gpu.sizes())
         out_gpu = kfunca.gemm(a_gpu, b_gpu, 1.0, 0.0)
         out = np.matmul(a, b)
-        assert(np.allclose(out, out_gpu.numpy()) == True)
+        assert_allclose(out, out_gpu)
 
 
 if __name__ == '__main__':

test/test_nn.py

@@ -2,6 +2,7 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
+from common import assert_allclose
 
 print(kfunca.__file__)
 
@@ -29,7 +30,7 @@ def test_causal_attention(self):
             k_ref = torch.from_numpy(k_)
             v_ref = torch.from_numpy(v_)
             out_ref = F.scaled_dot_product_attention(q_ref, k_ref, v_ref, is_causal=True).numpy()
-            assert(np.allclose(out, out_ref, rtol=1e-3, atol=1e-3) == True)
+            assert_allclose(out, out_ref)
 
 
 if __name__ == '__main__':

test/test_tensor.py

@@ -1,6 +1,7 @@
 import kfunca
 import torch
 import numpy as np
+from common import assert_allclose
 
 print(kfunca.__file__)
 
@@ -9,7 +10,7 @@ class TestTensorImpl(object):
     def test_tensor_impl(self):
         arr = np.random.uniform(-10, 10, size=(2, 3))
         arr_gpu = kfunca.from_numpy(arr, 0)
-        assert(np.allclose(arr, arr_gpu.numpy()) == True)
+        assert_allclose(arr, arr_gpu)
 
     def test_tensor_add(self):
         for shape in ((2,3), (1000), (12,11,3331)):
@@ -18,12 +19,12 @@ def test_tensor_add(self):
             arr_gpu = kfunca.from_numpy(arr, 0)
             arr_gpu_2 = arr_gpu + arr_gpu
             arr_gpu_2_cpu = arr_gpu_2.numpy()
-            assert(np.allclose(arr_2, arr_gpu_2_cpu) == True)
+            assert_allclose(arr_2, arr_gpu_2_cpu)
             arr1 = np.random.uniform(-10, 10, size=shape).astype(np.int32)
             arr2 = np.random.uniform(-10, 10, size=shape).astype(np.float32)
             out = arr1 + arr2
             out_gpu = kfunca.from_numpy(arr1, 0) + kfunca.from_numpy(arr2, 0)
-            assert(np.allclose(out, out_gpu.numpy()) == True)
+            assert_allclose(out, out_gpu)
 
     def test_inplace_op(self):
         shape1 = (5,7,11)
@@ -36,35 +37,35 @@ def test_inplace_op(self):
         arr1 += arr2
         arr1_gpu += arr2_gpu
         assert(addr1 == arr1_gpu.data_ptr())
-        assert(np.allclose(arr1, arr1_gpu.numpy()) == True)
+        assert_allclose(arr1, arr1_gpu)
         arr1 -= arr2
         arr1_gpu -= arr2_gpu
         assert(addr1 == arr1_gpu.data_ptr())
-        assert(np.allclose(arr1, arr1_gpu.numpy()) == True)
+        assert_allclose(arr1, arr1_gpu)
         arr1 *= arr2
         arr1_gpu *= arr2_gpu
         assert(addr1 == arr1_gpu.data_ptr())
-        assert(np.allclose(arr1, arr1_gpu.numpy()) == True)
+        assert_allclose(arr1, arr1_gpu)
         arr1 /= arr2
         arr1_gpu /= arr2_gpu
         assert(addr1 == arr1_gpu.data_ptr())
-        assert(np.allclose(arr1, arr1_gpu.numpy()) == True)
+        assert_allclose(arr1, arr1_gpu)
         arr1 += 2
         arr1_gpu += 2
         assert(addr1 == arr1_gpu.data_ptr())
-        assert(np.allclose(arr1, arr1_gpu.numpy()) == True)
+        assert_allclose(arr1, arr1_gpu)
         arr1 -= 3
         arr1_gpu -= 3
         assert(addr1 == arr1_gpu.data_ptr())
-        assert(np.allclose(arr1, arr1_gpu.numpy()) == True)
+        assert_allclose(arr1, arr1_gpu)
         arr1 *= 4
         arr1_gpu *= 4
         assert(addr1 == arr1_gpu.data_ptr())
-        assert(np.allclose(arr1, arr1_gpu.numpy()) == True)
+        assert_allclose(arr1, arr1_gpu)
         arr1 /= 5
         arr1_gpu /= 5
         assert(addr1 == arr1_gpu.data_ptr())
-        assert(np.allclose(arr1, arr1_gpu.numpy()) == True)
+        assert_allclose(arr1, arr1_gpu)
 
     def test_data_ptr(self):
         import copy
@@ -98,12 +99,12 @@ def test_broadcast_basic_binary(self):
                     arr2 = np.random.uniform(-10, 10, size=shape[1]).astype(np.float32)
                     out = eval("arr1 {} arr2".format(op))
                     out_gpu = eval("kfunca.from_numpy(arr1, 0) {} kfunca.from_numpy(arr2, 0)".format(op))
-                    assert(np.allclose(out, out_gpu.numpy()) == True)
+                    assert_allclose(out, out_gpu)
                     arr1 = np.random.uniform(-10, 10, size=shape[0]).astype(np.int32)
                     arr2 = np.random.uniform(-10, 10, size=shape[1]).astype(np.float32)
                     out = eval("arr1 {} arr2".format(op))
                     out_gpu = eval("kfunca.from_numpy(arr1, 0) {} kfunca.from_numpy(arr2, 0)".format(op))
-                    assert(np.allclose(out, out_gpu.numpy()) == True)
+                    assert_allclose(out, out_gpu)
 
     def test_reduce(self):
         for op in ['sum', 'mean']:
@@ -113,7 +114,7 @@ def test_reduce(self):
                 arr_sum = eval("np.{}(arr, axis=dim, keepdims=True)".format(op))
                 arr_gpu = kfunca.from_numpy(arr, 0)
                 arr_gpu_sum = eval("arr_gpu.{}(dim)".format(op))
-                assert(np.allclose(arr_sum, arr_gpu_sum.numpy(), rtol=1e-2, atol=1e-2) == True)
+                assert_allclose(arr_sum, arr_gpu_sum, atol=1e-2, rtol=1e-2)
 
     def test_mean_std(self):
         shape = (13, 325, 127)
@@ -125,8 +126,8 @@ def test_mean_std(self):
         var = ((arr_ - mean) * (arr_ - mean)).sum(dim)
         var = var / divisor
         mean_var = arr_.mean_var(dim, False)
-        assert(np.allclose(mean.numpy(), mean_var[0].numpy(), rtol=1e-2, atol=1e-2) == True)
-        assert(np.allclose(var.numpy(), mean_var[1].numpy(), rtol=1e-2, atol=1e-2) == True)
+        assert_allclose(mean, mean_var[0], atol=1e-2, rtol=1e-2)
+        assert_allclose(var, mean_var[1], atol=1e-2, rtol=1e-2)
         kfunca.memstat()
 
     def test_norm_stat(self):
@@ -140,23 +141,23 @@ def test_norm_stat(self):
             var = np.sum(var, axis=dim, keepdims=True)
             invstd = 1.0 / np.sqrt(var / divisor)
             mean_invstd = arr_.norm_stat(dim)
-            assert(np.allclose(mean, mean_invstd[0].numpy(), rtol=1e-3, atol=1e-3) == True)
-            assert(np.allclose(invstd, mean_invstd[1].numpy(), rtol=1e-3, atol=1e-3) == True)
+            assert_allclose(mean, mean_invstd[0])
+            assert_allclose(invstd, mean_invstd[1])
 
     def test_convert(self):
         arr = np.random.uniform(-10, 10, size=(2, 3))
         arr_gpu = kfunca.from_numpy(arr, 0)
         arr_gpu_half = arr_gpu.half()
         arr_gpu *= arr_gpu
         arr_gpu_half *= arr_gpu_half
-        assert(np.allclose(arr_gpu.numpy(), arr_gpu_half.float().numpy(), rtol=1e-3, atol=1e-3) == True)
+        assert_allclose(arr_gpu, arr_gpu_half.float())
 
     def test_permute(self):
         arr = np.random.uniform(-10, 10, size=(16, 8, 64, 11)) # 0,1,2,3
         arr_p = arr.transpose(2,1,0,3)
         arr_gpu = kfunca.from_numpy(arr, 0)
-        arr_gpu_p = arr_gpu.permute([2,1,0,3]).contiguous()
-        assert(np.allclose(arr_gpu_p.numpy(), arr_p, rtol=1e-3, atol=1e-3) == True)
+        arr_gpu_p = arr_gpu.permute(2,1,0,3).contiguous()
+        assert_allclose(arr_gpu_p, arr_p)
 
     def test_sort_small_slice(self):
         shapes = [
@@ -180,17 +181,17 @@ def test_sort_small_slice(self):
                         res, ind = torch.sort(arr_t, dim=dim, descending=descending, stable=True)
                         arr_gpu = kfunca.from_numpy(arr, 0)
                         res_gpu, ind_gpu = arr_gpu.sort(dim, descending)
-                        assert(np.allclose(res_gpu.numpy(), res.numpy(), rtol=1e-3, atol=1e-3) == True)
-                        assert(np.allclose(ind_gpu.numpy(), ind.numpy(), rtol=1e-3, atol=1e-3) == True)
+                        assert_allclose(res_gpu, res)
+                        assert_allclose(ind_gpu, ind)
 
     def test_sort_large_slice(self):
         arr = np.random.uniform(-1000, 1000, size=(4, 1024000)).astype(np.float32)
         res = np.sort(arr, axis=1)
         ind = np.argsort(arr, axis=1, kind='stable')
         arr_gpu = kfunca.from_numpy(arr, 0)
         res_gpu, ind_gpu = arr_gpu.sort(1, False)
-        assert(np.allclose(res_gpu.numpy(), res, rtol=1e-3, atol=1e-3) == True)
-        assert(np.allclose(ind_gpu.numpy(), ind, rtol=1e-3, atol=1e-3) == True)
+        assert_allclose(res_gpu, res)
+        assert_allclose(ind_gpu, ind)
 
     def test_topk_small(self):
         shapes = [
@@ -211,7 +212,7 @@ def test_topk_small(self):
                         res, ind = torch.topk(arr_t, k, dim=dim, largest=descending)
                         arr_gpu = kfunca.from_numpy(arr, 0)
                         res_gpu, ind_gpu = arr_gpu.topk(k, dim, descending)
-                        assert(np.allclose(res_gpu.numpy(), res.numpy(), rtol=1e-3, atol=1e-3) == True)
+                        assert_allclose(res_gpu, res)
 
     def test_topk_large(self):
         for k in [2049, 22223]:
@@ -220,15 +221,23 @@ def test_topk_large(self):
             res, ind = torch.topk(arr_t, k, dim=1, largest=True)
             arr_gpu = kfunca.from_numpy(arr, 0)
             res_gpu, ind_gpu = arr_gpu.topk(k, 1, True)
-            assert(np.allclose(res_gpu.numpy(), res.numpy(), rtol=1e-3, atol=1e-3) == True)
+            assert_allclose(res_gpu, res)
 
     def test_tensor_slice(self):
-        arr = np.random.uniform(-10000, 10000, size=(11, 155, 33)).astype(np.float32)
+        arr = np.random.uniform(-10000, 10000, size=(11, 155, 33, 5)).astype(np.float32)
         arr_t = torch.from_numpy(arr)
         arr_gpu = kfunca.from_numpy(arr, 0)
         arr_t_ = arr_t[3, 3:8, 4:11:2]
         arr_gpu_ = arr_gpu[3, 3:8, 4:11:2]
-        assert(np.allclose(arr_t_.numpy(), arr_gpu_.contiguous().numpy(), rtol=1e-3, atol=1e-3) == True)
+        assert_allclose(arr_t_, arr_gpu_.contiguous())
+
+    def test_view(self):
+        arr = np.random.uniform(-10000, 10000, size=(5,2,11,23)).astype(np.float32)
+        arr_t = torch.from_numpy(arr)
+        arr_gpu = kfunca.from_numpy(arr, 0)
+        arr_t = arr_t.view(5,-1,23).contiguous() + 1
+        arr_gpu = arr_gpu.view(5,-1,23).contiguous() + 1
+        assert_allclose(arr_t, arr_gpu)
 
 
 if __name__ == '__main__':