codegen_opencl.h 2.57 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

20 21 22
/*!
 *  Copyright (c) 2017 by Contributors
 * \file codegen_opencl.h
23
 * \brief Generate OpenCL device code.
24 25 26 27 28 29 30
 */
#ifndef TVM_CODEGEN_CODEGEN_OPENCL_H_
#define TVM_CODEGEN_CODEGEN_OPENCL_H_

#include <tvm/codegen.h>
#include <tvm/packed_func_ext.h>
#include <string>
31
#include "codegen_c.h"
32 33 34 35

namespace tvm {
namespace codegen {

36
class CodeGenOpenCL final : public CodeGenC {
37
 public:
38
  CodeGenOpenCL();
39
  void AddFunction(LoweredFunc f);
40 41
  std::string Finish();

42
  // override print thread tag.
43
  void InitFuncState(LoweredFunc f) final;
44
  void BindThreadIndex(const IterVar& iv) final;  // NOLINT(*)
45
  void PrintStorageScope(const std::string& scope, std::ostream& os) final; // NOLINT(*)
46
  void PrintStorageSync(const Call* op) final;  // NOLINT(*)
47
  void PrintType(Type t, std::ostream& os) final; // NOLINT(*)
48 49
  std::string GetVecLoad(Type t, const Variable* buffer,
                         Expr base) final;
50
  void PrintVecStore(const Variable* buffer,
51 52
                     Type t, Expr base,
                     const std::string& value) final;  // NOLINT(*)
53 54 55
  // the address of load/store
  void PrintVecAddr(const Variable* buffer, Type t,
                    Expr base, std::ostream& os);  // NOLINT(*)
56 57
  std::string CastFromTo(std::string value, Type from, Type target); // NOLINT(*)

58 59
  // overload visitor
  void VisitExpr_(const Broadcast* op, std::ostream& os) final; // NOLINT(*)
60 61
  void VisitExpr_(const Call* op, std::ostream& os) final; // NOLINT(*)
  void VisitExpr_(const Select* op, std::ostream& os) final; // NOLINT(*)
62
  void VisitExpr_(const FloatImm *op, std::ostream& os) final; // NOLINT(*)
63 64 65 66 67

 private:
  // whether enable fp16 and fp64 extension
  bool enable_fp16_{false};
  bool enable_fp64_{false};
68 69 70 71 72 73
};

}  // namespace codegen
}  // namespace tvm

#endif  // TVM_CODEGEN_CODEGEN_OPENCL_H_