[Relax] Move TIR backend to gpu_generic

python/tvm/relax/backend/dispatch_sampling.py

@@ -36,7 +36,7 @@ def visit_call_(self, call: relax.Call) -> relax.Expr:
             return super().visit_call_(call)
 
         if call.op.name == "relax.multinomial_from_uniform":
-            from tvm.relax.backend_tir import (  # pylint: disable=import-outside-toplevel
+            from tvm.relax.backend.gpu_generic import (  # pylint: disable=import-outside-toplevel
                 generic_get_sample_index,
                 gpu_multinomial_from_uniform,
             )

python/tvm/relax/backend/dispatch_sort_scan.py

@@ -141,7 +141,7 @@ def visit_call_(self, call: relax.Call) -> relax.Expr:
                 and call.op.name == "relax.cumsum"
                 and call.attrs.exclusive == 0
             ):
-                from tvm.relax.backend_tir import (  # pylint: disable=import-outside-toplevel
+                from tvm.relax.backend.gpu_generic import (  # pylint: disable=import-outside-toplevel
                     gpu_2d_continuous_cumsum,
                 )
 

python/tvm/relax/backend/gpu_generic/__init__.py

@@ -15,10 +15,12 @@
 # specific language governing permissions and limitations
 # under the License.
 """The Relax Metal backend compilation pipeline and other passes."""
+from .cumsum import gpu_2d_continuous_cumsum
 from .pipeline import (
+    dataflow_lower_passes,
     finalize_passes,
     get_default_pipeline,
     legalize_passes,
-    dataflow_lower_passes,
     library_dispatch_passes,
 )
+from .sampling import generic_get_sample_index, gpu_multinomial_from_uniform