Skip to content

Commit

Permalink
Adding unit tests for collapse with remap and convolution
Browse files Browse the repository at this point in the history
  • Loading branch information
luitjens committed Aug 25, 2022
1 parent 4a69c27 commit d3fa10c
Show file tree
Hide file tree
Showing 2 changed files with 153 additions and 6 deletions.
4 changes: 2 additions & 2 deletions include/matx/operators/collapse.h
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ namespace matx
// comptue size of collapsed dimension
size_ = 1;

// Collapse right-most dims
// Collapse left-most dims
#pragma unroll
for(int i = 0 ; i <= DIM; i++) {
size_ *= op_.Size(i);
Expand All @@ -85,7 +85,7 @@ namespace matx
auto ind = in[0];
#pragma unroll
for(int i = 0; i <= DIM; i++) {
index_t d = DIM - i;
int d = DIM - i;
out[d] = ind % op_.Size(d);
ind /= op_.Size(d);
}
Expand Down
155 changes: 151 additions & 4 deletions test/00_operators/OperatorTests.cu
Original file line number Diff line number Diff line change
Expand Up @@ -614,7 +614,7 @@ TYPED_TEST(OperatorTestsNumericNonComplex, CollapseOp)
for(int n = 0; n < N; n++) {
for(int m = 0; m < M; m++) {
for(int k = 0; k < K; k++) {
EXPECT_TRUE(tiv(n,m,k) == tov(n,m*K+k));
ASSERT_TRUE(tiv(n,m,k) == tov(n,m*K+k));
}
}
}
Expand All @@ -637,7 +637,7 @@ TYPED_TEST(OperatorTestsNumericNonComplex, CollapseOp)
for(int n = 0; n < N; n++) {
for(int m = 0; m < M; m++) {
for(int k = 0; k < K; k++) {
EXPECT_TRUE(tiv(n,m,k) == tov(n*M+m,k));
ASSERT_TRUE(tiv(n,m,k) == tov(n*M+m,k));
}
}
}
Expand All @@ -658,7 +658,7 @@ TYPED_TEST(OperatorTestsNumericNonComplex, CollapseOp)
for(int n = 0; n < N; n++) {
for(int m = 0; m < M; m++) {
for(int k = 0; k < K; k++) {
EXPECT_TRUE(tiv(n,m,k) == tov(n*M*K+m*K+k));
ASSERT_TRUE(tiv(n,m,k) == tov(n*M*K+m*K+k));
}
}
}
Expand All @@ -679,7 +679,7 @@ TYPED_TEST(OperatorTestsNumericNonComplex, CollapseOp)
for(int n = 0; n < N; n++) {
for(int m = 0; m < M; m++) {
for(int k = 0; k < K; k++) {
EXPECT_TRUE(tiv(n,m,k) == tov(n*M*K+m*K+k));
ASSERT_TRUE(tiv(n,m,k) == tov(n*M*K+m*K+k));
}
}
}
Expand Down Expand Up @@ -2618,6 +2618,153 @@ TEST(OperatorTests, Cast)
MATX_EXIT_HANDLER();
}

TEST(OperatorTestsAdvanced, AdvancedRemapOp)
{
typedef cuda::std::complex<float> complex;
MATX_ENTER_HANDLER();

int I = 4;
int J = 4;
int K = 14;
int L = 133;

int F = 4096;
int P = 288;

int M = 2;

auto idx = matx::make_tensor<int, 1>({M});

idx(0) = 1;
idx(1) = 3;

auto A = matx::make_tensor<complex, 4>({I, J, K, L});
//collapsed tensor
auto B = matx::make_tensor<complex, 2>({I * M * K, L});

auto index = [&] (int i, int j, int k, int l) {
return i * J * K * L +
j * K * L +
k * L +
l;
};
for (int i = 0; i < I ; i++) {
for (int j = 0; j < J ; j++) {
for (int k = 0; k < K ; k++) {
for (int l = 0; l < L ; l++) {
float val = (float)index(i,j,k,l);
A(i,j,k,l) = complex(val, val/100);
}
}
}
}

(B = 0).run();

auto rop = remap<1>(A, idx);
auto lop = lcollapse<2>(rop);

ASSERT_EQ(lop.Rank() , 2);
ASSERT_EQ(lop.Size(1) , A.Size(3));
ASSERT_EQ(lop.Size(0) , I * M * K);

(B = lop).run();

cudaDeviceSynchronize();

for (int i = 0; i < I; i++) {
for (int m = 0; m < M; m++) {
for (int k = 0; k < K; k++) {
for (int l = 0; l < L; l++) {
int j = idx(m);
int fidx = i * M * K + m * K + k;
float val = (float)index(i,j,k,l);
complex expected_val = complex(val,val/100);
complex a_val = A(i,j,k,l);
complex b_val = B(fidx, l);
complex lop_val = lop(fidx, l);
complex rop_val = rop(i, m, k, l);

// printf("fidx: %d, i: %d, j: %d, k: %d, l: %d, val: %f,%f\n", fidx, i, j, k, l, val, val/100);
// printf("a_val: %f, %f, rop_val: %f, %f, lop_val: %f, %f, b_val: %f, %f\n",
// a_val.real(), a_val.imag(),
// rop_val.real(), rop_val.imag(),
// lop_val.real(), lop_val.imag(),
// b_val.real(), b_val.imag());
ASSERT_EQ(a_val, expected_val);
ASSERT_EQ(rop_val, expected_val);
ASSERT_EQ(lop_val, expected_val);
ASSERT_EQ(b_val, expected_val);

ASSERT_EQ(B(fidx, l) , lop(fidx, l));
}
}
}
}


// convolution test
auto O1 = matx::make_tensor<complex, 4>({I, J, K, F + P + L - 1});
auto O2 = matx::make_tensor<complex, 4>({I, J, K, F + P + L - 1});
auto O3 = matx::make_tensor<complex, 4>({I, J, K, F + P + L - 1});
auto O4 = matx::make_tensor<complex, 4>({I, J, K, F + P + L - 1});

auto C = matx::make_tensor<complex, 3>({I, K, F + P});
//collapsed tensor
auto D = matx::make_tensor<complex, 2>({I * M * K, F + P});

auto indexc = [&] (int i, int j, int k) {
return i * C.Size(1) * C.Size(2) +
j * C.Size(2) +
k;
};

for (int i = 0; i < I ; i++) {
for (int j = 0; j < J ; j++) {
for (int k = 0; k < K ; k++) {
float val = (float) indexc(i,j,k);
C(i,j,k) = complex(val, val/100);
}
}
}

auto o1op = lcollapse<2>(remap<1>(O1, idx));
auto o2op = lcollapse<2>(remap<1>(O2, idx));
auto o3op = lcollapse<2>(remap<1>(O3, idx));
auto o4op = lcollapse<2>(remap<1>(O4, idx));

auto cop = C.Clone<4>({matxKeepDim, M, matxKeepDim, matxKeepDim});
auto rcop = lcollapse<2>(remap<1>(cop, idx));

(B = lop).run();
(D = rcop).run();

// two operators as input
matx::conv1d(o1op, lop, rcop, matx::matxConvCorrMode_t::MATX_C_MODE_FULL, 0);

// one tensor and one operators as input
matx::conv1d(o2op, B, rcop, matx::matxConvCorrMode_t::MATX_C_MODE_FULL, 0);

// one tensor and one operators as input
matx::conv1d(o3op, lop, D, matx::matxConvCorrMode_t::MATX_C_MODE_FULL, 0);

//two tensors as input
matx::conv1d(o4op, B, D, matx::matxConvCorrMode_t::MATX_C_MODE_FULL, 0);

cudaDeviceSynchronize();

for (int i = 0; i < o1op.Size(0); i++) {
for (int l = 0; l < o1op.Size(1); l++) {
ASSERT_EQ(o1op(i,l), o2op(i,l));
ASSERT_EQ(o2op(i,l), o3op(i,l));
ASSERT_EQ(o3op(i,l), o4op(i,l));
}
}

MATX_EXIT_HANDLER();
}


TYPED_TEST(OperatorTestsFloat, Print)
{
MATX_ENTER_HANDLER();
Expand Down

0 comments on commit d3fa10c

Please sign in to comment.