diff --git a/python/tvm/relax/backend/cpu_generic/pipeline.py b/python/tvm/relax/backend/cpu_generic/pipeline.py index dc078ee25d68..d0b819cea7f8 100644 --- a/python/tvm/relax/backend/cpu_generic/pipeline.py +++ b/python/tvm/relax/backend/cpu_generic/pipeline.py @@ -22,7 +22,10 @@ def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unused-argument """The default library dispatch passes for CPU backend.""" - return [] + return [ + relax.backend.DispatchSampling(), + relax.backend.DispatchSortScan(), + ] def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py index 68592d67f870..adc0f7ad8383 100644 --- a/python/tvm/relax/vm_build.py +++ b/python/tvm/relax/vm_build.py @@ -248,7 +248,21 @@ def _extract_attrs(mod: tvm.IRModule): if relax_pipeline is not None: if isinstance(relax_pipeline, str): - relax_pipeline = relax.get_pipeline(relax_pipeline) + # For GPU targets, prefer the target-specific pipeline which + # includes DLight scheduling. Without it, TIR functions generated + # from ops like Clip/ReLU6 lack thread bindings and fail + # VerifyMemory. CPU targets continue to use the generic pipeline + # since the CPU-specific pipeline applies fusion passes that can + # incorrectly remove call_pure_packed calls whose results are + # unused but whose side effects are relied upon. + _is_gpu = target is not None and "gpu" in target.keys + if relax_pipeline == "default" and _is_gpu: + try: + relax_pipeline = relax.get_default_pipeline(target) + except (ValueError, AttributeError): + relax_pipeline = relax.get_pipeline(relax_pipeline) + else: + relax_pipeline = relax.get_pipeline(relax_pipeline) if target is None: mod = relax_pipeline(mod) else: