dim_test.cu 3.4 KB
Newer Older
1
//  Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2 3 4 5 6 7 8 9 10 11 12 13
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
F
fengjiayi 已提交
14 15 16 17
#include <thrust/device_vector.h>
#include <sstream>

#include "gtest/gtest.h"
Y
Yi Wang 已提交
18
#include "paddle/fluid/framework/dim.h"
F
fengjiayi 已提交
19

20
__global__ void test(paddle::framework::Dim<2>* o) {
21
  o[0] = paddle::framework::make_dim(5, 6);
F
fengjiayi 已提交
22 23
}

Q
qijun 已提交
24
__global__ void dyn_idx_gpu(int64_t* o) {
25 26
  auto d = paddle::framework::make_dim(5, 6);
  o[0] = d[1];
F
fengjiayi 已提交
27 28 29
}

TEST(Dim, Equality) {
30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62
  // construct a Dim on the CPU
  auto a = paddle::framework::make_dim(3, 4);
  EXPECT_EQ(paddle::framework::get<0>(a), 3);
  EXPECT_EQ(paddle::framework::get<1>(a), 4);

  // construct a Dim on the GPU
  thrust::device_vector<paddle::framework::Dim<2>> t(2);
  test<<<1, 1>>>(thrust::raw_pointer_cast(t.data()));
  a = t[0];
  EXPECT_EQ(paddle::framework::get<0>(a), 5);
  EXPECT_EQ(paddle::framework::get<1>(a), 6);

  // linearization
  auto b = paddle::framework::make_dim(7, 8);
  EXPECT_EQ(paddle::framework::linearize(a, b), 83);

  // product
  EXPECT_EQ(paddle::framework::product(a), 30);

  // mutate a Dim
  paddle::framework::get<1>(b) = 10;
  EXPECT_EQ(paddle::framework::get<0>(b), 7);
  EXPECT_EQ(paddle::framework::get<1>(b), 10);

  // dynamic access
  paddle::framework::get(b, 0) = 8;
  b[1] = 11;
  EXPECT_EQ(paddle::framework::get<0>(b), 8);
  EXPECT_EQ(paddle::framework::get<1>(b), 11);
  EXPECT_EQ(paddle::framework::get(b, 0), 8);
  EXPECT_EQ(b[1], 11);

  // dynamic access on GPU
Q
qijun 已提交
63
  thrust::device_vector<int64_t> r(1);
64
  dyn_idx_gpu<<<1, 1>>>(thrust::raw_pointer_cast(r.data()));
Q
qijun 已提交
65
  int64_t res = r[0];
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84
  EXPECT_EQ(res, 6);

  // ex_prefix_mul
  paddle::framework::Dim<3> c =
      paddle::framework::ex_prefix_mul(paddle::framework::Dim<3>(3, 4, 5));
  EXPECT_EQ(paddle::framework::get<0>(c), 1);
  EXPECT_EQ(paddle::framework::get<1>(c), 3);
  EXPECT_EQ(paddle::framework::get<2>(c), 12);

  // generate from an index
  auto size = paddle::framework::make_dim(4, 5, 2);
  c = paddle::framework::Dim<3>(14, size);
  EXPECT_EQ(paddle::framework::get<0>(c), 2);
  EXPECT_EQ(paddle::framework::get<1>(c), 3);
  EXPECT_EQ(paddle::framework::get<2>(c), 0);
  c = paddle::framework::Dim<3>(25, size);
  EXPECT_EQ(paddle::framework::get<0>(c), 1);
  EXPECT_EQ(paddle::framework::get<1>(c), 1);
  EXPECT_EQ(paddle::framework::get<2>(c), 1);
F
fengjiayi 已提交
85 86 87
}

TEST(Dim, Bool) {
88 89 90 91 92 93 94 95 96 97 98 99
  auto a = paddle::framework::make_dim(3, 4);
  auto b = paddle::framework::make_dim(5, 6);
  auto c = paddle::framework::make_dim(3, 4);

  // in_bounds check
  EXPECT_TRUE(paddle::framework::contained(a, b));
  EXPECT_FALSE(paddle::framework::contained(b, a));

  // comparison
  EXPECT_TRUE(a == a);
  EXPECT_FALSE(a == b);
  EXPECT_TRUE(a == c);
F
fengjiayi 已提交
100 101 102
}

TEST(Dim, Print) {
103 104 105 106 107 108 109 110 111 112 113
  {
    std::stringstream ss;
    auto a = paddle::framework::make_dim(2, 3);
    ss << a;
    EXPECT_EQ(ss.str(), "2, 3");
  }
  {
    std::stringstream ss;
    ss << paddle::framework::make_dim(8);
    EXPECT_EQ(ss.str(), "8");
  }
F
fengjiayi 已提交
114
}