ptx_parser/parser/instruction/
cp_async_bulk_tensor.rs

1//! Original PTX specification:
2//!
3//! // global -> shared::cta
4//! cp.async.bulk.tensor.dim.dst.src{.load_mode}.completion_mechanism{.cta_group}{.level::cache_hint} [dstMem], [tensorMap, tensorCoords], [mbar]{, im2colInfo} {, cache-policy};
5//! .dst =                  { .shared::cta };
6//! .src =                  { .global };
7//! .dim =                  { .1d, .2d, .3d, .4d, .5d };
8//! .completion_mechanism = { .mbarrier::complete_tx::bytes };
9//! .cta_group =            { .cta_group::1, .cta_group::2 };
10//! .load_mode =            { .tile, .tile::gather4, .im2col, .im2col::w, .im2col::w::128 };
11//! .level::cache_hint =    { .L2::cache_hint };
12//! ----------------------------------------------------------------
13//! // global -> shared::cluster
14//! cp.async.bulk.tensor.dim.dst.src{.load_mode}.completion_mechanism{.multicast}{.cta_group}{.level::cache_hint} [dstMem], [tensorMap, tensorCoords], [mbar]{, im2colInfo} {, ctaMask} {, cache-policy};
15//! .dst =                  { .shared::cluster };
16//! .src =                  { .global };
17//! .dim =                  { .1d, .2d, .3d, .4d, .5d };
18//! .completion_mechanism = { .mbarrier::complete_tx::bytes };
19//! .cta_group =            { .cta_group::1, .cta_group::2 };
20//! .load_mode =            { .tile, .tile::gather4, .im2col, .im2col::w, .im2col::w::128 };
21//! .level::cache_hint =    { .L2::cache_hint };
22//! .multicast =            { .multicast::cluster  };
23//! ----------------------------------------------------------------
24//! // shared::cta -> global;
25//! cp.async.bulk.tensor.dim.dst.src{.load_mode}.completion_mechanism{.level::cache_hint} [tensorMap, tensorCoords], [srcMem] {, cache-policy};
26//! .dst =                  { .global };
27//! .src =                  { .shared::cta };
28//! .dim =                  { .1d, .2d, .3d, .4d, .5d };
29//! .completion_mechanism = { .bulk_group };
30//! .load_mode =            { .tile, .tile::scatter4, .im2col_no_offs };
31//! .level::cache_hint =    { .L2::cache_hint };
32
33#![allow(unused)]
34
35use crate::parser::{
36    PtxParseError, PtxParser, PtxTokenStream, Span,
37    util::{
38        between, comma_p, directive_p, exclamation_p, lbracket_p, lparen_p, map, minus_p, optional,
39        pipe_p, rbracket_p, rparen_p, semicolon_p, sep_by, string_p, try_map,
40    },
41};
42use crate::r#type::common::*;
43use crate::{alt, ok, seq_n};
44
45pub mod section_0 {
46    use super::*;
47    use crate::r#type::instruction::cp_async_bulk_tensor::section_0::*;
48
49    // ============================================================================
50    // Generated enum parsers
51    // ============================================================================
52
53    impl PtxParser for CompletionMechanism {
54        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
55            alt!(map(
56                string_p(".mbarrier::complete_tx::bytes"),
57                |_, _span| CompletionMechanism::MbarrierCompleteTxBytes
58            ))
59        }
60    }
61
62    impl PtxParser for CtaGroup {
63        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
64            alt!(
65                map(string_p(".cta_group::1"), |_, _span| CtaGroup::CtaGroup1),
66                map(string_p(".cta_group::2"), |_, _span| CtaGroup::CtaGroup2)
67            )
68        }
69    }
70
71    impl PtxParser for Dim {
72        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
73            alt!(
74                map(string_p(".1d"), |_, _span| Dim::_1d),
75                map(string_p(".2d"), |_, _span| Dim::_2d),
76                map(string_p(".3d"), |_, _span| Dim::_3d),
77                map(string_p(".4d"), |_, _span| Dim::_4d),
78                map(string_p(".5d"), |_, _span| Dim::_5d)
79            )
80        }
81    }
82
83    impl PtxParser for Dst {
84        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
85            alt!(map(string_p(".shared::cta"), |_, _span| Dst::SharedCta))
86        }
87    }
88
89    impl PtxParser for LevelCacheHint {
90        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
91            alt!(map(string_p(".L2::cache_hint"), |_, _span| {
92                LevelCacheHint::L2CacheHint
93            }))
94        }
95    }
96
97    impl PtxParser for LoadMode {
98        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
99            alt!(
100                map(string_p(".im2col::w::128"), |_, _span| LoadMode::Im2colW128),
101                map(string_p(".tile::gather4"), |_, _span| LoadMode::TileGather4),
102                map(string_p(".im2col::w"), |_, _span| LoadMode::Im2colW),
103                map(string_p(".im2col"), |_, _span| LoadMode::Im2col),
104                map(string_p(".tile"), |_, _span| LoadMode::Tile)
105            )
106        }
107    }
108
109    impl PtxParser for Src {
110        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
111            alt!(map(string_p(".global"), |_, _span| Src::Global))
112        }
113    }
114
115    impl PtxParser for CpAsyncBulkTensorDimDstSrcLoadModeCompletionMechanismCtaGroupLevelCacheHint {
116        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
117            try_map(
118                seq_n!(
119                    string_p("cp"),
120                    string_p(".async"),
121                    string_p(".bulk"),
122                    string_p(".tensor"),
123                    Dim::parse(),
124                    Dst::parse(),
125                    Src::parse(),
126                    optional(LoadMode::parse()),
127                    CompletionMechanism::parse(),
128                    optional(CtaGroup::parse()),
129                    optional(LevelCacheHint::parse()),
130                    AddressOperand::parse(),
131                    comma_p(),
132                    TexHandler2::parse(),
133                    comma_p(),
134                    AddressOperand::parse(),
135                    map(
136                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
137                        |value, _| value.map(|(_, operand)| operand)
138                    ),
139                    map(
140                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
141                        |value, _| value.map(|(_, operand)| operand)
142                    ),
143                    semicolon_p()
144                ),
145                |(
146                    _,
147                    async_,
148                    bulk,
149                    tensor,
150                    dim,
151                    dst,
152                    src,
153                    load_mode,
154                    completion_mechanism,
155                    cta_group,
156                    level_cache_hint,
157                    dstmem,
158                    _,
159                    tensormap,
160                    _,
161                    mbar,
162                    im2colinfo,
163                    cache_policy,
164                    _,
165                ),
166                 span| {
167                    ok!(CpAsyncBulkTensorDimDstSrcLoadModeCompletionMechanismCtaGroupLevelCacheHint {
168                        async_ = async_,
169                        bulk = bulk,
170                        tensor = tensor,
171                        dim = dim,
172                        dst = dst,
173                        src = src,
174                        load_mode = load_mode,
175                        completion_mechanism = completion_mechanism,
176                        cta_group = cta_group,
177                        level_cache_hint = level_cache_hint,
178                        dstmem = dstmem,
179                        tensormap = tensormap,
180                        mbar = mbar,
181                        im2colinfo = im2colinfo,
182                        cache_policy = cache_policy,
183
184                    })
185                },
186            )
187        }
188    }
189}
190
191pub mod section_1 {
192    use super::*;
193    use crate::r#type::instruction::cp_async_bulk_tensor::section_1::*;
194
195    // ============================================================================
196    // Generated enum parsers
197    // ============================================================================
198
199    impl PtxParser for CompletionMechanism {
200        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
201            alt!(map(
202                string_p(".mbarrier::complete_tx::bytes"),
203                |_, _span| CompletionMechanism::MbarrierCompleteTxBytes
204            ))
205        }
206    }
207
208    impl PtxParser for CtaGroup {
209        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
210            alt!(
211                map(string_p(".cta_group::1"), |_, _span| CtaGroup::CtaGroup1),
212                map(string_p(".cta_group::2"), |_, _span| CtaGroup::CtaGroup2)
213            )
214        }
215    }
216
217    impl PtxParser for Dim {
218        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
219            alt!(
220                map(string_p(".1d"), |_, _span| Dim::_1d),
221                map(string_p(".2d"), |_, _span| Dim::_2d),
222                map(string_p(".3d"), |_, _span| Dim::_3d),
223                map(string_p(".4d"), |_, _span| Dim::_4d),
224                map(string_p(".5d"), |_, _span| Dim::_5d)
225            )
226        }
227    }
228
229    impl PtxParser for Dst {
230        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
231            alt!(map(string_p(".shared::cluster"), |_, _span| {
232                Dst::SharedCluster
233            }))
234        }
235    }
236
237    impl PtxParser for LevelCacheHint {
238        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
239            alt!(map(string_p(".L2::cache_hint"), |_, _span| {
240                LevelCacheHint::L2CacheHint
241            }))
242        }
243    }
244
245    impl PtxParser for LoadMode {
246        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
247            alt!(
248                map(string_p(".im2col::w::128"), |_, _span| LoadMode::Im2colW128),
249                map(string_p(".tile::gather4"), |_, _span| LoadMode::TileGather4),
250                map(string_p(".im2col::w"), |_, _span| LoadMode::Im2colW),
251                map(string_p(".im2col"), |_, _span| LoadMode::Im2col),
252                map(string_p(".tile"), |_, _span| LoadMode::Tile)
253            )
254        }
255    }
256
257    impl PtxParser for Multicast {
258        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
259            alt!(map(string_p(".multicast::cluster"), |_, _span| {
260                Multicast::MulticastCluster
261            }))
262        }
263    }
264
265    impl PtxParser for Src {
266        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
267            alt!(map(string_p(".global"), |_, _span| Src::Global))
268        }
269    }
270
271    impl PtxParser
272        for CpAsyncBulkTensorDimDstSrcLoadModeCompletionMechanismMulticastCtaGroupLevelCacheHint
273    {
274        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
275            try_map(
276                seq_n!(
277                    string_p("cp"),
278                    string_p(".async"),
279                    string_p(".bulk"),
280                    string_p(".tensor"),
281                    Dim::parse(),
282                    Dst::parse(),
283                    Src::parse(),
284                    optional(LoadMode::parse()),
285                    CompletionMechanism::parse(),
286                    optional(Multicast::parse()),
287                    optional(CtaGroup::parse()),
288                    optional(LevelCacheHint::parse()),
289                    AddressOperand::parse(),
290                    comma_p(),
291                    TexHandler2::parse(),
292                    comma_p(),
293                    AddressOperand::parse(),
294                    map(
295                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
296                        |value, _| value.map(|(_, operand)| operand)
297                    ),
298                    map(
299                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
300                        |value, _| value.map(|(_, operand)| operand)
301                    ),
302                    map(
303                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
304                        |value, _| value.map(|(_, operand)| operand)
305                    ),
306                    semicolon_p()
307                ),
308                |(
309                    _,
310                    async_,
311                    bulk,
312                    tensor,
313                    dim,
314                    dst,
315                    src,
316                    load_mode,
317                    completion_mechanism,
318                    multicast,
319                    cta_group,
320                    level_cache_hint,
321                    dstmem,
322                    _,
323                    tensormap,
324                    _,
325                    mbar,
326                    im2colinfo,
327                    ctamask,
328                    cache_policy,
329                    _,
330                ),
331                 span| {
332                    ok!(CpAsyncBulkTensorDimDstSrcLoadModeCompletionMechanismMulticastCtaGroupLevelCacheHint {
333                        async_ = async_,
334                        bulk = bulk,
335                        tensor = tensor,
336                        dim = dim,
337                        dst = dst,
338                        src = src,
339                        load_mode = load_mode,
340                        completion_mechanism = completion_mechanism,
341                        multicast = multicast,
342                        cta_group = cta_group,
343                        level_cache_hint = level_cache_hint,
344                        dstmem = dstmem,
345                        tensormap = tensormap,
346                        mbar = mbar,
347                        im2colinfo = im2colinfo,
348                        ctamask = ctamask,
349                        cache_policy = cache_policy,
350
351                    })
352                },
353            )
354        }
355    }
356}
357
358pub mod section_2 {
359    use super::*;
360    use crate::r#type::instruction::cp_async_bulk_tensor::section_2::*;
361
362    // ============================================================================
363    // Generated enum parsers
364    // ============================================================================
365
366    impl PtxParser for CompletionMechanism {
367        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
368            alt!(map(string_p(".bulk_group"), |_, _span| {
369                CompletionMechanism::BulkGroup
370            }))
371        }
372    }
373
374    impl PtxParser for Dim {
375        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
376            alt!(
377                map(string_p(".1d"), |_, _span| Dim::_1d),
378                map(string_p(".2d"), |_, _span| Dim::_2d),
379                map(string_p(".3d"), |_, _span| Dim::_3d),
380                map(string_p(".4d"), |_, _span| Dim::_4d),
381                map(string_p(".5d"), |_, _span| Dim::_5d)
382            )
383        }
384    }
385
386    impl PtxParser for Dst {
387        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
388            alt!(map(string_p(".global"), |_, _span| Dst::Global))
389        }
390    }
391
392    impl PtxParser for LevelCacheHint {
393        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
394            alt!(map(string_p(".L2::cache_hint"), |_, _span| {
395                LevelCacheHint::L2CacheHint
396            }))
397        }
398    }
399
400    impl PtxParser for LoadMode {
401        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
402            alt!(
403                map(string_p(".tile::scatter4"), |_, _span| {
404                    LoadMode::TileScatter4
405                }),
406                map(string_p(".im2col_no_offs"), |_, _span| {
407                    LoadMode::Im2colNoOffs
408                }),
409                map(string_p(".tile"), |_, _span| LoadMode::Tile)
410            )
411        }
412    }
413
414    impl PtxParser for Src {
415        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
416            alt!(map(string_p(".shared::cta"), |_, _span| Src::SharedCta))
417        }
418    }
419
420    impl PtxParser for CpAsyncBulkTensorDimDstSrcLoadModeCompletionMechanismLevelCacheHint {
421        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
422            try_map(
423                seq_n!(
424                    string_p("cp"),
425                    string_p(".async"),
426                    string_p(".bulk"),
427                    string_p(".tensor"),
428                    Dim::parse(),
429                    Dst::parse(),
430                    Src::parse(),
431                    optional(LoadMode::parse()),
432                    CompletionMechanism::parse(),
433                    optional(LevelCacheHint::parse()),
434                    TexHandler2::parse(),
435                    comma_p(),
436                    AddressOperand::parse(),
437                    map(
438                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
439                        |value, _| value.map(|(_, operand)| operand)
440                    ),
441                    semicolon_p()
442                ),
443                |(
444                    _,
445                    async_,
446                    bulk,
447                    tensor,
448                    dim,
449                    dst,
450                    src,
451                    load_mode,
452                    completion_mechanism,
453                    level_cache_hint,
454                    tensormap,
455                    _,
456                    srcmem,
457                    cache_policy,
458                    _,
459                ),
460                 span| {
461                    ok!(CpAsyncBulkTensorDimDstSrcLoadModeCompletionMechanismLevelCacheHint {
462                        async_ = async_,
463                        bulk = bulk,
464                        tensor = tensor,
465                        dim = dim,
466                        dst = dst,
467                        src = src,
468                        load_mode = load_mode,
469                        completion_mechanism = completion_mechanism,
470                        level_cache_hint = level_cache_hint,
471                        tensormap = tensormap,
472                        srcmem = srcmem,
473                        cache_policy = cache_policy,
474
475                    })
476                },
477            )
478        }
479    }
480}
ptx_parser/parser/instruction/cp_async_bulk_tensor.rs

ptx_parser/parser/instruction/
cp_async_bulk_tensor.rs
