ready for release

Author: AlvinZhu

PKGBUILD

@@ -0,0 +1,18 @@
+ 
+# Maintainer: Alvin Zhu <[email protected]>
+pkgname=python-nbasm
+pkgver=11.6.2
+pkgrel=1
+pkgdesc="NB Assembler"
+arch=('any')
+url="https://github.com/AlvinZhu/NBAssembler"
+license=('GPL')
+groups=('aur-alvin')
+depends=('python')
+optdepends=("CUDA: NVIDIA's GPU programming toolkit")
+makedepends=('python-setuptools')
+
+package() {
+  cp -r ../nbas ../nbasm ../setup.py "$srcdir"
+  python setup.py install --root="$pkgdir/" --optimize=1
+}

nbas/__main__.py

@@ -145,7 +145,7 @@ def decompile_ptx(asm_path, ptx_path, define_list):
         print(ptx, end='')
 
 
-def assemble(asm_path, out_cubin_path, define_list, out_asm_path, sort_banks):
+def assemble(asm_path, out_cubin_path, define_list, out_asm_path, sort_banks, strip):
     if not out_cubin_path:
         out_cubin_path = 'out.cubin'
     cubin = Cubin()
@@ -202,69 +202,12 @@ def assemble(asm_path, out_cubin_path, define_list, out_asm_path, sort_banks):
         for constant in cubin.constant_dict.values():
             constant_asm += constant.print() + '\n'
         asm = header_asm + global_asm + constant_asm + kernel_asm
-        with open(out_asm_path, 'w') as f:
-            f.write(asm)
-
-
-def preprocess(asm_path, out_asm_path, define_list, strip):
-    cubin = Cubin()
-
-    define_dict = {}
-    for define in define_list:
-        if not define:
-            continue
-        d = define.split('=')
-        if len(d) < 2:
-            exec(f'{define} = True', define_dict)
-        else:
-            exec(f'{define}', define_dict)
 
-    cubin.load_asm(asm_path, define_dict)
-    header_asm = cubin.header.print() + '\n'
-    global_asm = ''
-    constant_asm = ''
-    kernel_asm = ''
+        if strip:
+            asm = strip_comment(asm)
 
-    consts = set()
-    globals_ = set()
-    for kernel in cubin.kernel_dict.values():
-        # Unmap global, const0, const3
-        cubin.unmap_constant3(kernel)
-        kernel.unmap_reg()
-        kernel.unmap_constant0()
-        kernel.unmap_jump()
-        kernel.unmap_global()
-        kernel.schedule()
-        kernel.sort_banks()
-        cubin.map_constant3(kernel)
-        kernel.map_global()
-        kernel.map_constant0()
-        kernel.map_jump(rel=True)
-        kernel.mark_const2()
-        kernel_asm += '\n' + kernel.print()
-        consts = consts.union(kernel.consts)
-        globals_ = globals_.union(kernel.globals)
-
-    for global_ in cubin.global_dict.values():
-        if global_.name in globals_:
-            global_asm += global_.print() + '\n'
-    for global_ in cubin.global_init_dict.values():
-        if global_.name in globals_:
-            global_asm += global_.print() + '\n'
-    for constant in cubin.constant_dict.values():
-        if constant.name in consts or 'ALL_CONST3' in consts:
-            constant_asm += constant.print() + '\n'
-
-    asm = header_asm + global_asm + constant_asm + kernel_asm
-
-    if strip:
-        asm = strip_comment(asm)
-
-    if out_asm_path:
         with open(out_asm_path, 'w') as f:
             f.write(asm)
-    else:
-        print(asm, end='')
 
 
 def test_cubin(cubin_path, kernel_names, global_only, check=False):
@@ -397,14 +340,8 @@ def main():
     parser_as.add_argument('-D', '--define', metavar='DEFINE', nargs='+', type=str, default='',
                            help='define variable for embedded python code')
     parser_as.add_argument('-d', '--debug', metavar='OUTPUT_ASM', type=str, default='', help='output asm for debug')
-    parser_as.add_argument('-s', '--sort', action='store_true', help='sort banks')
-
-    parser_pre = subparsers.add_parser('pre', help='preprocess asm')
-    parser_pre.add_argument('asm', help='input asm', metavar='ASM')
-    parser_pre.add_argument('-D', '--define', metavar='DEFINE', nargs='+', type=str, default='',
-                            help='define variable for embedded python code')
-    parser_pre.add_argument('-o', '--output', metavar='OUTPUT', type=str, default='', help='output asm file path')
-    parser_pre.add_argument('-s', '--strip', action='store_true', help='strip comment')
+    parser_as.add_argument('-b', '--bank', action='store_true', help='sort banks')
+    parser_as.add_argument('-s', '--strip', action='store_true', help='strip comment')
 
     parser_pdas = subparsers.add_parser('dcc', help='decompile asm to ptx')
     parser_pdas.add_argument('asm', help='input asm', metavar='ASM')
@@ -435,9 +372,7 @@ def main():
                     global_only=args.global_only, no_line_info=args.no_line_info)
     elif args.cmd == 'as':
         assemble(asm_path=args.asm, out_cubin_path=args.output, define_list=args.define, out_asm_path=args.debug,
-                 sort_banks=args.sort)
-    elif args.cmd == 'pre':
-        preprocess(asm_path=args.asm, out_asm_path=args.output, define_list=args.define, strip=args.strip)
+                 sort_banks=args.bank, strip=args.strip)
     elif args.cmd == 'dcc':
         decompile_ptx(asm_path=args.asm, ptx_path=args.output, define_list=args.define)
     elif args.cmd == 'test':

nbas/tool.py

@@ -38,7 +38,6 @@ def load_ptx_line_info(path):
 
 
 def disassemble_nv(binary, arch):
-    # 使用nvdisasm 反汇编
     fd, tmp_file = mkstemp()
     os.close(fd)
     with open(tmp_file, 'wb') as f:

test.py

@@ -3,7 +3,7 @@
 
 import os
 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
 import pycuda.driver as cuda
 import pycuda.tools
@@ -15,20 +15,23 @@
 mod = cuda.module_from_file('test.cubin')
 kernel = mod.get_function("kTest")
 
-a = np.array([0x70000]*THREADS).astype(np.uint32)
+a = np.array([0x3]*THREADS).astype(np.uint32)
 a_gpu, a_size = mod.get_global('a')
 
-b = np.array([0xf]*THREADS).astype(np.uint32)
+b = np.array([0x2]*THREADS).astype(np.uint32)
 b_gpu, b_size = mod.get_global('b')
 
-d = np.array([0]*THREADS).astype(np.uint32)
-d_gpu, d_size = mod.get_global('d')
+c = np.array([0]*THREADS).astype(np.uint32)
+c_gpu, c_size = mod.get_global('c')
+
+# d = np.array([0]*THREADS).astype(np.uint32)
+# d_gpu, d_size = mod.get_global('d')
 
 cuda.memcpy_htod(a_gpu, a)
 cuda.memcpy_htod(b_gpu, b)
 
 kernel(block=(1, 1, 1))
 
-cuda.memcpy_dtoh(d, d_gpu)
+cuda.memcpy_dtoh(c, c_gpu)
 
-print(f'Result = {np.unique(d)[0]:#0x}')
+print(f'Result = {np.unique(c)[0]:#0x}')