diff --git a/README.md b/README.md index 662f0ae..68b889b 100644 --- a/README.md +++ b/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: diff --git a/src/core/include/tensor.h b/src/core/include/tensor.h index ba3546e..3e6c876 100644 --- a/src/core/include/tensor.h +++ b/src/core/include/tensor.h @@ -223,11 +223,12 @@ class Tensor { Tensor &fill_(const any_t &value); int64_t offset(const std::vector &indices) const; Tensor contiguous() const; - Tensor as_strided(const std::vector sizes, const std::vector strides, int64_t storage_offset = 0) const; + Tensor as_strided(std::vector sizes, std::vector strides, int64_t storage_offset = 0) const; Tensor permute(const std::vector dims) const; Tensor slice(int64_t dim, std::optional start, std::optional 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 sizes) const; Tensor _half() const; Tensor _float() const; diff --git a/src/core/tensor.cpp b/src/core/tensor.cpp index 1ee3dbb..b2ef315 100644 --- a/src/core/tensor.cpp +++ b/src/core/tensor.cpp @@ -116,12 +116,21 @@ Tensor Tensor::contiguous() const { return gpu::clone(*this); } -Tensor Tensor::as_strided(const std::vector sizes, - const std::vector strides, - int64_t storage_offset) const { - // in-bounds check +Tensor Tensor::as_strided(std::vector sizes, + std::vector 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 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); } diff --git a/src/register.cpp b/src/register.cpp index c4f0f47..741a263 100644 --- a/src/register.cpp +++ b/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 dims; + for (auto arg : args) { + int64_t idx = arg.cast(); + dims.emplace_back(idx); + } + return self.permute(dims); + }) + .def("view", [](Tensor &self, py::args args) { + std::vector dims; + for (auto arg : args) { + int64_t idx = arg.cast(); + 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(key)) { auto t = key.cast(); - CHECK_FAIL(t.size() == self.dim()); + CHECK_FAIL(t.size() <= self.dim()); int dim = 0; for (auto item : t) { if (py::isinstance(item)) { @@ -132,7 +148,7 @@ PYBIND11_MODULE(kfunca, m) { output = output.slice(dim, start, end, step); dim++; } else if (py::isinstance(item)) { - size_t idx = item.cast(); + int64_t idx = item.cast(); 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(); + int64_t idx = key.cast(); output = output.select(0, idx); } return output; diff --git a/test/common.py b/test/common.py new file mode 100644 index 0000000..a561f76 --- /dev/null +++ b/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) diff --git a/test/test_gemm.py b/test/test_gemm.py index a8ce911..2f89f0e 100644 --- a/test/test_gemm.py +++ b/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__': diff --git a/test/test_nn.py b/test/test_nn.py index 7dcb862..ce41ba6 100644 --- a/test/test_nn.py +++ b/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__': diff --git a/test/test_tensor.py b/test/test_tensor.py index 12a49d7..84c5256 100644 --- a/test/test_tensor.py +++ b/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,8 +141,8 @@ 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)) @@ -149,14 +150,14 @@ def test_convert(self): 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,8 +181,8 @@ 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) @@ -189,8 +190,8 @@ def test_sort_large_slice(self): 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__':