Repository navigation
[Feat] Add support for IPC-based warp-level pull from remote tensor and refactor relevent tests - #14
Conversation
…naming to clarify remote copy operators' scope
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Tip For best results, initiate chat on the files or code changes. @chengyupku Hello! I see you've mentioned me. How can I help you with this PR? Would you like me to:
Please let me know what you'd like me to assist with! (( )) )_) |
|
@coderabbitai review this pr |
|
@Rachmanino I'll review this PR for you. This looks like a significant refactor that introduces IPC-based warp-level pull operations and splits the previous ✅ Actions performedReview triggered.
|
WalkthroughIntroduces warp-level distributed copy primitives by replacing remote_copy with push_warp and adding pull_warp across TL core, language API, and CUDA templates. Updates an example to use push_warp, adds a new pull_warp example, and removes legacy IPC tests and CUDA extension. Minor type hint added in distributed utils. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant App as Example Kernel (push_warp)
participant TL as TL Op (PushWarpOp)
participant MD as MetaData (optional)
participant GPU as GPU Kernel
participant PE as Remote PE
App->>TL: push_warp(src, dst, size, dst_pe?, unroll)
alt Symmetric (dst_pe set)
TL->>MD: Load remote addr via metadata
TL->>GPU: emit cp_unrolled<N,UF>(dst_addr_u64, src_ptr)
else UVA (no dst_pe)
TL->>GPU: emit cp_unrolled<N,UF>(dst_u64, src_ptr)
end
GPU-->>PE: Writes to remote memory (push)
sequenceDiagram
autonumber
participant App as Example Kernel (pull_warp)
participant TL as TL Op (PullWarpOp)
participant MD as MetaData (optional)
participant GPU as GPU Kernel
participant PE as Remote PE
App->>TL: pull_warp(src, dst, size, src_pe?, unroll)
alt Symmetric (src_pe set)
TL->>MD: Load remote addr via metadata
TL->>GPU: emit cp_unrolled<N,UF>(dst_ptr, src_addr_u64)
else UVA (no src_pe)
TL->>GPU: emit cp_unrolled<N,UF>(dst_ptr, src_u64)
end
GPU-->>PE: Reads from remote memory (pull)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/op/remote_copy.cc (1)
34-57: Incorrect IntImm extraction; will not compile.
args[4].as<IntImm>().value()->valueis invalid. Use IntImmNode and guard.- unroll_factor = args[4].as<IntImm>().value()->value; + const auto* imm = args[4].as<IntImmNode>(); + ICHECK(imm) << "unroll_factor must be an IntImm"; + unroll_factor = static_cast<int>(imm->value);
🧹 Nitpick comments (18)
tilelang/distributed/utils.py (1)
47-49: Type annotation addition looks good.Accurate return type; no runtime impact. Consider annotating get_local_ipc_handle/create_dist_tensor in a follow-up for consistency.
examples/distributed/ipc/example_push_warp.py (2)
1-1: Ensure NCCL_DEBUG is set before any NCCL usage.Setting the env var early is good. To be safest, set it before importing torch.distributed to avoid any lazy init edge cases on some stacks.
-import os +import os +os.environ['NCCL_DEBUG'] = 'WARN' # silence NCCL log (set before importing torch/distributed) import tilelang import tilelang.language as T @@ -tilelang.disable_cache() -os.environ['NCCL_DEBUG'] = 'WARN' # silence NCCL log +tilelang.disable_cache()Also applies to: 11-11
24-31: push_warp usage matches the new API, but example assumes 2 ranks.
dst_pe = rank[0] ^ 1only swaps 0↔1, 2↔3, etc. If users pass num_processes != 2, results aren’t obvious. Consider asserting num_ranks == 2 or generalizing neighbor selection.- #TODO(ipc): Add work partition among warps + # TODO(ipc): Add work partition among warps T.push_warp( src=T.address_of(src[bid * block_M, 0]), dst=T.address_of(dst[bid * block_M, 0]), size=block_M * N, - dst_pe=rank[0] ^ 1, + # NOTE: This example targets 2 ranks; for >2 ranks, choose an explicit dst_pe. + dst_pe=rank[0] ^ 1, unroll_factor=4)Option: add
assert num_rank == 2on the host before launching.tilelang/language/distributed/common.py (2)
8-13: Docstring wording: “Push to a remote buffer”, not “from”.Minor clarity fix to reflect semantics.
- """Push from a remote buffer with unrolled loop. + """Push to a remote buffer with an unrolled loop.
33-39: Docstring consistency for pull_warp.Tweak wording and parameter note for symmetry/UVA to mirror push_warp.
- """Pull from a remote buffer with unrolled loop. + """Pull from a remote buffer with an unrolled loop.Optional: mention “If src_pe is provided, src is symmetric; otherwise UVA” in one line to match push_warp’s style.
src/op/remote_copy.cc (3)
3-3: Update file brief to reflect both ops.- * \brief Push warp operator. + * \brief Warp-level remote copy operators (push_warp, pull_warp).
59-82: Validate extern call argument order matches tl_templates.You pass dst then src (symmetric or UVA). Ensure tl::cp_unrolled signature is (dst, src). If templates expect (src, dst), swap here.
109-133: Minor: consistent argument placement and comments.You prepend
dst_addrwith a comment “Always dst first in tl_templates”. Consider adding the same clarifying comment in PushWarpOp for parity.src/op/remote_copy.h (2)
21-39: PushWarpOp data shape and Clone are fine; consider documenting symmetry.Add a short comment on dst_pe-driven symmetry to aid readers.
7-8: Header guard name mismatch (optional).Guard uses BULK_COPY; file is remote_copy. Low impact but can reduce confusion.
-#ifndef TVM_TL_OP_BULK_COPY_H_ -#define TVM_TL_OP_BULK_COPY_H_ +#ifndef TVM_TL_OP_REMOTE_COPY_H_ +#define TVM_TL_OP_REMOTE_COPY_H_ @@ -#endif // TVM_TL_OP_BULK_COPY_H_ +#endif // TVM_TL_OP_REMOTE_COPY_H_Also applies to: 64-64
src/tl_templates/cuda/copy.h (4)
195-201: Add restrict to help alias analysis and keep signatures consistent across overloads.No behavior change; improves codegen and matches non-const dst correction.
-template <int N, int UNROLL_FACTOR, typename dtype_t> -TL_DEVICE void cp_unrolled(uint64_t dst_addr_uint64, - dtype_t const *const src_addr) { +template <int N, int UNROLL_FACTOR, typename dtype_t> +TL_DEVICE void cp_unrolled(uint64_t dst_addr_uint64, + dtype_t const *__restrict__ src_addr) { dtype_t *dst_addr = reinterpret_cast<dtype_t *>(dst_addr_uint64);
218-224: Align overload with restrict and keep dst non-const.-template <int N, int UNROLL_FACTOR, typename dtype_t> -TL_DEVICE void cp_unrolled(dtype_t *const dst_addr, uint64_t src_addr_uint64) { +template <int N, int UNROLL_FACTOR, typename dtype_t> +TL_DEVICE void cp_unrolled(dtype_t *__restrict__ const dst_addr, uint64_t src_addr_uint64) { const dtype_t *src_addr = reinterpret_cast<const dtype_t *>(src_addr_uint64);
135-141: Confirm 16-byte alignment for int4 vector LD/ST or provide a scalar fallback.v4.s32 requires 16B alignment. If T.address_of(...) can be misaligned, add alignment guarantees or branch to scalar path to avoid misaligned accesses.
I can add a guarded path using LD_FUNC/ST_FUNC with scalar types when (reinterpret_cast<uintptr_t>(ptr) & 0xF) != 0.
Also applies to: 164-167
75-90: VecInt<1>/<2> are currently unusable without int8_t/int16_t LD/ST specializations.Either remove them for now or add the missing specializations; otherwise they suggest support that doesn’t exist.
I can provide minimal, correct int16_t load/store PTX; int8_t store is trickier in inline PTX and may be deferred.
examples/distributed/ipc/example_pull_warp.py (4)
41-46: Ensure src is CUDA; NCCL all_gather requires CUDA tensors (and pull_warp needs device-accessible memory).Verify that create_tensor returns a CUDA tensor; otherwise all_gather will fail or pull_warp will dereference host memory.
Optionally enforce:
- src = create_tensor([M, N], torch.float32).random_() + src = create_tensor([M, N], torch.float32).random_() + assert src.is_cuda, "src must be a CUDA tensor"If create_tensor accepts a device arg, pass device="cuda" explicitly.
Also applies to: 57-59
47-50: Only rank 0 prints kernel source — good. Consider caching compilation across ranks.Compilation is per-process; if compilation time is noticeable, precompile on rank 0 and broadcast artifacts, or set persistent cache. Not blocking.
61-67: Tighten validation output for easier diffing.Large tensors printed fully can spam logs. Print summaries (shape, dtype, a few elements) on failure.
- print(f"dst_torch: {dst_torch}, dst: {dst}") + print(f"Mismatch on rank {local_rank}: dst_torch.mean={dst_torch.float().mean().item():.6f}, " + f"dst.mean={dst.float().mean().item():.6f}, max_abs_err={(dst_torch - dst).abs().max().item():.3e}")
14-22: Meta buffer shape matches (1, num_rank); confirm dtype/stride expectations.Buffer is reshaped to (1, num_local_ranks) and annotated as uint64 meta_data. Ensure create_dist_tensor returns uint64 device pointers and the kernel expects row-major contiguous layout.
I can wire up a small check to print buffer_ptrs_gpu.dtype/stride on rank 0 before launch.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (11)
examples/distributed/ipc/example_pull_warp.py(1 hunks)examples/distributed/ipc/example_push_warp.py(3 hunks)src/op/remote_copy.cc(4 hunks)src/op/remote_copy.h(3 hunks)src/tl_templates/cuda/copy.h(1 hunks)tilelang/distributed/testing/ipc/README.md(0 hunks)tilelang/distributed/testing/ipc/test_ipc.py(0 hunks)tilelang/distributed/testing/ipc/test_kernel/set_value.cu(0 hunks)tilelang/distributed/testing/ipc/test_kernel/setup.py(0 hunks)tilelang/distributed/utils.py(1 hunks)tilelang/language/distributed/common.py(2 hunks)
💤 Files with no reviewable changes (4)
- tilelang/distributed/testing/ipc/README.md
- tilelang/distributed/testing/ipc/test_kernel/setup.py
- tilelang/distributed/testing/ipc/test_ipc.py
- tilelang/distributed/testing/ipc/test_kernel/set_value.cu
🔇 Additional comments (3)
src/op/remote_copy.cc (1)
135-143: Operator registration count aligns with 5 inputs.Looks correct given (src, dst, size, pe, unroll). No action.
src/op/remote_copy.h (2)
3-3: Header brief matches the new scope.Good rename to warp-level ops.
41-59: PullWarpOp mirrors PushWarpOp appropriately.API parity looks good.
| with T.Kernel(T.ceildiv(M, block_M), threads=threads) as bid: | ||
| #TODO(ipc): Add work partition among warps | ||
| T.pull_warp( | ||
| src=T.address_of(src[bid * block_M, 0]), | ||
| dst=T.address_of(dst[bid * block_M, 0]), | ||
| size=block_M * N, | ||
| src_pe=rank[0] ^ 1, | ||
| unroll_factor=4) | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Guard pull_warp to a single warp or partition work across warps.
With threads=128, all 4 warps will issue identical T.pull_warp copies, wasting bandwidth. Quick fix: use a single warp per block; better: partition the [block_M*N] span among warps.
Quick fix:
- threads = 128
+ threads = 32 # one warp per block since pull_warp is warp-scopedIf you prefer multi-warp, I can draft a warp-chunked version that offsets src/dst by warp_id * (size / num_warps).
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| with T.Kernel(T.ceildiv(M, block_M), threads=threads) as bid: | |
| #TODO(ipc): Add work partition among warps | |
| T.pull_warp( | |
| src=T.address_of(src[bid * block_M, 0]), | |
| dst=T.address_of(dst[bid * block_M, 0]), | |
| size=block_M * N, | |
| src_pe=rank[0] ^ 1, | |
| unroll_factor=4) | |
| threads = 32 # one warp per block since pull_warp is warp-scoped |
🤖 Prompt for AI Agents
In examples/distributed/ipc/example_pull_warp.py around lines 23 to 31,
T.pull_warp is currently executed by all threads (all warps) causing duplicate
transfers; either restrict the call to a single warp per block or partition the
transfer across warps. Quick fix: guard the call so only one warp (e.g., warp_id
== 0) invokes T.pull_warp for the whole block. Better fix: compute per-warp_size
= ceil_div(block_M * N, num_warps), compute src_offset = warp_id * per-warp_size
and dst_offset similarly, and call T.pull_warp with size = min(per-warp_size,
remaining) and src/dst addresses offset by those bytes; ensure correct integer
division/ceil and handle remainder so total bytes equal block_M * N.
| PullWarpOp::PullWarpOp(Array<PrimExpr> args, BufferMap vmap) { | ||
| src_addr = args[0]; | ||
| dst_addr = args[1]; | ||
| ICHECK(src_addr.as<CallNode>()) << "src_addr must be a call node"; | ||
| ICHECK(src_addr.as<CallNode>()->op.same_as(builtin::address_of())) | ||
| << "src_addr must be address_of op"; | ||
| ICHECK(dst_addr.as<CallNode>()) << "dst_addr must be a call node"; | ||
| ICHECK(dst_addr.as<CallNode>()->op.same_as(builtin::address_of())) | ||
| << "dst_addr must be address_of op"; | ||
|
|
||
| src_offset = | ||
| get_offset(src_addr.as<CallNode>()->args[0].as<BufferLoadNode>()); | ||
| dst_offset = | ||
| get_offset(dst_addr.as<CallNode>()->args[0].as<BufferLoadNode>()); | ||
| src_buffer = src_addr.as<CallNode>()->args[0].as<BufferLoadNode>()->buffer; | ||
| dst_buffer = dst_addr.as<CallNode>()->args[0].as<BufferLoadNode>()->buffer; | ||
|
|
||
| copy_size = args[2]; | ||
| src_pe = args[3]; | ||
| unroll_factor = args[4].as<IntImm>().value()->value; | ||
| if (src_pe.defined()) { | ||
| is_symmetric = true; | ||
| } | ||
| } |
There was a problem hiding this comment.
Repeat the IntImm fix in PullWarpOp.
Same extraction bug.
- unroll_factor = args[4].as(IntImm)().value()->value;
+ const auto* imm = args[4].as<IntImmNode>();
+ ICHECK(imm) << "unroll_factor must be an IntImm";
+ unroll_factor = static_cast<int>(imm->value);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| PullWarpOp::PullWarpOp(Array<PrimExpr> args, BufferMap vmap) { | |
| src_addr = args[0]; | |
| dst_addr = args[1]; | |
| ICHECK(src_addr.as<CallNode>()) << "src_addr must be a call node"; | |
| ICHECK(src_addr.as<CallNode>()->op.same_as(builtin::address_of())) | |
| << "src_addr must be address_of op"; | |
| ICHECK(dst_addr.as<CallNode>()) << "dst_addr must be a call node"; | |
| ICHECK(dst_addr.as<CallNode>()->op.same_as(builtin::address_of())) | |
| << "dst_addr must be address_of op"; | |
| src_offset = | |
| get_offset(src_addr.as<CallNode>()->args[0].as<BufferLoadNode>()); | |
| dst_offset = | |
| get_offset(dst_addr.as<CallNode>()->args[0].as<BufferLoadNode>()); | |
| src_buffer = src_addr.as<CallNode>()->args[0].as<BufferLoadNode>()->buffer; | |
| dst_buffer = dst_addr.as<CallNode>()->args[0].as<BufferLoadNode>()->buffer; | |
| copy_size = args[2]; | |
| src_pe = args[3]; | |
| unroll_factor = args[4].as<IntImm>().value()->value; | |
| if (src_pe.defined()) { | |
| is_symmetric = true; | |
| } | |
| } | |
| copy_size = args[2]; | |
| src_pe = args[3]; | |
| const auto* imm = args[4].as<IntImmNode>(); | |
| ICHECK(imm) << "unroll_factor must be an IntImm"; | |
| unroll_factor = static_cast<int>(imm->value); | |
| if (src_pe.defined()) { | |
| is_symmetric = true; | |
| } | |
| } |
🤖 Prompt for AI Agents
In src/op/remote_copy.cc around lines 84-107, the extraction of unroll_factor
incorrectly uses args[4].as<IntImm>().value()->value (repeat of the earlier
IntImm bug); change it to cast to the IntImmNode and read its value like
args[4].as<IntImmNode>()->value, and add an ICHECK that args[4].as<IntImmNode>()
is non-null (or args[4].defined() and is<IntImmNode>()) before reading to avoid
null/deref errors; update any similar IntImm extractions in this block to follow
the same pattern.
| template <typename dtype_t> TL_DEVICE dtype_t ld_nc_global(const dtype_t *ptr) { | ||
| auto ret = ld_nc_global( | ||
| reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(ptr)); | ||
| return *reinterpret_cast<dtype_t *>(&ret); | ||
| } |
There was a problem hiding this comment.
Fix generic ld_nc_global recursion for 1- and 2-byte dtypes (potential infinite self-call).
With VecInt<1>/<2> present but no ld_nc_global specializations for int8_t/int16_t, the generic template calls itself recursively for those sizes. Either restrict supported sizes or add missing specializations.
Apply this diff to guard unsupported sizes now:
template <typename dtype_t> TL_DEVICE dtype_t ld_nc_global(const dtype_t *ptr) {
+ static_assert(sizeof(dtype_t) == 4 || sizeof(dtype_t) == 8 || sizeof(dtype_t) == 16,
+ "ld_nc_global: unsupported dtype size (1/2-byte types need explicit specialization).");
auto ret = ld_nc_global(
reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(ptr));
return *reinterpret_cast<dtype_t *>(&ret);
}Optional: I can also provide provably-correct int16_t/int8_t specializations if you intend to support those.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| template <typename dtype_t> TL_DEVICE dtype_t ld_nc_global(const dtype_t *ptr) { | |
| auto ret = ld_nc_global( | |
| reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(ptr)); | |
| return *reinterpret_cast<dtype_t *>(&ret); | |
| } | |
| template <typename dtype_t> TL_DEVICE dtype_t ld_nc_global(const dtype_t *ptr) { | |
| static_assert(sizeof(dtype_t) == 4 || sizeof(dtype_t) == 8 || sizeof(dtype_t) == 16, | |
| "ld_nc_global: unsupported dtype size (1/2-byte types need explicit specialization)."); | |
| auto ret = ld_nc_global( | |
| reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(ptr)); | |
| return *reinterpret_cast<dtype_t *>(&ret); | |
| } |
| TL_DEVICE void st_na_global(const dtype_t *ptr, const dtype_t &value) { | ||
| st_na_global( | ||
| reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(ptr), | ||
| *reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>( | ||
| &value)); | ||
| } |
There was a problem hiding this comment.
Prevent generic st_na_global from recursing for 1/2-byte dtypes and document supported sizes.
Same recursion risk as loads; add a size gate and avoid calling the same template on int8_t/int16_t.
Apply:
template <typename dtype_t>
TL_DEVICE void st_na_global(const dtype_t *ptr, const dtype_t &value) {
+ static_assert(sizeof(dtype_t) == 4 || sizeof(dtype_t) == 8 || sizeof(dtype_t) == 16,
+ "st_na_global: unsupported dtype size (1/2-byte types need explicit specialization).");
st_na_global(
reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(ptr),
*reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(
&value));
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| TL_DEVICE void st_na_global(const dtype_t *ptr, const dtype_t &value) { | |
| st_na_global( | |
| reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(ptr), | |
| *reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>( | |
| &value)); | |
| } | |
| TL_DEVICE void st_na_global(const dtype_t *ptr, const dtype_t &value) { | |
| static_assert(sizeof(dtype_t) == 4 || sizeof(dtype_t) == 8 || sizeof(dtype_t) == 16, | |
| "st_na_global: unsupported dtype size (1/2-byte types need explicit specialization)."); | |
| st_na_global( | |
| reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>(ptr), | |
| *reinterpret_cast<const typename VecInt<sizeof(dtype_t)>::vec_t *>( | |
| &value)); | |
| } |
🤖 Prompt for AI Agents
In src/tl_templates/cuda/copy.h around lines 144-149, the generic st_na_global
can recurse when dtype_t is 1- or 2-byte (e.g., int8_t/int16_t) because the
VecInt mapping yields a larger vec_t; add a size gate and documentation: guard
the reinterpret_cast/recursive call with an if constexpr (or static if) so it
only runs when sizeof(dtype_t) >= sizeof(typename
VecInt<sizeof(dtype_t)>::vec_t) (or equivalently when VecInt yields an
equal-or-smaller element type), otherwise perform a scalar-safe store path
(e.g., call the non-vector scalar st_na_global or use a memcpy/byte-wise store)
to avoid recursion; add a brief comment above the function listing supported
dtype sizes and that 1/2-byte types are handled by the scalar path.
| template <> TL_DEVICE void st_na_global(const int *ptr, const int &value) { | ||
| asm volatile(ST_NA_FUNC ".s32 [%0], %1;" ::"l"(ptr), "r"(value)); | ||
| } | ||
|
|
||
| template <> | ||
| TL_DEVICE void st_na_global(const int64_t *ptr, const int64_t& value) { | ||
| asm volatile(ST_NA_FUNC ".s64 [%0], %1;" ::"l"(ptr), "l"(value)); | ||
| TL_DEVICE void st_na_global(const int64_t *ptr, const int64_t &value) { | ||
| asm volatile(ST_NA_FUNC ".s64 [%0], %1;" ::"l"(ptr), "l"(value)); | ||
| } | ||
|
|
||
| template <> | ||
| TL_DEVICE void st_na_global(const float *ptr, const float& value) { | ||
| asm volatile(ST_NA_FUNC ".f32 [%0], %1;" ::"l"(ptr), "f"(value)); | ||
| template <> TL_DEVICE void st_na_global(const float *ptr, const float &value) { | ||
| asm volatile(ST_NA_FUNC ".f32 [%0], %1;" ::"l"(ptr), "f"(value)); | ||
| } | ||
|
|
||
| template <> | ||
| TL_DEVICE void st_na_global(const int4 *ptr, const int4& value) { | ||
| asm volatile(ST_NA_FUNC ".v4.s32 [%0], {%1, %2, %3, %4};" | ||
| ::"l"(ptr), "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w)); | ||
| template <> TL_DEVICE void st_na_global(const int4 *ptr, const int4 &value) { | ||
| asm volatile(ST_NA_FUNC ".v4.s32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), | ||
| "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w)); | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Mark stores as writing to memory to prevent reordering issues.
Inline PTX stores lack a "memory" clobber; the compiler may reorder surrounding memory ops across these stores. Add a memory clobber to all store specializations.
template <> TL_DEVICE void st_na_global(const int *ptr, const int &value) {
- asm volatile(ST_NA_FUNC ".s32 [%0], %1;" ::"l"(ptr), "r"(value));
+ asm volatile(ST_NA_FUNC ".s32 [%0], %1;" ::"l"(ptr), "r"(value) : "memory");
}
template <>
TL_DEVICE void st_na_global(const int64_t *ptr, const int64_t &value) {
- asm volatile(ST_NA_FUNC ".s64 [%0], %1;" ::"l"(ptr), "l"(value));
+ asm volatile(ST_NA_FUNC ".s64 [%0], %1;" ::"l"(ptr), "l"(value) : "memory");
}
template <> TL_DEVICE void st_na_global(const float *ptr, const float &value) {
- asm volatile(ST_NA_FUNC ".f32 [%0], %1;" ::"l"(ptr), "f"(value));
+ asm volatile(ST_NA_FUNC ".f32 [%0], %1;" ::"l"(ptr), "f"(value) : "memory");
}
template <> TL_DEVICE void st_na_global(const int4 *ptr, const int4 &value) {
- asm volatile(ST_NA_FUNC ".v4.s32 [%0], {%1, %2, %3, %4};" ::"l"(ptr),
- "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w));
+ asm volatile(ST_NA_FUNC ".v4.s32 [%0], {%1, %2, %3, %4};" ::"l"(ptr),
+ "r"(value.x), "r"(value.y), "r"(value.z), "r"(value.w) : "memory");
}🤖 Prompt for AI Agents
In src/tl_templates/cuda/copy.h around lines 151 to 167, the inline PTX store
specializations use asm volatile without a "memory" clobber allowing the
compiler to reorder surrounding memory accesses; update each asm volatile call
to include the "memory" clobber (i.e. add a clobber list containing "memory"
after the input operand section) so the compiler treats these stores as touching
memory and prevents reordering.
| template <int N, int UNROLL_FACTOR, typename dtype_t> | ||
| TL_DEVICE void cp_unrolled(dtype_t const* const dst_addr, dtype_t const* const src_addr) { | ||
| TL_DEVICE void cp_unrolled(dtype_t const *const dst_addr, | ||
| dtype_t const *const src_addr) { | ||
| int lane_id; | ||
| asm("mov.s32 %0, %laneid;" : "=r"(lane_id)); |
There was a problem hiding this comment.
Destination must not be const; current signature implies writing through a const pointer.
This can lead to UB and misoptimizations when combined with inline PTX stores.
-template <int N, int UNROLL_FACTOR, typename dtype_t>
-TL_DEVICE void cp_unrolled(dtype_t const *const dst_addr,
- dtype_t const *const src_addr) {
+template <int N, int UNROLL_FACTOR, typename dtype_t>
+TL_DEVICE void cp_unrolled(dtype_t *__restrict__ const dst_addr,
+ dtype_t const *__restrict__ src_addr) {📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| template <int N, int UNROLL_FACTOR, typename dtype_t> | |
| TL_DEVICE void cp_unrolled(dtype_t const* const dst_addr, dtype_t const* const src_addr) { | |
| TL_DEVICE void cp_unrolled(dtype_t const *const dst_addr, | |
| dtype_t const *const src_addr) { | |
| int lane_id; | |
| asm("mov.s32 %0, %laneid;" : "=r"(lane_id)); | |
| template <int N, int UNROLL_FACTOR, typename dtype_t> | |
| TL_DEVICE void cp_unrolled(dtype_t *__restrict__ const dst_addr, | |
| dtype_t const *__restrict__ src_addr) { | |
| int lane_id; | |
| asm("mov.s32 %0, %laneid;" : "=r"(lane_id)); | |
| // …rest of function… | |
| } |
🤖 Prompt for AI Agents
In src/tl_templates/cuda/copy.h around lines 172 to 176, the cp_unrolled
signature marks the destination pointer as pointing to const (dst_addr is
dtype_t const *const) which incorrectly implies writes through a const pointer
and can cause UB or misoptimizations with inline PTX stores; change the function
signature so the destination is non-const (e.g., dst_addr should be dtype_t
*const or dtype_t*), keep the src pointer const, and update any callers/types
accordingly so that stores to dst_addr are to a mutable pointer type.
This pull request introduces a significant refactor and extension of the distributed IPC (Inter-Process Communication) operators and examples, focusing on warp-level push and pull operations. The main changes include splitting and renaming of distributed copy operators, the addition of new kernel examples for both push and pull warp operations, and substantial code cleanup and modernization in the C++ CUDA templates.
Key changes:
Operator Refactor and Extension
RemoteCopyOpinto two distinct operators:PushWarpOpandPullWarpOp, each with their own implementation and registration. This clarifies the semantics and usage of distributed memory copy at the warp level. [1] [2] [3] [4]push_warpandpull_warpterminology, and updated the corresponding documentation comments. [1] [2] [3] [4]Example and Test Updates
example_pull_warp.pydemonstrating the usage of the newpull_warpoperator for distributed tensor copying.example_remote_copy.pytoexample_push_warp.pyand updated it to use the newpush_warpoperator instead of the deprecatedremote_copy. [1] [2] [3]CUDA Template Improvements
These changes collectively modernize and clarify the distributed IPC operator infrastructure, making it easier to extend and maintain, and provide clear, up-to-date examples for users.
Summary by CodeRabbit
New Features
Refactor
Documentation
Tests
Chores