ptx_parser/parser/instruction/
tcgen05_mma.rs

1//! Original PTX specification:
2//!
3//! // 1. Floating-point type without block scaling:
4//! tcgen05.mma.cta_group.kind   [d-tmem],  a-desc,  b-desc, idesc
5//! {, disable-output-lane }, enable-input-d {, scale-input-d};
6//! tcgen05.mma.cta_group.kind   [d-tmem], [a-tmem], b-desc, idesc
7//! {, disable-output-lane }, enable-input-d {, scale-input-d};
8//! .kind      = { .kind::f16, .kind::tf32, .kind::f8f6f4 };
9//! .cta_group = { .cta_group::1, .cta_group::2 };
10//! ------------------------------------------------------------------
11//! // 2. Floating-point type with block scaling:
12//! tcgen05.mma.cta_group.kind.block_scale{.scale_vectorsize}
13//! [d-tmem],  a-desc,  b-desc, idesc,
14//! [scale-A-tmem], [scale-B-tmem], enable-input-d;
15//! tcgen05.mma.cta_group.kind.block_scale{.scale_vectorsize}
16//! [d-tmem], [a-tmem], b-desc, idesc,
17//! [scale-A-tmem], [scale-B-tmem], enable-input-d;
18//! .kind = { .kind::mxf8f6f4, .kind::mxf4, .kind::mxf4nvf4 };
19//! .cta_group      = { .cta_group::1,   .cta_group::2 };
20//! .scale_vectorsize = { .scale_vec::1X, .scale_vec::2X, .scale_vec::4X, .block16, .block32 };
21//! ------------------------------------------------------------------
22//! // 3. Convolution MMA for floating-point type without block scaling:
23//! tcgen05.mma.cta_group.kind.collector_usage [d-tmem],  a-desc,  b-desc, idesc
24//! {, disable-output-lane }, enable-input-d {, scale-input-d};
25//! tcgen05.mma.cta_group.kind{.ashift}.collector_usage [d-tmem], [a-tmem], b-desc, idesc
26//! {, disable-output-lane }, enable-input-d {, scale-input-d};
27//! tcgen05.mma.cta_group.kind.ashift{.collector_usage} [d-tmem], [a-tmem], b-desc, idesc
28//! {, disable-output-lane }, enable-input-d {, scale-input-d};
29//! .kind      = { .kind::f16, .kind::tf32, .kind::f8f6f4 };
30//! .cta_group = { .cta_group::1,   .cta_group::2 };
31//! .collector_usage = { .collector::buffer::op };
32//! ::buffer         = { ::a };
33//! ::op             = { ::fill, ::use, ::lastuse, ::discard* };
34//! ------------------------------------------------------------------
35//! // 4. Activation Stationary MMA for floating-point type with block scaling:
36//! tcgen05.mma.cta_group.kind.block_scale{.scale_vectorsize}.collector_usage
37//! [d-tmem],  a-desc,  b-desc, idesc,
38//! [scale-A-tmem], [scale-B-tmem], enable-input-d;
39//! tcgen05.mma.cta_group.kind.block_scale{.scale_vectorsize}.collector_usage
40//! [d-tmem], [a-tmem], b-desc, idesc,
41//! [scale-A-tmem], [scale-B-tmem], enable-input-d;
42//! .cta_group       = { .cta_group::1,   .cta_group::2 };
43//! .scale_vectorsize  = { .scale_vec::1X, .scale_vec::2X, .scale_vec::4X, .block16, .block32 };
44//! .kind            = { .kind::mxf8f6f4, .kind::mxf4, .kind::mxf4nvf4 };
45//! .collector_usage = { .collector::buffer::op };
46//! ::buffer         = { ::a };
47//! ::op             = { ::fill, ::use, ::lastuse, ::discard* };
48//! ------------------------------------------------------------------
49//! // 5. Integer type:
50//! tcgen05.mma.cta_group.kind::i8  [d-tmem],  a-desc,  b-desc, idesc
51//! {, disable-output-lane }, enable-input-d;
52//! tcgen05.mma.cta_group.kind::i8  [d-tmem], [a-tmem], b-desc, idesc
53//! {, disable-output-lane }, enable-input-d;
54//! .cta_group = { .cta_group::1,   .cta_group::2  };
55//! ------------------------------------------------------------------
56//! // 6. Convolution MMA for integer type:
57//! tcgen05.mma.cta_group.kind::i8.collector_usage          [d-tmem],  a-desc,  b-desc, idesc
58//! {, disable-output-lane }, enable-input-d;
59//! tcgen05.mma.cta_group.kind::i8.ashift{.collector_usage} [d-tmem], [a-tmem], b-desc, idesc
60//! {, disable-output-lane }, enable-input-d;
61//! tcgen05.mma.cta_group.kind::i8{.ashift}.collector_usage [d-tmem], [a-tmem], b-desc, idesc
62//! {, disable-output-lane }, enable-input-d;
63//! .cta_group       = { .cta_group::1,   .cta_group::2  };
64//! .collector_usage = { .collector::buffer::op };
65//! ::buffer         = { ::a };
66//! ::op             = { ::fill, ::use, ::lastuse, ::discard* };
67
68#![allow(unused)]
69
70use crate::parser::{
71    PtxParseError, PtxParser, PtxTokenStream, Span,
72    util::{
73        between, comma_p, directive_p, exclamation_p, lbracket_p, lparen_p, map, minus_p, optional,
74        pipe_p, rbracket_p, rparen_p, semicolon_p, sep_by, string_p, try_map,
75    },
76};
77use crate::r#type::common::*;
78use crate::{alt, ok, seq_n};
79
80pub mod section_0 {
81    use super::*;
82    use crate::r#type::instruction::tcgen05_mma::section_0::*;
83
84    // ============================================================================
85    // Generated enum parsers
86    // ============================================================================
87
88    impl PtxParser for CtaGroup {
89        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
90            alt!(
91                map(string_p(".cta_group::1"), |_, _span| CtaGroup::CtaGroup1),
92                map(string_p(".cta_group::2"), |_, _span| CtaGroup::CtaGroup2)
93            )
94        }
95    }
96
97    impl PtxParser for Kind {
98        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
99            alt!(
100                map(string_p(".kind::f8f6f4"), |_, _span| Kind::KindF8f6f4),
101                map(string_p(".kind::tf32"), |_, _span| Kind::KindTf32),
102                map(string_p(".kind::f16"), |_, _span| Kind::KindF16)
103            )
104        }
105    }
106
107    impl PtxParser for Tcgen05MmaCtaGroupKind {
108        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
109            try_map(
110                seq_n!(
111                    string_p("tcgen05"),
112                    string_p(".mma"),
113                    CtaGroup::parse(),
114                    Kind::parse(),
115                    AddressOperand::parse(),
116                    comma_p(),
117                    GeneralOperand::parse(),
118                    comma_p(),
119                    GeneralOperand::parse(),
120                    comma_p(),
121                    GeneralOperand::parse(),
122                    map(
123                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
124                        |value, _| value.map(|(_, operand)| operand)
125                    ),
126                    comma_p(),
127                    GeneralOperand::parse(),
128                    map(
129                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
130                        |value, _| value.map(|(_, operand)| operand)
131                    ),
132                    semicolon_p()
133                ),
134                |(
135                    _,
136                    mma,
137                    cta_group,
138                    kind,
139                    d_tmem,
140                    _,
141                    a_desc,
142                    _,
143                    b_desc,
144                    _,
145                    idesc,
146                    disable_output_lane,
147                    _,
148                    enable_input_d,
149                    scale_input_d,
150                    _,
151                ),
152                 span| {
153                    ok!(Tcgen05MmaCtaGroupKind {
154                        mma = mma,
155                        cta_group = cta_group,
156                        kind = kind,
157                        d_tmem = d_tmem,
158                        a_desc = a_desc,
159                        b_desc = b_desc,
160                        idesc = idesc,
161                        disable_output_lane = disable_output_lane,
162                        enable_input_d = enable_input_d,
163                        scale_input_d = scale_input_d,
164
165                    })
166                },
167            )
168        }
169    }
170
171    impl PtxParser for Tcgen05MmaCtaGroupKind1 {
172        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
173            try_map(
174                seq_n!(
175                    string_p("tcgen05"),
176                    string_p(".mma"),
177                    CtaGroup::parse(),
178                    Kind::parse(),
179                    AddressOperand::parse(),
180                    comma_p(),
181                    AddressOperand::parse(),
182                    comma_p(),
183                    GeneralOperand::parse(),
184                    comma_p(),
185                    GeneralOperand::parse(),
186                    map(
187                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
188                        |value, _| value.map(|(_, operand)| operand)
189                    ),
190                    comma_p(),
191                    GeneralOperand::parse(),
192                    map(
193                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
194                        |value, _| value.map(|(_, operand)| operand)
195                    ),
196                    semicolon_p()
197                ),
198                |(
199                    _,
200                    mma,
201                    cta_group,
202                    kind,
203                    d_tmem,
204                    _,
205                    a_tmem,
206                    _,
207                    b_desc,
208                    _,
209                    idesc,
210                    disable_output_lane,
211                    _,
212                    enable_input_d,
213                    scale_input_d,
214                    _,
215                ),
216                 span| {
217                    ok!(Tcgen05MmaCtaGroupKind1 {
218                        mma = mma,
219                        cta_group = cta_group,
220                        kind = kind,
221                        d_tmem = d_tmem,
222                        a_tmem = a_tmem,
223                        b_desc = b_desc,
224                        idesc = idesc,
225                        disable_output_lane = disable_output_lane,
226                        enable_input_d = enable_input_d,
227                        scale_input_d = scale_input_d,
228
229                    })
230                },
231            )
232        }
233    }
234}
235
236pub mod section_1 {
237    use super::*;
238    use crate::r#type::instruction::tcgen05_mma::section_1::*;
239
240    // ============================================================================
241    // Generated enum parsers
242    // ============================================================================
243
244    impl PtxParser for CtaGroup {
245        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
246            alt!(
247                map(string_p(".cta_group::1"), |_, _span| CtaGroup::CtaGroup1),
248                map(string_p(".cta_group::2"), |_, _span| CtaGroup::CtaGroup2)
249            )
250        }
251    }
252
253    impl PtxParser for Kind {
254        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
255            alt!(
256                map(string_p(".kind::mxf8f6f4"), |_, _span| Kind::KindMxf8f6f4),
257                map(string_p(".kind::mxf4nvf4"), |_, _span| Kind::KindMxf4nvf4),
258                map(string_p(".kind::mxf4"), |_, _span| Kind::KindMxf4)
259            )
260        }
261    }
262
263    impl PtxParser for ScaleVectorsize {
264        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
265            alt!(
266                map(string_p(".scale_vec::1X"), |_, _span| {
267                    ScaleVectorsize::ScaleVec1x
268                }),
269                map(string_p(".scale_vec::2X"), |_, _span| {
270                    ScaleVectorsize::ScaleVec2x
271                }),
272                map(string_p(".scale_vec::4X"), |_, _span| {
273                    ScaleVectorsize::ScaleVec4x
274                }),
275                map(string_p(".block16"), |_, _span| ScaleVectorsize::Block16),
276                map(string_p(".block32"), |_, _span| ScaleVectorsize::Block32)
277            )
278        }
279    }
280
281    impl PtxParser for Tcgen05MmaCtaGroupKindBlockScaleScaleVectorsize {
282        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
283            try_map(
284                seq_n!(
285                    string_p("tcgen05"),
286                    string_p(".mma"),
287                    CtaGroup::parse(),
288                    Kind::parse(),
289                    string_p(".block_scale"),
290                    optional(ScaleVectorsize::parse()),
291                    AddressOperand::parse(),
292                    comma_p(),
293                    GeneralOperand::parse(),
294                    comma_p(),
295                    GeneralOperand::parse(),
296                    comma_p(),
297                    GeneralOperand::parse(),
298                    comma_p(),
299                    AddressOperand::parse(),
300                    comma_p(),
301                    AddressOperand::parse(),
302                    comma_p(),
303                    GeneralOperand::parse(),
304                    semicolon_p()
305                ),
306                |(
307                    _,
308                    mma,
309                    cta_group,
310                    kind,
311                    block_scale,
312                    scale_vectorsize,
313                    d_tmem,
314                    _,
315                    a_desc,
316                    _,
317                    b_desc,
318                    _,
319                    idesc,
320                    _,
321                    scale_a_tmem,
322                    _,
323                    scale_b_tmem,
324                    _,
325                    enable_input_d,
326                    _,
327                ),
328                 span| {
329                    ok!(Tcgen05MmaCtaGroupKindBlockScaleScaleVectorsize {
330                        mma = mma,
331                        cta_group = cta_group,
332                        kind = kind,
333                        block_scale = block_scale,
334                        scale_vectorsize = scale_vectorsize,
335                        d_tmem = d_tmem,
336                        a_desc = a_desc,
337                        b_desc = b_desc,
338                        idesc = idesc,
339                        scale_a_tmem = scale_a_tmem,
340                        scale_b_tmem = scale_b_tmem,
341                        enable_input_d = enable_input_d,
342
343                    })
344                },
345            )
346        }
347    }
348
349    impl PtxParser for Tcgen05MmaCtaGroupKindBlockScaleScaleVectorsize1 {
350        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
351            try_map(
352                seq_n!(
353                    string_p("tcgen05"),
354                    string_p(".mma"),
355                    CtaGroup::parse(),
356                    Kind::parse(),
357                    string_p(".block_scale"),
358                    optional(ScaleVectorsize::parse()),
359                    AddressOperand::parse(),
360                    comma_p(),
361                    AddressOperand::parse(),
362                    comma_p(),
363                    GeneralOperand::parse(),
364                    comma_p(),
365                    GeneralOperand::parse(),
366                    comma_p(),
367                    AddressOperand::parse(),
368                    comma_p(),
369                    AddressOperand::parse(),
370                    comma_p(),
371                    GeneralOperand::parse(),
372                    semicolon_p()
373                ),
374                |(
375                    _,
376                    mma,
377                    cta_group,
378                    kind,
379                    block_scale,
380                    scale_vectorsize,
381                    d_tmem,
382                    _,
383                    a_tmem,
384                    _,
385                    b_desc,
386                    _,
387                    idesc,
388                    _,
389                    scale_a_tmem,
390                    _,
391                    scale_b_tmem,
392                    _,
393                    enable_input_d,
394                    _,
395                ),
396                 span| {
397                    ok!(Tcgen05MmaCtaGroupKindBlockScaleScaleVectorsize1 {
398                        mma = mma,
399                        cta_group = cta_group,
400                        kind = kind,
401                        block_scale = block_scale,
402                        scale_vectorsize = scale_vectorsize,
403                        d_tmem = d_tmem,
404                        a_tmem = a_tmem,
405                        b_desc = b_desc,
406                        idesc = idesc,
407                        scale_a_tmem = scale_a_tmem,
408                        scale_b_tmem = scale_b_tmem,
409                        enable_input_d = enable_input_d,
410
411                    })
412                },
413            )
414        }
415    }
416}
417
418pub mod section_2 {
419    use super::*;
420    use crate::r#type::instruction::tcgen05_mma::section_2::*;
421
422    // ============================================================================
423    // Generated enum parsers
424    // ============================================================================
425
426    impl PtxParser for CollectorUsage {
427        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
428            alt!(map(
429                |stream| {
430                    stream.try_with_span(|stream| {
431                        stream.with_partial_token_mode(|stream| {
432                            stream.expect_string(".collector")?;
433                            let part0 = match stream.expect_strings(&["::a"])? {
434                                0 => Buffer::A,
435                                _ => unreachable!(),
436                            };
437                            let part1 = match stream.expect_strings(&[
438                                "::discard*",
439                                "::lastuse",
440                                "::fill",
441                                "::use",
442                            ])? {
443                                0 => Op::Discard,
444                                1 => Op::Lastuse,
445                                2 => Op::Fill,
446                                3 => Op::Use,
447                                _ => unreachable!(),
448                            };
449                            Ok(((), part0, part1))
450                        })
451                    })
452                },
453                |(collector, buffer, op), _span| CollectorUsage::CollectorBufferOp(
454                    collector, buffer, op
455                )
456            ))
457        }
458    }
459
460    impl PtxParser for CtaGroup {
461        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
462            alt!(
463                map(string_p(".cta_group::1"), |_, _span| CtaGroup::CtaGroup1),
464                map(string_p(".cta_group::2"), |_, _span| CtaGroup::CtaGroup2)
465            )
466        }
467    }
468
469    impl PtxParser for Kind {
470        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
471            alt!(
472                map(string_p(".kind::f8f6f4"), |_, _span| Kind::KindF8f6f4),
473                map(string_p(".kind::tf32"), |_, _span| Kind::KindTf32),
474                map(string_p(".kind::f16"), |_, _span| Kind::KindF16)
475            )
476        }
477    }
478
479    impl PtxParser for Tcgen05MmaCtaGroupKindCollectorUsage {
480        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
481            try_map(
482                seq_n!(
483                    string_p("tcgen05"),
484                    string_p(".mma"),
485                    CtaGroup::parse(),
486                    Kind::parse(),
487                    CollectorUsage::parse(),
488                    AddressOperand::parse(),
489                    comma_p(),
490                    GeneralOperand::parse(),
491                    comma_p(),
492                    GeneralOperand::parse(),
493                    comma_p(),
494                    GeneralOperand::parse(),
495                    map(
496                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
497                        |value, _| value.map(|(_, operand)| operand)
498                    ),
499                    comma_p(),
500                    GeneralOperand::parse(),
501                    map(
502                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
503                        |value, _| value.map(|(_, operand)| operand)
504                    ),
505                    semicolon_p()
506                ),
507                |(
508                    _,
509                    mma,
510                    cta_group,
511                    kind,
512                    collector_usage,
513                    d_tmem,
514                    _,
515                    a_desc,
516                    _,
517                    b_desc,
518                    _,
519                    idesc,
520                    disable_output_lane,
521                    _,
522                    enable_input_d,
523                    scale_input_d,
524                    _,
525                ),
526                 span| {
527                    ok!(Tcgen05MmaCtaGroupKindCollectorUsage {
528                        mma = mma,
529                        cta_group = cta_group,
530                        kind = kind,
531                        collector_usage = collector_usage,
532                        d_tmem = d_tmem,
533                        a_desc = a_desc,
534                        b_desc = b_desc,
535                        idesc = idesc,
536                        disable_output_lane = disable_output_lane,
537                        enable_input_d = enable_input_d,
538                        scale_input_d = scale_input_d,
539
540                    })
541                },
542            )
543        }
544    }
545
546    impl PtxParser for Tcgen05MmaCtaGroupKindAshiftCollectorUsage {
547        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
548            try_map(
549                seq_n!(
550                    string_p("tcgen05"),
551                    string_p(".mma"),
552                    CtaGroup::parse(),
553                    Kind::parse(),
554                    map(optional(string_p(".ashift")), |value, _| value.is_some()),
555                    CollectorUsage::parse(),
556                    AddressOperand::parse(),
557                    comma_p(),
558                    AddressOperand::parse(),
559                    comma_p(),
560                    GeneralOperand::parse(),
561                    comma_p(),
562                    GeneralOperand::parse(),
563                    map(
564                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
565                        |value, _| value.map(|(_, operand)| operand)
566                    ),
567                    comma_p(),
568                    GeneralOperand::parse(),
569                    map(
570                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
571                        |value, _| value.map(|(_, operand)| operand)
572                    ),
573                    semicolon_p()
574                ),
575                |(
576                    _,
577                    mma,
578                    cta_group,
579                    kind,
580                    ashift,
581                    collector_usage,
582                    d_tmem,
583                    _,
584                    a_tmem,
585                    _,
586                    b_desc,
587                    _,
588                    idesc,
589                    disable_output_lane,
590                    _,
591                    enable_input_d,
592                    scale_input_d,
593                    _,
594                ),
595                 span| {
596                    ok!(Tcgen05MmaCtaGroupKindAshiftCollectorUsage {
597                        mma = mma,
598                        cta_group = cta_group,
599                        kind = kind,
600                        ashift = ashift,
601                        collector_usage = collector_usage,
602                        d_tmem = d_tmem,
603                        a_tmem = a_tmem,
604                        b_desc = b_desc,
605                        idesc = idesc,
606                        disable_output_lane = disable_output_lane,
607                        enable_input_d = enable_input_d,
608                        scale_input_d = scale_input_d,
609
610                    })
611                },
612            )
613        }
614    }
615
616    impl PtxParser for Tcgen05MmaCtaGroupKindAshiftCollectorUsage1 {
617        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
618            try_map(
619                seq_n!(
620                    string_p("tcgen05"),
621                    string_p(".mma"),
622                    CtaGroup::parse(),
623                    Kind::parse(),
624                    string_p(".ashift"),
625                    optional(CollectorUsage::parse()),
626                    AddressOperand::parse(),
627                    comma_p(),
628                    AddressOperand::parse(),
629                    comma_p(),
630                    GeneralOperand::parse(),
631                    comma_p(),
632                    GeneralOperand::parse(),
633                    map(
634                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
635                        |value, _| value.map(|(_, operand)| operand)
636                    ),
637                    comma_p(),
638                    GeneralOperand::parse(),
639                    map(
640                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
641                        |value, _| value.map(|(_, operand)| operand)
642                    ),
643                    semicolon_p()
644                ),
645                |(
646                    _,
647                    mma,
648                    cta_group,
649                    kind,
650                    ashift,
651                    collector_usage,
652                    d_tmem,
653                    _,
654                    a_tmem,
655                    _,
656                    b_desc,
657                    _,
658                    idesc,
659                    disable_output_lane,
660                    _,
661                    enable_input_d,
662                    scale_input_d,
663                    _,
664                ),
665                 span| {
666                    ok!(Tcgen05MmaCtaGroupKindAshiftCollectorUsage1 {
667                        mma = mma,
668                        cta_group = cta_group,
669                        kind = kind,
670                        ashift = ashift,
671                        collector_usage = collector_usage,
672                        d_tmem = d_tmem,
673                        a_tmem = a_tmem,
674                        b_desc = b_desc,
675                        idesc = idesc,
676                        disable_output_lane = disable_output_lane,
677                        enable_input_d = enable_input_d,
678                        scale_input_d = scale_input_d,
679
680                    })
681                },
682            )
683        }
684    }
685}
686
687pub mod section_3 {
688    use super::*;
689    use crate::r#type::instruction::tcgen05_mma::section_3::*;
690
691    // ============================================================================
692    // Generated enum parsers
693    // ============================================================================
694
695    impl PtxParser for CollectorUsage {
696        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
697            alt!(map(
698                |stream| {
699                    stream.try_with_span(|stream| {
700                        stream.with_partial_token_mode(|stream| {
701                            stream.expect_string(".collector")?;
702                            let part0 = match stream.expect_strings(&["::a"])? {
703                                0 => Buffer::A,
704                                _ => unreachable!(),
705                            };
706                            let part1 = match stream.expect_strings(&[
707                                "::discard*",
708                                "::lastuse",
709                                "::fill",
710                                "::use",
711                            ])? {
712                                0 => Op::Discard,
713                                1 => Op::Lastuse,
714                                2 => Op::Fill,
715                                3 => Op::Use,
716                                _ => unreachable!(),
717                            };
718                            Ok(((), part0, part1))
719                        })
720                    })
721                },
722                |(collector, buffer, op), _span| CollectorUsage::CollectorBufferOp(
723                    collector, buffer, op
724                )
725            ))
726        }
727    }
728
729    impl PtxParser for CtaGroup {
730        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
731            alt!(
732                map(string_p(".cta_group::1"), |_, _span| CtaGroup::CtaGroup1),
733                map(string_p(".cta_group::2"), |_, _span| CtaGroup::CtaGroup2)
734            )
735        }
736    }
737
738    impl PtxParser for Kind {
739        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
740            alt!(
741                map(string_p(".kind::mxf8f6f4"), |_, _span| Kind::KindMxf8f6f4),
742                map(string_p(".kind::mxf4nvf4"), |_, _span| Kind::KindMxf4nvf4),
743                map(string_p(".kind::mxf4"), |_, _span| Kind::KindMxf4)
744            )
745        }
746    }
747
748    impl PtxParser for ScaleVectorsize {
749        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
750            alt!(
751                map(string_p(".scale_vec::1X"), |_, _span| {
752                    ScaleVectorsize::ScaleVec1x
753                }),
754                map(string_p(".scale_vec::2X"), |_, _span| {
755                    ScaleVectorsize::ScaleVec2x
756                }),
757                map(string_p(".scale_vec::4X"), |_, _span| {
758                    ScaleVectorsize::ScaleVec4x
759                }),
760                map(string_p(".block16"), |_, _span| ScaleVectorsize::Block16),
761                map(string_p(".block32"), |_, _span| ScaleVectorsize::Block32)
762            )
763        }
764    }
765
766    impl PtxParser for Tcgen05MmaCtaGroupKindBlockScaleScaleVectorsizeCollectorUsage {
767        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
768            try_map(
769                seq_n!(
770                    string_p("tcgen05"),
771                    string_p(".mma"),
772                    CtaGroup::parse(),
773                    Kind::parse(),
774                    string_p(".block_scale"),
775                    optional(ScaleVectorsize::parse()),
776                    CollectorUsage::parse(),
777                    AddressOperand::parse(),
778                    comma_p(),
779                    GeneralOperand::parse(),
780                    comma_p(),
781                    GeneralOperand::parse(),
782                    comma_p(),
783                    GeneralOperand::parse(),
784                    comma_p(),
785                    AddressOperand::parse(),
786                    comma_p(),
787                    AddressOperand::parse(),
788                    comma_p(),
789                    GeneralOperand::parse(),
790                    semicolon_p()
791                ),
792                |(
793                    _,
794                    mma,
795                    cta_group,
796                    kind,
797                    block_scale,
798                    scale_vectorsize,
799                    collector_usage,
800                    d_tmem,
801                    _,
802                    a_desc,
803                    _,
804                    b_desc,
805                    _,
806                    idesc,
807                    _,
808                    scale_a_tmem,
809                    _,
810                    scale_b_tmem,
811                    _,
812                    enable_input_d,
813                    _,
814                ),
815                 span| {
816                    ok!(Tcgen05MmaCtaGroupKindBlockScaleScaleVectorsizeCollectorUsage {
817                        mma = mma,
818                        cta_group = cta_group,
819                        kind = kind,
820                        block_scale = block_scale,
821                        scale_vectorsize = scale_vectorsize,
822                        collector_usage = collector_usage,
823                        d_tmem = d_tmem,
824                        a_desc = a_desc,
825                        b_desc = b_desc,
826                        idesc = idesc,
827                        scale_a_tmem = scale_a_tmem,
828                        scale_b_tmem = scale_b_tmem,
829                        enable_input_d = enable_input_d,
830
831                    })
832                },
833            )
834        }
835    }
836
837    impl PtxParser for Tcgen05MmaCtaGroupKindBlockScaleScaleVectorsizeCollectorUsage1 {
838        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
839            try_map(
840                seq_n!(
841                    string_p("tcgen05"),
842                    string_p(".mma"),
843                    CtaGroup::parse(),
844                    Kind::parse(),
845                    string_p(".block_scale"),
846                    optional(ScaleVectorsize::parse()),
847                    CollectorUsage::parse(),
848                    AddressOperand::parse(),
849                    comma_p(),
850                    AddressOperand::parse(),
851                    comma_p(),
852                    GeneralOperand::parse(),
853                    comma_p(),
854                    GeneralOperand::parse(),
855                    comma_p(),
856                    AddressOperand::parse(),
857                    comma_p(),
858                    AddressOperand::parse(),
859                    comma_p(),
860                    GeneralOperand::parse(),
861                    semicolon_p()
862                ),
863                |(
864                    _,
865                    mma,
866                    cta_group,
867                    kind,
868                    block_scale,
869                    scale_vectorsize,
870                    collector_usage,
871                    d_tmem,
872                    _,
873                    a_tmem,
874                    _,
875                    b_desc,
876                    _,
877                    idesc,
878                    _,
879                    scale_a_tmem,
880                    _,
881                    scale_b_tmem,
882                    _,
883                    enable_input_d,
884                    _,
885                ),
886                 span| {
887                    ok!(Tcgen05MmaCtaGroupKindBlockScaleScaleVectorsizeCollectorUsage1 {
888                        mma = mma,
889                        cta_group = cta_group,
890                        kind = kind,
891                        block_scale = block_scale,
892                        scale_vectorsize = scale_vectorsize,
893                        collector_usage = collector_usage,
894                        d_tmem = d_tmem,
895                        a_tmem = a_tmem,
896                        b_desc = b_desc,
897                        idesc = idesc,
898                        scale_a_tmem = scale_a_tmem,
899                        scale_b_tmem = scale_b_tmem,
900                        enable_input_d = enable_input_d,
901
902                    })
903                },
904            )
905        }
906    }
907}
908
909pub mod section_4 {
910    use super::*;
911    use crate::r#type::instruction::tcgen05_mma::section_4::*;
912
913    // ============================================================================
914    // Generated enum parsers
915    // ============================================================================
916
917    impl PtxParser for CtaGroup {
918        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
919            alt!(
920                map(string_p(".cta_group::1"), |_, _span| CtaGroup::CtaGroup1),
921                map(string_p(".cta_group::2"), |_, _span| CtaGroup::CtaGroup2)
922            )
923        }
924    }
925
926    impl PtxParser for Tcgen05MmaCtaGroupKindI8 {
927        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
928            try_map(
929                seq_n!(
930                    string_p("tcgen05"),
931                    string_p(".mma"),
932                    CtaGroup::parse(),
933                    string_p(".kind::i8"),
934                    AddressOperand::parse(),
935                    comma_p(),
936                    GeneralOperand::parse(),
937                    comma_p(),
938                    GeneralOperand::parse(),
939                    comma_p(),
940                    GeneralOperand::parse(),
941                    map(
942                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
943                        |value, _| value.map(|(_, operand)| operand)
944                    ),
945                    comma_p(),
946                    GeneralOperand::parse(),
947                    semicolon_p()
948                ),
949                |(
950                    _,
951                    mma,
952                    cta_group,
953                    kind_i8,
954                    d_tmem,
955                    _,
956                    a_desc,
957                    _,
958                    b_desc,
959                    _,
960                    idesc,
961                    disable_output_lane,
962                    _,
963                    enable_input_d,
964                    _,
965                ),
966                 span| {
967                    ok!(Tcgen05MmaCtaGroupKindI8 {
968                        mma = mma,
969                        cta_group = cta_group,
970                        kind_i8 = kind_i8,
971                        d_tmem = d_tmem,
972                        a_desc = a_desc,
973                        b_desc = b_desc,
974                        idesc = idesc,
975                        disable_output_lane = disable_output_lane,
976                        enable_input_d = enable_input_d,
977
978                    })
979                },
980            )
981        }
982    }
983
984    impl PtxParser for Tcgen05MmaCtaGroupKindI81 {
985        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
986            try_map(
987                seq_n!(
988                    string_p("tcgen05"),
989                    string_p(".mma"),
990                    CtaGroup::parse(),
991                    string_p(".kind::i8"),
992                    AddressOperand::parse(),
993                    comma_p(),
994                    AddressOperand::parse(),
995                    comma_p(),
996                    GeneralOperand::parse(),
997                    comma_p(),
998                    GeneralOperand::parse(),
999                    map(
1000                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
1001                        |value, _| value.map(|(_, operand)| operand)
1002                    ),
1003                    comma_p(),
1004                    GeneralOperand::parse(),
1005                    semicolon_p()
1006                ),
1007                |(
1008                    _,
1009                    mma,
1010                    cta_group,
1011                    kind_i8,
1012                    d_tmem,
1013                    _,
1014                    a_tmem,
1015                    _,
1016                    b_desc,
1017                    _,
1018                    idesc,
1019                    disable_output_lane,
1020                    _,
1021                    enable_input_d,
1022                    _,
1023                ),
1024                 span| {
1025                    ok!(Tcgen05MmaCtaGroupKindI81 {
1026                        mma = mma,
1027                        cta_group = cta_group,
1028                        kind_i8 = kind_i8,
1029                        d_tmem = d_tmem,
1030                        a_tmem = a_tmem,
1031                        b_desc = b_desc,
1032                        idesc = idesc,
1033                        disable_output_lane = disable_output_lane,
1034                        enable_input_d = enable_input_d,
1035
1036                    })
1037                },
1038            )
1039        }
1040    }
1041}
1042
1043pub mod section_5 {
1044    use super::*;
1045    use crate::r#type::instruction::tcgen05_mma::section_5::*;
1046
1047    // ============================================================================
1048    // Generated enum parsers
1049    // ============================================================================
1050
1051    impl PtxParser for CollectorUsage {
1052        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
1053            alt!(map(
1054                |stream| {
1055                    stream.try_with_span(|stream| {
1056                        stream.with_partial_token_mode(|stream| {
1057                            stream.expect_string(".collector")?;
1058                            let part0 = match stream.expect_strings(&["::a"])? {
1059                                0 => Buffer::A,
1060                                _ => unreachable!(),
1061                            };
1062                            let part1 = match stream.expect_strings(&[
1063                                "::discard*",
1064                                "::lastuse",
1065                                "::fill",
1066                                "::use",
1067                            ])? {
1068                                0 => Op::Discard,
1069                                1 => Op::Lastuse,
1070                                2 => Op::Fill,
1071                                3 => Op::Use,
1072                                _ => unreachable!(),
1073                            };
1074                            Ok(((), part0, part1))
1075                        })
1076                    })
1077                },
1078                |(collector, buffer, op), _span| CollectorUsage::CollectorBufferOp(
1079                    collector, buffer, op
1080                )
1081            ))
1082        }
1083    }
1084
1085    impl PtxParser for CtaGroup {
1086        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
1087            alt!(
1088                map(string_p(".cta_group::1"), |_, _span| CtaGroup::CtaGroup1),
1089                map(string_p(".cta_group::2"), |_, _span| CtaGroup::CtaGroup2)
1090            )
1091        }
1092    }
1093
1094    impl PtxParser for Tcgen05MmaCtaGroupKindI8CollectorUsage {
1095        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
1096            try_map(
1097                seq_n!(
1098                    string_p("tcgen05"),
1099                    string_p(".mma"),
1100                    CtaGroup::parse(),
1101                    string_p(".kind::i8"),
1102                    CollectorUsage::parse(),
1103                    AddressOperand::parse(),
1104                    comma_p(),
1105                    GeneralOperand::parse(),
1106                    comma_p(),
1107                    GeneralOperand::parse(),
1108                    comma_p(),
1109                    GeneralOperand::parse(),
1110                    map(
1111                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
1112                        |value, _| value.map(|(_, operand)| operand)
1113                    ),
1114                    comma_p(),
1115                    GeneralOperand::parse(),
1116                    semicolon_p()
1117                ),
1118                |(
1119                    _,
1120                    mma,
1121                    cta_group,
1122                    kind_i8,
1123                    collector_usage,
1124                    d_tmem,
1125                    _,
1126                    a_desc,
1127                    _,
1128                    b_desc,
1129                    _,
1130                    idesc,
1131                    disable_output_lane,
1132                    _,
1133                    enable_input_d,
1134                    _,
1135                ),
1136                 span| {
1137                    ok!(Tcgen05MmaCtaGroupKindI8CollectorUsage {
1138                        mma = mma,
1139                        cta_group = cta_group,
1140                        kind_i8 = kind_i8,
1141                        collector_usage = collector_usage,
1142                        d_tmem = d_tmem,
1143                        a_desc = a_desc,
1144                        b_desc = b_desc,
1145                        idesc = idesc,
1146                        disable_output_lane = disable_output_lane,
1147                        enable_input_d = enable_input_d,
1148
1149                    })
1150                },
1151            )
1152        }
1153    }
1154
1155    impl PtxParser for Tcgen05MmaCtaGroupKindI8AshiftCollectorUsage {
1156        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
1157            try_map(
1158                seq_n!(
1159                    string_p("tcgen05"),
1160                    string_p(".mma"),
1161                    CtaGroup::parse(),
1162                    string_p(".kind::i8"),
1163                    string_p(".ashift"),
1164                    optional(CollectorUsage::parse()),
1165                    AddressOperand::parse(),
1166                    comma_p(),
1167                    AddressOperand::parse(),
1168                    comma_p(),
1169                    GeneralOperand::parse(),
1170                    comma_p(),
1171                    GeneralOperand::parse(),
1172                    map(
1173                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
1174                        |value, _| value.map(|(_, operand)| operand)
1175                    ),
1176                    comma_p(),
1177                    GeneralOperand::parse(),
1178                    semicolon_p()
1179                ),
1180                |(
1181                    _,
1182                    mma,
1183                    cta_group,
1184                    kind_i8,
1185                    ashift,
1186                    collector_usage,
1187                    d_tmem,
1188                    _,
1189                    a_tmem,
1190                    _,
1191                    b_desc,
1192                    _,
1193                    idesc,
1194                    disable_output_lane,
1195                    _,
1196                    enable_input_d,
1197                    _,
1198                ),
1199                 span| {
1200                    ok!(Tcgen05MmaCtaGroupKindI8AshiftCollectorUsage {
1201                        mma = mma,
1202                        cta_group = cta_group,
1203                        kind_i8 = kind_i8,
1204                        ashift = ashift,
1205                        collector_usage = collector_usage,
1206                        d_tmem = d_tmem,
1207                        a_tmem = a_tmem,
1208                        b_desc = b_desc,
1209                        idesc = idesc,
1210                        disable_output_lane = disable_output_lane,
1211                        enable_input_d = enable_input_d,
1212
1213                    })
1214                },
1215            )
1216        }
1217    }
1218
1219    impl PtxParser for Tcgen05MmaCtaGroupKindI8AshiftCollectorUsage1 {
1220        fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
1221            try_map(
1222                seq_n!(
1223                    string_p("tcgen05"),
1224                    string_p(".mma"),
1225                    CtaGroup::parse(),
1226                    string_p(".kind::i8"),
1227                    map(optional(string_p(".ashift")), |value, _| value.is_some()),
1228                    CollectorUsage::parse(),
1229                    AddressOperand::parse(),
1230                    comma_p(),
1231                    AddressOperand::parse(),
1232                    comma_p(),
1233                    GeneralOperand::parse(),
1234                    comma_p(),
1235                    GeneralOperand::parse(),
1236                    map(
1237                        optional(seq_n!(comma_p(), GeneralOperand::parse())),
1238                        |value, _| value.map(|(_, operand)| operand)
1239                    ),
1240                    comma_p(),
1241                    GeneralOperand::parse(),
1242                    semicolon_p()
1243                ),
1244                |(
1245                    _,
1246                    mma,
1247                    cta_group,
1248                    kind_i8,
1249                    ashift,
1250                    collector_usage,
1251                    d_tmem,
1252                    _,
1253                    a_tmem,
1254                    _,
1255                    b_desc,
1256                    _,
1257                    idesc,
1258                    disable_output_lane,
1259                    _,
1260                    enable_input_d,
1261                    _,
1262                ),
1263                 span| {
1264                    ok!(Tcgen05MmaCtaGroupKindI8AshiftCollectorUsage1 {
1265                        mma = mma,
1266                        cta_group = cta_group,
1267                        kind_i8 = kind_i8,
1268                        ashift = ashift,
1269                        collector_usage = collector_usage,
1270                        d_tmem = d_tmem,
1271                        a_tmem = a_tmem,
1272                        b_desc = b_desc,
1273                        idesc = idesc,
1274                        disable_output_lane = disable_output_lane,
1275                        enable_input_d = enable_input_d,
1276
1277                    })
1278                },
1279            )
1280        }
1281    }
1282}
ptx_parser/parser/instruction/tcgen05_mma.rs

ptx_parser/parser/instruction/
tcgen05_mma.rs
