ksgpu 1.0.0__tar.gz

ksgpu-1.0.0/.gitignore

@@ -0,0 +1,2 @@
+*~
+*.o

ksgpu-1.0.0/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.1
+Name: ksgpu
+Version: 1.0.0
+Summary: A library of low-level utilities for cuda/cupy
+Requires-Dist: pybind11
+Requires-Dist: numpy
+Requires-Dist: cupy

ksgpu-1.0.0/README.md

@@ -0,0 +1,5 @@
+### GPU C++/CUDA core utils
+
+Some day there will be useful documentation here.
+
+Contact: Kendrick Smith <kmsmith@perimeterinstitute.ca>

ksgpu-1.0.0/generate_device_mma_hpp.py

@@ -0,0 +1,257 @@
+#!/usr/bin/env python3
+#
+# Reference:
+#   https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma
+#
+# Note: I'm no longer running this as part of the build process (in an automated way).
+#
+# Instead, I'm keeping the output file (include/ksgpu/device_mma.hpp) in git, and occasionally
+# updating by hand:
+#
+#   ./generate_device_mma_hpp.py > include/ksgpu/device_mma.hpp     
+
+
+class Argument:
+    def __init__(self, name, cuda_type, num_registers, is_const=False, pack=True, ptx_type=None, is_scalar=False, is_immediate=False):
+        self.name = name
+        self.cuda_type = cuda_type
+        self.num_registers = num_registers
+        self.is_const = is_const
+        self.pack = pack
+        self.ptx_type = ptx_type
+        self.is_scalar = is_scalar
+        self.is_immediate = is_immediate
+
+        if is_immediate:
+            assert is_const and is_scalar and (ptx_type is None)
+
+        if is_scalar:
+            assert (not pack) and (num_registers == 1)
+        
+
+    def make_template_arglist(self):
+        return [ f'{self.cuda_type} {self.name}' ] if self.is_immediate else [ ]
+
+    
+    def make_cuda_arglist(self):
+        amp = '' if self.is_const else '&'
+        cv = 'const ' if self.is_const else ''
+        
+        if self.is_immediate:
+            return [ ]
+        elif self.is_scalar:
+            return [ f'{self.cuda_type} {amp}{self.name}' ]
+        elif self.pack:
+            return [ f'{cv}{self.cuda_type} {self.name}[{self.num_registers}]' ]
+        else:
+            return [ f'{self.cuda_type} {cv}{self.name}{i}' for i in range(self.num_registers) ]
+
+
+    def make_ptx_argstr(self, base):
+        t = [ f'%{i}' for i in range(base, base + self.num_registers) ]
+        t = ', '.join(t)
+        if not self.is_scalar:
+            t =  '{' + t + '}'
+        return t
+
+
+    def make_constraint_str(self):
+        cv = 'const ' if self.is_const else ''
+        constraint = '"r"' if self.is_const else '"=r"'
+
+        if self.is_immediate:
+            return f'"n" ({self.name})'
+
+        if self.is_scalar:
+            return f'{constraint} ({self.name})'
+            
+        ret = [ ]
+        for i in range(self.num_registers):
+            v = f'{self.name}[{i}]' if self.pack else f'{self.name}{i}'
+            if self.ptx_type is not None:
+                v = f'*({cv}{self.ptx_type} *) &{v}'
+            ret.append(f'{constraint} ({v})')
+
+        return ', '.join(ret)
+
+
+####################################################################################################
+
+
+def emit_kernel(cuda_name, ptx_name, *args):
+    # 'ordered_metadata' was introduced in nvcc 12.5 (I think), and we accommodate it with a hack (see below).
+    # FIXME some day, when nvcc 12.4 is ancient history, this hack can be removed.
+    omstr = '::ordered_metadata'
+    omi = ptx_name.find(omstr)
+    
+    template_arglist = [ ]
+    cuda_arglist = [ ]
+    
+    for arg in args:
+        template_arglist += arg.make_template_arglist()
+        cuda_arglist += arg.make_cuda_arglist()
+
+    template_argstr = ', '.join(template_arglist)
+    cuda_argstr = ', '.join(cuda_arglist)
+    
+    print(f'')
+    print(f'// D = A*B + C')
+
+    if len(template_arglist) > 0:
+        print(f'template<{template_argstr}>')
+    
+    print(f'__device__ __forceinline__')
+    print(f'void {cuda_name}({cuda_argstr})')
+    print(f'{{')
+
+    if omi < 0:
+        print(f'    asm("{ptx_name} "')
+    else:
+        ptx_name2 = ptx_name[:omi] + ptx_name[omi+len(omstr):]
+        print(f'    asm(')
+        print(f'#if CUDART_VERSION >= 12050')
+        print(f'        "{ptx_name} "')
+        print(f'#else')
+        print(f'        "{ptx_name2} "')
+        print(f'#endif')
+
+    base = 0
+    for i,arg in enumerate(args):
+        s = f'"{arg.make_ptx_argstr(base)}'
+        t = ', "' if (i < len(args)-1) else ';" :'
+        base += arg.num_registers
+        print(f'        {s}{t}')
+
+    for i,arg in enumerate(args):
+        t = ''
+        if i == 0:
+            t = ' :'
+        elif i < len(args)-1:
+            t = ','
+
+        s = arg.make_constraint_str()
+        print(f'        {s}{t}')
+    
+    print(f'    );')
+    print(f'}}')
+    print(f'')
+
+
+####################################################################################################
+
+    
+def emit_dense_mma(cuda_name, ptx_name, cuda_type, dbits, sbits, m, n, k, ptx_type=None, s=1):
+    # Register counts
+    na = (m*k*s*sbits) // 1024
+    nb = (k*n*s*sbits) // 1024
+    nc = (m*n*s*dbits) // 1024
+
+    emit_kernel(
+        cuda_name,
+        ptx_name,
+        Argument('d', cuda_type, nc, ptx_type=ptx_type),
+        Argument('a', cuda_type, na, ptx_type=ptx_type, is_const=True),
+        Argument('b', cuda_type, nb, ptx_type=ptx_type, is_const=True),
+        Argument('c', cuda_type, nc, ptx_type=ptx_type, is_const=True)
+    )
+
+
+def emit_dense_f16_mma(m, n, k, s=1, layout=None):
+    """The 'layout' parameter is only used for m8n8k4, and is a string pair such as ('row','col')."""
+    
+    cuda_name = f'mma_f16_m{m}_n{n}_k{k}'
+    if layout is not None:
+        cuda_name = f'{cuda_name}_{layout[0][0]}{layout[1][0]}'
+
+    a = 'row' if (layout is None) else layout[0]
+    b = 'col' if (layout is None) else layout[1]
+    ptx_name = f'mma.sync.aligned.m{m}n{n}k{k}.{a}.{b}.f16.f16.f16.f16'
+    
+    emit_dense_mma(cuda_name, ptx_name, '__half2', 16, 16, m, n, k, ptx_type='uint', s=s)
+
+
+def emit_dense_int_mma(sbits, m, n, k):
+    typename = f's{sbits}' if (sbits > 1) else 'b1'
+    satfinite = '.satfinite' if (sbits > 1) else ''
+    suffix = '' if (sbits > 1) else '.and.popc'
+    cuda_name = f'mma_{typename}_m{m}_n{n}_k{k}'
+    ptx_name = f'mma.sync.aligned.m{m}n{n}k{k}.row.col{satfinite}.s32.{typename}.{typename}.s32{suffix}'
+    emit_dense_mma(cuda_name, ptx_name, 'int', 32, sbits, m, n, k)
+
+
+def emit_sparse_f16_mma(m, n, k):
+    cuda_name = f'mma_sp_f16_m{m}_n{n}_k{k}'
+    ptx_name = f'mma.sp::ordered_metadata.sync.aligned.m{m}n{n}k{k}.row.col.f16.f16.f16.f16'
+
+    # Register counts
+    na = (m*k) // 128
+    nb = (k*n) // 64
+    nc = (m*n) // 64
+
+    emit_kernel(
+        cuda_name,
+        ptx_name,
+        Argument('d', '__half2', nc, ptx_type='uint'),
+        Argument('a', '__half2', na, ptx_type='uint', is_const=True),
+        Argument('b', '__half2', nb, ptx_type='uint', is_const=True),
+        Argument('c', '__half2', nc, ptx_type='uint', is_const=True),
+        Argument('e', 'uint', 1, pack=False, is_scalar=True, is_const=True),
+        Argument('F', 'uint', 1, pack=False, is_scalar=True, is_const=True, is_immediate=True)
+    )
+
+    
+
+####################################################################################################
+
+    
+if __name__ == '__main__':
+    print(f'#ifndef _KSGPU_DEVICE_MMA_HPP')
+    print(f'#define _KSGPU_DEVICE_MMA_HPP')
+    print(f'')
+    print(f'// Autogenerated by generate_device_mma_hpp.py')
+    print(f'//')
+    print(f'// Reference for matrix shapes:')
+    print(f'//   https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-shape')
+    print(f'//')
+    print(f'// Reference for PTX instruction syntax:')
+    print(f'//   https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma')
+    print()
+    print('#include <cuda_fp16.h>')
+    print(f'')
+    print(f'namespace ksgpu {{')
+    print(f'')
+
+    # float16
+    emit_dense_f16_mma(16, 8, 8)
+    emit_dense_f16_mma(16, 8, 16)
+
+    # int4
+    emit_dense_int_mma(4, 8, 8, 32)
+    emit_dense_int_mma(4, 16, 8, 32)
+    emit_dense_int_mma(4, 16, 8, 64)
+
+    # int8
+    emit_dense_int_mma(8, 8, 8, 16)
+    emit_dense_int_mma(8, 16, 8, 16)
+    emit_dense_int_mma(8, 16, 8, 32)
+
+    # int1
+    emit_dense_int_mma(1, 8, 8, 128)
+
+    # sparse float16
+    emit_sparse_f16_mma(16, 8, 16)
+    emit_sparse_f16_mma(16, 8, 32)
+
+    # The PTX ISA includes f16 m8n8k4 MMAs.
+    # I tried generating wrappers for these, but timing showed that they were extremely slow.
+    # I assume these MMAs are legacy instructions which are emulated on Ampere.
+    # I left commented-out code here, and in ../generate_device_mma_hpp.py in case I ever want to revisit this.
+
+    # for a in ['row','col']:
+    #     for b in ['row','col']:
+    #         emit_dense_f16_mma(8, 8, 4, s=4, layout=(a,b))
+    
+    print(f'')
+    print(f'}} // namespace ksgpu')
+    print(f'')
+    print(f'#endif // _KSGPU_DEVICE_MMA_HPP')