Fix some failing tests

* vload/vstore updates that were missed previously
* int1 -> bool updates
* fix gpu target tests

Fixes a test and updates comments referencing old load/store api

Change-Id: I26a0c480d2dedee442ca0116909a7751d1dfa9ac

Author: lhutton1

src/script/printer/tir/buffer.cc

@@ -275,7 +275,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
           ExprDoc buffer = d->AsDoc<ExprDoc>(store->buffer, p->Attr("buffer"));
           ExprDoc value = d->AsDoc<ExprDoc>(store->value, p->Attr("value"));
 
-          // Use .store(...) syntax when there is a predicate
+          // Use .vstore(...) syntax when there is a predicate
           if (store->predicate.defined()) {
             ExprDoc indices = d->AsDoc<ExprDoc>(store->indices, p->Attr("indices"));
             ExprDoc predicate = d->AsDoc<ExprDoc>(store->predicate, p->Attr("predicate"));
@@ -293,7 +293,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
         "", [](tir::BufferLoad load, ObjectPath p, IRDocsifier d) -> Doc {
           ExprDoc buffer = d->AsDoc<ExprDoc>(load->buffer, p->Attr("buffer"));
 
-          // Use .load(...) syntax when there is a predicate
+          // Use .vload(...) syntax when there is a predicate
           if (load->predicate.defined()) {
             ExprDoc indices = d->AsDoc<ExprDoc>(load->indices, p->Attr("indices"));
             ExprDoc predicate = d->AsDoc<ExprDoc>(load->predicate, p->Attr("predicate"));

tests/python/codegen/test_target_codegen.py

@@ -21,7 +21,7 @@
 from tvm.script import tir as T
 
 
-@tvm.testing.exclude_targets("llvm")
+@tvm.testing.parametrize_targets("c")
 def test_buffer_store_predicate_not_supported(target):
     @T.prim_func
     def func(b: T.handle):
@@ -34,7 +34,25 @@ def func(b: T.handle):
             tvm.build(func)
 
 
-@tvm.testing.exclude_targets("llvm")
+@tvm.testing.parametrize_targets("cuda", "opencl", "metal", "rocm", "vulkan -from_device=0")
+def test_buffer_store_predicate_not_supported_gpu(target):
+    @T.prim_func
+    def func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (2, 3), "float32")
+        B = T.match_buffer(b, (6,), "float32")
+        T.func_attr({"global_symbol": "main"})
+        for i_0 in T.thread_binding(3, thread="threadIdx.x"):
+            B.vstore(
+                [T.Ramp(i_0, 1, 4)], T.Broadcast(1.0, 4), predicate=T.Broadcast(T.bool(True), 4)
+            )
+
+    err_msg = "Predicated buffer store is not supported."
+    with pytest.raises(tvm.TVMError, match=err_msg):
+        with tvm.target.Target(target):
+            tvm.build(func)
+
+
+@tvm.testing.parametrize_targets("c")
 def test_buffer_load_predicate_not_supported(target):
     @T.prim_func
     def func(a: T.handle, b: T.handle):
@@ -52,5 +70,23 @@ def func(a: T.handle, b: T.handle):
             tvm.build(func)
 
 
+@tvm.testing.parametrize_targets("cuda", "opencl", "metal", "rocm", "vulkan -from_device=0")
+def test_buffer_load_predicate_not_supported_gpu(target):
+    @T.prim_func
+    def func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (8,), "float32")
+        B = T.match_buffer(b, (8,), "float32")
+        for i_0 in T.thread_binding(3, thread="threadIdx.x"):
+            B.vstore(
+                [T.Ramp(0, 2, 4)],
+                A.vload([T.Ramp(i_0, 1, 4)], predicate=T.Broadcast(T.bool(True), 4)),
+            )
+
+    err_msg = "Predicated buffer load is not supported."
+    with pytest.raises(tvm.TVMError, match=err_msg):
+        with tvm.target.Target(target):
+            tvm.build(func)
+
+
 if __name__ == "__main__":
     tvm.testing.main()

tests/python/tir-transform/test_tir_transform_vectorize.py

@@ -186,10 +186,10 @@ def main(a: T.handle, n: T.int32, x: T.int32):
                     T.float32(1), extent
                 )
             else:
-                A.store(
-                    T.Broadcast(T.float32(2), T.vscale() * 4),
+                A.vstore(
                     [T.Ramp(0, 1, T.vscale() * 4)],
-                    predicate=T.get_active_lane_mask("int1xvscalex4", 0, n),
+                    T.Broadcast(T.float32(2), T.vscale() * 4),
+                    predicate=T.get_active_lane_mask("uint1xvscalex4", 0, n),
                 )
 
     with tvm.target.Target(target):

tests/python/tvmscript/test_tvmscript_ir_builder_tir.py

@@ -472,7 +472,7 @@ def test_ir_builder_tir_buffer_store_predicate():
     buffer_a = T.Buffer((30,), "float32")
     value = T.broadcast(0.11, T.vscale() * 4)
     index = T.ramp(0, 1, T.vscale() * 4)
-    predicate = T.broadcast(1, T.vscale() * 4)
+    predicate = T.broadcast(T.bool(True), T.vscale() * 4)
 
     with IRBuilder() as ib:
         T.buffer_store(buffer_a, value, [index], predicate)