131 lines
3.5 KiB
C++
131 lines
3.5 KiB
C++
/*
|
|
* Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
* All rights reserved.
|
|
*
|
|
* This source code is licensed under the BSD-style license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
#include <random>
|
|
#include <type_traits>
|
|
#include <vector>
|
|
|
|
#include <gtest/gtest.h>
|
|
|
|
#include "fbgemm/Utils.h"
|
|
|
|
using namespace std;
|
|
using namespace fbgemm;
|
|
|
|
template <typename T>
|
|
::testing::AssertionResult compare_tranpose_results(
|
|
vector<T> expected,
|
|
vector<T> acutal,
|
|
int m,
|
|
int n,
|
|
int ld_src,
|
|
int ld_dst) {
|
|
std::stringstream ss;
|
|
if (is_same<T, float>::value) {
|
|
ss << " float results ";
|
|
} else if (is_same<T, uint8_t>::value) {
|
|
ss << " i8 results ";
|
|
} else {
|
|
ss << " i16 results ";
|
|
}
|
|
ss << " mismatch at ";
|
|
bool match = true;
|
|
for (int i = 0; i < m; ++i) {
|
|
for (int j = 0; j < n; ++j) {
|
|
int exp = expected[i * ld_src + j];
|
|
int act = acutal[i + j * ld_dst];
|
|
if (exp != act) {
|
|
ss << "(" << i << ", " << j << "). ref " << exp << " actual " << act;
|
|
match = false;
|
|
}
|
|
}
|
|
}
|
|
if (match)
|
|
return ::testing::AssertionSuccess();
|
|
else
|
|
return ::testing::AssertionFailure() << "results differ: " << ss.str();
|
|
}
|
|
|
|
TEST(TransposeTest, TransposeTest) {
|
|
// Generate shapes to test
|
|
vector<tuple<int, int, int, int>> shapes;
|
|
uniform_int_distribution<int> dist(0, 32);
|
|
random_device r;
|
|
default_random_engine generator(r());
|
|
for (int i = 0; i < 1024; ++i) {
|
|
int m = dist(generator);
|
|
int n = dist(generator);
|
|
int ld_src = n + dist(generator);
|
|
int ld_dst = m + dist(generator);
|
|
shapes.push_back(make_tuple(m, n, ld_src, ld_dst));
|
|
}
|
|
for (int i = 0; i < 1024; ++i) {
|
|
int m = dist(generator);
|
|
int n = 2;
|
|
int ld_src = n;
|
|
int ld_dst = m + dist(generator);
|
|
shapes.push_back(make_tuple(m, n, ld_src, ld_dst));
|
|
}
|
|
for (int i = 0; i < 1024; ++i) {
|
|
int m = dist(generator);
|
|
int n = 4;
|
|
int ld_src = n;
|
|
int ld_dst = m + dist(generator);
|
|
shapes.push_back(make_tuple(m, n, ld_src, ld_dst));
|
|
}
|
|
for (int i = 0; i < 1024; ++i) {
|
|
int m = 2;
|
|
int n = dist(generator);
|
|
int ld_src = n + dist(generator);
|
|
int ld_dst = m;
|
|
shapes.push_back(make_tuple(m, n, ld_src, ld_dst));
|
|
}
|
|
for (int i = 0; i < 1024; ++i) {
|
|
int m = 4;
|
|
int n = dist(generator);
|
|
int ld_src = n + dist(generator);
|
|
int ld_dst = m;
|
|
shapes.push_back(make_tuple(m, n, ld_src, ld_dst));
|
|
}
|
|
|
|
for (const auto& shape : shapes) {
|
|
int m, n, ld_src, ld_dst;
|
|
tie(m, n, ld_src, ld_dst) = shape;
|
|
|
|
// float test
|
|
vector<float> a(m * ld_src);
|
|
vector<float> b(n * ld_dst);
|
|
generate(
|
|
a.begin(), a.end(), [&dist, &generator] { return dist(generator); });
|
|
|
|
transpose_simd(m, n, a.data(), ld_src, b.data(), ld_dst);
|
|
|
|
EXPECT_TRUE(compare_tranpose_results(a, b, m, n, ld_src, ld_dst));
|
|
|
|
// i8 test
|
|
vector<uint8_t> a_i8(m * ld_src);
|
|
vector<uint8_t> b_i8(n * ld_dst);
|
|
generate(a_i8.begin(), a_i8.end(), [&dist, &generator] {
|
|
return dist(generator);
|
|
});
|
|
|
|
transpose_simd(m, n, a_i8.data(), ld_src, b_i8.data(), ld_dst);
|
|
EXPECT_TRUE(compare_tranpose_results(a_i8, b_i8, m, n, ld_src, ld_dst));
|
|
|
|
// i16 test
|
|
vector<uint16_t> a_i16(m * ld_src);
|
|
vector<uint16_t> b_i16(n * ld_dst);
|
|
generate(a_i16.begin(), a_i16.end(), [&dist, &generator] {
|
|
return dist(generator);
|
|
});
|
|
|
|
transpose_simd(m, n, a_i16.data(), ld_src, b_i16.data(), ld_dst);
|
|
EXPECT_TRUE(compare_tranpose_results(a_i16, b_i16, m, n, ld_src, ld_dst));
|
|
}
|
|
}
|