From 383e8d01bff21a39793f64b10c543930213ac038 Mon Sep 17 00:00:00 2001
From: meshtag <prathameshtagore@gmail.com>
Date: Mon, 29 May 2023 10:01:14 +0530
Subject: [PATCH 1/2] Add tests for store, index and access

---
 tests/dialects/test_stencil.py | 80 ++++++++++++++++++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/tests/dialects/test_stencil.py b/tests/dialects/test_stencil.py
index 4fb70b97b4..5fd474c9d4 100644
--- a/tests/dialects/test_stencil.py
+++ b/tests/dialects/test_stencil.py
@@ -1,4 +1,5 @@
 import pytest
+from xdsl.dialects import builtin
 
 from xdsl.dialects.builtin import (
     AnyFloat,
@@ -19,10 +20,13 @@
     ReturnOp,
     ResultType,
     ApplyOp,
+    StoreOp,
     TempType,
     LoadOp,
     FieldType,
     IndexAttr,
+    IndexOp,
+    AccessOp,
 )
 from xdsl.dialects.stencil import CastOp
 from xdsl.ir import Block
@@ -498,3 +502,79 @@ def test_stencil_resulttype(float_type: AnyFloat):
 
     assert isinstance(stencil_resulttype, ResultType)
     assert stencil_resulttype.elem == float_type
+
+
+def test_stencil_store():
+    temp_type = TempType([5, 5], f32)
+    temp_type_ssa_val = TestSSAValue(temp_type)
+
+    field_type = FieldType([2, 2], f32)
+    field_type_ssa_val = TestSSAValue(field_type)
+
+    lb = IndexAttr.get(1, 1)
+    ub = IndexAttr.get(64, 64)
+
+    store = StoreOp.get(temp_type_ssa_val, field_type_ssa_val, lb, ub)
+
+    assert isinstance(store, StoreOp)
+    assert isinstance(store.field.typ, FieldType)
+    assert store.field.typ == field_type
+    assert isinstance(store.temp.typ, TempType)
+    assert store.temp.typ == temp_type
+    assert len(store.field.typ.shape) == 2
+    assert len(store.temp.typ.shape) == 2
+    assert store.lb is lb
+    assert store.ub is ub
+
+
+def test_stencil_store_load_overlap():
+    temp_type = TempType([5, 5], f32)
+    temp_type_ssa_val = TestSSAValue(temp_type)
+
+    field_type = FieldType([2, 2], f32)
+    field_type_ssa_val = TestSSAValue(field_type)
+
+    lb = IndexAttr.get(1, 1)
+    ub = IndexAttr.get(64, 64)
+
+    load = LoadOp.get(field_type_ssa_val, lb, ub)
+    store = StoreOp.get(temp_type_ssa_val, field_type_ssa_val, lb, ub)
+
+    with pytest.raises(VerifyException) as exc_info:
+        load.verify()
+    assert exc_info.value.args[0] == "Cannot Load and Store the same field!"
+
+    with pytest.raises(VerifyException) as exc_info:
+        store.verify()
+    assert exc_info.value.args[0] == "Cannot Load and Store the same field!"
+
+
+def test_stencil_index():
+    dim = IntAttr(10)
+    offset = IndexAttr.get(1)
+
+    index = IndexOp.build(
+        attributes={
+            "dim": dim,
+            "offset": offset,
+        },
+        result_types=[builtin.IndexType()],
+    )
+
+    assert isinstance(index, IndexOp)
+    assert index.dim is dim
+    assert index.offset is offset
+
+
+def test_stencil_access():
+    temp_type = TempType([5, 5], f32)
+    temp_type_ssa_val = TestSSAValue(temp_type)
+
+    offset = [1, 1]
+    offset_index_attr = IndexAttr.get(*offset)
+
+    access = AccessOp.get(temp_type_ssa_val, offset)
+
+    assert isinstance(access, AccessOp)
+    assert access.offset == offset_index_attr
+    assert access.temp.typ == temp_type

From f6ffb04d363207439d7e2330fcc57bbe0d869488 Mon Sep 17 00:00:00 2001
From: meshtag <prathameshtagore@gmail.com>
Date: Mon, 29 May 2023 10:03:11 +0530
Subject: [PATCH 2/2] Remove unnecessary include

---
 tests/dialects/test_stencil.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/dialects/test_stencil.py b/tests/dialects/test_stencil.py
index 5fd474c9d4..53a1aa9f78 100644
--- a/tests/dialects/test_stencil.py
+++ b/tests/dialects/test_stencil.py
@@ -1,5 +1,4 @@
 import pytest
-from xdsl.dialects import builtin
 
 from xdsl.dialects.builtin import (
     AnyFloat,
@@ -15,6 +14,7 @@
     i64,
     IntegerType,
     ArrayAttr,
+    IndexType,
 )
 from xdsl.dialects.experimental.stencil import (
     ReturnOp,
@@ -558,7 +558,7 @@ def test_stencil_index():
             "dim": dim,
             "offset": offset,
         },
-        result_types=[builtin.IndexType()],
+        result_types=[IndexType()],
     )
 
     assert isinstance(index, IndexOp)