Bug-fix] Fix tir allocation with multiple lanes

This PR stemmed from #6907
and it is fixing a small error in the getter and setter of a buffer for
the case where `t.lanes > 1`. I also added a test to stress the issue.

python/tvm/tir/ir_builder.py

@@ -103,7 +103,7 @@ def __getitem__(self, index):
         index = self._linear_index(index)
         if t.lanes > 1:
             base = index * t.lanes
-            index = _expr.Ramp(base, const(1, base.dtype), t.lanes)
+            index = _expr.Ramp(base, 1, t.lanes)
         return _expr.Load(self._content_type, self._buffer_var, index)
 
     def __setitem__(self, index, value):
@@ -116,7 +116,7 @@ def __setitem__(self, index, value):
         t = DataType(self._content_type)
         if t.lanes > 1:
             base = index * t.lanes
-            index = _expr.Ramp(base, const(1, base.dtype), t.lanes)
+            index = _expr.Ramp(base, 1, t.lanes)
         self._builder.emit(_stmt.Store(self._buffer_var, value, index))
 
 

tests/python/unittest/test_tir_ir_builder.py

@@ -20,6 +20,21 @@
 import tvm.testing
 
 
+def test_allocate_with_lanes():
+    ib = tvm.tir.ir_builder.create()
+    n = te.size_var("n")
+    A = ib.allocate("float32x4", n, name="A", scope="global")
+    B = ib.allocate("float32x4", n, name="A", scope="global")
+    A[0] = tvm.tir.const(0, "float32x4")
+    B[0] = A[0]
+    body = ib.get()
+    assert A == A
+    print(body)
+    assert isinstance(body, tvm.tir.AttrStmt)
+    body = body.body
+    assert isinstance(body, tvm.tir.Allocate)
+
+
 def test_for():
     ib = tvm.tir.ir_builder.create()
     n = te.size_var("n")
@@ -177,5 +192,6 @@ def check_target(target):
     test_prefetch()
     test_if()
     test_for()
+    test_allocate_with_lanes()
     test_cpu()
     test_gpu()