Address comments

- Correct doc strings
- Correct typo in error message
- Add some additional checks for BufferLoad

Change-Id: Ie25563d569c0ed729ac915a6ba3a724a9e191014

Author: lhutton1

include/tvm/tir/buffer.h

@@ -210,7 +210,7 @@ class Buffer : public ObjectRef {
    * \param begin The beginning index
    * \param dtype The data type to be loaded.
    * \param predicate A vector mask of boolean values indicating which lanes of a vector are to be
-   * stored. The number lanes of the mask must be equal to the number of lanes in value.
+   * loaded. The number lanes of the mask must be equal to the number of lanes in being loaded.
    */
   TVM_DLL PrimExpr vload(Array<PrimExpr> begin, DataType dtype,
                          Optional<PrimExpr> predicate = NullOpt) const;

python/tvm/tir/buffer.py

@@ -115,8 +115,7 @@ def vload(self, begin, dtype=None, predicate=None):
 
         predicate : Optional[PrimExpr]
             A vector mask of boolean values indicating which lanes of a vector are to be
-            stored. The number lanes of the mask must be equal to the number of lanes in
-            value.
+            loaded. The number lanes of the mask must be equal to the number of lanes being loaded.
 
         Returns
         -------

python/tvm/tir/expr.py

@@ -1100,8 +1100,7 @@ class BufferLoad(PrimExprWithOp):
 
     predicate : Optional[PrimExpr]
         A vector mask of boolean values indicating which lanes of a vector are to be
-        stored. The number lanes of the mask must be equal to the number of lanes in
-        value.
+        loaded. The number lanes of the mask must be equal to the number of lanes being loaded.
     """
 
     buffer: Buffer

src/target/llvm/codegen_llvm.h

@@ -331,8 +331,8 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
    * \param indices The indices at which the buffer is being accessed.
    *
    * \param predicate A vector mask of boolean values indicating which lanes of a
-   * vector are to be stored. The number lanes of the mask must be equal to the
-   * number of lanes in value.
+   * vector are to be accessed. The number lanes of the mask must be equal to the
+   * number of lanes being accessed.
    *
    * \param value_dtype The datatype to be read from (BufferLoad) or
    * written to (BufferStore) the buffer.

src/tir/ir/expr.cc

@@ -780,7 +780,21 @@ BufferLoad::BufferLoad(Buffer buffer, Array<PrimExpr> indices, Optional<PrimExpr
       << "-dimensional indices provided.";
 
   if (predicate.defined()) {
-    DataType predicate_element_dtype = predicate.value().dtype().element_of();
+    DataType predicate_dtype = predicate.value().dtype();
+
+    bool is_index_scalable = indices.empty() ? false : indices.back().dtype().is_scalable_vector();
+    bool is_predicate_scalable = predicate_dtype.is_scalable_vector();
+    ICHECK_EQ(is_index_scalable, is_predicate_scalable)
+        << "Predicate mask dtype and load indices must both be scalable.";
+
+    int index_lanes = indices.empty() ? 1 : indices.back().dtype().get_lanes_or_vscale_factor();
+    int predicate_lanes = predicate_dtype.get_lanes_or_vscale_factor();
+    ICHECK_EQ(index_lanes, predicate_lanes)
+        << "Got a predicate mask with " << predicate_lanes
+        << " lanes, but trying to load a vector with " << index_lanes
+        << " lanes. The number of lanes must match.";
+
+    DataType predicate_element_dtype = predicate_dtype.element_of();
     ICHECK(predicate_element_dtype.is_bool())
         << "Predicate mask elements must be boolean values, but got " << predicate_element_dtype
         << ".";

src/tir/transforms/inject_rolling_buffer.cc

@@ -257,8 +257,8 @@ class RollingBufferInjector : public StmtExprMutator {
           indices.push_back(index);
         }
       }
-      ICHECK(!op->predicate.defined())
-          << "Predicated buffer store is not current supported in the inject rolling buffer pass.";
+      ICHECK(!op->predicate.defined()) << "Predicated buffer store is not currently supported in "
+                                          "the inject rolling buffer pass.";
       Stmt buffer_store = BufferStore(op->buffer, op->value, indices, op->predicate, op->span);
       // Then wrap the BufferStores in some Ifs to avoid recomputing elements
       for (size_t i{0}; i < rolling_buffer_info.axis_iter_vars.size(); ++i) {

tests/python/tir-base/test_tir_nodes.py

@@ -514,6 +514,29 @@ def test_buffer_load_predicate_elements_invalid_type():
         tvm.tir.BufferLoad(b, [index], predicate)
 
 
+def test_buffer_store_predicate_invalid_scalability():
+    b = tvm.tir.decl_buffer((24,), "int32")
+    index = tvm.tir.expr.Ramp(0, 1, 4 * tvm.tir.vscale())
+    predicate = tvm.tir.expr.Broadcast(tvm.tir.IntImm("int1", 1), 4)
+
+    err_msg = "Predicate mask dtype and load indices must both be scalable."
+    with pytest.raises(tvm.TVMError, match=err_msg):
+        tvm.tir.BufferLoad(b, [index], predicate)
+
+
+def test_buffer_store_predicate_invalid_lanes():
+    b = tvm.tir.decl_buffer((24,), "int32")
+    index = tvm.tir.expr.Ramp(0, 1, 4 * tvm.tir.vscale())
+    predicate = tvm.tir.expr.Broadcast(tvm.tir.IntImm("int1", 1), 8 * tvm.tir.vscale())
+
+    err_msg = (
+        "Got a predicate mask with 8 lanes, but trying to load a "
+        "vector with 4 lanes. The number of lanes must match."
+    )
+    with pytest.raises(tvm.TVMError, match=err_msg):
+        tvm.tir.BufferLoad(b, [index], predicate)
+
+
 def test_scalable_vec_cast():
     b = tvm.tir.decl_buffer((24,), "float32")
     value = tvm.tir.expr.Broadcast(1, 12 * tvm.tir.vscale()).astype("float32xvscalex12")