builtin.cc 5.37 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
/*!
 * \file tl/op/builtin.cc
 * \brief Builtin intrinsics.
 *
 */

#include "builtin.h"

#include <tvm/tir/builtin.h>
#include <tvm/tir/op.h>
#include <tvm/tir/op_attr_types.h>

#include "../target/cuda.h"
14
#include "../target/utils.h"
15
16
17
18

namespace tvm {
namespace tl {

19
TVM_REGISTER_PASS_CONFIG_OPTION(kDebugMergeSharedMemoryAllocations, Bool);
20
TVM_REGISTER_PASS_CONFIG_OPTION(kDisableTMALower, Bool);
21
TVM_REGISTER_PASS_CONFIG_OPTION(kDisableSafeMemoryLegalize, Bool);
22
TVM_REGISTER_PASS_CONFIG_OPTION(kDisableWarpSpecialized, Bool);
23
TVM_REGISTER_PASS_CONFIG_OPTION(kConfigIndexBitwidth, Integer);
24
25
TVM_REGISTER_PASS_CONFIG_OPTION(kDisableDynamicTailSplit, Bool);
TVM_REGISTER_PASS_CONFIG_OPTION(kDynamicAlignment, Integer);
26
TVM_REGISTER_PASS_CONFIG_OPTION(kEnableAggressiveSharedMemoryMerge, Bool);
27
TVM_REGISTER_PASS_CONFIG_OPTION(kDisableFastMath, Bool);
28
TVM_REGISTER_PASS_CONFIG_OPTION(kPtxasRegisterUsageLevel, Integer);
29
TVM_REGISTER_PASS_CONFIG_OPTION(kEnablePTXASVerboseOutput, Bool);
30

31
32
33
34
35
36
37
#define TIR_DEFINE_TL_BUILTIN(OpName)                                          \
  const Op &OpName() {                                                         \
    static const Op &op = Op::Get("tl." #OpName);                              \
    return op;                                                                 \
  }                                                                            \
  TVM_REGISTER_OP("tl." #OpName)                                               \
      .set_attr<TScriptPrinterName>("TScriptPrinterName", #OpName)
38

39
TIR_DEFINE_TL_BUILTIN(create_list_of_mbarrier)
40
    .set_num_inputs(-1)
41
42
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
43

44
TIR_DEFINE_TL_BUILTIN(create_tma_descriptor)
45
    .set_num_inputs(-1)
46
47
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kPure));
48

49
TIR_DEFINE_TL_BUILTIN(create_tma_im2col_descriptor)
50
    .set_num_inputs(-1)
51
52
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kPure));
53

54
TIR_DEFINE_TL_BUILTIN(get_mbarrier)
55
    .set_num_inputs(1)
56
57
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kPure));
58

59
TIR_DEFINE_TL_BUILTIN(tma_load).set_num_inputs(-1).set_attr<TCallEffectKind>(
60
61
    "TCallEffectKind", Integer(CallEffectKind::kOpaque));

62
TIR_DEFINE_TL_BUILTIN(tma_load_im2col)
63
64
65
66
    .set_num_inputs(-1)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));

67
68
TIR_DEFINE_TL_BUILTIN(tma_store).set_num_inputs(-1).set_attr<TCallEffectKind>(
    "TCallEffectKind", Integer(CallEffectKind::kOpaque));
69

70
TIR_DEFINE_TL_BUILTIN(mbarrier_wait_parity)
71
    .set_num_inputs(2)
72
73
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
74

75
TIR_DEFINE_TL_BUILTIN(mbarrier_expect_tx)
76
    .set_num_inputs(2)
77
78
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
79

80
TIR_DEFINE_TL_BUILTIN(ptx_ldmatirx)
81
    .set_num_inputs(4)
82
83
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
84

85
TIR_DEFINE_TL_BUILTIN(ptx_stmatirx)
86
    .set_num_inputs(-1)
87
88
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
89

90
TIR_DEFINE_TL_BUILTIN(sync_thread_partial)
91
    .set_num_inputs(1)
92
93
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
94

95
TIR_DEFINE_TL_BUILTIN(fence_proxy_async)
96
    .set_num_inputs(0)
97
98
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
99

100
TIR_DEFINE_TL_BUILTIN(tma_store_arrive)
101
102
103
104
    .set_num_inputs(0)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));

105
TIR_DEFINE_TL_BUILTIN(tma_store_wait)
106
107
108
    .set_num_inputs(0)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
109
TIR_DEFINE_TL_BUILTIN(set_max_nreg)
110
    .set_num_inputs(2)
111
112
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
113

114
TIR_DEFINE_TL_BUILTIN(no_set_max_nreg)
115
116
117
118
    .set_num_inputs(0)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));

119
120
121
122
TIR_DEFINE_TL_BUILTIN(wait_wgmma)
    .set_num_inputs(1)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
123

124
TIR_DEFINE_TL_BUILTIN(pack_b16).set_num_inputs(2).set_attr<TCallEffectKind>(
125
    "TCallEffectKind", Integer(CallEffectKind::kPure));
126
127
128
129
130
131
132
133

TIR_DEFINE_TL_BUILTIN(sync_grid).set_num_inputs(0).set_attr<TCallEffectKind>(
    "TCallEffectKind", Integer(CallEffectKind::kOpaque));

TIR_DEFINE_TL_BUILTIN(loop_break)
    .set_num_inputs(0)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
134
135
136
137
138
139
140
141
142

TIR_DEFINE_TL_BUILTIN(tl_gemm).set_num_inputs(4).set_attr<TCallEffectKind>(
    "TCallEffectKind", Integer(CallEffectKind::kOpaque));

TIR_DEFINE_TL_BUILTIN(tl_gemm_sp)
    .set_num_inputs(5)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));

143
} // namespace tl
144
} // namespace tvm