ptx asm in llvm ir

I asked 7 month ago in r/llvm question how I can insert PTX asm right in LLVM IR and got exactly zero answers. So, as usual, I had to figure it out on my own (depressing little song "No Help is Coming" is playing in the background)

How inline PTX looks like in text form:
%7 = call i32 asm sideeffect "madc.hi.cc.u32 $0,$1,$2,$3;", "=r,r,r,r"(i32 %.sroa.018.0.extract.trunc, i32 %.sroa.282.0.extract.trunc, i32 0) #5, !srcloc !9 

So basically it is just call result-type asm with some arguments in parentheses (note that type of result $0 is i32 and it described after keyword call). If you need result of PTX instruction - just assign it to some variable. Official documentation says that #5 is attributes list - somewhere below it defined as
attributes #5 = { nounwind }
and !9 is metadata - is this case for debug info srcloc:
!9 = !{i32 46731}


Well, that was easy part of story - and now Something Completely Different (tm). LLVM IR is strictly typed (I would say - redundantly strictly), so types first time defined for each argument - like i32 for $1, $2 and $3. Second time - in string called operand constraint codes - in my case it is "=r,r,r,r". And official documentation blatantly lies about them. Let's check some source code - method getRegForInlineAsmConstraint in NVPTXISelLowering.cpp. As you can see it accepts following codes:

• b - 1bit, predicates
• c & h - 16bit, like (.b16 / .u16 / .s16)
• r & f - 32bit, like (.b32 / .u32 / .s32) and .f32 for f
• l, N, d - 64bit, (.b64 / .u64 / .s64) & .f64 for d
• q - 128bit since sm70+
• 0 - meaning is still unknown
  

Symbol '=' is so called Constraint Modifier:

• = Write-only output operand (overwrites previous contents)
• + Read-write operand (input and output tied to the same register)
• & Early-clobber operand (modified before inputs are consumed)
• ~ Clobber list marker (tells LLVM a register or memory/flags are modified implicitly
  

Yet another unpleasant discovery - you can freely swap order of operands - for example this variant is exactly the same as above one:
call i32 asm sideeffect "madc.hi.cc.u32 $0,$2,$1,$3;", "=r,r,r,r"(i32 %.sroa.282.0.extract.trunc, i32 %.sroa.018.0.extract.trunc, i32 0)

This makes the task of parsing & comparison of PTX instructions non-trivial - especially in complex cases like
%1 = call { i32, i32, i32, i32 } asm sideeffect "tex.grad.1d.v4.u32.f32 {$0, $1, $2, $3}, [$4, {$5}], {$6}, {$7};", "=r,=r,=r,=r,l,f,f,f"(i64 %tmp5, float %tmp6, float %tmp7, float %tmp8)

PTX from cicc

Once you understand how inline PTX is represented in LLVM IR, the next step is examining how nvidia’s own internal toolchain leverages it.

While doing some RE of nvidia's llvm-based back-end I dumped inline PTX instructions. Now when I have PTX parser the next logic step is try to parse PTX from cicc and for example try find some undocumented instruction/attributes (which nvidia uses for unfair competitive advantage). So I added to my parser option -r to dump instructions with unrecognized attributes, and also wrote little perl script to collect them. Then run whole pipe like

../ptx.parse/tp -r < ptx.txt | perl ../ptx.parse/ra.pl

And try to guess what happened? Yes - nvidia uses ~5-7% of instructions with undocumented attributes

shrinking holes in SASS allocated registers

In my previous post I discovered that SASS register allocation can have holes - non-contiguous register indices allocated to a function where intermediate registers like R57 remain unused despite R56 and R58 being live/allocated

To estimate size of problem I wrote some code to collect statistics about number (and share) of functions having holes in register allocations:

• libcublas.so.13.7.0.74.sm_90.cubin
  ; 34 holes in regs (784), 0.043367
; 3 holes in uregs (434), 0.006912
; 34 functions with holes (0.629630 from total)
• libcublas.so.13.7.0.608.sm_90.cubin
  ; 41 holes in regs (8546), 0.004798
; 15 holes in uregs (2784), 0.005388
; 39 functions with holes (0.102632 from total)
• libcublas.so.13.7.0.935.sm_90.cubin
  ; 19 holes in regs (5016), 0.003788
; 17 holes in uregs (2012), 0.008449
; 11 functions with holes (0.059783 from total)

As you can see holes occupy up to 4% of total registers (in average 0.3-0.5%) and it's very tempting to try reduce them. It would seem—what could be simpler? If we have something like that
; RHoles max 61: R56 R58

just remap in whole function R61 to R58 and R60 to R56, and then reduce EIATTR_REGCOUNT,right? Well, actually no

Cyclomatic complexity of SASS code

Many algorithms of my SASS optimizer work with code blocks. So I decided to collect some metrics about blocks, resources usage per block/function etc to find outliers (like too short blocks or blocks with anomaly high registers numbers). But before I present the results I should note that building of cyclomatic complexity for SASS is not easy task:

• approximately half of EIATTR attributes are undocumented. And yet, there are some very remarkable ones there - like EIATTR_COROUTINE_RESUME_ID_OFFSETS. Judging by the name they are clearly related to coroutines and so should be taken into account while carving code blocks. Unfortunately, there are no Cubin files in my collection that contain this attribute
• Predicated instructions. Well, this is not SASS-specific problem - for example old 32bit arm had them too. But for example how consider case with several instructions with predicates in the same block? Each of them can modify value of predicate register - then this will be another branch, right? So I just ignore predicates for now
• How to carve code blocks? For my needs, I require maximum-sized blocks to minimize the number of blocks used - so I ignore many cases like BSSY/BSYNC pairs. Sure you can modify my logic of CFG building and make it more similar to classic SSA - see function dg in dg2.pl and the preceding extensive commentary
  

All tests were conducted on libcublas.so.13.7.0.10 from CUDA SDK 13.4 for sm90 cubins, command line options for dg2.pl -gtT:

• -g to build CFG
• -t for registers/predicates tracking
• -T is new option added special to produce various useless metrics 

Count and length of blocks

parser of PTX instructions

A couple of facts to start things off

From official "Inline PTX Assembly in CUDA":

The compiler front end does not parse the asm() statement template string and does not know what it means or even whether it is valid PTX input

And second, less well-known one: order of instruction's attributes (except types of operand) is not important

The combination of these facts leads to stark conclusion - CUDA compiler front-ends totally ignore PTX inline asm and only PTXAS known how to parse them. For example cuKLEE does this wrong

So I made simple (and hopefully fast) parser of PTX instructions

Note: this is not full featured replacement of PTX parser. it is designed specifically to extract instruction attributes and determine the correct instruction form based on argument types and counts

For example for

cvt.bf16x2.e5m2x2.rn.relu.scaled::n2::ue8m0.satfinite d, a, scale-factor;

output will be something like

tail: d, a, scale-factor;
3 tail operands
--> cvt
 line 71: 01x E32Q16
--- types 2:
 bf16x2
 e5m2x2
--- attrs 4:
 1:5 satfinite
 2:7 scaled::n2::ue8m0
 1:0 relu
 3:3 rn 

optimization of SASS stall counts, part 2

In part 1 I suggested that "native" latency tables are too conservative and can be relaxed for some instructions. Indeed, let's look at couple of examples:

In c8.txt there is two delays for IMAD - with value 4 and IMAD.WIDE with value 9. In *_2.txt IMAD included in many groups but none reflect 'wide' form, like

 IMAD_OP = {IMAD,IMADfmalighter_pipe,IMAD32I,IMAD32Ifmalighter_pipe,
             IMUL,IMULfmalighter_pipe,IMUL32I,IMUL32Ifmalighter_pipe}

Corresponding row in RaW table looks like

IMAD_OP`{Rd @RdRange,Rd2 @Rd2Range} : 5 4 6 6 6 6 8 6 6 7 7 7 7 7 7 6 4

For what instructions such relaxation is possible? Well, FP instructions already cleanly separated right at ISA level - for FP64 we have DADD/DMUL/DFMA vs standard FP32 ops. So I patched only restricted set of integer instructions like IMAD/IMUL/IMNMX & SEL, then made binding of this method in Perl and ran tests

An unpleasant discovery awaited me - we can't safely patch delay for xxSETP instructions (ISETP/PSETP/UISETP). I don't know why - maybe due to the fact that predicates can be used to select every instruction for execution and so update requires some hardcore synchronization with instructions decoder/scheduler

results

As usually it depends from version of CUDA SDK, optimization options and your kernel. For FP intensive kernels speed-up is negligible like 0.06%

However for kernels with lots of integer arithmetic it can be much bigger - 0.2-0.3% 

new cmd line options for dg2.pl

• -R to apply delays relaxation
• -S to collect detailed statistics on instructions types distribution
  

optimization of SASS stall counts

Optimal instructions scheduling is NP-hard task. For this reason almost all compilers implement metaheuristic methods like list scheduling/Gibbons–Muchnick algorithm etc. ptxas is no exception - it also generates non-optimal scheduling, and this opens some opportunity for automatic optimization. Here I want to present scheduling model for SASS, tool for optimization of stall counts for binary CUBIN files, achieved results and possible direction for further improvements

First version

Has name dg.pl and used latency table extracted from ptxas with RE. Unfortunately it has at least 2 fatal flaws

1. This latency table is incomplete - for example some instructions like PRET/BFE/ICMP are missed. I tried to find similar table in more old ptxas versions - like 12, 11 and even 9 - it seems that they all incomplete
2. It takes into account only Read after Write joints and patched code crashed in random places, and on each launch in different places

After considerable and agonizing reflection, I concluded that it had used an incorrect model

so I decided to continue experiments with latency tables early extracted from nvdisasm

Latency tables

identification of const bank0 params

Official documentation doesn't disclose ConstBank0 (c[0x0]) memory layout used at the SASS level for kernel arguments and special registers (like %gridid, %nctaid)

So I've spent last week trying to solve this deceptively simple problem. Names of params are documented in official doc - seems that this time ptxas can't add something new to this list

Unfortunately I was unable to find inside ptxas some nice looking tables for pile of SM.  What other approaches can we use? As usually the first thought is to do some brute-force.

Brute-forcing

Lets write in plain PTX dummy function trash with u32 return value - something like

.visible .func (.param .u32 func_retval0) trash
{
  .reg .u32       %r<3>;
  mov.u32    %r0, %gridid;
  mov.u32    %r1, %nctaid.x;
  add.u32    %r0, %r0, %r1;
  st.param.u32 [func_retval0+0], %r0;
  ret;
}
The final st.param is very important bcs otherwise compiler will just eliminate whole code. Instead of gridid & nctaid.x we can substitute pair of special registers, compile with ptxas to specific SM and then parse output of nvdisasm/nvd/whatever can disasm SASS

Surprisingly, this stupid method worked very well, however there are holes in params. So it's time to check

CUDA runtime

I extracted them in December and now we can parse output of nvdisasm to find not identified yet offsets. The funny part is that official nvdisasm failed on several files, like sm54.elf

nvdisasm error   : Could not establish the target of this branch operation

or on sm23.elf

nvdisasm error   : Wrong Anti dependency order in function 'vfprintf_internal'
nvdisasm         .         @P1 LD.E.CG.64 R14, [R4], P0
nvdisasm         .          -- Anti(PRED,0),0*,0 -->
nvdisasm         .         @!P1 LEA.HI.X P0, R7, R12, RZ, R13

So I was forced to use my own nvd

It turned out that the parameter space is divided into two parts - there are block of parameters at offset 0x1860 (holding for example starting PC of kernel) used by kernel launch logic and CnpXXX functions

So now we know lots of offsets and their sizes. However to identify semantics of many found offsets we need debugger

cuda-gdb rushes to the rescue

I made fake PTX for each SM, patched it with my ced and inspected in debugger values with command $_cuda_const_bank(0, offset). Actually this was the most boring part of work and I still didn't recognized some fields. Also I don't have expensive monsters like sm100+ so I extracted only params from Maxwell till Hopper

 

Results

Source

I also add this code to my XS perl module and nvd, so output looks like

/*58*/  XMAD R02,R17,c[0][0x8],R02 ?trans1;
 ; cb0 param %ntid_x


Names starting with '%' were extracted with just disasm of dummy trash function

 

Happy hacking!