generics (#80)

Author: makslevental

examples/cuda_matmul_opt.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import contextlib
+import math
+
+import cupy as cp
+import mlir.extras.types as T
+import numpy as np
+from cupy.cuda import Module
+
+from mlir.extras.ast.canonicalize import canonicalize
+from mlir.extras.context import (
+    mlir_mod_ctx,
+    MLIRContext,
+)
+from mlir.extras.dialects.ext import arith, memref, gpu, scf
+from mlir.extras.dialects.ext.gpu import (
+    block_id,
+    thread_id,
+    block_dim,
+    get_compile_object_bytes,
+)
+from mlir.extras.dialects.ext.scf import range_
+from mlir.extras.runtime.passes import Pipeline, run_pipeline
+
+# noinspection PyUnresolvedReferences
+from mlir.extras.util import find_ops, enable_debug as enable_debug
+
+# just so it doesn't get DCE'd by black/reformat
+_ = memref
+
+
+def build_cuda_func(compiled_module, kernel_name="mat_product_kernel"):
+    ptx = get_compile_object_bytes(compiled_module)
+    mod = Module()
+    mod.load(ptx)
+    return mod.get_function(kernel_name)
+
+
+@contextlib.contextmanager
+def time_cuda():
+    start_gpu = cp.cuda.Event()
+    end_gpu = cp.cuda.Event()
+
+    start_gpu.record()
+    yield start_gpu, end_gpu
+    end_gpu.record()
+    end_gpu.synchronize()
+
+
+@gpu.func
+@canonicalize(using=(arith.canonicalizer, scf.canonicalizer))
+def mat_product_kernel[
+    M, K, N, dtype
+](A: T.memref(M, K, dtype), B: T.memref(K, N, dtype), C: T.memref(M, N, dtype)):
+    x = block_dim.x * block_id.x + thread_id.x
+    y = block_dim.y * block_id.y + thread_id.y
+
+    one = arith.constant(1.0, type=dtype)
+    tmp = arith.constant(0, type=dtype)
+    for k, tmp in range_(K, iter_args=[tmp]):
+        tmp += A[x, k] * B[k, y]
+        tmp = yield tmp
+    C[x, y] = tmp + one
+
+
+def main(ctx: MLIRContext, M, K, N, BLOCK_SIZE=32, repeat_times=50):
+    dtype = T.f32()
+    npy_dtype = np.float32
+
+    gpu.set_container_module(ctx.module)
+
+    @gpu.module("naive", ["#nvvm.target"])
+    def _():
+        mat_product_kernel[M, K, N, dtype].emit()
+
+    # print(ctx.module)
+    ctx.module.operation.verify()
+
+    compiled_module = run_pipeline(
+        ctx.module,
+        Pipeline().add_pass(
+            "gpu-lower-to-nvvm-pipeline",
+            # https://github.com/llvm/llvm-project/blob/ace69e6b942b8fa7e610d70be2a92e801ceea481/mlir/include/mlir/Dialect/GPU/Pipelines/Passes.h#L18
+            **{
+                "cubin-chip": "sm_80",
+                "cubin-features": "+ptx83",
+                "cubin-format": "isa",
+                "kernel-bare-ptr-calling-convention": "1",
+                # "cubin-format": "fatbin",
+                # "cubin-format": "bin",
+            },
+        ),
+    )
+    cuda_func = build_cuda_func(compiled_module)
+    # print(compiled_module)
+    # print_ptx(compiled_module)
+
+    A = np.random.randint(0, 10, (M, K)).astype(npy_dtype)
+    B = np.random.randint(0, 10, (K, N)).astype(npy_dtype)
+    C = np.zeros((M, N)).astype(npy_dtype)
+
+    dA = cp.asarray(A)
+    dB = cp.asarray(B)
+    dC = cp.asarray(C)
+
+    with time_cuda() as (start_gpu, end_gpu):
+        for _ in range(repeat_times):
+            cuda_func(
+                (math.ceil(M / BLOCK_SIZE), math.ceil(N / BLOCK_SIZE), 1),
+                (BLOCK_SIZE, BLOCK_SIZE, 1),
+                (dA.data.ptr, dB.data.ptr, dC.data.ptr),
+            )
+
+    t_gpu = cp.cuda.get_elapsed_time(start_gpu, end_gpu)
+
+    print(f"t_gpu={t_gpu / repeat_times:.6f} ms")
+
+    if not cp.array_equal(dC, dA @ dB + 1):
+        print(dA @ dB + 1)
+        print(dC)
+
+
+for s in [128, 256, 512, 1024]:
+    with (
+        mlir_mod_ctx() as ctx,
+        # enable_debug()
+    ):
+        main(ctx, s, s, s)

mlir/extras/ast/canonicalize.py

@@ -4,16 +4,17 @@
 import inspect
 import logging
 import types
+import warnings
 from abc import ABC, abstractmethod
 from dis import findlinestarts
 from opcode import opmap
 from types import CodeType
-from typing import List, Union
+from typing import List, Union, Sequence
 
 import astunparse
 from bytecode import ConcreteBytecode
 
-from ..ast.util import get_module_cst, copy_func
+from ..ast.util import get_module_cst
 
 logger = logging.getLogger(__name__)
 
@@ -59,28 +60,63 @@ def transform_func(f, *transformer_ctors: type(Transformer)):
     return module
 
 
+# TODO(max): unify with `replace_closure` in ast/utils.py
+def insert_closed_vars(f, module):
+    enclosing_mod = ast.FunctionDef(
+        name="enclosing_mod",
+        args=ast.arguments(
+            posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]
+        ),
+        body=[],
+        decorator_list=[],
+    )
+    for var in f.__code__.co_freevars:
+        enclosing_mod.body.append(
+            ast.Assign(
+                targets=[ast.Name(var, ctx=ast.Store())],
+                value=ast.Constant(None, kind="None"),
+            )
+        )
+    enclosing_mod.body.extend(module.body)
+    module.body = [enclosing_mod]
+    return module
+
+
+def find_func_in_code_object(co, func_name):
+    for c in co.co_consts:
+        if type(c) is CodeType:
+            if c.co_name == func_name:
+                return c
+            else:
+                f = find_func_in_code_object(c, func_name)
+                if f is not None:
+                    return f
+
+
 def transform_ast(
     f, transformers: List[Union[type(Transformer), type(StrictTransformer)]] = None
 ):
     if transformers is None:
         return f
 
     module = transform_func(f, *transformers)
+    if f.__closure__:
+        module = insert_closed_vars(f, module)
     module = ast.fix_missing_locations(module)
     module = ast.increment_lineno(module, f.__code__.co_firstlineno - 1)
     module_code_o = compile(module, f.__code__.co_filename, "exec")
-    new_f_code_o = next(
-        c
-        for c in module_code_o.co_consts
-        if type(c) is CodeType and c.co_name == f.__name__
-    )
+    new_f_code_o = find_func_in_code_object(module_code_o, f.__name__)
     n_lines = len(inspect.getsource(f).splitlines())
     line_starts = list(findlinestarts(new_f_code_o))
-    assert (
+    if (
         max([l for _, l in line_starts]) - min([l for _, l in line_starts]) + 1
-        <= n_lines
-    ), f"something went wrong with the line numbers for the rewritten/canonicalized function"
-    return copy_func(f, new_f_code_o)
+        > n_lines
+    ) or (f.__code__.co_firstlineno != min([l for _, l in line_starts])):
+        warnings.warn(
+            "something went wrong with the line numbers for the rewritten/canonicalized function"
+        )
+    f.__code__ = new_f_code_o
+    return f
 
 
 # this is like this because i couldn't figure out how to subclass
@@ -117,7 +153,8 @@ def patch_bytecode(f, patchers: List[type(BytecodePatcher)] = None):
     for patcher in patchers:
         code = patcher(context).patch_bytecode(code, f)
 
-    return copy_func(f, code.to_code())
+    f.__code__ = code.to_code()
+    return f
 
 
 class Canonicalizer(ABC):
@@ -132,10 +169,18 @@ def bytecode_patchers(self) -> List[BytecodePatcher]:
         pass
 
 
-def canonicalize(*, using: Canonicalizer):
+def canonicalize(*, using: Union[Canonicalizer, Sequence[Canonicalizer]]):
+    if not isinstance(using, Sequence):
+        using = [using]
+    cst_transformers = []
+    bytecode_patchers = []
+    for u in using:
+        cst_transformers.extend(u.cst_transformers)
+        bytecode_patchers.extend(u.bytecode_patchers)
+
     def wrapper(f):
-        f = transform_ast(f, using.cst_transformers)
-        f = patch_bytecode(f, using.bytecode_patchers)
+        f = transform_ast(f, cst_transformers)
+        f = patch_bytecode(f, bytecode_patchers)
         return f
 
     return wrapper

mlir/extras/ast/util.py

@@ -2,7 +2,13 @@
 import functools
 import inspect
 import types
+from itertools import dropwhile
+from opcode import opmap
 from textwrap import dedent
+from typing import Dict
+
+from bytecode import ConcreteBytecode
+from cloudpickle import cloudpickle
 
 
 def set_lineno(node, n=1):
@@ -26,8 +32,8 @@ def ast_call(name, args=None, keywords=None):
 
 
 def get_module_cst(f):
-    f_src = dedent(inspect.getsource(f))
-    # tree = cst.parse_module(f_src)
+    lines, _lnum = inspect.getsourcelines(f)
+    f_src = dedent("".join(list(dropwhile(lambda l: l.startswith("@"), lines))))
     tree = ast.parse(f_src)
     assert isinstance(
         tree.body[0], ast.FunctionDef
@@ -43,21 +49,89 @@ def bind(func, instance, as_name=None):
     return bound_method
 
 
-def copy_func(f, new_code):
-    """Based on http://stackoverflow.com/a/6528148/190597 (Glenn Maynard)"""
+# based on https://github.com/python/cpython/blob/6078f2033ea15a16cf52fe8d644a95a3be72d2e3/Tools/build/deepfreeze.py#L48
+def get_localsplus_name_to_idx(code: types.CodeType):
+    localsplus = code.co_varnames + code.co_cellvars + code.co_freevars
+    return localsplus, {v: i for i, v in enumerate(localsplus)}
+
+
+class _empty_cell_value:
+    """Sentinel for empty closures."""
+
+    @classmethod
+    def __reduce__(cls):
+        return cls.__name__
+
+
+_empty_cell_value = _empty_cell_value()
+
+
+# based on https://github.com/cloudpipe/cloudpickle/blob/f111f7ab6d302e9b1e2a568d0e4c574895db6a6e/cloudpickle/cloudpickle.py#L513
+def make_empty_cell():
+    if False:
+        # trick the compiler into creating an empty cell in our lambda
+        cell = None
+        raise AssertionError("this route should not be executed")
+
+    return (lambda: cell).__closure__[0]
+
+
+def make_cell(value=_empty_cell_value):
+    cell = make_empty_cell()
+    if value is not _empty_cell_value:
+        cell.cell_contents = value
+    return cell
+
+
+# based on https://github.com/python/cpython/blob/a4b44d39cd6941cc03590fee7538776728bdfd0a/Lib/test/test_code.py#L197
+def replace_closure(code, new_closure: Dict):
+    COPY_FREE_VARS = opmap["COPY_FREE_VARS"]
+    LOAD_DEREF = opmap["LOAD_DEREF"]
+
+    # get the orig localplus that will be loaded from by the orig bytecode LOAD_DEREF arg_i
+    localsplus, localsplus_name_to_idx = get_localsplus_name_to_idx(code)
+
+    # closure vars go into co_freevars
+    new_code = code.replace(co_freevars=tuple(new_closure.keys()))
+    # closure is a tuple of cells
+    closure = tuple(
+        make_cell(v) if not isinstance(v, types.CellType) else v
+        for v in new_closure.values()
+    )
+
+    new_code = ConcreteBytecode.from_code(new_code)
+    # update how many closure vars are loaded from frame
+    # see https://github.com/python/cpython/blob/6078f2033ea15a16cf52fe8d644a95a3be72d2e3/Python/bytecodes.c#L1571
+    assert new_code[0].opcode == COPY_FREE_VARS
+    new_code[0].arg = len(closure)
+
+    # map orig localsplus arg_i to new localplus position/arg_i
+    new_localsplus = new_code.varnames + new_code.cellvars + new_code.freevars
+    new_localsplus_name_to_idx = {v: i for i, v in enumerate(new_localsplus)}
+    for c in new_code:
+        if c.opcode == LOAD_DEREF and c.arg < len(localsplus):
+            c.arg = new_localsplus_name_to_idx[localsplus[c.arg]]
+    new_code = new_code.to_code()
+
+    return new_code, closure
+
+
+# Based on http://stackoverflow.com/a/6528148/190597 (Glenn Maynard);
+# potentially more complete approach https://stackoverflow.com/a/56901529/9045206
+def copy_func(f, new_closure: Dict = None):
+    if new_closure is not None:
+        code, closure = replace_closure(f.__code__, new_closure)
+    else:
+        code, closure = f.__code__, f.__closure__
+
     g = types.FunctionType(
-        code=new_code,
-        globals={
-            **f.__globals__,
-            **{
-                fr: f.__closure__[i].cell_contents
-                for i, fr in enumerate(f.__code__.co_freevars)
-            },
-        },
+        code=code,
+        globals=f.__globals__,
         name=f.__name__,
         argdefs=f.__defaults__,
-        # TODO(max): ValueError: foo requires closure of length 0, not 1
-        # closure=f.__closure__,
+        # see https://github.com/cloudpipe/cloudpickle/blob/f111f7ab6d302e9b1e2a568d0e4c574895db6a6e/cloudpickle/cloudpickle.py#L813
+        # for how this trick is accomplished (dill and pickle both fail to pickle eg generic typevars)
+        closure=cloudpickle.loads(cloudpickle.dumps(closure)),
     )
     g.__kwdefaults__ = f.__kwdefaults__
     g.__dict__.update(f.__dict__)
@@ -66,3 +140,12 @@ def copy_func(f, new_code):
     if inspect.ismethod(f):
         g = bind(g, f.__self__)
     return g
+
+
+def append_hidden_node(node_body, new_node):
+    last_statement = node_body[-1]
+    new_node = ast.fix_missing_locations(
+        set_lineno(new_node, last_statement.end_lineno)
+    )
+    node_body.append(new_node)
+    return node_body