Add llvm, transform, gpu dialects (#209)

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

Co-authored-by: Maksim Levental <[email protected]>

Author: ashermancinelli

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

@@ -0,0 +1,41 @@
+# 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 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,70 @@
+# 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
+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,24 @@
+# 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 ..util import get_user_code_loc
+from ... import ir
+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 *
+
+_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(v16, A, B, C, opsel=opsel, loc=loc, ip=ip).result

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

@@ -375,8 +375,6 @@ 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(),
         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: