This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 2393db61bb [Relax] Return values a match_cast reads from a CUDA graph
capture (#20485)
2393db61bb is described below
commit 2393db61bbd0100f048b1a6d3c2519fb70914564
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Sep 29 19:05:42 2026 -0700
[Relax] Return values a match_cast reads from a CUDA graph capture (#20485)
RewriteCUDAGraph adds a lifted value to the capture function's outputs
when a binding outside the region uses it. The planner does that for
var, call, tuple and tuple_get_item bindings but never overrode the
match_cast visitor, so a match_cast over a lifted value kept reading a
var that no longer exists in the original function. Compiling
Phi-3.5-vision for CUDA in MLC LLM hits this in the image embedding,
where a static reshape is captured and then match_cast to the symbolic
crop grid, and VM codegen fails with `Var lv383 is not defined`.
1. Add a `MatchCastNode` visitor to the planner that marks the value as
an output of the region that lifted it
2. Add a test with a captured reshape followed by a match_cast and a
launch that depends on it
---
src/relax/transform/rewrite_cuda_graph.cc | 11 +++
.../relax/test_transform_rewrite_cuda_graph.py | 87 ++++++++++++++++++++++
2 files changed, 98 insertions(+)
diff --git a/src/relax/transform/rewrite_cuda_graph.cc
b/src/relax/transform/rewrite_cuda_graph.cc
index ab654772e4..3074510af7 100644
--- a/src/relax/transform/rewrite_cuda_graph.cc
+++ b/src/relax/transform/rewrite_cuda_graph.cc
@@ -456,6 +456,17 @@ class CUDAGraphRewritePlanner : public ExprVisitor {
AddStaticBinding(binding, false);
}
+ void VisitBinding_(const MatchCastNode* binding) final {
+ // A match_cast stays in the original function even when its value was
lifted into the
+ // current capture region, so the region has to return that value.
+ if (const auto* var = binding->value.as<VarNode>()) {
+ if (auto it = binding_to_region_.find(var); it !=
binding_to_region_.end()) {
+ it->second->MarkOutput(var);
+ }
+ }
+ ExprVisitor::VisitBinding_(binding);
+ }
+
void VisitBinding_(const VarBindingNode* binding, const TupleNode* tuple)
final {
std::vector<const VarNode*> args;
std::vector<PrimVar> tir_vars;
diff --git a/tests/python/relax/test_transform_rewrite_cuda_graph.py
b/tests/python/relax/test_transform_rewrite_cuda_graph.py
index f7af28d745..cf953f6f28 100644
--- a/tests/python/relax/test_transform_rewrite_cuda_graph.py
+++ b/tests/python/relax/test_transform_rewrite_cuda_graph.py
@@ -1190,3 +1190,90 @@ def test_static_input_with_symbolic_shape():
if __name__ == "__main__":
tvm.testing.main()
+
+
+def test_match_cast_of_captured_value():
+ m = T.dynamic("m")
+
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor((8,), "float16")):
+ R.func_attr({"relax.force_pure": True, "num_input": 1})
+ storage1 = R.memory.alloc_storage(R.shape([8]), 0, "global",
"float16")
+ alloc1 = R.memory.alloc_tensor(storage1, 0, R.shape([8]),
"float16")
+ _ = R.call_packed("dummy", x, alloc1, ty_args=(R.Tuple,))
+ storage2 = R.memory.alloc_storage(R.shape([8]), 0, "global",
"float16")
+ alloc2 = R.memory.alloc_tensor(storage2, 0, R.shape([8]),
"float16")
+ _1 = R.call_packed("dummy", alloc1, alloc2, ty_args=(R.Tuple,))
+ lv = R.reshape(alloc2, R.shape([2, 4]))
+ lv1 = R.match_cast(lv, R.Tensor((m, 4), "float16"))
+ storage3 = R.memory.alloc_storage(R.shape([8]), 0, "global",
"float16")
+ alloc3 = R.memory.alloc_tensor(storage3, 0, R.shape([8]),
"float16")
+ _2 = R.call_packed("dummy", lv1, alloc3, ty_args=(R.Tuple,))
+ gv = (alloc3,)
+ return gv
+
+ @I.ir_module
+ class Expected:
+ @R.function(private=True)
+ def cuda_graph_alloc() -> R.Tuple(R.Any, R.Any):
+ R.func_attr({"relax.force_pure": True})
+ storage1: R.Any = R.memory.alloc_storage(
+ R.shape([8]), R.prim_value(0), R.str("global"),
R.dtype("float16")
+ )
+ storage2: R.Any = R.memory.alloc_storage(
+ R.shape([8]), R.prim_value(0), R.str("global"),
R.dtype("float16")
+ )
+ gv: R.Tuple(R.Any, R.Any) = storage1, storage2
+ return gv
+
+ @R.function(private=True)
+ def main_cuda_graph_capture(
+ alloc1: R.Tensor((8,), dtype="float16"), alloc2: R.Tensor((8,),
dtype="float16")
+ ) -> R.Tuple(R.Tensor((2, 4), dtype="float16")):
+ R.func_attr({"relax.force_pure": True})
+ R.call_packed("dummy", alloc1, alloc2, ty_args=(R.Tuple,))
+ lv: R.Tensor((2, 4), dtype="float16") = R.reshape(alloc2,
R.shape([2, 4]))
+ gv: R.Tuple(R.Tensor((2, 4), dtype="float16")) = (lv,)
+ return gv
+
+ @R.function
+ def main(x: R.Tensor((8,), dtype="float16")) -> R.Tuple(R.Tensor((8,),
dtype="float16")):
+ R.func_attr({"num_input": 1, "relax.force_pure": True})
+ cls = Expected
+ gv: R.Tuple(R.Any, R.Any) = R.call_builtin_with_ctx(
+ "vm.builtin.cuda_graph.get_cached_alloc",
+ (cls.cuda_graph_alloc, R.prim_value(0)),
+ ty_args=(R.Tuple(R.Any, R.Any),),
+ )
+ storage1: R.Any = gv[0]
+ alloc1: R.Tensor((8,), dtype="float16") = R.memory.alloc_tensor(
+ storage1, R.prim_value(0), R.shape([8]), R.dtype("float16")
+ )
+ R.call_packed("dummy", x, alloc1, ty_args=(R.Tuple,))
+ storage2: R.Any = gv[1]
+ alloc2: R.Tensor((8,), dtype="float16") = R.memory.alloc_tensor(
+ storage2, R.prim_value(0), R.shape([8]), R.dtype("float16")
+ )
+ gv1: R.Tuple(R.Tensor((2, 4), dtype="float16")) =
R.call_builtin_with_ctx(
+ "vm.builtin.cuda_graph.run_or_capture",
+ (cls.main_cuda_graph_capture, (alloc1, alloc2),
R.prim_value(0)),
+ ty_args=(R.Tuple(R.Tensor((2, 4), dtype="float16")),),
+ )
+ lv: R.Tensor((2, 4), dtype="float16") = gv1[0]
+ lv1: R.Tensor((m, 4), dtype="float16") = R.match_cast(
+ lv, R.Tensor((m, 4), dtype="float16")
+ )
+ storage3: R.Any = R.memory.alloc_storage(
+ R.shape([8]), R.prim_value(0), R.str("global"),
R.dtype("float16")
+ )
+ alloc3: R.Tensor((8,), dtype="float16") = R.memory.alloc_tensor(
+ storage3, R.prim_value(0), R.shape([8]), R.dtype("float16")
+ )
+ R.call_packed("dummy", lv1, alloc3, ty_args=(R.Tuple,))
+ gv_1: R.Tuple(R.Tensor((8,), dtype="float16")) = (alloc3,)
+ return gv_1
+
+ after = relax.transform.RewriteCUDAGraph()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
