Fix memref (#76)

Author: makslevental

mlir/extras/dialects/ext/memref.py

@@ -108,7 +108,7 @@ def __getitem__(self, idx: tuple) -> "MemRef":
         if idx is None:
             return expand_shape(self, (0,), loc=loc)
 
-        idx = list((idx,) if isinstance(idx, int) else idx)
+        idx = list((idx,) if isinstance(idx, (int, slice)) else idx)
         for i, d in enumerate(idx):
             if isinstance(d, int):
                 idx[i] = constant(d, index=True, loc=loc)
@@ -291,6 +291,9 @@ def global_(
         sym_name = _get_sym_name(
             previous_frame, check_func_call="memref\\.global_|global_"
         )
+        assert (
+            sym_name is not None
+        ), "couldn't automatically find sym_name in previous frame"
     if loc is None:
         loc = get_user_code_loc()
     if initial_value is None:

tests/test_memref.py

@@ -56,6 +56,24 @@ def test_simple_literal_indexing(ctx: MLIRContext):
     filecheck(correct, ctx.module)
 
 
+def test_simple_slicing(ctx: MLIRContext):
+    mem = alloc(10, T.i32())
+
+    w = mem[5:]
+    w = mem[:5]
+
+    correct = dedent(
+        """\
+    module {
+      %alloc = memref.alloc() : memref<10xi32>
+      %subview = memref.subview %alloc[5] [5] [1] : memref<10xi32> to memref<5xi32, strided<[1], offset: 5>>
+      %subview_0 = memref.subview %alloc[0] [5] [1] : memref<10xi32> to memref<5xi32>
+    }
+    """
+    )
+    filecheck(correct, ctx.module)
+
+
 def test_simple_literal_indexing_alloca(ctx: MLIRContext):
     @alloca_scope([])
     def demo_scope2():