backport generics

Author: makslevental

examples/cuda_matmul_opt.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 
-import ast
 import contextlib
 import math
-import re
 
 import cupy as cp
 import mlir.extras.types as T
@@ -20,8 +18,8 @@
     block_id,
     thread_id,
     block_dim,
+    get_compile_object_bytes,
 )
-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
 
@@ -33,9 +31,7 @@
 
 
 def build_cuda_func(compiled_module, kernel_name="mat_product_kernel"):
-    ptx = get_ptx(compiled_module)
-    ptx = re.sub(r"\\(\w\w)", lambda m: r"\x" + m.groups(0)[0].lower(), ptx)
-    ptx = ast.literal_eval(rf"b'{ptx}'")
+    ptx = get_compile_object_bytes(compiled_module)
     mod = Module()
     mod.load(ptx)
     return mod.get_function(kernel_name)

mlir/extras/dialects/ext/func.py

@@ -2,7 +2,7 @@
 import sys
 from dataclasses import dataclass
 from functools import update_wrapper
-from typing import Optional
+from typing import Optional, List, Union, TypeVar
 
 from ...ast.util import copy_func
 from ...meta import op_region_builder
@@ -138,14 +138,16 @@ def __init__(
         func_op_ctor,
         return_op_ctor,
         call_op_ctor,
+        *,
         return_types=None,
         sym_visibility=None,
         arg_attrs=None,
         res_attrs=None,
         func_attrs=None,
+        generics: List[Union[TypeVar, ReifiedTypeParams]] = None,
+        qualname=None,
         loc=None,
         ip=None,
-        qualname=None,
     ):
         assert inspect.isfunction(body_builder), body_builder
         assert inspect.isclass(func_op_ctor), func_op_ctor
@@ -159,6 +161,7 @@ def __init__(
         self.call_op_ctor = call_op_ctor
         self.arg_attrs = arg_attrs
         self.res_attrs = res_attrs
+        self.generics = generics
         self.loc = loc
         self.ip = ip
         self._func_op = None
@@ -203,11 +206,8 @@ def emit(self, *call_args, decl=False, force=False) -> FuncOp:
             if len(call_args) == 0:
                 input_types = self.input_types[:]
                 locals = {}
-                if (
-                    hasattr(self.body_builder, "__type_params__")
-                    and self.body_builder.__type_params__
-                ):
-                    for t in self.body_builder.__type_params__:
+                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
@@ -267,17 +267,14 @@ def __call__(self, *call_args):
         return call(self.emit(*call_args), call_args)
 
     def __getitem__(self, item):
-        if (
-            not hasattr(self.body_builder, "__type_params__")
-            or not self.body_builder.__type_params__
-        ):
+        if self.generics is None:
             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"
+        # 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(body_builder.__type_params__):
+        for i, t in enumerate(self.generics):
             if t.__bound__ is not None:
                 r = ReifiedTypeParams(t.__name__, t.__bound__)
             else:
@@ -287,21 +284,22 @@ def __getitem__(self, item):
                 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)
+        generics = 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,
+            return_types=self.return_types,
+            sym_visibility=self.sym_visibility,
+            arg_attrs=self.arg_attrs,
+            res_attrs=self.res_attrs,
+            func_attrs=self.func_attrs,
+            generics=generics,
+            qualname=self.qualname,
+            loc=self.loc,
+            ip=self.ip,
         )
 
 
@@ -314,11 +312,20 @@ def func(
     res_attrs=None,
     func_attrs=None,
     emit=False,
+    generics=None,
     loc=None,
     ip=None,
 ) -> FuncBase:
     if loc is None:
         loc = get_user_code_loc()
+    if hasattr(f, "__type_params__") and f.__type_params__:
+        assert generics is None, "generics XOR type params"
+        generics = f.__type_params__
+    if generics is not None:
+        for i, g in enumerate(generics):
+            if isinstance(g, str):
+                generics[i] = TypeVar(g)
+
     func_ = FuncBase(
         body_builder=f,
         func_op_ctor=FuncOp.__base__,
@@ -328,6 +335,7 @@ def func(
         arg_attrs=arg_attrs,
         res_attrs=res_attrs,
         func_attrs=func_attrs,
+        generics=generics,
         loc=loc,
         ip=ip,
     )

mlir/extras/dialects/ext/gpu.py

@@ -13,6 +13,7 @@
     _get_previous_frame_idents,
     get_user_code_loc,
     make_maybe_no_args_decorator,
+    find_ops,
 )
 from ....dialects._gpu_ops_gen import _Dialect
 from ....dialects._ods_common import (
@@ -414,12 +415,16 @@ def func(
     res_attrs=None,
     func_attrs=None,
     emit=False,
+    generics=None,
     loc=None,
     ip=None,
     emit_grid=False,
 ) -> Grid:
     if loc is None:
         loc = get_user_code_loc()
+    if hasattr(f, "__type_params__") and f.__type_params__:
+        assert generics is None, "generics XOR type params"
+        generics = f.__type_params__
     func_ = GPUFunc(
         body_builder=f,
         func_op_ctor=GPUFuncOp,
@@ -429,6 +434,7 @@ def func(
         arg_attrs=arg_attrs,
         res_attrs=res_attrs,
         func_attrs=func_attrs,
+        generics=generics,
         loc=loc,
         ip=ip,
     )
@@ -544,3 +550,9 @@ def memcpy(dst, src, async_dependencies=None, *, loc=None, ip=None):
         loc=loc,
         ip=ip,
     )
+
+
+def get_compile_object_bytes(compiled_module):
+    binary = find_ops(compiled_module, lambda o: isinstance(o, BinaryOp), single=True)
+    objects = list(map(ObjectAttr, binary.objects))
+    return objects[-1].object

tests/test_func.py

@@ -1,14 +1,17 @@
 import inspect
 import sys
 from textwrap import dedent
+from typing import TypeVar
 
 import pytest
 
 import mlir.extras.types as T
+
+from mlir.extras.ast.canonicalize import canonicalize
 from mlir.extras.context import mlir_mod_ctx
 from mlir.extras.dialects.ext.arith import constant
 from mlir.extras.dialects.ext.func import func
-from mlir.extras.dialects.ext import linalg
+from mlir.extras.dialects.ext import linalg, arith, scf, memref
 
 # noinspection PyUnresolvedReferences
 from mlir.extras.testing import mlir_ctx as ctx, filecheck, MLIRContext
@@ -40,8 +43,7 @@ def demo_fun1():
 
 
 def test_declare_byte_rep(ctx: MLIRContext):
-    def demo_fun1():
-        ...
+    def demo_fun1(): ...
 
     if sys.version_info.minor == 12:
         assert demo_fun1.__code__.co_code == b"\x97\x00y\x00"
@@ -55,20 +57,16 @@ def demo_fun1():
 
 def test_declare(ctx: MLIRContext):
     @func
-    def demo_fun1() -> T.i32():
-        ...
+    def demo_fun1() -> T.i32(): ...
 
     @func
-    def demo_fun2() -> (T.i32(), T.i32()):
-        ...
+    def demo_fun2() -> (T.i32(), T.i32()): ...
 
     @func
-    def demo_fun3(x: T.i32()) -> (T.i32(), T.i32()):
-        ...
+    def demo_fun3(x: T.i32()) -> (T.i32(), T.i32()): ...
 
     @func
-    def demo_fun4(x: T.i32(), y: T.i32()) -> (T.i32(), T.i32()):
-        ...
+    def demo_fun4(x: T.i32(), y: T.i32()) -> (T.i32(), T.i32()): ...
 
     demo_fun1()
     demo_fun2()
@@ -197,3 +195,99 @@ def test_func_no_context_2(ctx: MLIRContext):
     """
     )
     filecheck(correct, ctx.module)
+
+
+def test_generics_just_args(ctx: MLIRContext):
+    @func(generics=["M", "K", "N", "dtype"])
+    def mat_product_kernel(
+        A: "T.memref(M, K, dtype)",
+        B: "T.memref(K, N, dtype)",
+        C: "T.memref(M, N, dtype)",
+    ):
+        one = arith.constant(1.0)
+
+    mat_product_kernel[32, 32, 32, T.i32()].emit()
+    correct = dedent(
+        """\
+    module {
+      func.func @mat_product_kernel(%arg0: memref<32x32xi32>, %arg1: memref<32x32xi32>, %arg2: memref<32x32xi32>) {
+        %cst = arith.constant 1.000000e+00 : f32
+        return
+      }
+    }
+    """
+    )
+    filecheck(correct, ctx.module)
+
+
+def test_generics_closure(ctx: MLIRContext):
+    dtype = None
+
+    @func(generics=["M", "K", "N", "dtype"])
+    def mat_product_kernel(
+        A: "T.memref(M, K, dtype)",
+        B: "T.memref(K, N, dtype)",
+        C: "T.memref(M, N, dtype)",
+    ):
+        one = arith.constant(1, dtype)
+
+    mat_product_kernel[32, 32, 32, T.i32()].emit()
+    correct = dedent(
+        """\
+    module {
+      func.func @mat_product_kernel(%arg0: memref<32x32xi32>, %arg1: memref<32x32xi32>, %arg2: memref<32x32xi32>) {
+        %c1_i32 = arith.constant 1 : i32
+        return
+      }
+    }
+    """
+    )
+    filecheck(correct, ctx.module)
+
+
+def test_generics_with_canonicalizations(ctx: MLIRContext):
+    dtype = None
+    K = None
+
+    @func(generics=["M", "K", "N", "dtype"])
+    @canonicalize(using=(arith.canonicalizer, scf.canonicalizer))
+    def mat_product_kernel(
+        A: "T.memref(M, K, dtype)",
+        B: "T.memref(K, N, dtype)",
+        C: "T.memref(M, N, dtype)",
+    ):
+        x = arith.constant(1, index=True)
+        y = arith.constant(1, index=True)
+        one = arith.constant(1.0, type=dtype)
+        tmp = arith.constant(0, type=dtype)
+        for k, tmp in scf.range_(K, iter_args=[tmp]):
+            tmp += A[x, k] * B[k, y]
+            tmp = yield tmp
+        C[x, y] = tmp + one
+
+    mat_product_kernel[32, 32, 32, T.f32()].emit()
+    correct = dedent(
+        """\
+    module {
+      func.func @mat_product_kernel(%arg0: memref<32x32xf32>, %arg1: memref<32x32xf32>, %arg2: memref<32x32xf32>) {
+        %c1 = arith.constant 1 : index
+        %c1_0 = arith.constant 1 : index
+        %cst = arith.constant 1.000000e+00 : f32
+        %cst_1 = arith.constant 0.000000e+00 : f32
+        %c0 = arith.constant 0 : index
+        %c32 = arith.constant 32 : index
+        %c1_2 = arith.constant 1 : index
+        %0 = scf.for %arg3 = %c0 to %c32 step %c1_2 iter_args(%arg4 = %cst_1) -> (f32) {
+          %2 = memref.load %arg0[%c1, %arg3] : memref<32x32xf32>
+          %3 = memref.load %arg1[%arg3, %c1_0] : memref<32x32xf32>
+          %4 = math.fma %2, %3, %arg4 : f32
+          scf.yield %4 : f32
+        }
+        %1 = arith.addf %0, %cst : f32
+        memref.store %1, %arg2[%c1, %c1_0] : memref<32x32xf32>
+        return
+      }
+    }
+    """
+    )
+    filecheck(correct, ctx.module)