add more memref tests

Author: makslevental

examples/flash_attention.py

@@ -104,7 +104,7 @@ def manual_attn(q, k, v):
     return y
 
 
-rank_reduce = memref.MemRef.rank_reduce
+rank_reduce = memref.rank_reduce
 
 
 # https://github.com/tspeterkim/flash-attention-minimal/blob/main/flash.cu

mlir/extras/dialects/ext/memref.py

@@ -170,6 +170,9 @@ def store(
     return get_op_result_or_op_results(StoreOp(value, memref, indices, loc=loc, ip=ip))
 
 
+rank_reduce = object()
+
+
 @register_value_caster(MemRefType.static_typeid)
 @ShapedValue
 class MemRef(Value):
@@ -179,8 +182,6 @@ def __str__(self):
     def __repr__(self):
         return str(self)
 
-    rank_reduce = object()
-
     def __getitem__(self, idx: tuple) -> "MemRef":
         loc = get_user_code_loc()
 
@@ -195,9 +196,9 @@ def __getitem__(self, idx: tuple) -> "MemRef":
             return expand_shape(self, (0,), loc=loc)
 
         idx = list((idx,) if isinstance(idx, (int, Scalar, slice)) else idx)
-        rank_reduce = MemRef.rank_reduce in idx
-        if rank_reduce:
-            idx.remove(MemRef.rank_reduce)
+        should_rank_reduce = rank_reduce in idx
+        if should_rank_reduce:
+            idx.remove(rank_reduce)
 
         for i, d in enumerate(idx):
             # TODO(max): rethink this since subview and etc probably take constant attributes?
@@ -207,7 +208,7 @@ def __getitem__(self, idx: tuple) -> "MemRef":
         if all(isinstance(d, Scalar) for d in idx) and len(idx) == len(self.shape):
             return load(self, idx, loc=loc)
         else:
-            return _subview(self, tuple(idx), rank_reduce=rank_reduce, loc=loc)
+            return _subview(self, tuple(idx), rank_reduce=should_rank_reduce, loc=loc)
 
     def __setitem__(self, idx, val):
         loc = get_user_code_loc()

tests/test_memref.py

@@ -17,6 +17,7 @@
     alloca_scope,
     alloca_scope_return,
     global_,
+    rank_reduce,
 )
 from mlir.extras.dialects.ext.scf import (
     range_,
@@ -41,6 +42,10 @@ def test_simple_literal_indexing(ctx: MLIRContext):
     w = mem[2, 4, 6, 8]
     assert isinstance(w, Scalar)
 
+    two = constant(1) * 2
+    w = mem[two, 4, 6, 8]
+    mem[two, 4, 6, 8] = w
+
     correct = dedent(
         """\
     module {
@@ -50,6 +55,19 @@ def test_simple_literal_indexing(ctx: MLIRContext):
       %c6 = arith.constant 6 : index
       %c8 = arith.constant 8 : index
       %0 = memref.load %alloc[%c2, %c4, %c6, %c8] : memref<10x22x333x4444xi32>
+      %c1_i32 = arith.constant 1 : i32
+      %c2_i32 = arith.constant 2 : i32
+      %1 = arith.muli %c1_i32, %c2_i32 : i32
+      %c4_0 = arith.constant 4 : index
+      %c6_1 = arith.constant 6 : index
+      %c8_2 = arith.constant 8 : index
+      %2 = arith.index_cast %1 : i32 to index
+      %3 = memref.load %alloc[%2, %c4_0, %c6_1, %c8_2] : memref<10x22x333x4444xi32>
+      %c4_3 = arith.constant 4 : index
+      %c6_4 = arith.constant 6 : index
+      %c8_5 = arith.constant 8 : index
+      %4 = arith.index_cast %1 : i32 to index
+      memref.store %3, %alloc[%4, %c4_3, %c6_4, %c8_5] : memref<10x22x333x4444xi32>
     }
     """
     )
@@ -62,8 +80,8 @@ def test_simple_slicing(ctx: MLIRContext):
     w = mem[5:]
     w = mem[:5]
 
-    one = constant(1, index=True) * 2
-    w = mem[one:]
+    two = constant(1, index=True) * 2
+    w = mem[two:]
 
     correct = dedent(
         """\
@@ -153,9 +171,9 @@ def test_ellipsis_and_full_slice_plus_coordinate_1(ctx: MLIRContext):
     w = mem[1, :, ...]
     w = mem[1, :, :, ...]
 
-    one = constant(1, index=True) * 2
-    w = mem[one, :, :, ...]
-    w = mem[one:, :, :, ...]
+    two = constant(1, index=True) * 2
+    w = mem[two, :, :, ...]
+    w = mem[two:, :, :, ...]
 
     correct = dedent(
         f"""\
@@ -183,36 +201,43 @@ def test_ellipsis_and_full_slice_plus_coordinate_1(ctx: MLIRContext):
         w = mem[1, :, :, :, :]
     except IndexError as e:
         assert (
-                str(e)
-                == "Too many indices for shaped type with rank: 5 non-None/Ellipsis indices for dim 4."
+            str(e)
+            == "Too many indices for shaped type with rank: 5 non-None/Ellipsis indices for dim 4."
         )
 
 
-
 def test_ellipsis_and_full_slice_plus_coordinate_2(ctx: MLIRContext):
     sizes = (10, 22, 333, 4444)
     dtype_size_in_bytes = np.int32().dtype.itemsize
     golden_mem = np.zeros(sizes, dtype=np.int32)
     golden_w_1 = golden_mem[1:2, :]
+    golden_w_1_rank_reduce = golden_mem[1, :]
     golden_w_2 = golden_mem[1:2, :, :]
     golden_w_3 = golden_mem[1:2, :, :, :]
     golden_w_4 = golden_mem[:, 1:2]
     golden_w_5 = golden_mem[:, :, 1:2]
 
     golden_w_1_strides = (np.array(golden_w_1.strides) // dtype_size_in_bytes).tolist()
+    golden_w_1_rank_reduce_strides = (
+        np.array(golden_w_1_rank_reduce.strides) // dtype_size_in_bytes
+    ).tolist()
     golden_w_2_strides = (np.array(golden_w_2.strides) // dtype_size_in_bytes).tolist()
     golden_w_3_strides = (np.array(golden_w_3.strides) // dtype_size_in_bytes).tolist()
     golden_w_4_strides = (np.array(golden_w_4.strides) // dtype_size_in_bytes).tolist()
     golden_w_5_strides = (np.array(golden_w_5.strides) // dtype_size_in_bytes).tolist()
 
     golden_w_1_offset = get_np_view_offset(golden_w_1) // dtype_size_in_bytes
+    golden_w_1_rank_reduce_offset = (
+        get_np_view_offset(golden_w_1_rank_reduce) // dtype_size_in_bytes
+    )
     golden_w_2_offset = get_np_view_offset(golden_w_2) // dtype_size_in_bytes
     golden_w_3_offset = get_np_view_offset(golden_w_3) // dtype_size_in_bytes
     golden_w_4_offset = get_np_view_offset(golden_w_4) // dtype_size_in_bytes
     golden_w_5_offset = get_np_view_offset(golden_w_5) // dtype_size_in_bytes
 
     mem = alloc(sizes, T.i32())
     w = mem[1, :]
+    w = mem[1, :, rank_reduce]
     w = mem[1, :, :]
     w = mem[1, :, :, :]
     w = mem[:, 1]
@@ -224,13 +249,15 @@ def test_ellipsis_and_full_slice_plus_coordinate_2(ctx: MLIRContext):
       %c1 = arith.constant 1 : index
       %subview = memref.subview %alloc[1, 0, 0, 0] [1, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<1x22x333x4444xi32, strided<{golden_w_1_strides}, offset: {golden_w_1_offset}>>
       %c1_0 = arith.constant 1 : index
-      %subview_2 = memref.subview %alloc[1, 0, 0, 0] [1, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<1x22x333x4444xi32, strided<{golden_w_2_strides}, offset: {golden_w_2_offset}>>
+      %subview_1 = memref.subview %alloc[1, 0, 0, 0] [1, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<22x333x4444xi32, strided<{golden_w_1_rank_reduce_strides}, offset: {golden_w_1_rank_reduce_offset}>>
       %c1_2 = arith.constant 1 : index
-      %subview_3 = memref.subview %alloc[1, 0, 0, 0] [1, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<1x22x333x4444xi32, strided<{golden_w_3_strides}, offset: {golden_w_3_offset}>>
+      %subview_3 = memref.subview %alloc[1, 0, 0, 0] [1, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<1x22x333x4444xi32, strided<{golden_w_2_strides}, offset: {golden_w_2_offset}>>
       %c1_4 = arith.constant 1 : index
-      %subview_5 = memref.subview %alloc[0, 1, 0, 0] [10, 1, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x1x333x4444xi32, strided<{golden_w_4_strides}, offset: {golden_w_4_offset}>>
+      %subview_5 = memref.subview %alloc[1, 0, 0, 0] [1, 22, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<1x22x333x4444xi32, strided<{golden_w_3_strides}, offset: {golden_w_3_offset}>>
       %c1_6 = arith.constant 1 : index
-      %subview_7 = memref.subview %alloc[0, 0, 1, 0] [10, 22, 1, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x1x4444xi32, strided<{golden_w_5_strides}, offset: {golden_w_5_offset}>>
+      %subview_7 = memref.subview %alloc[0, 1, 0, 0] [10, 1, 333, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x1x333x4444xi32, strided<{golden_w_4_strides}, offset: {golden_w_4_offset}>>
+      %c1_8 = arith.constant 1 : index
+      %subview_9 = memref.subview %alloc[0, 0, 1, 0] [10, 22, 1, 4444] [1, 1, 1, 1] : memref<10x22x333x4444xi32> to memref<10x22x1x4444xi32, strided<{golden_w_5_strides}, offset: {golden_w_5_offset}>>
     }}
     """
     )
@@ -688,6 +715,9 @@ def test_memref_view(ctx: MLIRContext):
     ab_buffer = alloc(((m * k + k * n) * byte_width_dtype,), T.i8())
     a_buffer = memref.view(ab_buffer, (m, k), dtype=dtype)
     b_buffer = memref.view(ab_buffer, (k, n), dtype=dtype, shift=m * k)
+    two = constant(1) * 2
+    # TODO(max): should the type here also contain the offset...?
+    c_buffer = memref.view(ab_buffer, (k, n), dtype=dtype, shift=m * k + two)
 
     correct = dedent(
         """\
@@ -697,6 +727,15 @@ def test_memref_view(ctx: MLIRContext):
       %view = memref.view %alloc[%c0][] : memref<2048xi8> to memref<16x16xf32>
       %c1024 = arith.constant 1024 : index
       %view_0 = memref.view %alloc[%c1024][] : memref<2048xi8> to memref<16x16xf32>
+      %c1_i32 = arith.constant 1 : i32
+      %c2_i32 = arith.constant 2 : i32
+      %0 = arith.muli %c1_i32, %c2_i32 : i32
+      %c256_i32 = arith.constant 256 : i32
+      %1 = arith.addi %c256_i32, %0 : i32
+      %c4_i32 = arith.constant 4 : i32
+      %2 = arith.muli %1, %c4_i32 : i32
+      %3 = arith.index_cast %2 : i32 to index
+      %view_1 = memref.view %alloc[%3][] : memref<2048xi8> to memref<16x16xf32>
     }
     """
     )