helper function for transposing host tensor

library/include/ck/library/utility/host_tensor.hpp

@@ -433,6 +433,29 @@ struct Tensor
         return mData[mDesc.GetOffsetFromMultiIndex(idx)];
     }
 
+    Tensor<T> Transpose(std::vector<size_t> axes = {})
+    {
+        if(axes.empty())
+        {
+            axes.resize(this->GetNumOfDimension());
+            std::iota(axes.rbegin(), axes.rend(), 0);
+        }
+        if(axes.size() != mDesc.GetNumOfDimension())
+        {
+            throw std::runtime_error(
+                "Tensor::Transpose(): size of axes must match tensor dimension");
+        }
+        std::vector<size_t> tlengths, tstrides;
+        for(const auto& axis : axes)
+        {
+            tlengths.push_back(GetLengths()[axis]);
+            tstrides.push_back(GetStrides()[axis]);
+        }
+        Tensor<T> ret(*this);
+        ret.mDesc = HostTensorDescriptor(tlengths, tstrides);
+        return ret;
+    }
+
     typename Data::iterator begin() { return mData.begin(); }
 
     typename Data::iterator end() { return mData.end(); }

test/CMakeLists.txt

@@ -58,3 +58,4 @@ add_subdirectory(batchnorm)
 if(GPU_TARGETS MATCHES "gfx1100")
     add_subdirectory(wmma_op)
 endif()
+add_subdirectory(host_tensor)

test/host_tensor/CMakeLists.txt

+add_gtest_executable(test_host_tensor test_host_tensor.cpp)
+target_link_libraries(test_host_tensor PRIVATE utility)
\ No newline at end of file

test/host_tensor/test_host_tensor.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iostream>
+#include <vector>
+#include <gtest/gtest.h>
+
+#include "ck/ck.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+
+using namespace ck;
+
+TEST(HostTensorTranspose, TestBadArugment)
+{
+    Tensor<float> tensor({13, 7});
+
+    EXPECT_THROW(tensor.Transpose({0}), std::runtime_error);
+    EXPECT_THROW(tensor.Transpose({0, 1, 2}), std::runtime_error);
+}
+
+TEST(HostTensorTranspose, Test2D)
+{
+    std::vector<size_t> lengths  = {13, 7};
+    std::vector<size_t> tlengths = {7, 13};
+
+    Tensor<float> tensor(lengths);
+    tensor(0, 0) = 0.f;
+    tensor(3, 4) = 34.f;
+
+    EXPECT_EQ(tensor.GetLengths(), lengths);
+    EXPECT_EQ(tensor(0, 0), 0.f);
+    EXPECT_EQ(tensor(3, 4), 34.f);
+    EXPECT_EQ(tensor(4, 3), 0.f);
+
+    EXPECT_EQ(tensor.Transpose().GetLengths(), tlengths);
+    EXPECT_EQ(tensor.Transpose()(0, 0), 0.f);
+    EXPECT_EQ(tensor.Transpose()(4, 3), 34.f);
+    EXPECT_EQ(tensor.Transpose()(3, 4), 0.f);
+}
+
+TEST(HostTensorTranspose, Test3D)
+{
+    std::vector<size_t> lengths  = {13, 7, 5};
+    std::vector<size_t> tlengths = {5, 7, 13};
+
+    Tensor<float> tensor(lengths);
+    tensor(0, 0, 0) = 0.f;
+    tensor(3, 4, 2) = 342.f;
+
+    EXPECT_EQ(tensor.GetLengths(), lengths);
+    EXPECT_EQ(tensor(0, 0, 0), 0.f);
+    EXPECT_EQ(tensor(3, 4, 2), 342.f);
+    EXPECT_EQ(tensor(4, 3, 2), 0.f);
+
+    EXPECT_EQ(tensor.Transpose().GetLengths(), tlengths);
+    EXPECT_EQ(tensor.Transpose()(0, 0, 0), 0.f);
+    EXPECT_EQ(tensor.Transpose()(2, 4, 3), 342.f);
+    EXPECT_EQ(tensor.Transpose()(2, 3, 4), 0.f);
+}
+
+TEST(HostTensorTranspose, Test3D_021)
+{
+    std::vector<size_t> lengths  = {13, 7, 5};
+    std::vector<size_t> tlengths = {13, 5, 7};
+
+    Tensor<float> tensor(lengths);
+    tensor(0, 0, 0) = 0.f;
+    tensor(3, 4, 2) = 342.f;
+
+    EXPECT_EQ(tensor.GetLengths(), lengths);
+    EXPECT_EQ(tensor(0, 0, 0), 0.f);
+    EXPECT_EQ(tensor(3, 4, 2), 342.f);
+    EXPECT_EQ(tensor(4, 3, 2), 0.f);
+
+    // transpose last two dimensions
+    EXPECT_EQ(tensor.Transpose({0, 2, 1}).GetLengths(), tlengths);
+    EXPECT_EQ(tensor.Transpose({0, 2, 1})(0, 0, 0), 0.f);
+    EXPECT_EQ(tensor.Transpose({0, 2, 1})(2, 4, 3), 0.f);
+    EXPECT_EQ(tensor.Transpose({0, 2, 1})(3, 2, 4), 342.f);
+    EXPECT_EQ(tensor.Transpose({0, 2, 1})(2, 3, 4), 0.f);
+
+    // transpose last two dimensions back again
+    EXPECT_EQ(tensor.Transpose({0, 2, 1}).Transpose({0, 2, 1}).GetLengths(), lengths);
+    EXPECT_EQ(tensor.Transpose({0, 2, 1}).Transpose({0, 2, 1})(3, 4, 2), 342.f);
+}
+
+TEST(HostTensorTranspose, TestNonpacked2D)
+{
+    std::vector<size_t> lengths  = {13, 7};
+    std::vector<size_t> strides  = {100, 1};
+    std::vector<size_t> tlengths = {7, 13};
+
+    Tensor<float> tensor(lengths, strides);
+    tensor(0, 0) = 0.f;
+    tensor(3, 4) = 34.f;
+
+    EXPECT_EQ(tensor.GetLengths(), lengths);
+    EXPECT_EQ(tensor(0, 0), 0.f);
+    EXPECT_EQ(tensor(3, 4), 34.f);
+    EXPECT_EQ(tensor(4, 3), 0.f);
+
+    EXPECT_EQ(tensor.Transpose().GetLengths(), tlengths);
+    EXPECT_EQ(tensor.Transpose()(0, 0), 0.f);
+    EXPECT_EQ(tensor.Transpose()(4, 3), 34.f);
+    EXPECT_EQ(tensor.Transpose()(3, 4), 0.f);
+}