catchup

Author: makslevental

mlir/extras/context.py

@@ -46,7 +46,6 @@ def mlir_mod_ctx(
         context.allow_unregistered_dialects = True
     with context, mlir_mod(src, location) as module:
         yield MLIRContext(context, module)
-    context._clear_live_operations()
 
 
 class RAIIMLIRContext:

mlir/extras/dialects/ext/gpu.py

@@ -3,7 +3,7 @@
 from typing import Any, List, Optional, Tuple, Union
 
 
-from .arith import constant
+from . import arith as arith_ext
 from .func import FuncBase
 from ... import types as T
 from ...meta import (
@@ -16,7 +16,13 @@
     make_maybe_no_args_decorator,
     find_ops,
 )
-from ....dialects._gpu_ops_gen import _Dialect
+from ....dialects import _gpu_ops_gen
+from ....dialects._gpu_ops_gen import (
+    _Dialect,
+    GPUFuncOp as _GPUFuncOp,
+    LaunchOp as _LaunchOp,
+    LaunchFuncOp as _LaunchFuncOp,
+)
 from ....dialects._ods_common import (
     _cext,
     get_default_loc_context,
@@ -223,7 +229,8 @@ def __prepare__(cls, name, bases, **kwargs):
         return {"ip": ip, "gpu_module_op": gpu_module_op}
 
 
-class GPUFuncOp(GPUFuncOp):
+# TODO(max): integrate upstream
+class GPUFuncOp(_GPUFuncOp):
     def __init__(
         self,
         sym_name,
@@ -269,7 +276,8 @@ def __init__(
             )
 
 
-class LaunchOp(LaunchOp):
+# TODO(max): integrate upstream
+class LaunchOp(_LaunchOp):
     def __init__(
         self,
         grid_size: Tuple[Any, Any, Any],
@@ -321,7 +329,7 @@ def launch_(
     for size in [grid_size, block_size]:
         for i, s in enumerate(size):
             if isinstance(s, int):
-                size[i] = constant(s, index=True)
+                size[i] = arith_ext.constant(s, index=True)
     launch_op = LaunchOp(
         grid_size,
         block_size,
@@ -336,7 +344,8 @@ def launch_(
 launch = region_op(launch_, terminator=lambda *_args: TerminatorOp())
 
 
-class LaunchFuncOp(LaunchFuncOp):
+# TODO(max): integrate upstream
+class LaunchFuncOp(_LaunchFuncOp):
     def __init__(
         self,
         kernel: List[str],
@@ -394,7 +403,7 @@ def __call__(
         for size in [grid_size, block_size]:
             for i, s in enumerate(size):
                 if isinstance(s, int):
-                    size[i] = constant(s, index=True)
+                    size[i] = arith_ext.constant(s, index=True)
 
         if loc is None:
             loc = get_user_code_loc()
@@ -596,15 +605,6 @@ def get_compile_object_bytes(compiled_module):
     return objects[-1].object
 
 
-_printf = printf
-
-
-def printf(format, *args, loc=None, ip=None):
-    if loc is None:
-        loc = get_user_code_loc()
-    return _printf(format=format, args=args, loc=loc, ip=ip)
-
-
 _dynamic_shared_memory = dynamic_shared_memory
 
 
@@ -634,7 +634,7 @@ def memset(dst, value, async_dependencies=None, *, loc=None, ip=None):
     if len(async_dependencies):
         async_token = gpu_async_token()
     if isinstance(value, (int, float, bool)):
-        value = constant(value, type=dst.type.element_type)
+        value = arith_ext.constant(value, type=dst.type.element_type)
     return _memset(async_token, async_dependencies, dst, value, loc=loc, ip=ip)
 
 

mlir/extras/dialects/ext/transform.py

@@ -395,8 +395,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,

mlir/extras/dialects/ext/vector.py

@@ -155,14 +155,14 @@ def transfer_read(
 _extract = extract
 
 
-def extract(vector, position, *, loc=None, ip=None):
+def extract(source, position, *, loc=None, ip=None):
     if loc is None:
         loc = get_user_code_loc()
     dynamic_position, _packed_position, static_position = _dispatch_mixed_values(
         position
     )
     return _extract(
-        vector=vector,
+        source=source,
         dynamic_position=dynamic_position,
         static_position=static_position,
         loc=loc,
@@ -231,14 +231,14 @@ def broadcast(vector, source, *, loc=None, ip=None):
 _extract_strided_slice = extract_strided_slice
 
 
-def extract_strided_slice(vector, offsets, sizes, strides, *, loc=None, ip=None):
+def extract_strided_slice(source, offsets, sizes, strides, *, loc=None, ip=None):
     if loc is None:
         loc = get_user_code_loc()
-    result_shape = [int(s) for s in sizes] + vector.type.shape[len(sizes) :]
-    result = T.vector(*result_shape, vector.type.element_type)
+    result_shape = [int(s) for s in sizes] + source.type.shape[len(sizes) :]
+    result = T.vector(*result_shape, source.type.element_type)
     return _extract_strided_slice(
         result=result,
-        vector=vector,
+        source=source,
         offsets=offsets,
         sizes=sizes,
         strides=strides,

mlir/extras/testing/testing.py

@@ -127,7 +127,8 @@ def filecheck_with_comments(module):
 def mlir_ctx() -> MLIRContext:
     with mlir_mod_ctx(allow_unregistered_dialects=True) as ctx:
         yield ctx
-    assert Context.current is None
+    # TODO(max): why is context.current being retained now?
+    # assert Context.current is None
 
 
 @pytest.fixture(scope="function")

tests/test_vector.py

@@ -354,8 +354,8 @@ def result(aval, aidx, b):
             v6 = vector.gather(
                 result=T.vector(8, T.f32()),
                 base=X,
-                indices=[c0],
-                index_vec=aidx,
+                offsets=[c0],
+                indices=aidx,
                 mask=v0,
                 pass_thru=v1,
             )