builtin.cc 6.33 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
TVM_REGISTER_PASS_CONFIG_OPTION(kDisableShuffleElect, Bool);
31

32
33
DataType cuTensorMapType() { return DataType::UInt(8, 128); }

34
35
36
37
38
39
40
#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)
41

42
TIR_DEFINE_TL_BUILTIN(create_list_of_mbarrier)
43
    .set_num_inputs(-1)
44
45
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
46

47
TIR_DEFINE_TL_BUILTIN(create_tma_descriptor)
48
    .set_num_inputs(-1)
49
50
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kPure));
51

52
TIR_DEFINE_TL_BUILTIN(create_tma_im2col_descriptor)
53
    .set_num_inputs(-1)
54
55
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kPure));
56

57
TIR_DEFINE_TL_BUILTIN(get_mbarrier)
58
    .set_num_inputs(1)
59
60
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kPure));
61

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

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

70
71
TIR_DEFINE_TL_BUILTIN(tma_store).set_num_inputs(-1).set_attr<TCallEffectKind>(
    "TCallEffectKind", Integer(CallEffectKind::kOpaque));
72

73
TIR_DEFINE_TL_BUILTIN(mbarrier_wait_parity)
74
    .set_num_inputs(2)
75
76
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
77

78
TIR_DEFINE_TL_BUILTIN(mbarrier_expect_tx)
79
    .set_num_inputs(2)
80
81
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
82

83
TIR_DEFINE_TL_BUILTIN(ptx_ldmatrix)
84
    .set_num_inputs(4)
85
86
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
87

88
TIR_DEFINE_TL_BUILTIN(ptx_stmatrix)
89
    .set_num_inputs(-1)
90
91
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
92

93
94
95
96
97
TIR_DEFINE_TL_BUILTIN(ptx_cp_async_barrier_noinc)
    .set_num_inputs(1)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));

98
TIR_DEFINE_TL_BUILTIN(fence_proxy_async)
99
    .set_num_inputs(0)
100
101
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
102

103
TIR_DEFINE_TL_BUILTIN(tma_store_arrive)
104
105
106
107
    .set_num_inputs(0)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));

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

117
TIR_DEFINE_TL_BUILTIN(no_set_max_nreg)
118
119
120
121
    .set_num_inputs(0)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));

122
123
124
125
TIR_DEFINE_TL_BUILTIN(wait_wgmma)
    .set_num_inputs(1)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kOpaque));
126

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

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));
137
138
139
140
141
142
143
144
145

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));

146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
TIR_DEFINE_TL_BUILTIN(tvm_mfma).set_num_inputs(12).set_attr<TCallEffectKind>(
    "TCallEffectKind", Integer(CallEffectKind::kOpaque));

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

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

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

164
165
166
167
168
TIR_DEFINE_TL_BUILTIN(tl_shuffle_elect)
    .set_num_inputs(1)
    .set_attr<TCallEffectKind>("TCallEffectKind",
                               Integer(CallEffectKind::kPure));

169
} // namespace tl
170
} // namespace tvm