Add llvm, xform, gpu dialects

Adds tests and dialects for rocdl, llvm, gpu and transform.

Author: ashermancinelli

projects/eudsl-python-extras/mlir/extras/dialects/gpu.py

@@ -0,0 +1,38 @@
+from typing import Union, Optional
+
+from ...util import infer_mlir_type
+
+# noinspection PyUnresolvedReferences
+from ....dialects.llvm import *
+from ....ir import Type, Value, IntegerAttr, FloatAttr
+from ....dialects._ods_common import get_op_result_or_op_results
+
+ValueRef = Value
+
+
+def llvm_ptr_t():
+    return Type.parse("!llvm.ptr")
+
+
+try:
+    from . import amdgcn
+except ImportError:
+    pass
+
+
+def mlir_constant(
+    value: Union[int, float, bool], type: Optional[Type] = None, *, loc=None, ip=None
+) -> Value:
+    if type is None:
+        type = infer_mlir_type(value, vector=False)
+
+    if isinstance(value, int):
+        value = IntegerAttr.get(type, value)
+    elif isinstance(value, float):
+        value = FloatAttr.get(type, value)
+    else:
+        raise NotImplementedError(f"{value} is not a valid type")
+
+    return get_op_result_or_op_results(
+        ConstantOp(res=value.type, value=value, loc=loc, ip=ip)
+    )

projects/eudsl-python-extras/mlir/extras/dialects/llvm/__init__.py

@@ -0,0 +1,67 @@
+import warnings
+
+# noinspection PyUnresolvedReferences
+from ....dialects.llvm import *
+from ....ir import Type, F16Type, F32Type, F64Type, BF16Type, IntegerType
+
+try:
+    from llvm import intrinsic_is_overloaded, intrinsic_get_type, print_type_to_string
+    from llvm import types_
+    from llvm.context import context as llvm_context
+except ImportError:
+    warnings.warn(
+        "llvm bindings not installed; call_intrinsic won't work without supplying return type explicitly"
+    )
+
+
+def mlir_type_to_llvm_type(mlir_type, llvm_ctx):
+    if F16Type.isinstance(mlir_type):
+        return types_.half_type_in_context(llvm_ctx)
+    if F32Type.isinstance(mlir_type):
+        return types_.float_type_in_context(llvm_ctx)
+    if F64Type.isinstance(mlir_type):
+        return types_.double_type_in_context(llvm_ctx)
+    if BF16Type.isinstance(mlir_type):
+        return types_.b_float_type_in_context(llvm_ctx)
+    if IntegerType.isinstance(mlir_type):
+        return types_.int_type_in_context(llvm_ctx, mlir_type.width)
+
+    raise NotImplementedError(f"{mlir_type} is not supported")
+
+
+def llvm_type_str_to_mlir_type(llvm_type: str):
+    if llvm_type.startswith("<"):
+        return Type.parse(f"vector{llvm_type}")
+    if llvm_type == "float":
+        return F32Type.get()
+    raise NotImplementedError(f"{llvm_type} is not supported")
+
+
+_call_intrinsic = call_intrinsic
+
+
+def call_intrinsic(*args, **kwargs):
+    intr_id = kwargs.pop("intr_id")
+    intr_name = kwargs.pop("intr_name")
+    mlir_ret_type = kwargs.pop("return_type", None)
+    if mlir_ret_type:
+        return _call_intrinsic(mlir_ret_type, intr_name, args, [], [])
+
+    is_overloaded = kwargs.pop("is_overloaded", None)
+    if is_overloaded is None:
+        is_overloaded = intrinsic_is_overloaded(intr_id)
+    with llvm_context() as ctx:
+        types = []
+        if is_overloaded:
+            types = [mlir_type_to_llvm_type(a.type, ctx.context) for a in args]
+        intr_decl_fn_ty = intrinsic_get_type(ctx.context, intr_id, types)
+
+    ret_type_str = print_type_to_string(intr_decl_fn_ty).split(" (")[0].strip()
+    mlir_ret_type = None
+    if ret_type_str:
+        mlir_ret_type = llvm_type_str_to_mlir_type(ret_type_str)
+
+    return _call_intrinsic(mlir_ret_type, intr_name, args, [], [])
+
+
+call_intrinsic_ = call_intrinsic

projects/eudsl-python-extras/mlir/extras/dialects/llvm/util.py

@@ -0,0 +1,69 @@
+# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+from . import arith
+from ..util import get_user_code_loc
+
+from ...dialects._ods_common import (
+    _dispatch_mixed_values,
+    _cext,
+    get_op_results_or_values,
+    get_default_loc_context,
+    get_op_result_or_op_results,
+    get_default_loc_context,
+    segmented_accessor,
+)
+
+# noinspection PyUnresolvedReferences
+from ...dialects.rocdl import *
+from ...dialects._rocdl_ops_gen import _Dialect
+from ... import ir
+
+
+@_cext.register_operation(_Dialect, replace=True)
+class WMMA_F16_16X16X16_F16(ir.OpView):
+    OPERATION_NAME = "rocdl.wmma.f16.16x16x16.f16"
+
+    _ODS_REGIONS = (0, True)
+
+    def __init__(self, res, args, *, loc=None, ip=None):
+        operands = []
+        results = []
+        attributes = {}
+        regions = None
+        operands.extend(get_op_results_or_values(args))
+        _ods_context = get_default_loc_context(loc)
+        results.append(res)
+        _ods_successors = None
+        super().__init__(
+            self.OPERATION_NAME,
+            self._ODS_REGIONS,
+            self._ODS_OPERAND_SEGMENTS,
+            self._ODS_RESULT_SEGMENTS,
+            attributes=attributes,
+            results=results,
+            operands=operands,
+            successors=_ods_successors,
+            regions=regions,
+            loc=loc,
+            ip=ip,
+        )
+
+    @property
+    def args(self):
+        _ods_variadic_group_length = len(self.operation.operands) - 1 + 1
+        return self.operation.operands[0 : 0 + _ods_variadic_group_length]
+
+    @property
+    def res(self):
+        return self.operation.results[0]
+
+
+wmma_f16_16x16x16_f16_ = wmma_f16_16x16x16_f16
+
+
+def wmma_f16_16x16x16_f16(A, B, C, *, opsel=False, loc=None, ip=None) -> ir.Value:
+    v16 = ir.VectorType.get((16,), ir.F16Type.get())
+    return wmma_f16_16x16x16_f16_(
+        res=v16, A=A, B=B, C=C, opsel=opsel, loc=loc, ip=ip
+    ).result

projects/eudsl-python-extras/mlir/extras/dialects/rocdl.py

@@ -375,8 +375,8 @@ def _structured_bufferize_to_allocation(
             memory_space = StringAttr.get(memory_space)
 
     return __structured_bufferize_to_allocation(
-        allocated_buffer=transform_any_value_t(),
-        new_ops=transform_any_op_t(),
+        # allocated_buffer=transform_any_value_t(),
+        # new_ops=transform_any_op_t(),
         target=target,
         memory_space=memory_space,
         memcpy_op=memcpy_op,

projects/eudsl-python-extras/mlir/extras/dialects/transform.py

@@ -72,7 +72,13 @@ def filecheck(correct: str, module):
 
     correct = "\n".join(filter(None, correct.splitlines()))
     correct = dedent(correct)
-    correct_with_checks = main(correct).replace("CHECK:", "CHECK-NEXT:")
+    correct_with_checks = main(correct).strip().splitlines()
+    correct_with_checks = "\n".join(
+        [
+            (line.replace("CHECK:", "CHECK-NEXT:") if i > 0 else line)
+            for i, line in enumerate(correct_with_checks)
+        ]
+    )
 
     filecheck_path = get_filecheck_path()
     with tempfile.NamedTemporaryFile() as tmp:
@@ -81,13 +87,17 @@ def filecheck(correct: str, module):
         p = Popen([filecheck_path, tmp.name], stdout=PIPE, stdin=PIPE, stderr=PIPE)
         out, err = map(lambda o: o.decode(), p.communicate(input=op.encode()))
         if p.returncode:
+            breakpoint()
+            if "error: " in err:
+                raise RuntimeError(err)
             diff = list(
                 difflib.unified_diff(
                     op.splitlines(),  # to this
                     correct.splitlines(),  # delta from this
                     lineterm="",
                 )
             )
+            breakpoint()
             diff.insert(1, "delta from module to correct")
             print("lit report:", err, file=sys.stderr)
             raise ValueError("\n" + "\n".join(diff))