infershape_launchers_test.cc 3.2 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// 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.

#include <gtest/gtest.h>

17 18 19
#include "paddle/infrt/kernel/phi/infershaped/infershaped_kernel_launcher.h"
#include "paddle/infrt/kernel/phi/infershaped/infershaped_kernel_launchers.h"
#include "paddle/infrt/kernel/phi/infershaped/infershaped_utils.h"
20 21 22 23
#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/common/place.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/meta_tensor.h"
24 25

namespace infrt {
26
namespace kernel {
27

28
namespace {
29 30 31
static void ElementwiseAddTest(const ::Tensor& a,
                               const ::Tensor& b,
                               ::Tensor* c);
32 33 34 35 36 37 38 39
}

TEST(utils, registry) {
  constexpr uint8_t count =
      InferShapeHelper<decltype(&ElementwiseAddTest)>::count;
  CHECK_EQ(count, 2U);
}

40
class FancyAllocator : public ::phi::Allocator {
W
Wilber 已提交
41
 public:
42
  static void Delete(::phi::Allocation* allocation) {
W
Wilber 已提交
43 44 45 46 47
    ::operator delete(allocation->ptr());
  }

  AllocationPtr Allocate(size_t bytes_size) override {
    void* data = ::operator new(bytes_size);
48 49
    auto* allocation =
        new ::phi::Allocation(data, bytes_size, ::phi::CPUPlace());
W
Wilber 已提交
50 51 52 53
    return AllocationPtr(allocation, Delete);
  }
};

54 55
TEST(ElementwiseAdd, launcher_registry) {
  host_context::KernelRegistry registry;
56
  RegisterInferShapeLaunchers(&registry);
W
Wilber 已提交
57
  ASSERT_GE(registry.size(), 1UL);
58
  auto creator = registry.GetKernel("phi_cpu.add.float32.any");
W
Wilber 已提交
59

60 61 62 63 64
  const ::phi::DDim dims({1, 2});
  const ::phi::DataType dtype{::phi::DataType::FLOAT32};
  const ::phi::DataLayout layout{::phi::DataLayout::NHWC};
  const ::phi::LoD lod{};
  ::phi::DenseTensorMeta meta(dtype, dims, layout, lod);
W
Wilber 已提交
65

66
  auto fancy_allocator = std::unique_ptr<::phi::Allocator>(new FancyAllocator);
W
Wilber 已提交
67
  auto* alloc = fancy_allocator.get();
68

69 70 71
  ::Tensor a(alloc, meta);
  ::Tensor b(alloc, meta);
  ::Tensor c(alloc, meta);
W
Wilber 已提交
72

73
  auto place = ::phi::CPUPlace();
W
Wilber 已提交
74 75 76 77 78 79 80 81
  float* a_data = a.mutable_data<float>(place);
  float* b_data = b.mutable_data<float>(place);
  float* c_data = c.mutable_data<float>(place);
  for (size_t i = 0; i < 2; ++i) {
    a_data[i] = 1.f;
    b_data[i] = 2.f;
  }

82
  ::phi::CPUContext context;
W
Wilber 已提交
83
  context.SetAllocator(alloc);
84 85

  host_context::KernelFrameBuilder kernel_frame_builder;
W
Wilber 已提交
86
  kernel_frame_builder.AddArgument(new host_context::Value(std::move(context)));
87 88 89
  kernel_frame_builder.AddArgument(new host_context::Value(std::move(a)));
  kernel_frame_builder.AddArgument(new host_context::Value(std::move(b)));
  kernel_frame_builder.SetResults({new host_context::Value(std::move(c))});
W
Wilber 已提交
90

91
  creator(&kernel_frame_builder);
W
Wilber 已提交
92 93 94 95

  for (size_t i = 0; i < 2; ++i) {
    CHECK_EQ(c_data[i], 3.f);
  }
96 97
}

98
}  // namespace kernel
99
}  // namespace infrt