operator.h 2.69 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
/*!
 * \file tl/op/op.h
 * \brief Tile library operations.
 *
 */

#ifndef TVM_TL_OP_OP_H_
#define TVM_TL_OP_OP_H_

#include <tvm/arith/analyzer.h>
#include <tvm/ir/op.h>
#include <tvm/target/target.h>
#include <tvm/tir/buffer.h>
14
#include <tvm/tir/op.h>
15
#include <tvm/tir/op_attr_types.h>
16
#include <tvm/tir/stmt.h>
17
18
19
20
21
22
23
24
25
26
27
28

#include "../layout/layout.h"

namespace tvm {
namespace tl {

using namespace tir;

using AddWorkspaceCallback = std::function<PrimExpr(int, DataType)>;
using LayoutMap = Map<Buffer, Layout>;
using BufferMap = Map<Var, Buffer>;

29
enum class InferLevel : uint8_t {
30
31
32
33
34
35
36
  kFree = 0,
  kCommon = 1,
  kStrict = 2,
};

struct LowerArgs {
  Target target;
37
  Range thread_bounds;
38
39
40
41
  Var thread_var;
  AddWorkspaceCallback AddWorkspace;
  LayoutMap layout_map;
  Map<Buffer, Buffer> buffer_remap;
42
  Array<Var> buffer_var_gemm;
43
44
45
46
};

struct LayoutInferArgs {
  Target target;
47
  Range thread_bounds;
48
  LayoutMap layout_map;
49
50
  arith::Analyzer *analyzer;
  bool buffer_oob = false;
51
52
53
  Map<Buffer, Buffer> buffer_remap;
};

54
class TileOperator;
55

56
57
class TileOperatorNode : public Object {
public:
58
  virtual Stmt Lower(const LowerArgs &T, arith::Analyzer *analyzer) const = 0;
59

60
  virtual LayoutMap InferLayout(const LayoutInferArgs &T,
61
                                InferLevel level) const = 0;
62

63
64
  virtual TileOperator Clone() const = 0;

65
  TVM_FFI_DECLARE_OBJECT_INFO("tl.TileOperator", TileOperatorNode, Object);
66
};
67

68
class TileOperator : public ObjectRef {
69
public:
70
71
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TileOperator, ObjectRef,
                                             TileOperatorNode);
72
73
};

74
Var GetVarFromAccessPtr(const PrimExpr &expr);
75

76
77
78
TileOperator ParseOperator(Call call, BufferMap vmap);
TileOperator ParseOperator(Stmt stmt, BufferMap vmap);

79
80
using OpBuilderFunc =
    ffi::TypedFunction<TileOperator(Array<PrimExpr>, BufferMap)>;
81
82
83
84
85
86
87
88
89

#define TIR_REGISTER_TL_OP(Entry, OpName)                                      \
  const Op &Entry::Get() {                                                     \
    static const Op &op = Op::Get("tl." #OpName);                              \
    return op;                                                                 \
  }                                                                            \
  TVM_REGISTER_OP("tl." #OpName)                                               \
      .set_attr<TScriptPrinterName>("TScriptPrinterName", #OpName)             \
      .set_attr<OpBuilderFunc>("TLOpBuilder",                                  \
90
91
                               [](Array<PrimExpr> args, BufferMap vmap) {      \
                                 return Entry(args, vmap);                     \
92
93
                               })

94
95
} // namespace tl
} // namespace tvm
96

97
#endif // TVM_TL_OP_OP_H_