[Test] Update test cases for bit slicing (cornell-zhang#443)

chhzh123 · Mar 16, 2022 · c8d9ac7 · c8d9ac7
1 parent 6ea2e58
commit c8d9ac7
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 12 deletions.
diff --git a/tests/mlir/test_codegen_vhls.py b/tests/mlir/test_codegen_vhls.py
@@ -25,7 +25,7 @@ def test_set_bit():
     harness.test_set_bit(target, "[4] = 1")
 
 def test_set_slice():
-    harness.test_set_slice(target, "A[0](4, 1) = 1")
+    harness.test_set_slice(target, "(4, 1) = 1")
 
 def test_pack():
 

diff --git a/tests/mlir/test_dsl_basic.py b/tests/mlir/test_dsl_basic.py
@@ -457,7 +457,7 @@ def test_get_slice_expr():
     hcl.init()
 
     def kernel(A):
-        return hcl.compute(A.shape, lambda x: (A[x] + 1)[2:0])
+        return hcl.compute(A.shape, lambda x: (A[x] + 1)[0:2])
 
     A = hcl.placeholder((10,))
     s = hcl.create_schedule(A, kernel)
@@ -479,7 +479,7 @@ def test_get_slice_tensor():
     hcl.init()
 
     def kernel(A):
-        return hcl.compute(A.shape, lambda x: A[x][2:0])
+        return hcl.compute(A.shape, lambda x: A[x][0:2])
 
     A = hcl.placeholder((10,))
     s = hcl.create_schedule(A, kernel)
@@ -501,7 +501,7 @@ def test_get_slice_tensor_reverse():
     hcl.init()
 
     def kernel(A):
-        return hcl.compute(A.shape, lambda x: A[x][0:8])
+        return hcl.compute(A.shape, lambda x: (A[x][0:8]).reverse())
 
     A = hcl.placeholder((10,))
     s = hcl.create_schedule(A, kernel)
@@ -531,7 +531,7 @@ def test_set_slice_expr():
 
     def kernel(A, B):
         with hcl.for_(0, 10) as i:
-            (B[i]+1)[2:0] = A[i]
+            (B[i]+1)[0:2] = A[i]
 
     A = hcl.placeholder((10,))
     B = hcl.placeholder((10,))
@@ -548,7 +548,7 @@ def test_set_slice_tensor():
 
     def kernel(A, B):
         with hcl.for_(0, 10) as i:
-            B[i][2:0] = A[i]
+            B[i][0:2] = A[i]
 
     A = hcl.placeholder((10,))
     B = hcl.placeholder((10,))
@@ -572,7 +572,7 @@ def test_set_slice_tensor_reverse():
 
     def kernel(A, B):
         with hcl.for_(0, 10) as i:
-            B[i][0:8] = A[i]
+            B[i][0:8] = A[i].reverse()
 
     A = hcl.placeholder((10,))
     B = hcl.placeholder((10,))
@@ -601,7 +601,7 @@ def test_slice_op():
     hcl.init()
 
     def kernel(A):
-        return hcl.compute(A.shape, lambda x: A[x][8:0] + A[x][16:8])
+        return hcl.compute(A.shape, lambda x: A[x][0:8] + A[x][8:16])
 
     A = hcl.placeholder((10,))
     s = hcl.create_schedule(A, kernel)

diff --git a/tests/mlir/test_simplify.py b/tests/mlir/test_simplify.py
@@ -6,14 +6,13 @@ def test_remove_single_loop():
     b = hcl.compute(a.shape, lambda x: a[x] + 1)
     s = hcl.create_schedule([a, b])
     ir = hcl.lower(s)
-    assert "for (x, 0, 1)" not in str(ir)
+    assert "0 to 1" not in str(ir)
 
 def test_simplify_slice():
     hcl.init()
     A = hcl.placeholder((10,), "A")
     def kernel(A):
-        with hcl.Stage():
-            A[5][2:2] = 4
+        A[5][2:2] = 4
     s = hcl.create_schedule(A, kernel)
     ir = hcl.lower(s)
-    assert "2:2" not in str(ir)
+    assert "set_slice" not in str(ir)