gen_base.h 2.0 KB
Newer Older
T
tensor-tang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
/* Copyright (c) 2018 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. */

#pragma once

#include <gflags/gflags.h>
T
tensor-tang 已提交
18
#include <memory>  // for shared_ptr
T
tensor-tang 已提交
19
#include "paddle/fluid/operators/jit/kernel_base.h"
T
tensor-tang 已提交
20 21 22 23 24 25
#include "paddle/fluid/platform/macros.h"

DECLARE_bool(dump_jitcode);

namespace paddle {
namespace operators {
T
tensor-tang 已提交
26
namespace jit {
T
tensor-tang 已提交
27 28 29 30

// TODO(TJ): make these functions as virtual of a class

// Every JitCode should estimate the code size itself
T
tensor-tang 已提交
31
template <KernelType KT, typename T, typename Attr>
T
tensor-tang 已提交
32 33 34 35 36 37 38 39 40 41 42 43 44 45
size_t CodeSize(Attr attr) {
  return 4096;
}

// Every JitCode should have a condition when to use this JitCode
template <KernelType KT, typename T, typename Attr>
bool UseJitCode(Attr attr) {
  return false;
}

// Every JitCode should have a method to get the key from attribution
template <typename Attr>
size_t GetKey(Attr attr);

T
tensor-tang 已提交
46
class GenBase : public Kernel {
T
tensor-tang 已提交
47 48 49
 public:
  virtual const char* name() const = 0;
  virtual const unsigned char* getCodeInternal() = 0;
T
tensor-tang 已提交
50
  virtual size_t getSize() const = 0;
T
tensor-tang 已提交
51 52 53 54 55 56 57 58 59 60
  template <typename FUNC>
  const FUNC getCode() {
    const unsigned char* code = this->getCodeInternal();
    if (FLAGS_dump_jitcode) {
      this->dumpCode(code);
    }
    return reinterpret_cast<const FUNC>(code);
  }

 protected:
T
tensor-tang 已提交
61
  void dumpCode(const unsigned char* code) const;
T
tensor-tang 已提交
62 63
};

T
tensor-tang 已提交
64
template <KernelType KT, typename T, typename Attr>
T
tensor-tang 已提交
65
std::unique_ptr<GenBase> CreateJitCode(Attr attr);
T
tensor-tang 已提交
66

T
tensor-tang 已提交
67
}  // namespace jit
T
tensor-tang 已提交
68 69
}  // namespace operators
}  // namespace paddle