Skip to content

Commit 0eef4e9

Browse files
janssen-vjanssen-v
committed
feat: Add Ascend RDMA support
Add Ascend ACLSHMEM wrappers for RMA, put-signal, wait, quiet, and fence operations with dtype suffix mapping. Expose ACL SHMEM signal/compare constants and wire Ascend SHMEM extern handling.
1 parent 4a2767a commit 0eef4e9

9 files changed

Lines changed: 797 additions & 9 deletions

File tree

‎include/TritonDistributed/Conversion/TritonDistributedToHIVM/TritonDistributedToHIVMPass.h‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,6 @@ static const llvm::SmallVector<std::string> mixTypeList = {
3737
"aclshmem_barrier_all"};
3838
static const llvm::SmallVector<std::string> vecTypeList = {
3939
"aclshmemx_barrier_all_vec"};
40-
static const llvm::SmallVector<std::string> sideEffectSymbols = {
41-
"aclshmemx_barrier_all_vec", "aclshmem_barrier_all", "aclshmem_int8_p",
42-
"aclshmem_int16_p", "aclshmem_int32_p", "aclshmem_int64_p"};
4340
static const llvm::DenseMap<llvm::StringRef, hivm::PIPE> pipeMap;
4441
static const int distributedDialectPrefixLen = 12;
4542
void populateDistributedOpToHIVMPatterns(RewritePatternSet &patterns,

‎lib/Conversion/TritonDistributedToHIVM/ASCEND/DistributedOpToHIVM.cpp‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
#include "mlir/IR/SymbolTable.h"
3232
#include "mlir/IR/Value.h"
3333
#include "mlir/IR/ValueRange.h"
34+
#include "mlir/Interfaces/SideEffectInterfaces.h"
3435
#include "triton/Dialect/Triton/IR/Types.h"
3536
#include "llvm/ADT/DenseMap.h"
3637
#include "llvm/ADT/STLExtras.h"
@@ -122,7 +123,7 @@ class DistributedOpToHIVM : public OpRewritePattern<DistOp> {
122123
if (auto sym = op->template getAttrOfType<StringAttr>("symbol")) {
123124
symbolName = sym.str();
124125
}
125-
bool hasSideEffect = true;
126+
bool hasSideEffect = !mlir::isMemoryEffectFree(op.getOperation());
126127

127128
llvm::TypeSwitch<Operation *, void>(op)
128129
.Case([&](distributed::SymmAtOp) {
@@ -159,9 +160,6 @@ class DistributedOpToHIVM : public OpRewritePattern<DistOp> {
159160
symbolName = op->getName().getStringRef().drop_front(
160161
ASCEND::distributedDialectPrefixLen);
161162
}
162-
if (!llvm::is_contained(ASCEND::sideEffectSymbols, symbolName)) {
163-
hasSideEffect = false;
164-
}
165163
std::string customName = customPrefix + "." + symbolName;
166164
ValueRange operands = op->getOperands();
167165
if (symbolName == "aclshmem_my_pe" || symbolName == "aclshmem_n_pes") {

‎python/triton_dist/jit.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,10 @@ def get_shmem_extern_lib() -> Dict[str, str]:
142142
extern_libs = {"libshmem": str(mxshmem_lib)}
143143
return extern_libs
144144

145+
elif is_ascend():
146+
extern_libs = {"libshmem": str("")}
147+
return extern_libs
148+
145149
else:
146150
raise NotImplementedError("Unsupported device type to get shmem bitcode lib path.")
147151

‎python/triton_dist/language/extra/ascend/libaclshmem_device.py‎

Lines changed: 317 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,57 @@
55
from triton_dist.language.core import extern_call
66

77
pi_u64_t = tl.core.pointer_type(tl.core.dtype("uint64"))
8+
pi_i32_t = tl.core.pointer_type(tl.core.dtype("int32"))
89

910
void_ptr = core.pointer_type(core.void)
1011

12+
# The subset of dtypes supported by Ascend SHMEM RMA ops.
13+
RMA_DTYPES = [
14+
"char",
15+
"fp16",
16+
"fp32",
17+
"int8",
18+
"int16",
19+
"int32",
20+
"int64",
21+
"uint8",
22+
"uint16",
23+
"uint32",
24+
"uint64",
25+
"bf16",
26+
]
27+
28+
DTYPE_TO_KERNEL_SUFFIX = {
29+
"char": "char",
30+
"fp16": "half",
31+
"fp32": "float",
32+
"int8": "int8",
33+
"int16": "int16",
34+
"int32": "int32",
35+
"int64": "int64",
36+
"uint8": "uint8",
37+
"uint16": "uint16",
38+
"uint32": "uint32",
39+
"uint64": "uint64",
40+
"bf16": "bfloat16",
41+
}
42+
43+
44+
def _pointer_type_hash(self):
45+
return hash((self.name, self.element_ty, "tt_ptr"))
46+
47+
48+
def patch_hash_method_for_pointer_type():
49+
elem_dtype_list = (tl.core.dtype.SINT_TYPES + tl.core.dtype.UINT_TYPES + tl.core.dtype.FP_TYPES +
50+
tl.core.dtype.OTHER_TYPES)
51+
for elem_dtype in elem_dtype_list:
52+
ptr_ty = type(tl.core.pointer_type(tl.core.dtype(elem_dtype)))
53+
ptr_ty.__hash__ = _pointer_type_hash
54+
55+
56+
# Keep pointer_type hashable for Triton extern signature dict keys.
57+
patch_hash_method_for_pointer_type()
58+
1159

1260
@core.extern
1361
def my_pe(_semantic=None):
@@ -37,6 +85,55 @@ def n_pes(_semantic=None):
3785
)
3886

3987

88+
@core.extern
89+
def team_my_pe(team, _semantic=None):
90+
return extern_call(
91+
"libshmem_device",
92+
"",
93+
[tl.cast(team, tl.int32, _semantic=_semantic)],
94+
{
95+
(tl.int32, ): ("aclshmem_team_my_pe", core.dtype("int32")),
96+
},
97+
is_pure=True,
98+
_semantic=_semantic,
99+
)
100+
101+
102+
@core.extern
103+
def team_n_pes(team, _semantic=None):
104+
return extern_call(
105+
"libshmem_device",
106+
"",
107+
[tl.cast(team, tl.int32, _semantic=_semantic)],
108+
{
109+
(tl.int32, ): ("aclshmem_team_n_pes", core.dtype("int32")),
110+
},
111+
is_pure=True,
112+
_semantic=_semantic,
113+
)
114+
115+
116+
@core.extern
117+
def team_translate_pe(src_team, pe_in_src_team, dest_team, _semantic=None):
118+
return extern_call(
119+
"libshmem_device",
120+
"",
121+
[
122+
tl.cast(src_team, tl.int32, _semantic=_semantic),
123+
tl.cast(pe_in_src_team, tl.int32, _semantic=_semantic),
124+
tl.cast(dest_team, tl.int32, _semantic=_semantic),
125+
],
126+
{
127+
(tl.int32, tl.int32, tl.int32): (
128+
"aclshmem_team_translate_pe",
129+
core.dtype("int32"),
130+
),
131+
},
132+
is_pure=True,
133+
_semantic=_semantic,
134+
)
135+
136+
40137
@core.extern
41138
def int_p(dest, value, pe, _semantic=None):
42139
# force have a return value, even not used.
@@ -100,3 +197,223 @@ def barrier_all_vec(_semantic=None):
100197
is_pure=False,
101198
_semantic=_semantic,
102199
)
200+
201+
202+
@core.extern
203+
def barrier(team, _semantic=None):
204+
return extern_call(
205+
"libshmem_device",
206+
"",
207+
[tl.cast(team, tl.int32, _semantic=_semantic)],
208+
{
209+
(tl.int32, ): ("aclshmem_barrier", ()),
210+
},
211+
is_pure=False,
212+
_semantic=_semantic,
213+
)
214+
215+
216+
@core.extern
217+
def barrier_vec(team, _semantic=None):
218+
return extern_call(
219+
"libshmem_device",
220+
"",
221+
[tl.cast(team, tl.int32, _semantic=_semantic)],
222+
{
223+
(tl.int32, ): ("aclshmemx_barrier_vec", ()),
224+
},
225+
is_pure=False,
226+
_semantic=_semantic,
227+
)
228+
229+
230+
@core.extern
231+
def _rma_mem_impl(dest, source, elem_size, pe, op: core.constexpr, _semantic=None):
232+
return extern_call(
233+
"libshmem_device",
234+
"",
235+
[
236+
dest,
237+
source,
238+
tl.cast(elem_size, tl.uint32, _semantic=_semantic),
239+
tl.cast(pe, tl.int32, _semantic=_semantic),
240+
],
241+
{(
242+
core.pointer_type(core.dtype(core_dtype)),
243+
core.pointer_type(core.dtype(core_dtype)),
244+
tl.uint32,
245+
tl.int32,
246+
): (
247+
f"aclshmem_{DTYPE_TO_KERNEL_SUFFIX[core_dtype]}_{op.value}",
248+
(),
249+
)
250+
for core_dtype in RMA_DTYPES},
251+
is_pure=False,
252+
_semantic=_semantic,
253+
)
254+
255+
256+
@core.extern
257+
def getmem(dest, source, bytes, pe, _semantic=None):
258+
return _rma_mem_impl(dest, source, bytes, pe, core.constexpr("getmem"), _semantic=_semantic)
259+
260+
261+
@core.extern
262+
def putmem(dest, source, bytes, pe, _semantic=None):
263+
return _rma_mem_impl(dest, source, bytes, pe, core.constexpr("putmem"), _semantic=_semantic)
264+
265+
266+
@core.extern
267+
def getmem_nbi(dest, source, bytes, pe, qp_id=0, _semantic=None):
268+
return _rma_mem_impl(dest, source, bytes, pe, core.constexpr("getmem_nbi"), _semantic=_semantic)
269+
270+
271+
@core.extern
272+
def putmem_nbi(dest, source, bytes, pe, qp_id=0, _semantic=None):
273+
return _rma_mem_impl(dest, source, bytes, pe, core.constexpr("putmem_nbi"), _semantic=_semantic)
274+
275+
276+
@core.extern
277+
def _putmem_signal_impl(
278+
dest,
279+
source,
280+
nbytes,
281+
sig_addr,
282+
signal,
283+
sig_op,
284+
pe,
285+
op: core.constexpr,
286+
_semantic=None,
287+
):
288+
return extern_call(
289+
"libshmem_device",
290+
"",
291+
[
292+
dest,
293+
source,
294+
tl.cast(nbytes, tl.uint32, _semantic=_semantic),
295+
sig_addr,
296+
tl.cast(signal, tl.int32, _semantic=_semantic),
297+
tl.cast(sig_op, tl.int32, _semantic=_semantic),
298+
tl.cast(pe, tl.int32, _semantic=_semantic),
299+
],
300+
{(
301+
core.pointer_type(core.dtype(core_dtype)),
302+
core.pointer_type(core.dtype(core_dtype)),
303+
tl.uint32,
304+
core.pointer_type(core.dtype("int32")),
305+
tl.int32,
306+
tl.int32,
307+
tl.int32,
308+
): (
309+
f"aclshmem_{DTYPE_TO_KERNEL_SUFFIX[core_dtype]}_{op.value}",
310+
(),
311+
)
312+
for core_dtype in RMA_DTYPES},
313+
is_pure=False,
314+
_semantic=_semantic,
315+
)
316+
317+
318+
@core.extern
319+
def putmem_signal(dest, source, nbytes, sig_addr, signal, sig_op, pe, _semantic=None):
320+
return _putmem_signal_impl(
321+
dest,
322+
source,
323+
nbytes,
324+
sig_addr,
325+
signal,
326+
sig_op,
327+
pe,
328+
core.constexpr("putmem_signal"),
329+
_semantic=_semantic,
330+
)
331+
332+
333+
@core.extern
334+
def putmem_signal_nbi(dest, source, nbytes, sig_addr, signal, sig_op, pe, _semantic=None):
335+
return _putmem_signal_impl(
336+
dest,
337+
source,
338+
nbytes,
339+
sig_addr,
340+
signal,
341+
sig_op,
342+
pe,
343+
core.constexpr("putmem_signal_nbi"),
344+
_semantic=_semantic,
345+
)
346+
347+
348+
@core.extern
349+
def signal_op(sig_addr, signal, sig_op, pe, _semantic=None):
350+
return extern_call(
351+
"libshmem_device",
352+
"",
353+
[
354+
sig_addr,
355+
tl.cast(signal, tl.int32, _semantic=_semantic),
356+
tl.cast(sig_op, tl.int32, _semantic=_semantic),
357+
tl.cast(pe, tl.int32, _semantic=_semantic),
358+
],
359+
{
360+
(pi_i32_t, tl.int32, tl.int32, tl.int32): (
361+
"aclshmemx_signal_op",
362+
(),
363+
),
364+
},
365+
is_pure=False,
366+
_semantic=_semantic,
367+
)
368+
369+
370+
@core.extern
371+
def signal_wait_until(sig_addr, cmp_, cmp_val, _builder=None):
372+
return extern_call(
373+
"libshmem_device",
374+
"",
375+
[
376+
sig_addr,
377+
tl.cast(cmp_, tl.int32, _builder=_builder),
378+
tl.cast(cmp_val, tl.int32, _builder=_builder),
379+
],
380+
{(
381+
core.pointer_type(core.dtype(core_dtype)),
382+
tl.int32,
383+
tl.int32,
384+
): (
385+
f"aclshmem_wait_until_{DTYPE_TO_KERNEL_SUFFIX[core_dtype]}",
386+
tl.int32,
387+
)
388+
for core_dtype in RMA_DTYPES},
389+
is_pure=False,
390+
_builder=_builder,
391+
)
392+
393+
394+
@core.extern
395+
def quiet(_semantic=None):
396+
return extern_call(
397+
"libshmem_device",
398+
"",
399+
[],
400+
{
401+
(): ("aclshmem_quiet", ()),
402+
},
403+
is_pure=False,
404+
_semantic=_semantic,
405+
)
406+
407+
408+
@core.extern
409+
def fence(_semantic=None):
410+
return extern_call(
411+
"libshmem_device",
412+
"",
413+
[],
414+
{
415+
(): ("aclshmem_fence", ()),
416+
},
417+
is_pure=False,
418+
_semantic=_semantic,
419+
)

0 commit comments

Comments
 (0)