cuda_opt_v4_tensor_cores (#87)

Author: makslevental

examples/cuda_matmul_opt.py

@@ -4,17 +4,15 @@
 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, Sequence
 
 import astunparse
 from bytecode import ConcreteBytecode
 
-from ..ast.util import get_module_cst, set_lineno
+from ..ast.util import get_module_cst, set_lineno, find_func_in_code_object
 
 logger = logging.getLogger(__name__)
 
@@ -86,17 +84,6 @@ def insert_closed_vars(f, module):
     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
 ):

mlir/extras/ast/canonicalize.py

@@ -2,13 +2,14 @@
 import functools
 import inspect
 import types
-from itertools import dropwhile
 from opcode import opmap
 from textwrap import dedent
-from typing import Dict
+from typing import Dict, TypeVar
 
 from bytecode import ConcreteBytecode
 from cloudpickle import cloudpickle
+from einspect import ptr
+from einspect.structs import PyVarObject, PyObject
 
 
 def set_lineno(node, n=1):
@@ -101,7 +102,7 @@ def replace_closure(code, new_closure: Dict):
     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
+    assert new_code[0].opcode == COPY_FREE_VARS, f"{new_code[0].opcode=}"
     new_code[0].arg = len(closure)
 
     # map orig localsplus arg_i to new localplus position/arg_i
@@ -121,16 +122,17 @@ 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__
+        # 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(f.__closure__))
+        code = f.__code__
 
     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)),
+        closure=closure,
     )
     g.__kwdefaults__ = f.__kwdefaults__
     g.__dict__.update(f.__dict__)
@@ -151,3 +153,24 @@ def append_hidden_node(node_body, new_node):
     )
     node_body.append(new_node)
     return node_body
+
+
+def find_func_in_code_object(co, func_name):
+    for c in co.co_consts:
+        if type(c) is types.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
+
+
+_T = TypeVar("_T")
+
+
+# https://github.com/python/cpython/blob/809aa9a682fc865f7502e7421da0a74d204aab6d/Objects/typevarobject.c#L29
+class PyTypeVarObject(PyVarObject[_T, None, None]):
+    name: ptr[PyObject]
+    # not sure why but this is the only thing that works but that's fine because it's the only thing we need
+    bound: ptr[PyObject]

mlir/extras/ast/util.py

@@ -1,11 +1,13 @@
+import copy
 import inspect
 import sys
 from dataclasses import dataclass
 from functools import update_wrapper
 from typing import Optional, List, Union, TypeVar
 
-from ...ast.util import copy_func
+from ...ast.util import copy_func, PyTypeVarObject
 from ...meta import op_region_builder
+from ... import types as T
 from ...util import get_user_code_loc, make_maybe_no_args_decorator
 from ....dialects._ods_common import get_op_result_or_op_results
 from ....dialects.func import *
@@ -105,17 +107,17 @@ def prep_func_types(sig, return_types):
         return_types = [return_types]
     return_types = list(return_types)
     assert all(
-        isinstance(r, Type) for r in return_types
-    ), f"all return types must be mlir types {return_types=}"
+        isinstance(r, (str, Type, TypeVar)) or isalambda(r) for r in return_types
+    ), f"all return types must be mlir types or strings or TypeVars or lambdas {return_types=}"
 
     input_types = [
         p.annotation
         for p in sig.parameters.values()
         if not p.annotation is inspect.Signature.empty
     ]
     assert all(
-        isinstance(r, (str, Type)) or isalambda(r) for r in input_types
-    ), f"all input types must be mlir types {input_types=}"
+        isinstance(r, (str, Type, TypeVar)) or isalambda(r) for r in input_types
+    ), f"all input types must be mlir types or strings or TypeVars or lambdas {input_types=}"
     user_loc = get_user_code_loc()
     # If ir.Context is none (like for deferred func emit)
     if user_loc is None:
@@ -205,13 +207,15 @@ 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[:]
-                locals = {}
+                locals = {"T": T}
                 if self.generics is not None:
                     for t in self.generics:
                         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, TypeVar):
+                        v = v.__name__
                     if isinstance(v, str):
                         input_types[i] = Type(
                             eval(v, self.body_builder.__globals__, locals)
@@ -274,12 +278,38 @@ def __getitem__(self, item):
         # this also copies the function so that the original body_builder remains "generic" (via its closure)
         body_builder = copy_func(self.body_builder)
         reified_type_params = []
-        for i, t in enumerate(self.generics):
-            if t.__bound__ is not None:
-                r = ReifiedTypeParams(t.__name__, t.__bound__)
+        # dumb but whatever
+        already_reified_type_params = {}
+        generics = copy.deepcopy(self.generics)
+        for i, t in enumerate(generics):
+            if sys.version_info >= (3, 12):
+                type_var_bound = PyTypeVarObject.try_from(t).bound
+            else:
+                type_var_bound = t.__bound__
+            if type_var_bound:
+                # before 3.12 typevar was just a python class
+                # https://github.com/python/cpython/blob/3.11/Lib/typing.py#L966
+                if sys.version_info < (3, 12):
+                    type_var_bound = lambda: type_var_bound
+                else:
+                    type_var_bound = type_var_bound.contents.into_object()
+                    cvrs = inspect.getclosurevars(type_var_bound).nonlocals
+                    if len(cvrs):
+                        for k, v in cvrs.items():
+                            if not isinstance(v, TypeVar):
+                                continue
+                            if k not in already_reified_type_params:
+                                raise RuntimeError(
+                                    f"typevar {k} not reified prior to evaluating dependent typevar {t}"
+                                )
+                            cvrs[k] = already_reified_type_params[k]
+                        type_var_bound = copy_func(type_var_bound, cvrs)
+                r = ReifiedTypeParams(t.__name__, type_var_bound())
             else:
                 r = ReifiedTypeParams(t.__name__, item[i])
+
             reified_type_params.append(r)
+            already_reified_type_params[r.name] = r.val
 
             if t.__name__ in body_builder.__globals__:
                 body_builder.__globals__[t.__name__] = r.val
@@ -290,8 +320,6 @@ def __getitem__(self, item):
                 ), "typevars don't match"
                 body_builder.__closure__[free_i].cell_contents = r.val
 
-        generics = reified_type_params
-
         return FuncBase(
             body_builder,
             self.func_op_ctor,
@@ -302,7 +330,7 @@ def __getitem__(self, item):
             arg_attrs=self.arg_attrs,
             res_attrs=self.res_attrs,
             func_attrs=self.func_attrs,
-            generics=generics,
+            generics=reified_type_params,
             qualname=self.qualname,
             loc=self.loc,
             ip=self.ip,

mlir/extras/dialects/ext/func.py

@@ -2,7 +2,6 @@
 from functools import partial
 from typing import Any, List, Optional, Tuple, Union
 
-from mlir.dialects._gpu_enum_gen import AddressSpace
 
 from .arith import constant
 from .func import FuncBase
@@ -119,8 +118,8 @@ def get_device_mapping_array_attr(
     return ArrayAttr.get(mapping, context=context)
 
 
-def gpu_attr(mnemonic, mapping_id_enum: MappingId):
-    return Attribute.parse(f"#gpu.{mnemonic}<{mapping_id_enum}>")
+def gpu_attr(mnemonic, attr_value):
+    return Attribute.parse(f"#gpu.{mnemonic}<{attr_value}>")
 
 
 def thread_attr(thread):