Lunderberg commented on code in PR #16250:
URL: https://github.com/apache/tvm/pull/16250#discussion_r1439641672
##########
tests/python/tir-transform/test_tir_transform_split_host_device.py:
##########
@@ -273,5 +273,44 @@ def main_kernel():
return mod
+def test_dynamic_launch_thread():
+ """Dynamic T.launch_thread may depend on host-side variable
+
+ A dynamic parameter for `T.launch_thread` may have an extent that
+ is computed using variables outside of the `T.target` section.
+
+ This is a regression test to catch a previous failure mode, in
+ which SplitHostDevice generated output with undefined variables,
+ if the only use of a variable occurred in the extent of a
+ `T.launch_thread` statement.
+
+ While the lowering pass `LowerDeviceKernelLaunch` will hoist the
+ computation of the extent from the device kernel to the host
+ function, the IRModule must be well-defined at all stages of
+ lowering. Even if a variable is only used as part of a thread
+ extent, `SplitHostDevice` should treat it as a kernel parameter, to
+ provide a definition of the variable within the TIR device kernel.
+ """
+
+ @I.ir_module
+ class before:
+ @T.prim_func
+ def default_function(var_A: T.handle, var_B: T.handle, seq_len:
T.int32):
+ T.func_attr({"target": T.target("cuda")})
+
+ A = T.match_buffer(var_A, [seq_len], "int32")
+ B = T.match_buffer(var_B, [seq_len], "int32")
+
+ num_blocks: T.int32 = (seq_len + 127) // 128
+ with T.attr(T.target("cuda"), "target", 0):
+ blockIdx_x = T.launch_thread("blockIdx.x", num_blocks)
+ threadIdx_x = T.launch_thread("threadIdx.x", 128)
+ if blockIdx_x * 128 + threadIdx_x < seq_len:
+ B[blockIdx_x * 128 + threadIdx_x] = A[blockIdx_x * 128 +
threadIdx_x]
+
+ after = tvm.tir.transform.SplitHostDevice()(before)
Review Comment:
I didn't initially, as I've been trying to be more careful about adding
tests on structural equal, for two main reasons. First, they tend to be very
fragile to upstream changes (e.g. breakage due to improved simplification).
Second, they are difficult for a reader to identify the behavior being tested
(e.g. which changes should be made in the expected output, and which indicate a
bug). Since this test is primarily to ensure that there are no undefined
variables in the split device/host functions, that was the functionality tested.
Added for now, though I'll need to think on whether it should be there for
the long-term.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
