do closures right-er

Author: makslevental

examples/cuda_matmul_opt.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
+
 import ast
+import contextlib
 import math
 import re
-import time
 
 import cupy as cp
 import mlir.extras.types as T
@@ -20,7 +21,7 @@
     thread_id,
     block_dim,
 )
-from mlir.extras.dialects.ext.nvgpu import get_ptx, print_ptx
+from mlir.extras.dialects.ext.nvgpu import get_ptx
 from mlir.extras.dialects.ext.scf import range_
 from mlir.extras.runtime.passes import Pipeline, run_pipeline
 
@@ -40,12 +41,22 @@ def build_cuda_func(compiled_module, kernel_name="mat_product_kernel"):
     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)):
-    M, K, N, dtype
     x = block_dim.x * block_id.x + thread_id.x
     y = block_dim.y * block_id.y + thread_id.y
 
@@ -57,9 +68,7 @@ def mat_product_kernel[
     C[x, y] = tmp + one
 
 
-def main(ctx: MLIRContext):
-    M, K, N = 2048, 2048, 2048
-    BLOCK_SIZE = 32
+def main(ctx: MLIRContext, M, K, N, BLOCK_SIZE=32, repeat_times=50):
     dtype = T.f32()
     npy_dtype = np.float32
 
@@ -87,6 +96,7 @@ def _():
             },
         ),
     )
+    cuda_func = build_cuda_func(compiled_module)
     # print(compiled_module)
     # print_ptx(compiled_module)
 
@@ -98,34 +108,26 @@ def _():
     dB = cp.asarray(B)
     dC = cp.asarray(C)
 
-    start_gpu = cp.cuda.Event()
-    end_gpu = cp.cuda.Event()
-
-    cuda_func = build_cuda_func(compiled_module)
-    start_gpu.record()
-    start_cpu = time.perf_counter()
-    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),
-    )
-    end_cpu = time.perf_counter()
-    end_gpu.record()
-    end_gpu.synchronize()
+    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)
-    t_cpu = end_cpu - start_cpu
 
-    print(f"{t_gpu=}ms")
-    print(f"t_cpu={t_cpu / 1000}ms")
+    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)
 
 
-with (
-    mlir_mod_ctx() as ctx,
-    # enable_debug()
-):
-    main(ctx)
+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

@@ -60,6 +60,7 @@ 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",

mlir/extras/ast/util.py

@@ -1,7 +1,14 @@
 import ast
+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):
@@ -42,6 +49,99 @@ def bind(func, instance, as_name=None):
     return bound_method
 
 
+# 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=code,
+        globals=f.__globals__,
+        name=f.__name__,
+        argdefs=f.__defaults__,
+        # 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__)
+    g = functools.update_wrapper(g, f)
+
+    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(

mlir/extras/dialects/ext/func.py

@@ -1,7 +1,10 @@
 import inspect
 import sys
+from dataclasses import dataclass
 from functools import update_wrapper
+from typing import Optional
 
+from ...ast.util import copy_func
 from ...meta import op_region_builder
 from ...util import get_user_code_loc, make_maybe_no_args_decorator
 from ....dialects._ods_common import get_op_result_or_op_results
@@ -122,6 +125,12 @@ def prep_func_types(sig, return_types):
     return input_types, return_types, user_locs
 
 
+@dataclass
+class ReifiedTypeParams:
+    name: str
+    val: object
+
+
 class FuncBase:
     def __init__(
         self,
@@ -193,18 +202,20 @@ def emit(self, *call_args, decl=False, force=False) -> FuncOp:
         if self._func_op is None or decl or force:
             if len(call_args) == 0:
                 input_types = self.input_types[:]
-                env = dict(self.body_builder.__globals__)
-                if self.body_builder.__closure__:
-                    closure = dict(
-                        zip(
-                            self.body_builder.__code__.co_freevars,
-                            [c.cell_contents for c in self.body_builder.__closure__],
-                        )
-                    )
-                    env.update(closure)
+                locals = {}
+                if (
+                    hasattr(self.body_builder, "__type_params__")
+                    and self.body_builder.__type_params__
+                ):
+                    for t in self.body_builder.__type_params__:
+                        if not isinstance(t, ReifiedTypeParams):
+                            raise RuntimeError(f"{t=} must reified")
+                        locals[t.name] = t.val
                 for i, v in enumerate(input_types):
                     if isinstance(v, str):
-                        input_types[i] = Type(eval(v, env))
+                        input_types[i] = Type(
+                            eval(v, self.body_builder.__globals__, locals)
+                        )
                     elif isalambda(v):
                         input_types[i] = v()
             else:
@@ -256,24 +267,42 @@ def __call__(self, *call_args):
         return call(self.emit(*call_args), call_args)
 
     def __getitem__(self, item):
-        if self.body_builder.__code__.co_freevars and self.body_builder.__closure__:
-            closure = {
-                k: v
-                for k, v in zip(
-                    self.body_builder.__code__.co_freevars,
-                    self.body_builder.__closure__,
-                )
-                if v.cell_contents in self.body_builder.__type_params__
-            }
-
-            for i, t in enumerate(self.body_builder.__type_params__):
-                if t.__bound__ is not None:
-                    v = t.__bound__
-                else:
-                    v = item[i]
-                closure[t.__name__].cell_contents = v
-
-        return self
+        if (
+            not hasattr(self.body_builder, "__type_params__")
+            or not self.body_builder.__type_params__
+        ):
+            raise RuntimeError(
+                "using a generic call requires the func be generic (i.e., have type_params)"
+            )
+        # this also copies the function so that the original body_builder remains "generic"
+        body_builder = copy_func(self.body_builder)
+        reified_type_params = []
+        for i, t in enumerate(body_builder.__type_params__):
+            if t.__bound__ is not None:
+                r = ReifiedTypeParams(t.__name__, t.__bound__)
+            else:
+                r = ReifiedTypeParams(t.__name__, item[i])
+            reified_type_params.append(r)
+            if r.name in body_builder.__code__.co_freevars:
+                free_i = body_builder.__code__.co_freevars.index(r.name)
+                body_builder.__closure__[free_i].cell_contents = r.val
+
+        body_builder.__type_params__ = tuple(reified_type_params)
+
+        return FuncBase(
+            body_builder,
+            self.func_op_ctor,
+            self.return_op_ctor,
+            self.call_op_ctor,
+            self.return_types,
+            self.sym_visibility,
+            self.arg_attrs,
+            self.res_attrs,
+            self.func_attrs,
+            self.loc,
+            self.ip,
+            self.qualname,
+        )
 
 
 @make_maybe_no_args_decorator

requirements.txt

@@ -4,4 +4,5 @@ black
 bytecode
 inflection
 numpy
-astunparse
+astunparse
+cloudpickle