├── .github └── workflows │ └── python-package.yml ├── .gitignore ├── LICENSE ├── MANIFEST.in ├── README.md ├── nasmcompiler.py ├── setup.py ├── src └── pymult.asm ├── test.py └── winnasmcompiler.py /.github/workflows/python-package.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests and lint with a variety of Python versions 2 | # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions 3 | 4 | name: Python package 5 | 6 | on: 7 | push: 8 | branches: [ master ] 9 | pull_request: 10 | branches: [ master ] 11 | 12 | jobs: 13 | build-linux: 14 | runs-on: ubuntu-latest 15 | strategy: 16 | matrix: 17 | python-version: [3.8] 18 | 19 | steps: 20 | - name: Install NASM 21 | run: | 22 | sudo apt-get install -y nasm 23 | - uses: actions/checkout@v2 24 | - name: Set up Python ${{ matrix.python-version }} 25 | uses: actions/setup-python@v2 26 | with: 27 | python-version: ${{ matrix.python-version }} 28 | - name: Install dependencies 29 | run: | 30 | python -m pip install --upgrade pip pytest 31 | python setup.py install 32 | - name: Test with pytest 33 | run: | 34 | python -X dev -m pytest test.py 35 | build-mac: 36 | runs-on: macos-latest 37 | strategy: 38 | matrix: 39 | python-version: [3.7, 3.8] 40 | steps: 41 | - name: Install NASM 42 | run: | 43 | brew install nasm 44 | - uses: actions/checkout@v2 45 | - name: Set up Python ${{ matrix.python-version }} 46 | uses: actions/setup-python@v2 47 | with: 48 | python-version: ${{ matrix.python-version }} 49 | - name: Install dependencies 50 | run: | 51 | python -m pip install --upgrade pip pytest 52 | python setup.py install 53 | - name: Test with pytest 54 | run: | 55 | python -X dev -m pytest test.py 56 | 57 | build-windows: 58 | runs-on: windows-latest 59 | strategy: 60 | matrix: 61 | python-version: [3.8] 62 | steps: 63 | - name: Install NASM 64 | run: | 65 | choco install nasm 66 | - name: Add NASM to path 67 | run: echo '::add-path::c:\\Program Files\\NASM' 68 | - name: Add VC to path 69 | run: echo '::add-path::C:\\Program Files (x86)\\Microsoft Visual Studio 14.0\\VC\\bin' 70 | - uses: actions/checkout@v2 71 | - name: Set up Python ${{ matrix.python-version }} 72 | uses: actions/setup-python@v2 73 | with: 74 | python-version: ${{ matrix.python-version }} 75 | - name: Install dependencies 76 | run: | 77 | python -m pip install --upgrade pip pytest 78 | python setup.py install 79 | - name: Test with pytest 80 | run: | 81 | python -X dev -m pytest test.py 82 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /.idea/ 2 | /cmake-build-debug/ 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 2-Clause License 2 | 3 | Copyright (c) 2020, Anthony Shaw 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | 1. Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | 2. Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 17 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 18 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 19 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 20 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 21 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 22 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 23 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 24 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include nasmcompiler.py README.md 2 | include src/pymult.asm -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Building Python Extension Modules in Assembly 2 | 3 | ![GitHub Actions](https://github.com/tonybaloney/python-assembly-poc/workflows/Python%20package/badge.svg) 4 | 5 | 6 | This repository is a proof-of-concept to demonstrate how you can create a Python Extension in 7 | 100% assembly. 8 | 9 | Demonstrates: 10 | 11 | - How to write a Python module in pure assembly 12 | - How to write a function in pure assembly and call it from Python with Python objects 13 | - How to call the C API to create a PyObject and parse PyTuple (arguments) into raw pointers 14 | - How to pass data back into Python 15 | - How to register a module from assembly 16 | - How to create a method definition in assembly 17 | - How to write back to the Python stack using the dynamic module loader 18 | - How to package a NASM/Assembly Python extension with distutils 19 | 20 | The simple proof-of-concept function takes 2 parameters, 21 | 22 | ```default 23 | >>> import pymult 24 | >>> pymult.multiply(2, 4) 25 | 8 26 | ``` 27 | 28 | ## But, Why? 29 | 30 | Just because it can be done. 31 | 32 | Also, I want to see if some AVX/AVX2 instructions (high-performance matrix multiplication especially) can be used 33 | directly from Python. 34 | -------------------------------------------------------------------------------- /nasmcompiler.py: -------------------------------------------------------------------------------- 1 | """ 2 | Distutils doesn't support nasm, so this is a custom compiler for NASM 3 | """ 4 | 5 | from distutils.unixccompiler import UnixCCompiler 6 | from distutils.sysconfig import get_config_var 7 | import sys 8 | 9 | 10 | class NasmCompiler(UnixCCompiler) : 11 | compiler_type = 'nasm' 12 | src_extensions = ['.asm'] 13 | obj_extension = '.obj' 14 | language_map = {".asm" : "asm",} 15 | language_order = ["asm"] 16 | executables = {'preprocessor' : None, 17 | 'compiler' : ["nasm"], 18 | 'compiler_so' : ["nasm"], 19 | 'compiler_cxx' : ["nasm"], 20 | 'linker_so' : ["cc", "-shared"], 21 | 'linker_exe' : ["cc", "-shared"], 22 | 'archiver' : ["ar", "-cr"], 23 | 'ranlib' : None, 24 | } 25 | def __init__ (self, 26 | verbose=0, 27 | dry_run=0, 28 | force=0): 29 | 30 | UnixCCompiler.__init__ (self, verbose, dry_run, force) 31 | self.set_executable("compiler", "nasm") 32 | 33 | def _is_gcc(self, compiler_name): 34 | return False # ¯\_(ツ)_/¯ 35 | 36 | def _get_cc_args(self, pp_opts, debug, before): 37 | if sys.platform == 'darwin': 38 | # Fix the symbols on macOS 39 | cc_args = pp_opts + ["-f macho64","-DNOPIE","--prefix=_"] 40 | else: 41 | # Use 64-bit elf format for Linux 42 | cc_args = pp_opts + ["-f elf64"] 43 | if debug: 44 | # Debug symbols from NASM 45 | cc_args[:0] = ['-g'] 46 | if before: 47 | cc_args[:0] = before 48 | return cc_args 49 | 50 | def link(self, target_desc, objects, 51 | output_filename, output_dir=None, libraries=None, 52 | library_dirs=None, runtime_library_dirs=None, 53 | export_symbols=None, debug=0, extra_preargs=None, 54 | extra_postargs=None, build_temp=None, target_lang=None): 55 | # Make sure libpython gets linked 56 | if not self.runtime_library_dirs: 57 | self.runtime_library_dirs.append(get_config_var('LIBDIR')) 58 | if not self.libraries: 59 | libraries = ["python" + get_config_var("LDVERSION")] 60 | if not extra_preargs: 61 | extra_preargs = [] 62 | 63 | return super().link(target_desc, objects, 64 | output_filename, output_dir, libraries, 65 | library_dirs, runtime_library_dirs, 66 | export_symbols, debug, extra_preargs, 67 | extra_postargs, build_temp, target_lang) 68 | 69 | def runtime_library_dir_option(self, dir): 70 | if sys.platform == "darwin": 71 | return "-L" + dir 72 | else: 73 | return "-Wl,-R" + dir 74 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- encoding: utf-8 -*- 3 | 4 | import io 5 | import os 6 | from glob import glob 7 | from os.path import basename 8 | from os.path import dirname 9 | from os.path import join 10 | from os.path import relpath 11 | from os.path import splitext 12 | import sys 13 | from setuptools import Extension 14 | from setuptools import find_packages 15 | from setuptools import setup 16 | from setuptools.command.build_ext import build_ext 17 | from distutils.sysconfig import customize_compiler 18 | 19 | 20 | class NasmBuildCommand(build_ext): 21 | description = "build NASM extensions (compile/link to build directory)" 22 | 23 | def run(self): 24 | try: 25 | if self.distribution.has_c_libraries(): 26 | build_clib = self.get_finalized_command('build_clib') 27 | self.libraries.extend(build_clib.get_library_names() or []) 28 | self.library_dirs.append(build_clib.build_clib) 29 | if sys.platform == "win32": 30 | from winnasmcompiler import WinNasmCompiler 31 | self.compiler = WinNasmCompiler(verbose=True) 32 | else: 33 | from nasmcompiler import NasmCompiler 34 | self.compiler = NasmCompiler(verbose=True) 35 | customize_compiler(self.compiler) 36 | 37 | 38 | if self.include_dirs is not None: 39 | self.compiler.set_include_dirs(self.include_dirs) 40 | 41 | if self.define is not None: 42 | # 'define' option is a list of (name,value) tuples 43 | for (name, value) in self.define: 44 | self.compiler.define_macro(name, value) 45 | if self.undef is not None: 46 | for macro in self.undef: 47 | self.compiler.undefine_macro(macro) 48 | if self.libraries is not None: 49 | self.compiler.set_libraries(self.libraries) 50 | if self.library_dirs is not None: 51 | self.compiler.set_library_dirs(self.library_dirs) 52 | if self.rpath is not None: 53 | self.compiler.set_runtime_library_dirs(self.rpath) 54 | if self.link_objects is not None: 55 | self.compiler.set_link_objects(self.link_objects) 56 | 57 | self.build_extensions() 58 | except Exception as e: 59 | import traceback 60 | exc_type, exc_value, exc_traceback = sys.exc_info() 61 | print("*** print_tb:") 62 | traceback.print_tb(exc_traceback, file=sys.stdout) 63 | self._unavailable(e) 64 | self.extensions = [] # avoid copying missing files (it would fail). 65 | 66 | def _unavailable(self, e): 67 | print('*' * 80) 68 | print('''WARNING: 69 | 70 | An optional code optimization (C extension) could not be compiled. 71 | 72 | Optimizations for this package will not be available! 73 | ''') 74 | 75 | print('CAUSE:') 76 | print('') 77 | print(' ' + repr(e)) 78 | print('*' * 80) 79 | 80 | 81 | def read(*names, **kwargs): 82 | with io.open( 83 | join(dirname(__file__), *names), 84 | encoding=kwargs.get('encoding', 'utf8') 85 | ) as fh: 86 | return fh.read() 87 | 88 | 89 | setup( 90 | name='pymult', 91 | version='0.0.2', 92 | license='BSD-2-Clause', 93 | description='An example Python package written in x86-64 assembly.', 94 | long_description= read('README.md'), 95 | author='Anthony Shaw', 96 | author_email='anthonyshaw@apache.org', 97 | url='https://github.com/tonybaloney/python-assembly-poc', 98 | packages=find_packages('src'), 99 | package_dir={'': 'src'}, 100 | py_modules=[splitext(basename(path))[0] for path in glob('src/*.py')], 101 | include_package_data=True, 102 | zip_safe=False, 103 | classifiers=[ 104 | # complete classifier list: http://pypi.python.org/pypi?%3Aaction=list_classifiers 105 | 'Development Status :: 3 - Alpha', 106 | 'Intended Audience :: Developers', 107 | 'License :: OSI Approved :: BSD License', 108 | 'Operating System :: Unix', 109 | 'Operating System :: POSIX', 110 | 'Programming Language :: Python', 111 | 'Programming Language :: Python :: 3.7', 112 | 'Programming Language :: Python :: 3.8', 113 | 'Topic :: Utilities', 114 | ], 115 | python_requires='>=3.7', 116 | install_requires=[ 117 | ], 118 | extras_require={ 119 | }, 120 | cmdclass={'build_ext': NasmBuildCommand}, 121 | ext_modules=[ 122 | Extension( 123 | splitext(relpath(path, 'src').replace(os.sep, '.'))[0], 124 | sources=[path], 125 | extra_compile_args=[], 126 | extra_link_args=[], 127 | include_dirs=[dirname(path)] 128 | ) 129 | for root, _, _ in os.walk('src') 130 | for path in glob(join(root, '*.asm')) 131 | ], 132 | ) 133 | -------------------------------------------------------------------------------- /src/pymult.asm: -------------------------------------------------------------------------------- 1 | default rel 2 | bits 64 3 | %ifdef NOPIE 4 | %define PYARG_PARSETUPLE PyArg_ParseTuple 5 | %define PYLONG_FROMLONG PyLong_FromLong 6 | %define PYMODULE_CREATE2 PyModule_Create2 7 | %else 8 | %define PYARG_PARSETUPLE PyArg_ParseTuple wrt ..plt 9 | %define PYLONG_FROMLONG PyLong_FromLong wrt ..plt 10 | %define PYMODULE_CREATE2 PyModule_Create2 wrt ..plt 11 | %endif 12 | section .data 13 | modulename db "pymult", 0 14 | docstring db "Simple Multiplication function", 0 15 | 16 | struc moduledef 17 | ;pyobject header 18 | m_object_head_size: resq 1 19 | m_object_head_type: resq 1 20 | ;pymoduledef_base 21 | m_init: resq 1 22 | m_index: resq 1 23 | m_copy: resq 1 24 | ;moduledef 25 | m_name: resq 1 26 | m_doc: resq 1 27 | m_size: resq 1 28 | m_methods: resq 1 29 | m_slots: resq 1 30 | m_traverse: resq 1 31 | m_clear: resq 1 32 | m_free: resq 1 33 | endstruc 34 | 35 | struc methoddef 36 | ml_name: resq 1 37 | ml_meth: resq 1 38 | ml_flags: resd 1 39 | ml_doc: resq 1 40 | ml_term: resq 1 41 | ml_term2: resq 1 42 | endstruc 43 | 44 | section .bss 45 | section .text 46 | 47 | global PyInit_pymult 48 | global PyMult_multiply 49 | 50 | PyMult_multiply: 51 | ; 52 | ; pymult.multiply (a, b) 53 | ; Multiplies a and b 54 | ; Returns value as PyLong(PyObject*) 55 | extern PyLong_FromLong 56 | extern PyLong_AsLong 57 | extern PyArg_ParseTuple 58 | section .data 59 | parseStr db "LL", 0 ; convert arguments to Long, Long 60 | section .bss 61 | result resq 1 ; long result 62 | x resq 1 ; long input 63 | y resq 1 ; long input 64 | section .text 65 | push rbp ; preserve stack pointer 66 | mov rbp, rsp 67 | push rbx 68 | sub rsp, 0x18 69 | 70 | mov rdi, rsi ; args 71 | lea rsi, [parseStr] ; Parse args to LL 72 | xor ebx, ebx ; clear the ebx 73 | lea rdx, [x] ; set the address of x as the 3rd arg 74 | lea rcx, [y] ; set the address of y as the 4th arg 75 | 76 | xor eax, eax ; clear eax 77 | call PYARG_PARSETUPLE ; Parse Args via C-API 78 | 79 | test eax, eax ; if PyArg_ParseTuple is NULL, exit with error 80 | je badinput 81 | 82 | mov rax, [x] ; multiply x and y 83 | imul qword[y] 84 | mov [result], rax 85 | 86 | mov edi, [result] ; convert result to PyLong 87 | call PYLONG_FROMLONG 88 | 89 | mov rsp, rbp ; reinit stack pointer 90 | pop rbp 91 | ret 92 | 93 | badinput: 94 | mov rax, rbx 95 | add rsp, 0x18 96 | pop rbx 97 | pop rbp 98 | ret 99 | 100 | PyInit_pymult: 101 | extern PyModule_Create2 102 | section .data 103 | method1name db "multiply", 0 104 | method1doc db "Multiply two values", 0 105 | 106 | _method1def: 107 | istruc methoddef 108 | at ml_name, dq method1name 109 | at ml_meth, dq PyMult_multiply 110 | at ml_flags, dd 0x0001 ; METH_VARARGS 111 | at ml_doc, dq 0x0 112 | at ml_term, dq 0x0 ; Method defs are terminated by two NULL values, 113 | at ml_term2, dq 0x0 ; equivalent to qword[0x0], qword[0x0] 114 | iend 115 | _moduledef: 116 | istruc moduledef 117 | at m_object_head_size, dq 1 118 | at m_object_head_type, dq 0x0 ; null 119 | at m_init, dq 0x0 ; null 120 | at m_index, dq 0 ; zero 121 | at m_copy, dq 0x0 ; null 122 | at m_name, dq modulename 123 | at m_doc, dq docstring 124 | at m_size, dq 2 125 | at m_methods, dq _method1def 126 | at m_slots, dq 0 ; null- no slots 127 | at m_traverse, dq 0 ; null 128 | at m_clear, dq 0 ; null - no custom clear 129 | at m_free, dq 0 ; null - no custom free() 130 | iend 131 | section .text 132 | push rbp ; preserve stack pointer 133 | mov rbp, rsp 134 | 135 | lea rdi, [_moduledef] ; load module def 136 | mov esi, 0x3f5 ; 1033 - module_api_version 137 | call PYMODULE_CREATE2 ; create module, leave return value in register as return result 138 | 139 | mov rsp, rbp ; reinit stack pointer 140 | pop rbp 141 | ret -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | from pymult import multiply 2 | 3 | 4 | def test_basic_multiplication(): 5 | assert multiply(2, 4) == 8 6 | -------------------------------------------------------------------------------- /winnasmcompiler.py: -------------------------------------------------------------------------------- 1 | """ 2 | Distutils doesn't support nasm, so this is a custom compiler for NASM 3 | """ 4 | from distutils.errors import DistutilsExecError, CompileError 5 | 6 | from distutils.msvc9compiler import MSVCCompiler 7 | import os 8 | 9 | 10 | class WinNasmCompiler(MSVCCompiler): 11 | compiler_type = 'winnasm' 12 | src_extensions = ['.asm'] 13 | 14 | def __init__ (self, 15 | verbose=0, 16 | dry_run=0, 17 | force=0): 18 | 19 | MSVCCompiler.__init__(self, verbose, dry_run, force) 20 | 21 | def initialize(self, plat_name=None): 22 | self.cc = self.find_exe("nasm.exe") 23 | self.linker = self.find_exe("link.exe") 24 | self.lib = self.find_exe("lib.exe") 25 | self.rc = self.find_exe("rc.exe") # resource compiler 26 | self.mc = self.find_exe("mc.exe") # message compiler 27 | 28 | self.compile_options = ["-f win64", "-DWINDOWS", "-DNOPIE"] 29 | self.compile_options_debug = ["-f win64", "-g", "-DWINDOWS", "-DNOPIE"] 30 | self.ldflags_shared = ['/DLL', '/nologo', '/INCREMENTAL:NO', '/NOENTRY'] 31 | self.ldflags_shared_debug = ['/DLL', '/nologo', '/INCREMENTAL:NO', '/DEBUG', '/NOENTRY'] 32 | self.ldflags_static = ['/nologo'] 33 | self.initialized = True 34 | 35 | def link(self, target_desc, objects, 36 | output_filename, output_dir=None, libraries=None, 37 | library_dirs=None, runtime_library_dirs=None, 38 | export_symbols=None, debug=0, extra_preargs=None, 39 | extra_postargs=None, build_temp=None, target_lang=None): 40 | 41 | return super().link(target_desc, objects, 42 | output_filename, output_dir, libraries, 43 | library_dirs, runtime_library_dirs, 44 | export_symbols, debug, extra_preargs, 45 | extra_postargs, build_temp, target_lang) 46 | 47 | def compile(self, sources, 48 | output_dir=None, macros=None, include_dirs=None, debug=0, 49 | extra_preargs=None, extra_postargs=None, depends=None): 50 | 51 | if not self.initialized: 52 | self.initialize() 53 | compile_info = self._setup_compile(output_dir, macros, include_dirs, 54 | sources, depends, extra_postargs) 55 | macros, objects, extra_postargs, pp_opts, build = compile_info 56 | 57 | compile_opts = extra_preargs or [] 58 | 59 | if debug: 60 | compile_opts.extend(self.compile_options_debug) 61 | else: 62 | compile_opts.extend(self.compile_options) 63 | 64 | for obj in objects: 65 | try: 66 | src, ext = build[obj] 67 | except KeyError: 68 | continue 69 | if debug: 70 | # pass the full pathname to MSVC in debug mode, 71 | # this allows the debugger to find the source file 72 | # without asking the user to browse for it 73 | src = os.path.abspath(src) 74 | 75 | input_opt = src 76 | output_opt = "-o" + obj 77 | try: 78 | self.spawn([self.cc] + compile_opts + pp_opts + 79 | [input_opt, output_opt] + 80 | extra_postargs) 81 | except DistutilsExecError as msg: 82 | raise CompileError(msg) 83 | 84 | return objects 85 | --------------------------------------------------------------------------------