memref.global wrapper (#64)

Author: makslevental

mlir/extras/context.py

@@ -1,4 +1,5 @@
 import contextlib
+import warnings
 from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass
 from typing import Optional
@@ -58,7 +59,7 @@ def __del__(self):
         self.context.__exit__(None, None, None)
         # i guess the extension gets destroyed before this object sometimes?
         if ir is not None:
-            assert ir.Context.current is None, str(ir.Context.current)
+            assert ir.Context is not self.context
 
 
 class ExplicitlyManagedModule:

mlir/extras/dialects/ext/memref.py

@@ -1,16 +1,24 @@
+import inspect
 from functools import cached_property, reduce
 from typing import Tuple, Sequence, Union
 
+import numpy as np
+
 from .arith import Scalar, constant
 from .tensor import _indices_to_indexer, compute_result_shape_reassoc_list
 from ... import types as T
 from ...meta import region_op
-from ...util import get_user_code_loc, _unpack_sizes_element_type
+from ...util import (
+    get_user_code_loc,
+    _unpack_sizes_element_type,
+    _get_sym_name,
+    infer_mlir_type,
+)
 from ...._mlir_libs._mlir import register_value_caster
 from ....dialects import memref, arith
 from ....dialects._ods_common import get_op_result_or_op_results
 from ....dialects.memref import *
-from ....ir import Type, Value, MemRefType, ShapedType
+from ....ir import Type, Value, MemRefType, ShapedType, DenseElementsAttr
 
 S = ShapedType.get_dynamic_size()
 
@@ -286,3 +294,44 @@ def dim(source, index, *, loc=None, ip=None):
     if isinstance(index, int):
         index = constant(index, index=True)
     return _dim(source=source, index=index, loc=loc, ip=ip)
+
+
+def global_(
+    initial_value=None,
+    sym_name=None,
+    type_=None,
+    sym_visibility="private",
+    constant=None,
+    alignment=None,
+    loc=None,
+    ip=None,
+):
+    if sym_name is None:
+        previous_frame = inspect.currentframe().f_back
+        sym_name = _get_sym_name(
+            previous_frame, check_func_call="memref\.global_|global_"
+        )
+    if loc is None:
+        loc = get_user_code_loc()
+    if initial_value is None:
+        assert type_ is not None
+    else:
+        assert isinstance(initial_value, np.ndarray)
+        type_ = infer_mlir_type(initial_value, memref=True)
+        initial_value = DenseElementsAttr.get(
+            initial_value,
+            type=type_.element_type,
+            context=None,
+        )
+        constant = True
+
+    return memref.global_(
+        sym_name,
+        type_,
+        sym_visibility=sym_visibility,
+        initial_value=initial_value,
+        constant=constant,
+        alignment=alignment,
+        loc=loc,
+        ip=ip,
+    ).opview

mlir/extras/util.py

@@ -2,6 +2,7 @@
 import ctypes
 import inspect
 import platform
+import re
 import sys
 import warnings
 from dataclasses import dataclass
@@ -23,6 +24,7 @@
     Location,
     OpView,
     Operation,
+    MemRefType,
     RankedTensorType,
     Value,
     _GlobalDebug,
@@ -165,7 +167,7 @@ def mlir_type_to_ctype(mlir_type):
 
 
 def infer_mlir_type(
-    py_val: Union[int, float, bool, np.ndarray]
+    py_val: Union[int, float, bool, np.ndarray], memref=False
 ) -> Union[IntegerType, F32Type, F64Type, RankedTensorType]:
     """Infer MLIR type (`ir.Type`) from supported python values.
 
@@ -202,7 +204,10 @@ def infer_mlir_type(
             return T.f64()
     elif isinstance(py_val, np.ndarray):
         dtype = np_dtype_to_mlir_type(py_val.dtype.type)
-        return RankedTensorType.get(py_val.shape, dtype)
+        if memref:
+            return MemRefType.get(py_val.shape, dtype)
+        else:
+            return RankedTensorType.get(py_val.shape, dtype)
     else:
         raise NotImplementedError(
             f"Unsupported Python value {py_val=} with type {type(py_val)}"
@@ -356,8 +361,8 @@ def _get_sym_name(previous_frame, check_func_call=None):
             src_lines = src_file.readlines()
             src_line = src_lines[previous_frame.f_lineno - 1].strip()
             ident, func_call = map(lambda x: x.strip(), src_line.split("=", maxsplit=1))
-            if check_func_call is None:
-                assert check_func_call in func_call
+            if check_func_call is not None:
+                assert re.match(check_func_call, func_call)
         return ident
     except:
         return None

tests/test_memref.py

@@ -1,29 +1,31 @@
-import numpy as np
+import platform
 import re
 from textwrap import dedent
 
+import numpy as np
 import pytest
-from mlir.ir import MLIRError, Type
 
 import mlir.extras.types as T
+from mlir.dialects.memref import subview
 from mlir.extras.ast.canonicalize import canonicalize
+from mlir.extras.dialects.ext import memref
 from mlir.extras.dialects.ext.arith import Scalar, constant
 from mlir.extras.dialects.ext.memref import (
     alloc,
     alloca,
-    S,
     alloca_scope,
     alloca_scope_return,
+    global_,
 )
 from mlir.extras.dialects.ext.scf import (
     range_,
     yield_,
     canonicalizer,
 )
-from mlir.dialects.memref import subview
 
 # noinspection PyUnresolvedReferences
 from mlir.extras.testing import mlir_ctx as ctx, filecheck, MLIRContext
+from mlir.ir import MLIRError, Type
 
 # needed since the fix isn't defined here nor conftest.py
 pytest.mark.usefixtures("ctx")
@@ -581,3 +583,63 @@ def tenfoo():
     )
 
     filecheck(correct, ctx.module)
+
+
+@pytest.mark.skipif(
+    platform.system() == "Windows",
+    reason="On windows int64 is inferred to be i64 ",
+)
+def test_memref_global_windows(ctx: MLIRContext):
+    k = 32
+    weight1 = global_(np.ones((k,), dtype=np.int32))
+    weight2 = global_(np.ones((k,), dtype=np.int64))
+    weight3 = global_(np.ones((k,), dtype=np.float32))
+    weight4 = global_(np.ones((k,), dtype=np.float64))
+    weight5 = memref.global_(np.ones((k,), dtype=np.int16))
+    weight6 = memref.global_(np.ones((k,), dtype=np.float16))
+    print(ctx.module)
+
+    correct = dedent(
+        """\
+    module {
+      memref.global "private" constant @weight1 : memref<32xi32> = dense<1>
+      memref.global "private" constant @weight2 : memref<32xi64> = dense<1>
+      memref.global "private" constant @weight3 : memref<32xf32> = dense<1.000000e+00>
+      memref.global "private" constant @weight4 : memref<32xf64> = dense<1.000000e+00>
+      memref.global "private" constant @weight5 : memref<32xi16> = dense<1>
+      memref.global "private" constant @weight6 : memref<32xf16> = dense<1.000000e+00>
+    }
+    """
+    )
+
+    filecheck(correct, ctx.module)
+
+
+@pytest.mark.skipif(
+    platform.system() != "Windows",
+    reason="On linux/mac int64 is inferred to be index (through np.longlong)",
+)
+def test_memref_global_non_windows(ctx: MLIRContext):
+    k = 32
+    weight1 = global_(np.ones((k,), dtype=np.int32))
+    weight2 = global_(np.ones((k,), dtype=np.int64))
+    weight3 = global_(np.ones((k,), dtype=np.float32))
+    weight4 = global_(np.ones((k,), dtype=np.float64))
+    weight5 = memref.global_(np.ones((k,), dtype=np.int16))
+    weight6 = memref.global_(np.ones((k,), dtype=np.float16))
+    print(ctx.module)
+
+    correct = dedent(
+        """\
+    module {
+      memref.global "private" constant @weight1 : memref<32xi32> = dense<1>
+      memref.global "private" constant @weight2 : memref<32xindex> = dense<1>
+      memref.global "private" constant @weight3 : memref<32xf32> = dense<1.000000e+00>
+      memref.global "private" constant @weight4 : memref<32xf64> = dense<1.000000e+00>
+      memref.global "private" constant @weight5 : memref<32xi16> = dense<1>
+      memref.global "private" constant @weight6 : memref<32xf16> = dense<1.000000e+00>
+    }
+    """
+    )
+
+    filecheck(correct, ctx.module)