[GitHub] [tvm] Lunderberg commented on a diff in pull request #11532: [Bugfix][TIR] Handle boolean tensor properly in Flatten-Buffer

[GitBox <gi...@apache.org>]

Lunderberg commented on code in PR #11532:
URL: https://github.com/apache/tvm/pull/11532#discussion_r887974135


##########
tests/python/unittest/test_tir_transform_flatten_buffer.py:
##########
@@ -268,6 +268,33 @@ def annotated_loops(a: T.handle) -> None:
         A[i] = 0.0
 
 
+@T.prim_func
+def boolean_handling_before(a: T.Buffer[10, "bool"], b: T.Buffer[10, "bool"]) -> None:
+    for i0 in T.serial(10):
+        with T.block("b"):
+            T.reads(a[i0])
+            T.writes(b[i0])
+            b[i0] = a[i0]
+
+
+@T.prim_func
+def boolean_handling_after(a: T.Buffer[10, "int8"], b: T.Buffer[10, "int8"]) -> None:
+    T.preflattened_buffer(a, [10], dtype="bool", data=a.data)
+    T.preflattened_buffer(b, [10], dtype="bool", data=b.data)
+    # body
+    for i0 in T.serial(10):
+        b[i0] = T.cast(T.cast(a[i0], "bool"), "int8")
+
+
+@T.prim_func
+def boolean_handle_after(a: T.Buffer[10, "int8"], b: T.Buffer[10, "int8"]) -> None:
+    T.preflattened_buffer(a, [10], dtype="bool", data=a.data)
+    T.preflattened_buffer(b, [10], dtype="bool", data=b.data)
+    # body
+    for i0 in T.serial(10):
+        b[i0] = T.cast(T.cast(a[i0], "bool"), "int8")

Review Comment:
   I think it could be simplified, and would need to be done in the mutator for `BufferStoreNode`.  In the general case, the `int8` to `bool` to `int8` sequence of casts can't be simplified out as a no-op, because it does reduce all nonzero values to one.  In this case, since `b` also represents a boolean, it would be allowed, so that would be where this access type could be checked for.
   
   Long-term, I'd like to move all of the boolean-specific handling into a dedicated lowering pass.  As it is, there's quite a lot of special cases in multiple different locations which all need to know that boolean arrays are backed by int8 arrays.  (And even that should be target-specific, as not all Vulkan implementations support int8 datatypes.)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org