Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 9db74c7

Browse files
spectrometerHBHspectrometerHBH
authored
[TIRx] Update scoped ops and CUDA launch bounds (#19677)
## Summary - replace the block-structured TIRx exec-scope surface with scope-qualified `Tx.<scope>.<op>` namespaces and migrate call sites - split TIRx op namespaces and remove the unused dynamic generic-op fallback - add explicit CUDA launch bounds plumbing through TIRx attrs and split-host-device lowering ## Validation - `git diff --check apache/main..HEAD` - `pre-commit run --from-ref apache/main --to-ref HEAD`
1 parent 4d9d129 commit 9db74c7

180 files changed

Lines changed: 11458 additions & 11919 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎include/tvm/tirx/builtin.h‎

Lines changed: 29 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ TVM_DLL const Op& large_uint_imm();
134134
* (i.e., round(x.1) = x and round (x.5) = x+1)
135135
*/
136136
TVM_DLL const Op& q_multiply_shift();
137+
TVM_DLL const Op& q_multiply_shift_per_axis();
137138

138139
/*!
139140
* \brief Returns the address of an element in the buffer (see pseudocode below).
@@ -499,6 +500,11 @@ TVM_DLL const Op& tvm_call_trace_packed_lowered();
499500
*/
500501
TVM_DLL const Op& tvm_storage_sync();
501502

503+
/*!
504+
* \brief Marker where a transform should replace generated kernel initialization.
505+
*/
506+
TVM_DLL const Op& tvm_kernel_replace_point();
507+
502508
/*!
503509
* \brief See pseudo code
504510
*
@@ -912,6 +918,11 @@ TVM_DLL const Op& cuda_atomic_add();
912918
*/
913919
TVM_DLL const Op& cuda_thread_fence();
914920

921+
/*!
922+
* \brief tvm intrinsic for cuda warpgroup sync instruction
923+
*/
924+
TVM_DLL const Op& cuda_warpgroup_sync();
925+
915926
/*!
916927
* \brief Warp-level butterfly shuffle-XOR reduction.
917928
*
@@ -952,6 +963,11 @@ TVM_DLL const Op& cuda_cta_sync();
952963
*/
953964
TVM_DLL const Op& cuda_grid_sync();
954965

966+
/*!
967+
* \brief tvm intrinsic for cuda cluster-wide sync instruction
968+
*/
969+
TVM_DLL const Op& cuda_cluster_sync();
970+
955971
/*!
956972
* \brief tvm intrinsic that returns ``cooperative_groups::thread_rank()``
957973
* for the enclosing CTA (linear thread index within the block).
@@ -1053,25 +1069,19 @@ TVM_DLL const Op& ptx_reduce3_max_f32();
10531069
*/
10541070
TVM_DLL const Op& ptx_reduce3_min_f32();
10551071

1056-
/*!
1057-
* \brief tvm intrinsic for PTX packed add instruction (sm_100a+)
1058-
*/
1059-
TVM_DLL const Op& ptx_add_packed_f32x2();
1060-
1061-
/*!
1062-
* \brief tvm intrinsic for PTX packed subtract instruction (sm_100a+)
1063-
*/
1064-
TVM_DLL const Op& ptx_sub_packed_f32x2();
1065-
1066-
/*!
1067-
* \brief tvm intrinsic for PTX packed multiply instruction (sm_100a+)
1068-
*/
1069-
TVM_DLL const Op& ptx_mul_packed_f32x2();
1070-
1071-
/*!
1072-
* \brief tvm intrinsic for PTX packed FMA instruction (sm_100a+)
1073-
*/
1074-
TVM_DLL const Op& ptx_fma_packed_f32x2();
1072+
TVM_DLL const Op& ptx_add_f32();
1073+
TVM_DLL const Op& ptx_add_f32x2();
1074+
TVM_DLL const Op& ptx_add_f64();
1075+
TVM_DLL const Op& ptx_sub_f32();
1076+
TVM_DLL const Op& ptx_sub_f32x2();
1077+
TVM_DLL const Op& ptx_sub_f64();
1078+
TVM_DLL const Op& ptx_mul_f32();
1079+
TVM_DLL const Op& ptx_mul_f32x2();
1080+
TVM_DLL const Op& ptx_mul_f64();
1081+
TVM_DLL const Op& ptx_fma_f32();
1082+
TVM_DLL const Op& ptx_fma_f32x2();
1083+
TVM_DLL const Op& ptx_fma_f64();
1084+
TVM_DLL const Op& ptx_max_f32();
10751085

10761086
} // namespace builtin
10771087
} // namespace tirx

‎include/tvm/tirx/exec_context.h‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -136,9 +136,6 @@ struct ExecContext {
136136
/*! \brief Apply modulo filter on a factorized CTA axis such as cbx/cby/cbz. */
137137
bool WithCtaAxisModulo(const std::string& axis, int64_t modulus, int64_t residue,
138138
ExecContext* out, std::string* err) const;
139-
140-
/*! \brief Apply scope_switch; A preserved, split recomputed for new scope_kind. */
141-
bool WithScopeSwitch(ScopeKind new_scope_kind, ExecContext* out, std::string* err) const;
142139
};
143140

144141
/*!

‎include/tvm/tirx/exec_scope.h‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,11 +35,12 @@ namespace tvm {
3535
namespace tirx {
3636

3737
/*!
38-
* \brief The target execution scope kind of an ExecScopeStmt.
38+
* \brief The target execution scope kind of a tile primitive call.
3939
*
40-
* Replaces the string-keyed name of ExecScope. One value per user-facing
41-
* `with T.<kind>():` construct. Ordered from coarsest to finest; smaller
42-
* integer = wider scope, so ``ScopeKindHigher`` is a plain ``<``.
40+
* Identifies the granularity at which an op executes (the per-call
41+
* ``scope`` on a ``TilePrimitiveCall``, e.g. ``Tx.warp.copy(...)``).
42+
* Ordered from coarsest to finest; smaller integer = wider scope, so
43+
* ``ScopeKindHigher`` is a plain ``<``.
4344
*/
4445
enum class ScopeKind : int {
4546
kCluster = 2,

‎include/tvm/tirx/function.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,13 @@ namespace attr {
305305
*/
306306
constexpr const char* kKernelLaunchParams = "tirx.kernel_launch_params";
307307

308+
/*!
309+
* \brief CUDA launch bound minimum CTAs per SM.
310+
*
311+
* Type: IntImm
312+
*/
313+
constexpr const char* kLaunchBoundsMinBlocksPerSM = "tirx.launch_bounds_min_blocks_per_sm";
314+
308315
/*!
309316
* \brief Whether to set noalias rule on the function arguments.
310317
*

‎include/tvm/tirx/op.h‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
#include <tvm/ir/type.h>
3434
#include <tvm/tirx/builtin.h>
3535
#include <tvm/tirx/expr.h>
36+
#include <tvm/tirx/op_attr_types.h>
3637
#include <tvm/tirx/stmt.h>
3738
#include <tvm/tirx/target_builtin/cuda.h>
3839
#include <tvm/tirx/target_builtin/trn.h>
@@ -43,8 +44,10 @@
4344

4445
namespace tvm {
4546

46-
#define TVM_TIR_REGISTER_OP(OpName) \
47-
TVM_REGISTER_OP("tirx." OpName).set_attr<TScriptPrinterName>("TScriptPrinterName", OpName)
47+
#define TVM_TIR_REGISTER_OP(OpName) \
48+
TVM_REGISTER_OP("tirx." OpName) \
49+
.set_attr<TScriptPrinterName>("TScriptPrinterName", OpName) \
50+
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"), /*plevel=*/1)
4851

4952
#define TVM_TIRX_REGISTER_OP(OpName) TVM_TIR_REGISTER_OP(OpName)
5053

‎include/tvm/tirx/op_attr_types.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,23 @@ enum class ScriptDtypePrintLocation : int {
8282

8383
using TScriptDtypePrintLocation = int64_t;
8484

85+
/*!
86+
* \brief Broad TIRx op category.
87+
*
88+
* Expected values:
89+
* - "builtin"
90+
* - "tile_primitive"
91+
* - "device_intrin"
92+
*/
93+
using TIRxOpCategory = ffi::String;
94+
95+
/*!
96+
* \brief Device intrinsic namespace.
97+
*
98+
* Expected values include "cuda", "ptx", "nvshmem", "nki", and "metal".
99+
*/
100+
using TDeviceIntrinsicNamespace = ffi::String;
101+
85102
/*!
86103
* \brief The effect type of the call.
87104
*/

‎include/tvm/tirx/script/builder/frame.h‎

Lines changed: 0 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -247,52 +247,6 @@ class BlockInitFrame : public TIRFrame {
247247
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BlockInitFrame, TIRFrame, BlockInitFrameNode);
248248
};
249249

250-
/*!
251-
* \brief A frame that represents an execution scope (e.g. cta, warp, thread).
252-
*
253-
* When exiting this frame, it produces an ExecScopeStmt wrapping the body.
254-
* This is the new IR pattern, replacing the old pattern of storing exec_scope on SBlock.
255-
*
256-
* \sa ExecScopeFrame
257-
*/
258-
class ExecScopeFrameNode : public TIRFrameNode {
259-
public:
260-
/*! \brief The execution scope (always plain kind; no slice). */
261-
ffi::Optional<tvm::tirx::ExecScope> exec_scope;
262-
/*! \brief Optional surface-syntax guards for ``with Tx.scope(cond)``. */
263-
ffi::Array<PrimExpr> guards;
264-
265-
static void RegisterReflection() {
266-
namespace refl = tvm::ffi::reflection;
267-
refl::ObjectDef<ExecScopeFrameNode>()
268-
.def_ro("exec_scope", &ExecScopeFrameNode::exec_scope)
269-
.def_ro("guards", &ExecScopeFrameNode::guards);
270-
}
271-
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ExecScopeFrame", ExecScopeFrameNode,
272-
TIRFrameNode);
273-
274-
public:
275-
/*!
276-
* \brief The method called when exiting RAII scope.
277-
* \sa tvm::support::With
278-
*/
279-
void ExitWithScope() final;
280-
};
281-
282-
/*!
283-
* \brief Managed reference to ExecScopeFrameNode.
284-
*
285-
* \sa ExecScopeFrameNode
286-
*/
287-
class ExecScopeFrame : public TIRFrame {
288-
public:
289-
explicit ExecScopeFrame(ffi::ObjectPtr<ExecScopeFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
290-
TVM_FFI_ICHECK(data != nullptr);
291-
data_ = std::move(data);
292-
}
293-
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExecScopeFrame, TIRFrame, ExecScopeFrameNode);
294-
};
295-
296250
/*!
297251
* \brief A frame that represents the for loop.
298252
*

‎include/tvm/tirx/script/builder/ir.h‎

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -139,21 +139,6 @@ SBlockFrame Block(ffi::String name, bool no_realize = false, ffi::String exec_sc
139139

140140
void TilePrimitiveCall(tvm::tirx::TilePrimitiveCall op_call);
141141

142-
/*!
143-
* \brief Create an ExecScopeFrame for execution scope contexts.
144-
* \param exec_scope_name The name of the execution scope (e.g. "cta", "warp").
145-
* \return The ExecScopeFrame.
146-
*/
147-
ExecScopeFrame ExecScopeBlock(ffi::String exec_scope_name,
148-
ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
149-
150-
ExecScopeFrame Kernel(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
151-
ExecScopeFrame Cluster(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
152-
ExecScopeFrame WarpGroup(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
153-
ExecScopeFrame CTA(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
154-
ExecScopeFrame Warp(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
155-
ExecScopeFrame Thread(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
156-
157142
ffi::Array<tvm::tirx::Var> KernelId(ffi::Array<PrimExpr> extents, ffi::String parent);
158143

159144
ffi::Array<tvm::tirx::Var> CtaId(ffi::Array<PrimExpr> extents, ffi::String parent);

‎include/tvm/tirx/stmt.h‎

Lines changed: 1 addition & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -952,53 +952,11 @@ class SBlockRealize : public Stmt {
952952
TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockRealizeNode);
953953
};
954954

955-
/*!
956-
* \brief A statement that annotates the execution scope for its body.
957-
*
958-
* ExecScopeStmt represents a hardware execution scope (e.g. cta, warp, thread)
959-
* that wraps a body statement. This decouples the execution scope concept from
960-
* SBlock, making the IR structure cleaner.
961-
*
962-
* Example:
963-
* \code
964-
* with T.cta():
965-
* ...
966-
* \endcode
967-
*/
968-
class ExecScopeStmtNode : public StmtNode {
969-
public:
970-
/*! \brief The execution scope. */
971-
ExecScope exec_scope;
972-
/*! \brief The body statement under this execution scope. */
973-
Stmt body;
974-
975-
static void RegisterReflection() {
976-
namespace refl = tvm::ffi::reflection;
977-
refl::ObjectDef<ExecScopeStmtNode>()
978-
.def_ro("exec_scope", &ExecScopeStmtNode::exec_scope)
979-
.def_ro("body", &ExecScopeStmtNode::body);
980-
}
981-
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ExecScopeStmt", ExecScopeStmtNode, StmtNode);
982-
};
983-
984-
/*!
985-
* \brief Managed reference to ExecScopeStmtNode.
986-
* \sa ExecScopeStmtNode
987-
*/
988-
class ExecScopeStmt : public Stmt {
989-
public:
990-
TVM_DLL ExecScopeStmt(ExecScope exec_scope, Stmt body, Span span = Span());
991-
992-
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExecScopeStmt, Stmt, ExecScopeStmtNode);
993-
TVM_DEFINE_OBJECT_REF_COW_METHOD(ExecScopeStmtNode);
994-
};
995-
996955
/*!
997956
* \brief Standalone statement that declares a scope-id binding (e.g. cta_id,
998957
* warp_id, lane_id). Carries a ``ScopeIdDef`` value.
999958
*
1000-
* Unlike legacy ``ExecScopeStmt::scope_id_def`` (an array payload), each
1001-
* declaration is a flat stmt within the device-region body. The declared
959+
* Each declaration is a flat stmt within the device-region body. The declared
1002960
* ``Var``\ s are visible in subsequent stmts in the same enclosing scope
1003961
* (the AttrStmt ``kDeviceEntry`` body), analogous to ``BindNode``.
1004962
*/

‎include/tvm/tirx/stmt_functor.h‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,6 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
100100
virtual R VisitStmt_(const EvaluateNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
101101
virtual R VisitStmt_(const SBlockNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
102102
virtual R VisitStmt_(const SBlockRealizeNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
103-
virtual R VisitStmt_(const ExecScopeStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
104103
virtual R VisitStmt_(const ScopeIdDefStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
105104
virtual R VisitStmt_(const tirx::TilePrimitiveCallNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
106105
virtual R VisitStmtDefault_(const ffi::Object* op, Args...) {
@@ -127,7 +126,6 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
127126
IR_STMT_FUNCTOR_DISPATCH(BufferStoreNode);
128127
IR_STMT_FUNCTOR_DISPATCH(SBlockNode);
129128
IR_STMT_FUNCTOR_DISPATCH(SBlockRealizeNode);
130-
IR_STMT_FUNCTOR_DISPATCH(ExecScopeStmtNode);
131129
IR_STMT_FUNCTOR_DISPATCH(ScopeIdDefStmtNode);
132130
IR_STMT_FUNCTOR_DISPATCH(tirx::TilePrimitiveCallNode);
133131
vtable.Finalize();
@@ -185,7 +183,6 @@ class TVM_DLL StmtVisitor : protected StmtFunctor<void(const Stmt&)> {
185183
void VisitStmt_(const EvaluateNode* op) override;
186184
void VisitStmt_(const SBlockNode* op) override;
187185
void VisitStmt_(const SBlockRealizeNode* op) override;
188-
void VisitStmt_(const ExecScopeStmtNode* op) override;
189186
void VisitStmt_(const ScopeIdDefStmtNode* op) override;
190187
void VisitStmt_(const tirx::TilePrimitiveCallNode* op) override;
191188
};
@@ -304,7 +301,6 @@ class TVM_DLL StmtMutator : protected StmtFunctor<Stmt(const Stmt&)> {
304301
Stmt VisitStmt_(const EvaluateNode* op) override;
305302
Stmt VisitStmt_(const SBlockNode* op) override;
306303
Stmt VisitStmt_(const SBlockRealizeNode* op) override;
307-
Stmt VisitStmt_(const ExecScopeStmtNode* op) override;
308304
Stmt VisitStmt_(const ScopeIdDefStmtNode* op) override;
309305
Stmt VisitStmt_(const tirx::TilePrimitiveCallNode* op) override;
310306
/*!

0 commit comments

Comments
 (0)