Closed MellowArtisan closed 9 hours ago
Hi all, I found the tvm.transform.Sequential([relax.transform.RealizeVDevice()])(mod)
can remove the "R.hint_on_device()", however, using relax.transform.RealizeVDevice()(mod)
directly cannot remove it.
@Lunderberg @tqchen Why do the different usages for the same pass (i.e., RealizeVDevice) give different optimization results?
Hmm. It looks like it's even weirder. It looks like only the first use of RealizeVDevice
produces the correct output. The first execution correctly removes the R.hint_on_device
, but the second execution does not.
relax.transform.RealizeVDevice()(Module).show(name='First')
relax.transform.RealizeVDevice()(Module).show(name='Second')
Aha! The problem is that HintOnDeviceRemover
(the first step of RealizeVDevice
) is mutating the relax expression in-place, which is not legal. As a result, expressions that are in the input Module
are being mutated to have different StructInfo
. The second time that RealizeVDevice
is applied, its input has been mutated to already include vdevice
annotations.
## Running these commands
Module["foo"].show(name="Before")
relax.transform.RealizeVDevice()(Module)
Module["foo"].show(name="OrigAfter")
## Produces the following output
@R.function
def Before(
x: R.Tensor((2, 3), dtype="float32"),
y: R.Tensor((2, 3), dtype="float32"),
z: R.Tensor((2, 3), dtype="float32"),
) -> R.Tensor((2, 3), dtype="float32"):
R.func_attr({"global_symbol": "foo"})
with R.dataflow():
lv0: R.Tensor((2, 3), dtype="float32") = R.hint_on_device(y, R.device(dev_type=1, dev_id=0))
R.output(lv0)
return lv0
@R.function
def OrigAfter(
x: R.Tensor((2, 3), dtype="float32"),
y: R.Tensor((2, 3), dtype="float32", vdevice="llvm:-1"),
z: R.Tensor((2, 3), dtype="float32"),
) -> R.Tensor((2, 3), dtype="float32"):
R.func_attr({"global_symbol": "foo"})
with R.dataflow():
lv0: R.Tensor((2, 3), dtype="float32", vdevice="llvm:-1") = R.hint_on_device(
y, R.device(dev_type=1, dev_id=0)
)
R.output(lv0)
return lv0
The input module should never be modified when running any IRModule
transform, so this definitely narrows the bug down to the RealizeVDevice
implementation.
@Cookiee235 Can you verify the fix implemented in #17213? It removes the in-place mutation from RealizeVDevice
, and should resolve this issue.
Edit: Whoops, meant @MellowArtisan . With multiple issues/PRs in-flight, I got them mixed up.
@Cookiee235 Can you verify the fix implemented in #17213? It removes the in-place mutation from
RealizeVDevice
, and should resolve this issue.
@Lunderberg Yes! This PR fixed the bug. Thanks for your efforts!
Actual behavior
Environment
TVM: 0.17.dev0
Steps to reproduce
cc @junrushao