codegen_opencl.h 1.56 KB
Newer Older
1 2 3
/*!
 *  Copyright (c) 2017 by Contributors
 * \file codegen_opencl.h
4
 * \brief Generate OpenCL device code.
5 6 7 8 9 10 11
 */
#ifndef TVM_CODEGEN_CODEGEN_OPENCL_H_
#define TVM_CODEGEN_CODEGEN_OPENCL_H_

#include <tvm/codegen.h>
#include <tvm/packed_func_ext.h>
#include <string>
12
#include "codegen_c.h"
13 14 15 16

namespace tvm {
namespace codegen {

17
class CodeGenOpenCL final : public CodeGenC {
18
 public:
19
  CodeGenOpenCL();
20
  void AddFunction(LoweredFunc f);
21 22
  std::string Finish();

23
  // override print thread tag.
24
  void InitFuncState(LoweredFunc f) final;
25
  void BindThreadIndex(const IterVar& iv) final;  // NOLINT(*)
26
  void PrintStorageScope(const std::string& scope, std::ostream& os) final; // NOLINT(*)
27
  void PrintStorageSync(const Call* op) final;  // NOLINT(*)
28
  void PrintType(Type t, std::ostream& os) final; // NOLINT(*)
29 30
  std::string GetVecLoad(Type t, const Variable* buffer,
                         Expr base) final;
31
  void PrintVecStore(const Variable* buffer,
32 33
                     Type t, Expr base,
                     const std::string& value) final;  // NOLINT(*)
34 35 36
  // the address of load/store
  void PrintVecAddr(const Variable* buffer, Type t,
                    Expr base, std::ostream& os);  // NOLINT(*)
37 38
  std::string CastFromTo(std::string value, Type from, Type target); // NOLINT(*)

39 40
  // overload visitor
  void VisitExpr_(const Broadcast* op, std::ostream& os) final; // NOLINT(*)
41 42 43 44 45

 private:
  // whether enable fp16 and fp64 extension
  bool enable_fp16_{false};
  bool enable_fp64_{false};
46 47 48 49 50 51
};

}  // namespace codegen
}  // namespace tvm

#endif  // TVM_CODEGEN_CODEGEN_OPENCL_H_