Skip to content

GitLab

"docs/vscode:/vscode.git/clone" did not exist on "427f12569bfc2e21e51ab9586832ae19ee4453fb"
Unverified Commit 6664d170 authored by Wenhao Xie's avatar Wenhao Xie Committed by GitHub
Browse files

[Enhancement] Add eviction policy support for TMA operations, enhance CUDA... 

[Enhancement] Add eviction policy support for TMA operations, enhance CUDA codegen, and introduce new pass config (#690)

* Enhance TMA and barrier handling in CUDA code generation

- Updated `CodeGenTileLangCUDA` to support eviction policies for TMA operations, allowing for more flexible memory management.
- Introduced a new `CacheHintSm90` enum to define eviction strategies in `copy_sm90.h`.
- Modified TMA load/store functions to accept eviction policies, improving performance on different architectures.
- Enhanced `TmaBarrierCollector` and `TmaBarrierRewriter` to account for SIMT copies, ensuring correct barrier insertion.
- Refactored thread synchronization logic to utilize barrier IDs, improving the efficiency of partial thread synchronization.
- Updated Python interface for `copy` and `c2d_im2col` to include optional eviction policy parameters, enhancing usability.

* update shuffle and elect optimization

* fix bug

* fix bug

* fix potential bug

* lint fix

* lint fix

* update shuffle_elect template

* fix bug

* fix bug

* fix template

* lint and fix

* fix typo
parent fe70549f
Show whitespace changes
Inline Side-by-side
Showing with 734 additions and 201 deletions
src/op/builtin.cc
View file @ 6664d170
......@@ -27,6 +27,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION(kEnableAggressiveSharedMemoryMerge, Bool);
TVM_REGISTER_PASS_CONFIG_OPTION(kDisableFastMath, Bool);
TVM_REGISTER_PASS_CONFIG_OPTION(kPtxasRegisterUsageLevel, Integer);
TVM_REGISTER_PASS_CONFIG_OPTION(kEnablePTXASVerboseOutput, Bool);
TVM_REGISTER_PASS_CONFIG_OPTION(kDisableShuffleElect, Bool);
#define TIR_DEFINE_TL_BUILTIN(OpName) \
const Op &OpName() { \
......@@ -88,7 +89,7 @@ TIR_DEFINE_TL_BUILTIN(ptx_stmatirx)
Integer(CallEffectKind::kOpaque));
TIR_DEFINE_TL_BUILTIN(sync_thread_partial)
.set_num_inputs(1)
.set_num_inputs(2)
.set_attr<TCallEffectKind>("TCallEffectKind",
Integer(CallEffectKind::kOpaque));
......@@ -140,5 +141,10 @@ TIR_DEFINE_TL_BUILTIN(tl_gemm_sp)
.set_attr<TCallEffectKind>("TCallEffectKind",
Integer(CallEffectKind::kOpaque));
TIR_DEFINE_TL_BUILTIN(tl_shuffle_elect)
.set_num_inputs(1)
.set_attr<TCallEffectKind>("TCallEffectKind",
Integer(CallEffectKind::kPure));
} // namespace tl
} // namespace tvm
src/op/builtin.h
View file @ 6664d170
......@@ -32,7 +32,7 @@ static constexpr const char *kPtxasRegisterUsageLevel =
"tl.ptxas_register_usage_level";
static constexpr const char *kEnablePTXASVerboseOutput =
"tl.enable_ptxas_verbose_output";
static constexpr const char *kDisableShuffleElect = "tl.disable_shuffle_elect";
/*!
* \brief Whether to disable dynamic tail split
*
......@@ -294,6 +294,13 @@ TVM_DLL const Op &tl_gemm();
*/
TVM_DLL const Op &tl_gemm_sp();
/*!
* \brief tilelang intrinsic for shuffle elect.
*
* This op is used to represent a shuffle elect operation in tilelang.
*/
TVM_DLL const Op &tl_shuffle_elect();
} // namespace tl
} // namespace tvm
......
src/op/bulk_copy.cc
View file @ 6664d170
......@@ -297,7 +297,7 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
Call(DataType::Handle(), create_tma_descriptor(), desc.EncodeCallArgs());
Array<PrimExpr> args;
args.reserve(desc.rank + 3);
args.reserve(desc.rank + 4);
args.push_back(create_descriptor);
if (is_load)
args.push_back(0); // mbarrier id placeholder
......@@ -319,6 +319,7 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
global_coords.Set(0, global_coords[0] + instruction_dim * loop_var);
for (auto coord : global_coords)
args.push_back(coord);
args.push_back(this->eviction_policy);
tma_copy = For(loop_var, 0, loop_extent, ForKind::kUnrolled,
Evaluate(Call(DataType::Handle(), op, args)));
} else {
......@@ -327,6 +328,7 @@ Stmt Copy::LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const {
args.push_back(shared_addr);
for (auto coord : global_coords)
args.push_back(coord);
args.push_back(this->eviction_policy);
tma_copy = Evaluate(Call(DataType::Handle(), op, args));
}
tma_copy = IfThenElse(EQ(T.thread_var, T.thread_bounds->min), tma_copy);
......@@ -368,6 +370,7 @@ Conv2DIm2ColOp::Conv2DIm2ColOp(Array<PrimExpr> args, BufferMap vmap) {
stride = args[5].as<IntImm>().value()->value;
dilation = args[6].as<IntImm>().value()->value;
padding = args[7].as<IntImm>().value()->value;
eviction_policy = args[8].as<IntImm>().value()->value;
}
Stmt Conv2DIm2ColOp::Lower(const LowerArgs &T,
......@@ -477,7 +480,7 @@ Stmt Conv2DIm2ColOp::Lower(const LowerArgs &T,
FloorDiv(nhw_step * desc.smem_box_pixel, w_dim * h_dim));
Array<PrimExpr> args;
args.reserve(desc.rank * 2 + 1);
args.reserve(desc.rank * 2 + 2);
args.push_back(create_desc);
args.push_back(0); // mbar placeholder
auto dst_buffer = T.buffer_remap.count(dst) ? T.buffer_remap[dst] : dst;
......@@ -487,7 +490,7 @@ Stmt Conv2DIm2ColOp::Lower(const LowerArgs &T,
args.push_back(coord);
for (auto offset : image_offset)
args.push_back(offset);
args.push_back(this->eviction_policy);
Stmt tma_copy =
IfThenElse(EQ(T.thread_var, T.thread_bounds->min),
Evaluate(Call(DataType::Handle(), tma_load_im2col(), args)));
......@@ -522,7 +525,7 @@ Array<PrimExpr> TMAIm2ColDesc::EncodeCallArgs() const {
}
TIR_REGISTER_TL_OP(Conv2DIm2ColOp, c2d_im2col)
.set_num_inputs(8)
.set_num_inputs(9)
.set_attr<TCallEffectKind>("TCallEffectKind",
Integer(CallEffectKind::kOpaque));
......
src/op/bulk_copy.h
View file @ 6664d170
......@@ -57,7 +57,7 @@ public:
private:
Buffer src, dst;
int stride, padding, dilation, kernel;
int stride, padding, dilation, kernel, eviction_policy;
PrimExpr nhw_step, c_step;
};
......
src/op/elem.cc
View file @ 6664d170
......@@ -45,6 +45,9 @@ Copy::Copy(Array<PrimExpr> args, BufferMap vmap) : args_(args) {
auto disable_tma = Downcast<Bool>(args[3]);
this->disable_tma = disable_tma;
}
if (args.size() >= 5) {
this->eviction_policy = args[4].as<IntImmNode>()->value;
}
}
Array<IterVar> Copy::MakeIterVars() const {
......@@ -477,7 +480,7 @@ Stmt Fill::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const {
}
TIR_REGISTER_TL_OP(Copy, copy)
.set_num_inputs(3)
.set_num_inputs(4)
.set_attr<TCallEffectKind>("TCallEffectKind",
Integer(CallEffectKind::kOpaque));
......
src/op/elem.h
View file @ 6664d170
......@@ -58,6 +58,8 @@ protected:
Bool disable_tma = Bool(false);
std::unique_ptr<ParallelOp> par_op_;
int eviction_policy;
};
class Fill : public Operator {
......
src/target/codegen_cuda.cc
View file @ 6664d170
......@@ -926,13 +926,14 @@ std::string CodeGenTileLangCUDA::GetBufferRef(DataType t,
}
void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
auto print_extern_call_stmt = [&](std::string name, size_t offset = 0) {
auto print_extern_call_stmt = [&](std::string name, size_t start = 0,
size_t end = 0) {
// Cache context into a private ss, otherwise the let node may generate
// within the function call arguments.
std::ostringstream ss;
for (size_t i = offset; i < op->args.size(); i++) {
if (i > offset)
for (size_t i = start; i < op->args.size() - end; i++) {
if (i > start)
ss << ", ";
ss << this->PrintExpr(op->args[i]);
}
......@@ -990,13 +991,16 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
} else if (op->op.same_as(tl::mbarrier_wait_parity())) {
print_extern_call_stmt("tl::mbarrier_wait");
} else if (op->op.same_as(tl::sync_thread_partial())) {
print_extern_call_stmt("tl::syncthreads_partial");
print_extern_call_stmt("cutlass::arch::NamedBarrier::sync");
} else if (op->op.same_as(tl::no_set_max_nreg())) {
return;
} else if (op->op.same_as(tl::tma_load())) {
std::ostringstream ss;
ICHECK_GE(op->args.size(), 2);
ss << "tl::tma_load(";
auto eviction_policy =
this->eviction_policy_names_
[op->args[op->args.size() - 1].as<IntImmNode>()->value];
ss << "tl::tma_load<tl::CacheHintSm90::" << eviction_policy << ">(";
auto desc = op->args[0];
ss << this->PrintExpr(desc) << ", ";
if (const IntImmNode *imm = op->args[1].as<IntImmNode>()) {
......@@ -1004,7 +1008,7 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
} else {
ss << this->PrintExpr(op->args[1]) << ", ";
}
for (size_t i = 2; i < op->args.size(); i++) {
for (size_t i = 2; i < op->args.size() - 1; i++) {
if (i > 2)
ss << ", ";
ss << this->PrintExpr(op->args[i]);
......@@ -1013,9 +1017,19 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
this->PrintIndent();
this->stream << ss.str();
} else if (op->op.same_as(tl::tma_load_im2col())) {
print_extern_call_stmt("tl::tma_load_im2col");
std::stringstream ss;
ss << "tl::tma_load_im2col<tl::CacheHintSm90::"
<< this->eviction_policy_names_
[op->args[op->args.size() - 1].as<IntImmNode>()->value]
<< ">";
print_extern_call_stmt(ss.str(), 0, 1);
} else if (op->op.same_as(tl::tma_store())) {
print_extern_call_stmt("tl::tma_store");
std::stringstream ss;
ss << "tl::tma_store<tl::CacheHintSm90::"
<< this->eviction_policy_names_
[op->args[op->args.size() - 1].as<IntImmNode>()->value]
<< ">";
print_extern_call_stmt(ss.str(), 0, 1);
} else if (op->op.same_as(tl::ptx_ldmatirx())) {
int trans = Downcast<IntImm>(op->args[0])->value;
int num = Downcast<IntImm>(op->args[1])->value;
......@@ -1537,6 +1551,8 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
enable_sparse_gemm_ = true;
this->PrintCallExtern(GetType(GetRef<PrimExpr>(op)), op_instance->value,
op->args, true, os);
} else if (op->op.same_as(tl::tl_shuffle_elect())) {
os << "tl::tl_shuffle_elect<" << PrintExpr(op->args[0]) << ">()";
} else {
CodeGenC::VisitExpr_(op, os);
}
......
src/target/codegen_cuda.h
View file @ 6664d170
......@@ -126,6 +126,8 @@ private:
int32_t GetWmmaFragmentSize(const std::string &scope, const VarNode *variable,
int32_t size);
std::vector<std::string> eviction_policy_names_ = {
"EVICT_NORMAL", "EVICT_FIRST", "EVICT_LAST"};
std::unordered_set<std::string> bf16_supported_ops_ = {
"bf1622float2", "bf1622int16", "float22bf162", "bf162bf162"};
};
......
src/tl_templates/cuda/common.h
View file @ 6664d170
......@@ -241,4 +241,13 @@ TL_DEVICE void __sync_thread_partial() {
asm volatile("bar.sync %0, %1;" : : "r"(barrier_id), "r"(thread_count));
}
template <int thread_extent> TL_DEVICE bool tl_shuffle_elect() {
if constexpr (thread_extent == 0) {
return cutlass::canonical_warp_idx_sync() == 0 && cute::elect_one_sync();
}
return __shfl_sync(0xffffffff, (threadIdx.x / 32) % (thread_extent / 32),
0) == 0 &&
cute::elect_one_sync();
}
} // namespace tl
src/tl_templates/cuda/copy_sm90.h
View file @ 6664d170
......@@ -7,6 +7,11 @@
#include "common.h"
namespace tl {
enum class CacheHintSm90 : uint64_t {
EVICT_NORMAL = 0x1000000000000000,
EVICT_FIRST = 0x12F0000000000000,
EVICT_LAST = 0x14F0000000000000,
};
TL_DEVICE void tma_load(void *smem_ptr, void *gmem_ptr, uint64_t &smem_mbar,
uint32_t size) {
......@@ -30,11 +35,13 @@ TL_DEVICE void tma_load_multicast(void *smem_ptr, void *gmem_ptr,
:);
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
void const *const smem_ptr, int32_t const &crd0) {
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile("cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::"
"complete_tx::bytes"
" [%0], [%1, {%3}], [%2];"
......@@ -42,14 +49,25 @@ TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0)
: "memory");
} else {
asm volatile("cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::"
"complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3}], [%2], %4;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "l"(cache_hint)
: "memory");
}
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
void const *const smem_ptr, int32_t const &crd0,
int32_t const &crd1) {
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::"
"complete_tx::bytes"
" [%0], [%1, {%3, %4}], [%2];"
......@@ -57,14 +75,25 @@ TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1)
: "memory");
} else {
asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::"
"complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4}], [%2], %5;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "l"(cache_hint)
: "memory");
}
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
void const *const smem_ptr, int32_t const &crd0,
int32_t const &crd1, int32_t const &crd2) {
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile("cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::"
"complete_tx::bytes"
" [%0], [%1, {%3, %4, %5}], [%2];"
......@@ -72,8 +101,17 @@ TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2)
: "memory");
} else {
asm volatile("cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::"
"complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5}], [%2], %6;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2), "l"(cache_hint)
: "memory");
}
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
void const *const smem_ptr, int32_t const &crd0,
int32_t const &crd1, int32_t const &crd2,
......@@ -81,6 +119,7 @@ TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile("cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::"
"complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6}], [%2];"
......@@ -88,8 +127,18 @@ TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3)
: "memory");
} else {
asm volatile("cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::"
"complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6}], [%2], %7;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "l"(cache_hint)
: "memory");
}
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
void const *const smem_ptr, int32_t const &crd0,
int32_t const &crd1, int32_t const &crd2,
......@@ -97,6 +146,7 @@ TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile("cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::"
"complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2];"
......@@ -104,8 +154,19 @@ TL_DEVICE void tma_load(const CUtensorMap &descriptor, uint64_t &smem_mbar,
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4)
: "memory");
} else {
asm volatile("cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::"
"complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4),
"l"(cache_hint)
: "memory");
}
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_load_im2col(const CUtensorMap &descriptor,
uint64_t &smem_mbar, void const *const smem_ptr,
int32_t const &coord_c, int32_t const &coord_w,
......@@ -115,64 +176,99 @@ TL_DEVICE void tma_load_im2col(const CUtensorMap &descriptor,
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_mbar = smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
asm volatile("cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier:"
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile(
"cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier:"
":complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6}], [%2], {%7, %8};"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n),
"h"(offset_w), "h"(offset_h)
"r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n), "h"(offset_w),
"h"(offset_h)
: "memory");
}
TL_DEVICE void tma_store(void *dst_gmem_ptr, void *smem_ptr, uint32_t size) {
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
} else {
asm volatile(
"cp.async.bulk.global.shared::cta.bulk_group [%1], [%0], %2; \n" ::"r"(
smem_int_ptr),
"l"(dst_gmem_ptr), "r"(size)
:);
"cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier:"
":complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6}], [%2], {%7, %8}, %9;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n), "h"(offset_w),
"h"(offset_h), "l"(cache_hint)
: "memory");
}
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_store(const CUtensorMap &descriptor,
void const *const smem_ptr, int32_t const &crd0) {
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
asm volatile(
"cp.async.bulk.tensor.1d.global.shared::cta.bulk_group [%0, {%2}], [%1];"
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile("cp.async.bulk.tensor.1d.global.shared::cta.bulk_group [%0, "
"{%2}], [%1];"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0)
: "memory");
} else {
asm volatile("cp.async.bulk.tensor.1d.global.shared::cta.bulk_group "
"::cache_hint [%0, {%2}], [%1], %3;"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0),
"l"(cache_hint)
: "memory");
}
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_store(const CUtensorMap &descriptor,
void const *const smem_ptr, int32_t const &crd0,
int32_t const &crd1) {
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile("cp.async.bulk.tensor.2d.global.shared::cta.bulk_group [%0, "
"{%2, %3}], [%1];"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1)
: "memory");
} else {
asm volatile("cp.async.bulk.tensor.2d.global.shared::cta.bulk_group "
"::cache_hint [%0, {%2, %3}], [%1], %4;"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
"l"(cache_hint)
: "memory");
}
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_store(const CUtensorMap &descriptor,
void const *const smem_ptr, int32_t const &crd0,
int32_t const &crd1, int32_t const &crd2) {
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile("cp.async.bulk.tensor.3d.global.shared::cta.bulk_group [%0, "
"{%2, %3, %4}], [%1];"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
"r"(crd2)
: "memory");
} else {
asm volatile("cp.async.bulk.tensor.3d.global.shared::cta.bulk_group "
"::cache_hint [%0, {%2, %3, %4}], [%1], %5;"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
"r"(crd2), "l"(cache_hint)
: "memory");
}
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_store(const CUtensorMap &descriptor,
void const *const smem_ptr, int32_t const &crd0,
int32_t const &crd1, int32_t const &crd2,
......@@ -180,14 +276,24 @@ TL_DEVICE void tma_store(const CUtensorMap &descriptor,
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile("cp.async.bulk.tensor.4d.global.shared::cta.bulk_group [%0, "
"{%2, %3, %4, %5}], [%1];"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
"r"(crd2), "r"(crd3)
: "memory");
} else {
asm volatile("cp.async.bulk.tensor.4d.global.shared::cta.bulk_group "
"::cache_hint [%0, {%2, %3, %4, %5}], [%1], %6;"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
"r"(crd2), "r"(crd3), "l"(cache_hint)
: "memory");
}
}
template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL>
TL_DEVICE void tma_store(const CUtensorMap &descriptor,
void const *const smem_ptr, int32_t const &crd0,
int32_t const &crd1, int32_t const &crd2,
......@@ -195,12 +301,21 @@ TL_DEVICE void tma_store(const CUtensorMap &descriptor,
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr);
if constexpr (cache_hint == CacheHintSm90::EVICT_NORMAL) {
asm volatile("cp.async.bulk.tensor.5d.global.shared::cta.bulk_group [%0, "
"{%2, %3, %4, %5, %6}], [%1];"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
"r"(crd2), "r"(crd3), "r"(crd4)
: "memory");
} else {
asm volatile("cp.async.bulk.tensor.5d.global.shared::cta.bulk_group "
"::cache_hint [%0, {%2, %3, %4, %5, %6}], [%1], %7;"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd0), "r"(crd1),
"r"(crd2), "r"(crd3), "r"(crd4), "l"(cache_hint)
: "memory");
}
}
TL_DEVICE void prefetch_tma_descriptor(const CUtensorMap &descriptor) {
......@@ -215,13 +330,52 @@ TL_DEVICE void mbarrier_init(uint64_t &smem_barrier, uint32_t arrive_count) {
: "r"(arrive_count), "r"(smem_int_ptr));
}
TL_DEVICE uint32_t mbarrier_try_wait(uint64_t &smem_barrier, int phase_bit) {
uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
uint32_t waitComplete;
asm volatile("{\n\t"
".reg .pred P1; \n\t"
"mbarrier.try_wait.parity.shared.b64 P1, [%1], %2; \n\t"
"selp.b32 %0, 1, 0, P1; \n\t"
"}"
: "=r"(waitComplete)
: "r"(smem_int_ptr), "r"(phase_bit));
return waitComplete;
}
TL_DEVICE void mbarrier_wait(uint64_t &smem_barrier, int phase_bit) {
if (mbarrier_try_wait(smem_barrier, phase_bit) == 0) {
uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
asm volatile("{\n"
// Arbitrarily large timer value after which try-wait expires and re-tries.
uint32_t ticks = 0x989680;
asm volatile("{\n\t"
".reg .pred P1; \n\t"
"LAB_WAIT: \n\t"
"mbarrier.try_wait.parity.shared.b64 P1, [%0], %1, %2; \n\t"
"@P1 bra DONE; \n\t"
"bra LAB_WAIT; \n\t"
"DONE: \n\t"
"}"
:
: "r"(smem_int_ptr), "r"(phase_bit), "r"(ticks));
}
}
TL_DEVICE void mbarrier_test_wait(uint64_t &smem_barrier, int phase_bit) {
uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
asm volatile(
"{\n"
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.try_wait.parity.shared.b64 P1, [%0], %1;\n"
"@!P1 bra.uni LAB_WAIT;\n"
"mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"nanosleep.u32 5;\n" // wait a few nanoseconds on pre-Hopper architectures
// to save instruction issue slots
"bra.uni LAB_WAIT;\n"
"DONE:\n"
"}\n" ::"r"(smem_int_ptr),
"r"(phase_bit));
}
......@@ -231,6 +385,20 @@ TL_DEVICE void mbarrier_arrive(uint64_t &smem_barrier) {
asm volatile("mbarrier.arrive.shared.b64 _, [%0];" : : "r"(smem_int_ptr));
}
TL_DEVICE void mbarrier_arrive(uint64_t &smem_barrier, int cta_id,
uint32_t pred) {
uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
if (pred) {
asm volatile("{\n\t"
".reg .b32 remAddr32;\n\t"
"mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t"
"mbarrier.arrive.shared::cluster.b64 _, [remAddr32];\n\t"
"}"
:
: "r"(smem_int_ptr), "r"(cta_id));
}
}
TL_DEVICE void mbarrier_expect_tx(uint64_t &smem_barrier,
uint32_t transaction_bytes) {
uint32_t smem_int_ptr = smem_ptr_to_uint(&smem_barrier);
......
src/transform/inject_tma_barrier.cc
View file @ 6664d170
......@@ -122,8 +122,11 @@ private:
Stmt VisitStmt_(const IfThenElseNode *op) {
// Check if this is the TMA block
const EQNode *eq = op->condition.as<EQNode>();
if (eq != nullptr) {
bool flag = false;
if (op->condition.as<CallNode>()) {
flag = op->condition.as<CallNode>()->op.same_as(tl_shuffle_elect());
}
if (op->condition.as<EQNode>() || flag) {
Stmt ret = IRMutatorWithAnalyzer::VisitStmt_(op);
if (visited_tma_load_) {
......@@ -164,6 +167,9 @@ private:
class TmaBarrierCollector : public IRVisitorWithAnalyzer {
public:
TmaBarrierCollector(Map<Var, Buffer> buffer_data_to_buffer)
: buffer_data_to_buffer_(std::move(buffer_data_to_buffer)) {}
Map<ObjectRef, PrimExpr> tma_op_to_barrier_id() {
return tma_op_to_barrier_id_;
}
......@@ -222,7 +228,128 @@ private:
std::vector<Call> pending_tma_ops_;
Map<ObjectRef, PrimExpr> tma_op_to_barrier_id_;
Map<PrimExpr, IntImm> barrier_id_to_range_;
Map<Var, Buffer> buffer_data_to_buffer_;
};
class TmaSequenceCollector : public IRVisitorWithAnalyzer {
public:
TmaSequenceCollector(Map<ObjectRef, PrimExpr> tma_op_to_barrier_id)
: tma_op_to_barrier_id_(std::move(tma_op_to_barrier_id)) {}
std::vector<bool> GetSequence() {
std::vector<bool> clear_zero_list(expect_tx_count_, false);
int zero_idx = -1;
int zero_count = 0;
for (auto v : sequence) {
if (v == 0) {
zero_count += 1;
zero_idx += 1;
} else {
if (zero_count == 1) {
clear_zero_list[zero_idx] = expect_[zero_idx] && !has_simt_copy_;
if (clear_zero_list[zero_idx] == false) {
int begin = int_sets_[zero_idx].min().as<IntImmNode>()->value;
int end = int_sets_[zero_idx].max().as<IntImmNode>()->value;
for (int i = begin; i <= end; ++i) {
restore_barrier_ids_.push_back(i);
}
}
} else {
for (int i{zero_idx}; i > zero_idx - zero_count; --i) {
int begin = int_sets_[i].min().as<IntImmNode>()->value;
int end = int_sets_[i].max().as<IntImmNode>()->value;
for (int i = begin; i <= end; ++i) {
restore_barrier_ids_.push_back(i);
}
}
}
zero_count = 0;
}
}
return clear_zero_list;
}
std::vector<int> GetRestoreBarrierIds() { return restore_barrier_ids_; }
void VisitStmt_(const ForNode *op) final {
var_int_set_.Set(op->loop_var,
arith::IntSet::FromMinExtent(op->min, op->extent));
IRVisitorWithAnalyzer::VisitStmt_(op);
}
void VisitExpr_(const CallNode *op) final {
if (op->op.same_as(mbarrier_expect_tx())) {
PrimExpr e =
tma_op_to_barrier_id_[GetRef<Call>(op)].as<CallNode>()->args[0];
auto int_set = arith::EvalSet(e, var_int_set_);
expect_.push_back(if_depth_ == 1);
sequence.push_back(0);
int_sets_.push_back(int_set);
expect_tx_count_ += 1;
} else if (op->op.same_as(builtin::ptx_arrive_barrier())) {
sequence.push_back(1);
} else if (op->op.same_as(builtin::ptx_cp_async_barrier())) {
has_simt_copy_ = true;
}
IRVisitorWithAnalyzer::VisitExpr_(op);
}
void VisitStmt_(const IfThenElseNode *op) final {
if_depth_ += 1;
IRVisitorWithAnalyzer::VisitStmt(op->then_case);
if (op->else_case) {
IRVisitorWithAnalyzer::VisitStmt(op->else_case.value());
}
if_depth_ -= 1;
}
std::vector<int> sequence;
int expect_tx_count_{0};
std::vector<bool> expect_;
bool has_simt_copy_{false};
std::vector<int> restore_barrier_ids_;
int if_depth_{0};
Map<ObjectRef, PrimExpr> tma_op_to_barrier_id_;
arith::Analyzer *analyzer_;
Map<Var, arith::IntSet> var_int_set_;
std::vector<arith::IntSet> int_sets_;
};
class BarrierCreationRewriter : public StmtExprMutator {
public:
BarrierCreationRewriter(std::vector<int> restore_barrier_ids,
PrimExpr producer_thread_extent)
: restore_barrier_ids_(std::move(restore_barrier_ids)),
producer_thread_extent_(producer_thread_extent) {}
PrimExpr VisitExpr_(const CallNode *op) {
if (op->op.same_as(create_list_of_mbarrier())) {
std::vector<bool> tmp_(op->args.size(), false);
Array<PrimExpr> new_args;
for (auto &id : restore_barrier_ids_) {
tmp_[id] = true;
}
for (size_t i{0}; i < op->args.size(); ++i) {
if (tmp_[i]) {
new_args.push_back(producer_thread_extent_);
} else {
new_args.push_back(op->args[i]);
}
}
return Call(op->dtype, op->op, new_args);
} else {
return StmtExprMutator::VisitExpr_(op);
}
}
std::vector<int> restore_barrier_ids_;
PrimExpr producer_thread_extent_;
};
// we trust mbarrier_wait_parity to be correct
class TmaBarrierRewriter : public IRMutatorWithAnalyzer {
public:
......@@ -236,8 +363,12 @@ public:
has_create_list_of_mbarrier_(has_create_list_of_mbarrier) {}
static PrimFunc Rewrite(PrimFunc f, arith::Analyzer *analyzer) {
auto buffer_lca = DetectBufferAccessLCA(f);
Map<Var, Buffer> buffer_data_to_buffer_;
for (auto [buffer, _] : buffer_lca)
buffer_data_to_buffer_.Set(buffer->data, buffer);
f = TmaExpectTxRewriter::Rewrite(f, analyzer);
TmaBarrierCollector collector;
TmaBarrierCollector collector(buffer_data_to_buffer_);
collector(f->body);
bool has_create_list_of_mbarrier = false;
PostOrderVisit(f->body, [&](const ObjectRef &node) {
......@@ -253,6 +384,9 @@ public:
collector.barrier_id_to_range(),
has_create_list_of_mbarrier);
f.CopyOnWrite()->body = rewriter(f->body);
auto barrier_creation_rewriter = BarrierCreationRewriter(
rewriter.restore_barrier_ids_, rewriter.producer_thread_extent_);
f.CopyOnWrite()->body = barrier_creation_rewriter(f->body);
return f;
}
......@@ -266,6 +400,42 @@ private:
return IRMutatorWithAnalyzer::VisitStmt_(op);
}
Stmt VisitStmt_(const IfThenElseNode *op) {
if (first_if) {
if (op->condition.as<GENode>()) {
producer_thread_extent_ =
thread_var_->dom->extent - op->condition.as<GENode>()->b;
}
TmaSequenceCollector collector(tma_op_to_barrier_id_);
collector(op->then_case);
clear_expect_list_ = collector.GetSequence();
restore_barrier_ids_ = collector.GetRestoreBarrierIds();
first_if = false;
is_producer_ = true;
auto then_case = StmtExprMutator::VisitStmt(op->then_case);
is_producer_ = false;
Stmt else_case;
if (op->else_case.defined())
else_case = StmtExprMutator::VisitStmt(op->else_case.value());
return IfThenElse(op->condition, then_case, else_case);
}
return StmtExprMutator::VisitStmt_(op);
}
Stmt VisitStmt_(const AttrStmtNode *op) final {
if (op->attr_key == "kWarpSpecializationScope") {
has_warp_specialization_ = true;
first_if = true;
} else if (op->attr_key == tir::attr::thread_extent &&
Downcast<IterVar>(op->node)->thread_tag == "threadIdx.x") {
thread_var_ = Downcast<IterVar>(op->node);
}
return IRMutatorWithAnalyzer::VisitStmt_(op);
}
PrimExpr VisitExpr_(const CallNode *op) {
if (op->op.same_as(tma_load())) {
// check this must be in the tma_op_to_barrier_id_
......@@ -281,6 +451,22 @@ private:
auto barrier_id = tma_op_to_barrier_id_[GetRef<Call>(op)];
auto new_args = op->args;
new_args.Set(0, barrier_id);
if (!has_warp_specialization_)
clear_arrive_ = false;
else
clear_arrive_ = clear_expect_list_[cur_expect_idx_++];
if (clear_arrive_) {
return Call(op->dtype, builtin::ptx_arrive_barrier_expect_tx(),
new_args);
}
return Call(op->dtype, op->op, new_args);
} else if (op->op.same_as(builtin::ptx_arrive_barrier())) {
if (clear_arrive_) {
clear_arrive_ = false;
return 0;
}
// by default, all threads must wait.
auto new_args = op->args;
return Call(op->dtype, op->op, new_args);
}
return IRMutatorWithAnalyzer::VisitExpr_(op);
......@@ -288,6 +474,13 @@ private:
Map<ObjectRef, PrimExpr> tma_op_to_barrier_id_;
Map<PrimExpr, IntImm> barrier_id_to_range_;
bool has_create_list_of_mbarrier_;
bool clear_arrive_{false};
bool first_if{false}, has_warp_specialization_{false}, is_producer_{false};
IterVar thread_var_;
int tma_expect_tx_{0}, cur_expect_idx_{0};
std::vector<bool> clear_expect_list_;
std::vector<int> restore_barrier_ids_;
PrimExpr producer_thread_extent_;
};
tvm::transform::Pass InjectTmaBarrier() {
......
src/transform/lower_hopper_intrin.cc
View file @ 6664d170
......@@ -6,6 +6,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/builtin.h>
#include <tvm/tir/op.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/transform.h>
......@@ -21,9 +22,9 @@ using namespace tir;
#if (CUDA_MAJOR_VERSION >= 12)
class LowerHopperIntrin : public StmtExprMutator {
public:
static PrimFunc Substitute(PrimFunc &f) {
static PrimFunc Substitute(PrimFunc &f, bool disable_shuffle_elect) {
PrimFuncNode *fptr = f.CopyOnWrite();
LowerHopperIntrin substituter;
LowerHopperIntrin substituter(disable_shuffle_elect);
fptr->body = substituter.VisitStmt(f->body);
Map<String, Array<PrimExpr>> init_desc_arg_map;
for (auto [call, var] : substituter.desc_map_) {
......@@ -73,10 +74,15 @@ public:
auto stmts = prefetch_calls_;
stmts.insert(stmts.end(), init_mbarrier_calls_.begin(),
init_mbarrier_calls_.end());
auto init_stmt =
IfThenElse(EQ(iv->var, IntImm(iv->var->dtype, 0)),
PrimExpr condition;
if (!disable_shuffle_elect_) {
condition = Call(DataType::Bool(), tl_shuffle_elect(), {0});
} else {
condition = EQ(iv->var, 0);
}
auto stmt_ = IfThenElse(condition,
stmts.size() > 1 ? SeqStmt(stmts) : stmts[0]);
stmt_seq.push_back(init_stmt);
stmt_seq.push_back(stmt_);
if (!init_mbarrier_calls_.empty()) {
Stmt mem_sync =
Evaluate(Call(DataType::Handle(), builtin::tvm_storage_sync(),
......@@ -121,14 +127,6 @@ public:
{mbarrier, call->args[i]})));
}
return 0;
} else if (call->op.same_as(sync_thread_partial())) {
int barrier_id = init_mbarrier_calls_.size();
PrimExpr mbarrier =
Call(DataType::Handle(), get_mbarrier(), {barrier_id});
init_mbarrier_calls_.push_back(Evaluate(
Call(DataType::Handle(), builtin::ptx_init_barrier_thread_count(),
{mbarrier, call->args[0]})));
return Call(DataType::Handle(), sync_thread_partial(), {mbarrier});
} else {
return StmtExprMutator::VisitExpr_(call);
}
......@@ -138,14 +136,18 @@ private:
Array<Stmt> prefetch_calls_;
Array<Stmt> init_mbarrier_calls_;
std::unordered_map<Call, Var, StructuralHash, ExprDeepEqual> desc_map_;
LowerHopperIntrin() = default;
LowerHopperIntrin(bool disable_shuffle_elect)
: disable_shuffle_elect_(disable_shuffle_elect) {}
bool disable_shuffle_elect_;
};
using namespace tir::transform;
tvm::transform::Pass LowerHopperIntrin() {
auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
return LowerHopperIntrin::Substitute(f);
bool disable_shuffle_elect =
ctx->GetConfig<Bool>(kDisableShuffleElect, Bool(false)).value();
return LowerHopperIntrin::Substitute(f, disable_shuffle_elect);
};
return CreatePrimFuncPass(pass_func, 0, "tl.LowerHopperIntrin", {});
}
......
src/transform/merge_if_stmt.cc
View file @ 6664d170
......@@ -44,11 +44,13 @@ private:
continue;
} else {
if (!current_if_bodies.empty()) {
new_seq.push_back(IfThenElse(current_condition,
auto if_stmt =
IfThenElse(current_condition,
current_if_bodies.size() == 1
? current_if_bodies[0]
: SeqStmt(current_if_bodies),
Stmt()));
: this->VisitStmt(SeqStmt(current_if_bodies)),
Stmt());
new_seq.push_back(if_stmt);
current_if_bodies.clear();
}
......@@ -60,11 +62,13 @@ private:
}
if (!current_if_bodies.empty()) {
new_seq.push_back(IfThenElse(current_condition,
auto if_stmt =
IfThenElse(current_condition,
current_if_bodies.size() == 1
? current_if_bodies[0]
: SeqStmt(current_if_bodies),
Stmt()));
: this->VisitStmt(SeqStmt(current_if_bodies)),
Stmt());
new_seq.push_back(if_stmt);
current_condition = PrimExpr();
current_if_bodies.clear();
}
......@@ -73,11 +77,13 @@ private:
}
if (!current_if_bodies.empty()) {
new_seq.push_back(IfThenElse(current_condition,
auto if_stmt =
IfThenElse(current_condition,
current_if_bodies.size() == 1
? current_if_bodies[0]
: SeqStmt(current_if_bodies),
Stmt()));
: this->VisitStmt(SeqStmt(current_if_bodies)),
Stmt());
new_seq.push_back(if_stmt);
}
return new_seq.size() == 1 ? new_seq[0] : SeqStmt(new_seq);
......
src/transform/thread_partial_sync.cc
View file @ 6664d170
......@@ -29,7 +29,8 @@ public:
// The syncs inserted before each statement
std::unordered_set<const Object *> syncs_inserted_;
std::unordered_map<const Object *, int> partial_syncs_inserted_;
std::unordered_map<const Object *, std::tuple<int, int>>
partial_syncs_inserted_;
protected:
bool Enabled(const VarNode *buf, const StorageScope &scope) const final {
......@@ -257,17 +258,24 @@ private:
scope_.push_back(std::vector<StmtEntry>());
num_partial_threads_ = partitions[0];
barrier_id_ += 1;
this->VisitStmt(body->then_case);
StmtEntry s;
s.stmt = op;
s.access = Summarize(std::move(scope_.back()), nullptr);
scope_.pop_back();
if (!has_sync_)
barrier_id_ -= 1;
has_sync_ = false;
num_partial_threads_ = partitions[1];
scope_.push_back(std::vector<StmtEntry>());
barrier_id_ += 1;
VisitStmt(body->else_case.value());
auto v = Summarize(std::move(scope_.back()), nullptr);
scope_.pop_back();
if (!has_sync_)
barrier_id_ -= 1;
has_sync_ = false;
s.access.insert(s.access.end(), v.begin(), v.end());
num_partial_threads_ = std::nullopt;
......@@ -281,10 +289,12 @@ private:
// condition";
if (syncs_inserted_.count(obj))
return;
if (num_partial_threads_.defined()) {
if (num_partial_threads_.defined() && barrier_id_ >= 0 &&
barrier_id_ < 16) {
syncs_inserted_.insert(obj);
partial_syncs_inserted_[obj] =
static_cast<int>(num_partial_threads_.value()->value);
partial_syncs_inserted_[obj] = std::make_tuple(
static_cast<int>(num_partial_threads_.value()->value), barrier_id_);
has_sync_ = true;
} else {
syncs_inserted_.insert(obj);
}
......@@ -294,6 +304,8 @@ private:
Optional<IntImm> num_partial_threads_;
// synchronization scope
StorageScope sync_scope_;
int barrier_id_{-1};
bool has_sync_{false};
};
// There are cases where necessary syncthreads is not inserted by
......@@ -318,7 +330,7 @@ class ThreadPartialSyncInserter : public StmtExprMutator {
public:
ThreadPartialSyncInserter(
StorageScope sync_scope, const std::unordered_set<const Object *> &syncs,
std::unordered_map<const Object *, int> partial_syncs)
std::unordered_map<const Object *, std::tuple<int, int>> partial_syncs)
: sync_scope_(sync_scope), syncs_(syncs), partial_syncs_(partial_syncs) {}
Stmt VisitStmt(const Stmt &stmt) final {
......@@ -329,8 +341,10 @@ public:
if (partial_syncs_.count(stmt.get())) {
auto iter = partial_syncs_.find(stmt.get());
ICHECK(sync_scope_.rank == StorageRank::kShared);
barrier = Evaluate(
Call(DataType::Int(32), tl::sync_thread_partial(), {iter->second}));
int num_threads, barrier_id;
std::tie(num_threads, barrier_id) = iter->second;
barrier = Evaluate(Call(DataType::Int(32), tl::sync_thread_partial(),
{num_threads, barrier_id}));
} else {
return StmtExprMutator::VisitStmt(stmt);
}
......@@ -347,7 +361,8 @@ private:
// data structure.
StorageScope sync_scope_;
const std::unordered_set<const Object *> &syncs_;
const std::unordered_map<const Object *, int> &partial_syncs_;
const std::unordered_map<const Object *, std::tuple<int, int>>
&partial_syncs_;
};
Stmt TileLangThreadPartialSync(Stmt stmt, std::string storage_scope) {
......
src/transform/warp_specialized_rewriter.cc
View file @ 6664d170
......@@ -242,10 +242,15 @@ static PrimExpr makeGetBarrier(PrimExpr barrier_id) {
return Call(DataType::Handle(), get_mbarrier(), {barrier_id});
}
static Stmt makeArriveBarrier(PrimExpr barrier_id) {
auto call = Call(DataType::Handle(), builtin::ptx_arrive_barrier(),
{makeGetBarrier(barrier_id)});
return Evaluate(call);
static Stmt makeArriveBarrier(PrimExpr barrier_id, int cta_id = -1,
PrimExpr pred = 1) {
Array<PrimExpr> args = {makeGetBarrier(barrier_id)};
if (cta_id != -1) {
args.push_back(cta_id);
args.push_back(pred);
}
return Evaluate(
Call(DataType::Handle(), builtin::ptx_arrive_barrier(), args));
}
static Stmt makeCpAsyncBarrier(PrimExpr barrier_id) {
......@@ -318,14 +323,18 @@ private:
class ThreadIdxRewriter : public StmtExprMutator {
public:
static Stmt Rewrite(Stmt stmt, Var thread_var, PrimExpr replaced) {
auto rewriter = ThreadIdxRewriter(thread_var, replaced);
static Stmt Rewrite(Stmt stmt, Var thread_var, PrimExpr replaced,
PrimExpr thread_extent, bool do_shuffle = false) {
auto rewriter =
ThreadIdxRewriter(thread_var, replaced, thread_extent, do_shuffle);
return rewriter(stmt);
}
private:
ThreadIdxRewriter(Var thread_var, PrimExpr replaced)
: thread_var_(thread_var), replaced_(replaced) {}
ThreadIdxRewriter(Var thread_var, PrimExpr replaced, PrimExpr thread_extent,
bool do_shuffle)
: thread_var_(thread_var), replaced_(replaced),
thread_extent_(thread_extent), do_shuffle_(do_shuffle) {}
PrimExpr VisitExpr_(const VarNode *var) final {
if (var == thread_var_.get()) {
......@@ -335,8 +344,34 @@ private:
}
}
Stmt VisitStmt_(const IfThenElseNode *op) final {
auto f_uses_thread_index = [=](const tvm::tir::VarNode *parameter) {
return parameter == thread_var_.get();
};
maybe_thread_opt_ = false;
if (!op->else_case.defined() && op->condition.as<EQNode>() &&
UsesVar(op->condition, f_uses_thread_index) &&
!(UsesVar(op->then_case, f_uses_thread_index))) {
auto eq_op = Downcast<EQ>(op->condition);
if (eq_op->a.as<VarNode>() == thread_var_.get() ||
eq_op->b.as<VarNode>() == thread_var_.get()) {
maybe_thread_opt_ = true;
}
maybe_thread_opt_ = do_shuffle_ && maybe_thread_opt_;
}
if (maybe_thread_opt_)
return IfThenElse(
Call(DataType::Bool(), tl_shuffle_elect(), {thread_extent_}),
StmtExprMutator::VisitStmt(op->then_case), std::nullopt);
else
return StmtExprMutator::VisitStmt_(op);
}
Var thread_var_;
PrimExpr replaced_;
PrimExpr thread_extent_;
bool maybe_thread_opt_ = false;
bool do_shuffle_;
};
Block MakeGroupBlock(const Stmt &stmt,
......@@ -497,6 +532,41 @@ private:
PipelineInfo pipeline_info_;
};
class WgMMACollector : public StmtExprVisitor {
public:
WgMMACollector() = default;
void VisitExpr_(const CallNode *op) final {
if (op->op.same_as(tl_gemm()) || op->op.same_as(tl_gemm_sp())) {
auto op_name = std::string(op->args[0].as<StringImmNode>()->value);
if (has_wgmma_) {
has_wgmma_ =
op_name.find("false") == std::string::npos && !in_if_scope_;
}
}
StmtExprVisitor::VisitExpr_(op);
}
void VisitStmt_(const IfThenElseNode *op) final {
in_if_scope_ = true;
StmtExprVisitor::VisitStmt(op->then_case);
if (op->else_case.defined()) {
StmtExprVisitor::VisitStmt(op->else_case.value());
}
in_if_scope_ = false;
}
static bool HasWgMMA(Stmt stmt) {
auto collector = WgMMACollector();
collector(stmt);
return collector.has_wgmma_;
}
bool has_wgmma_{true};
bool in_if_scope_{false};
};
class WSCodeEmitter : public StmtMutator {
public:
WSCodeEmitter(bool is_emitting_producer, IterVar thread_iv,
......@@ -507,6 +577,10 @@ public:
buffer_data_to_buffer_(buffer_data_to_buffer), marker_(marker),
thread_var_(thread_iv->var), mbarrier_only_(mbarrier_only) {}
bool onlyHasWgMMA() const { return only_has_wgmma_; }
bool hasSimtCopy() const { return has_simt_copy_; }
private:
template <typename NodeType> Stmt FilterByRole(const NodeType *op) {
Role role = marker_.GetRole(op);
......@@ -542,6 +616,9 @@ private:
op->seq.Map([&](Stmt stmt) { return VisitStmt(stmt); });
auto map = ExtractSyncPattern(op->seq);
only_has_wgmma_ = WgMMACollector::HasWgMMA(SeqStmt(op->seq));
/*
std::cout << "Print ExtractSyncPattern" << std::endl;
for (int i = 0; i < static_cast<int>(op->seq.size()); i++) {
......@@ -594,8 +671,9 @@ private:
MbarrierRewriter::Rewrite(seq_transformed[i], release_barrier_id);
collector.Collect(stmt);
block_stmt.push_back(stmt);
if (collector.HasSimtCopy() > 0) {
if (collector.HasSimtCopy()) {
block_stmt.push_back(makeCpAsyncBarrier(release_barrier_id));
has_simt_copy_ = true;
}
if (map.release_after[i][j]) {
block_stmt.push_back(makeArriveBarrier(release_barrier_id));
......@@ -630,6 +708,10 @@ private:
int pattern_idx = map.release[i][j];
PrimExpr release_barrier_id =
stage_ + num_barriers_ + num_stages_ * pattern_idx;
if (only_has_wgmma_)
block_stmt.push_back(makeArriveBarrier(
release_barrier_id, 0, EQ(FloorMod(thread_var_, 128), 0)));
else
block_stmt.push_back(makeArriveBarrier(release_barrier_id));
for (int s = 0; s < num_stages_; s++) {
released_barrier_.insert(s + num_barriers_ +
......@@ -982,6 +1064,8 @@ private:
bool mbarrier_only_ = false;
PipelineInfo pipeline_info_;
friend class WarpSpecializedRewriter;
bool only_has_wgmma_ = false;
bool has_simt_copy_ = false;
};
class SetMaxNRegCollector : public StmtExprVisitor {
......@@ -1022,9 +1106,12 @@ private:
class WarpSpecializedRewriter : public StmtExprMutator {
public:
WarpSpecializedRewriter(bool disable_warp_specialized)
: disable_warp_specialized_(disable_warp_specialized) {}
static PrimFunc Substitute(PrimFunc f, bool disable_warp_specialized) {
WarpSpecializedRewriter(bool disable_warp_specialized,
bool disable_shuffle_elect)
: disable_warp_specialized_(disable_warp_specialized),
disable_shuffle_elect_(disable_shuffle_elect) {}
static PrimFunc Substitute(PrimFunc f, bool disable_warp_specialized,
bool disable_shuffle_elect) {
// Check if function only uses threadIdx.x before proceeding
if (!ThreadTagChecker::HasOnlyThreadIdxX(f)) {
LOG(WARNING) << "WarpSpecialize will be disabled because the program "
......@@ -1035,7 +1122,8 @@ public:
return f;
}
auto T = WarpSpecializedRewriter(disable_warp_specialized);
auto T = WarpSpecializedRewriter(disable_warp_specialized,
disable_shuffle_elect);
T.nreg_ = SetMaxNRegCollector::Collect(f);
T.buffer_lca_ = DetectBufferAccessLCA(f);
for (auto [buffer, _] : T.buffer_lca_)
......@@ -1085,7 +1173,7 @@ private:
ICHECK(thread_tag == "threadIdx.x") << "Only support threadIdx.x";
Var thread_iv = Downcast<Var>(for_node->loop_var);
Stmt new_body =
ThreadIdxRewriter::Rewrite(for_node->body, thread_iv, thread_iv_);
ThreadIdxRewriter::Rewrite(for_node->body, thread_iv, thread_iv_, 0);
return new_body;
}
return for_node;
......@@ -1128,6 +1216,7 @@ private:
WSCodeEmitter consumer(false, thread_iv_, buffer_data_to_buffer_, marker);
Stmt producer_code = producer(block->body);
Stmt consumer_code = consumer(block->body);
bool only_has_wgmma = consumer.onlyHasWgMMA();
PrimExpr consumer_thread_extent = thread_iv_->dom->extent;
PrimExpr producer_thread_extent = thread_iv_->dom->extent;
// Need one warp-group for bulk-copy only case
......@@ -1150,10 +1239,15 @@ private:
producer_code = SeqStmt({dec_reg_stmt, producer_code});
consumer_code = SeqStmt({inc_reg_stmt, consumer_code});
producer_code =
ThreadIdxRewriter::Rewrite(producer_code, thread_iv_->var,
thread_iv_->var - consumer_thread_extent);
updated_thread_extent_ = consumer_thread_extent + producer_thread_extent;
producer_code = ThreadIdxRewriter::Rewrite(
producer_code, thread_iv_->var,
thread_iv_->var - consumer_thread_extent, producer_thread_extent,
!disable_shuffle_elect_);
consumer_code = ThreadIdxRewriter::Rewrite(
consumer_code, thread_iv_->var, thread_iv_->var, consumer_thread_extent,
!disable_shuffle_elect_);
need_update_thread_extent_ = true;
ICHECK(producer.num_barriers_ == consumer.num_barriers_)
......@@ -1162,9 +1256,11 @@ private:
Array<PrimExpr> barrier_num_threads;
barrier_num_threads.reserve(num_barriers);
for (int i = 0; i < num_barriers; i++) {
PrimExpr arrive_thread_count = producer.released_barrier_.count(i)
? producer_thread_extent
: consumer_thread_extent;
PrimExpr arrive_thread_count =
producer.released_barrier_.count(i)
? (producer.hasSimtCopy() ? producer_thread_extent : 1)
: (only_has_wgmma ? FloorDiv(consumer_thread_extent, 128)
: consumer_thread_extent);
barrier_num_threads.push_back(arrive_thread_count);
}
......@@ -1191,6 +1287,7 @@ private:
Optional<PrimExpr> updated_thread_extent_;
bool need_update_thread_extent_ = false;
bool disable_warp_specialized_ = false;
bool disable_shuffle_elect_ = false;
Array<IntImm> nreg_;
};
......@@ -1257,10 +1354,13 @@ tvm::transform::Pass WarpSpecialized() {
auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
bool disable_warp_specialized =
ctx->GetConfig<Bool>(kDisableWarpSpecialized, Bool(false)).value();
bool disable_shuffle_elect =
ctx->GetConfig<Bool>(kDisableShuffleElect, Bool(false)).value();
bool warp_specialized = WarpSpecializedDetector::Detect(f->body);
if (!warp_specialized) {
return WarpSpecializedRewriter::Substitute(f, disable_warp_specialized);
return WarpSpecializedRewriter::Substitute(f, disable_warp_specialized,
disable_shuffle_elect);
}
return f;
};
......
tilelang/language/copy.py
View file @ 6664d170
"""The language interface for tl programs."""
from typing import Union, List, Optional
from typing import Union, List, Optional, Literal
from tilelang import language as T
from tilelang.utils.language import get_buffer_region_from_load
from tvm import ir, tir
......@@ -81,12 +81,11 @@ def buffer_region_to_tile_region(buffer_region: tir.BufferRegion, access_type: s
return region(T.BufferLoad(buffer_region.buffer, mins), access_type, *region_extents)
def copy(
src: Union[tir.Buffer, tir.BufferLoad, tir.BufferRegion],
def copy(src: Union[tir.Buffer, tir.BufferLoad, tir.BufferRegion],
dst: Union[tir.Buffer, tir.BufferLoad],
coalesced_width: Optional[int] = None,
disable_tma: bool = False,
):
eviction_policy: Optional[Literal["evict_normal", "evict_first", "evict_last"]] = None):
"""Copy data between memory regions.
Args:
......@@ -145,12 +144,15 @@ def copy(
if coalesced_width is None:
coalesced_width = -1 # PrimExpr can not be None
if eviction_policy is None:
eviction_policy = 0
else:
eviction_policy = {"evict_normal": 0, "evict_first": 1, "evict_last": 2}[eviction_policy]
return tir.call_intrin("handle", tir.op.Op.get("tl.copy"), src, dst, coalesced_width,
disable_tma)
disable_tma, eviction_policy)
def c2d_im2col(
img: tir.Buffer,
def c2d_im2col(img: tir.Buffer,
col: tir.Buffer,
nhw_step: tir.PrimExpr,
c_step: tir.PrimExpr,
......@@ -158,7 +160,8 @@ def c2d_im2col(
stride: int,
dilation: int,
pad: int,
):
eviction_policy: Optional[Literal["evict_normal", "evict_first",
"evict_last"]] = None):
"""Perform im2col transformation for 2D convolution.
Args:
......@@ -174,15 +177,10 @@ def c2d_im2col(
Returns:
tir.Call: A handle to the im2col operation
"""
return tir.call_intrin(
"handle",
tir.op.Op.get("tl.c2d_im2col"),
img.access_ptr("r"),
col.access_ptr("w"),
nhw_step,
c_step,
kernel,
stride,
dilation,
pad,
)
if eviction_policy is None:
eviction_policy = 0
else:
eviction_policy = {"evict_normal": 0, "evict_first": 1, "evict_last": 2}[eviction_policy]
return tir.call_intrin("handle", tir.op.Op.get("tl.c2d_im2col"), img.access_ptr("r"),
col.access_ptr("w"), nhw_step, c_step, kernel, stride, dilation, pad,
eviction_policy)
tilelang/transform/pass_config.py
View file @ 6664d170
......@@ -43,6 +43,9 @@ class PassConfigKey(str, Enum):
TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE = "tl.enable_aggressive_shared_memory_merge"
"""Enable aggressive merge of shared memory allocations. Default: False"""
TL_DISABLE_SHUFFLE_ELECT = "tl.disable_shuffle_elect"
"""Disable shuffle election optimization. Default: False"""
# TIR related configs
TIR_ENABLE_EQUIV_TERMS_IN_CSE = "tir.enable_equiv_terms_in_cse_tir"
"""Enable equivalent terms in TIR Common Subexpression Elimination. Default: True"""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment