This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx/standalone-package-cleanup in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 78d2f04415d988a3f82da73dc1e4ee252d033371 Author: Tianqi Chen <[email protected]> AuthorDate: Thu Oct 15 09:05:08 2026 +0000 [FIX][Docs] Resolve extracted TIRx API namespaces Point existing API references and tutorial imports at the extracted core package while retaining parent Relax and scheduling namespaces. Include the core Python path in documentation builds and update the existing alias resolution maps. --- docs/conf.py | 15 +++--- docs/deep_dive/relax/tutorials/relax_creation.py | 8 +-- .../relax/tutorials/relax_transformation.py | 10 ++-- .../tensor_ir/tutorials/dlight_gpu_scheduling.py | 11 ++-- .../deep_dive/tensor_ir/tutorials/meta_schedule.py | 8 +-- docs/deep_dive/tensor_ir/tutorials/tir_creation.py | 19 +++---- .../tensor_ir/tutorials/tir_transformation.py | 19 +++---- docs/get_started/tutorials/ir_module.py | 17 +++--- docs/get_started/tutorials/quick_start.py | 9 ++-- docs/how_to/tutorials/bring_your_own_codegen.py | 55 ++++++++++---------- docs/how_to/tutorials/cross_compilation_and_rpc.py | 15 +++--- docs/how_to/tutorials/customize_opt.py | 26 ++++++---- docs/how_to/tutorials/e2e_opt_model.py | 9 ++-- .../how_to/tutorials/export_and_load_executable.py | 11 ++-- docs/how_to/tutorials/import_model.py | 7 +-- .../tutorials/mix_python_and_tvm_with_pymodule.py | 22 ++++---- docs/how_to/tutorials/optimize_llm.py | 31 ++++++----- docs/reference/api/python/codegen.rst | 6 +-- docs/reference/api/python/error.rst | 6 +-- docs/reference/api/python/instrument.rst | 6 +-- docs/reference/api/python/ir.rst | 6 +-- docs/reference/api/python/runtime/runtime.rst | 6 +-- docs/reference/api/python/script/ir_builder.rst | 10 ++-- docs/reference/api/python/script/parser.rst | 44 ++++++++-------- docs/reference/api/python/script/printer.rst | 6 +-- docs/reference/api/python/script/script.rst | 40 +++++++-------- docs/reference/api/python/support.rst | 60 +++++++++++----------- docs/reference/api/python/sym.rst | 6 +-- docs/reference/api/python/target.rst | 12 ++--- docs/reference/api/python/testing.rst | 6 +-- docs/reference/api/python/transform.rst | 6 +-- docs/tirx/api/analysis.rst | 2 +- docs/tirx/api/backend.rst | 16 +++--- docs/tirx/api/compilation.rst | 12 ++--- docs/tirx/api/cuda.rst | 28 +++++----- docs/tirx/api/cuda_compile.rst | 10 ++-- docs/tirx/api/execution.rst | 12 ++--- docs/tirx/api/layout.rst | 22 ++++---- docs/tirx/api/ptx.rst | 8 +-- docs/tirx/api/script.rst | 26 +++++----- docs/tirx/api/tirx.rst | 2 +- docs/tirx/api/trainium.rst | 12 ++--- docs/tirx/api/transform.rst | 2 +- 43 files changed, 344 insertions(+), 320 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 599e03ca38..361356a95e 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -51,6 +51,7 @@ else: tvm_path = Path(os.pardir) sys.path.insert(0, str(tvm_path.resolve() / "python")) +sys.path.insert(0, str(tvm_path.resolve() / "tirx" / "python")) sys.path.insert(0, str(tvm_path.resolve() / "docs")) # -- General configuration ------------------------------------------------ @@ -315,7 +316,7 @@ language = None # directories to ignore when looking for source files. exclude_patterns = ["_build", "_staging"] -# The TIRx API pages autodoc modules (tvm.tirx, tvm.backend.cuda) that re-export +# The TIRx API pages autodoc modules (tvm_tirx.tirx, tvm_tirx.backend.cuda) that re-export # common IR types (PrimExpr, Op, ...) from several modules, which makes a handful # of autodoc'd cross references ambiguous. Silence that specific, benign category. suppress_warnings = ["ref.python"] @@ -549,8 +550,8 @@ autodoc_default_options = { # Maps the original namespace to list of potential modules # that we can import alias from. tvm_alias_check_map = { - "tvm.te": ["tvm.tirx"], - "tvm.tirx": ["tvm.ir", "tvm.runtime"], + "tvm.te": ["tvm_tirx.tirx"], + "tvm_tirx.tirx": ["tvm_tirx.ir", "tvm_tirx.runtime"], } @@ -594,15 +595,15 @@ def update_alias_docstring(name, obj, lines): tvm_class_name_rewrite_map = { - "tvm.tirx": ["Var", "Call"], + "tvm_tirx.tirx": ["Var", "Call"], "tvm.relax": ["Var", "Call", "StringImm"], "tvm.relax.frontend.nn": ["Module"], } # When documenting modules under these prefixes, prefer types from the mapped module -# to resolve ambiguous cross-references (e.g. Var exists in both tvm.tirx and tvm.relax). +# to resolve ambiguous cross-references (e.g. Var exists in both tvm_tirx.tirx and tvm.relax). tvm_module_type_preference = { - "tvm.s_tir": "tvm.tirx", + "tvm.s_tir": "tvm_tirx.tirx", } @@ -715,7 +716,7 @@ def _patch_python_domain_find_obj(): return best_matches # Check module type preference for cross-module resolution - # (e.g. tvm.s_tir.analysis uses types from tvm.tirx). + # (e.g. tvm.s_tir.analysis uses types from tvm_tirx.tirx). for prefix, preferred_mod in tvm_module_type_preference.items(): if modname.startswith(prefix): preferred = [m for m in matches if m[0].startswith(preferred_mod + ".")] diff --git a/docs/deep_dive/relax/tutorials/relax_creation.py b/docs/deep_dive/relax/tutorials/relax_creation.py index 3d884e7942..166026596b 100644 --- a/docs/deep_dive/relax/tutorials/relax_creation.py +++ b/docs/deep_dive/relax/tutorials/relax_creation.py @@ -38,10 +38,10 @@ and relax NNModule API. from tvm import relax, topi from tvm.relax import ExternFunc as _ExternFunc -from tvm.script import ir as I -from tvm.script import relax as R -from tvm.script import s_tir as Ts -from tvm.script import tirx as T +from tvm.relax import script as R +from tvm.s_tir import script as Ts +from tvm_tirx.script import ir as I +from tvm_tirx.script import tirx as T n = T.dynamic("n", "int64") diff --git a/docs/deep_dive/relax/tutorials/relax_transformation.py b/docs/deep_dive/relax/tutorials/relax_transformation.py index 8abaa2d75e..49ff316d19 100644 --- a/docs/deep_dive/relax/tutorials/relax_transformation.py +++ b/docs/deep_dive/relax/tutorials/relax_transformation.py @@ -31,8 +31,10 @@ for optimizing and integrating with hardware backends. # the :ref:`previous section <relax-creation>`. import tvm -from tvm import IRModule, relax +import tvm_tirx +from tvm import relax from tvm.relax.frontend import nn +from tvm_tirx import IRModule class NNModule(nn.Module): @@ -73,7 +75,7 @@ mod.show() # in ML compilers. Note that in relax, fusion optimizations are done with the collaboration of # a set of passes. We can apply them in a sequence. -mod = tvm.ir.transform.Sequential( +mod = tvm_tirx.ir.transform.Sequential( [ tvm.relax.transform.AnnotateTIROpPattern(), tvm.relax.transform.FuseOps(), @@ -115,9 +117,9 @@ class ReluRewriter(PyExprMutator): # Then we can write a pass to apply the mutator to the whole module. [email protected]_pass(opt_level=0, name="ReluToGelu") +@tvm_tirx.ir.transform.module_pass(opt_level=0, name="ReluToGelu") class ReluToGelu: # pylint: disable=too-few-public-methods - def transform_module(self, mod: IRModule, _ctx: tvm.ir.transform.PassContext) -> IRModule: + def transform_module(self, mod: IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> IRModule: """IRModule-level transformation""" rewriter = ReluRewriter(mod) for g_var, func in mod.functions_items(): diff --git a/docs/deep_dive/tensor_ir/tutorials/dlight_gpu_scheduling.py b/docs/deep_dive/tensor_ir/tutorials/dlight_gpu_scheduling.py index 37c6f2d61c..1a2c2ea7cf 100644 --- a/docs/deep_dive/tensor_ir/tutorials/dlight_gpu_scheduling.py +++ b/docs/deep_dive/tensor_ir/tutorials/dlight_gpu_scheduling.py @@ -46,10 +46,11 @@ quality, and how to write custom rules. # rules: ``Linear`` layers produce GEMM (matrix multiplication) kernels, ``LayerNorm`` # produces a general-reduction kernel, and ``ReLU`` is a simple elementwise op. -import tvm -from tvm import relax, tirx +import tvm_tirx +from tvm import relax from tvm.relax.frontend import nn from tvm.s_tir import dlight as dl +from tvm_tirx import tirx class DemoModel(nn.Module): @@ -70,8 +71,8 @@ mod, params = DemoModel().export_tvm({"forward": {"x": nn.spec.Tensor((1, 768), ###################################################################### # Legalize Relax operators into TIR functions so that DLight has concrete kernels to schedule. -device = tvm.cuda(0) -target = tvm.target.Target.from_device(device) +device = tvm_tirx.cuda(0) +target = tvm_tirx.target.Target.from_device(device) with target: mod = relax.get_pipeline("zero")(mod) @@ -263,7 +264,7 @@ from tvm.s_tir.dlight.base.schedule_rule import ScheduleRule @ScheduleRule.from_callable("MyTileAndBind") -def my_tile_and_bind(func: tirx.Function, target: tvm.target.Target, tunable: bool): +def my_tile_and_bind(func: tirx.Function, target: tvm_tirx.target.Target, tunable: bool): """A minimal rule: for single-block injective functions, tile and bind to GPU threads.""" if not isinstance(func, tirx.Function): return None diff --git a/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py b/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py index 584a05bcf2..030e96cbb3 100644 --- a/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py +++ b/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py @@ -78,7 +78,7 @@ the space of possible schedules. import os import tempfile -import tvm +import tvm_tirx from tvm import relax from tvm.relax.frontend import nn @@ -100,8 +100,8 @@ class DemoModel(nn.Module): input_shape = (1, 784) mod, params = DemoModel().export_tvm({"forward": {"x": nn.spec.Tensor(input_shape, "float32")}}) -device = tvm.cuda(0) -target = tvm.target.Target.from_device(device) +device = tvm_tirx.cuda(0) +target = tvm_tirx.target.Target.from_device(device) ###################################################################### # User-Facing Entry Points @@ -129,7 +129,7 @@ target = tvm.target.Target.from_device(device) if os.getenv("CI", "") != "true": with target, tempfile.TemporaryDirectory() as tmp_dir: - tuned_mod = tvm.ir.transform.Sequential( + tuned_mod = tvm_tirx.ir.transform.Sequential( [ relax.get_pipeline("zero"), relax.transform.MetaScheduleTuneTIR( diff --git a/docs/deep_dive/tensor_ir/tutorials/tir_creation.py b/docs/deep_dive/tensor_ir/tutorials/tir_creation.py index a96f0cb45f..15ab2f0a43 100644 --- a/docs/deep_dive/tensor_ir/tutorials/tir_creation.py +++ b/docs/deep_dive/tensor_ir/tutorials/tir_creation.py @@ -56,9 +56,10 @@ import numpy as np import tvm_ffi import tvm -from tvm.script import ir as I -from tvm.script import s_tir as Ts -from tvm.script import tirx as T +import tvm_tirx +from tvm.s_tir import script as Ts +from tvm_tirx.script import ir as I +from tvm_tirx.script import tirx as T @I.ir_module @@ -205,10 +206,10 @@ class DynamicShapeModule: # Now let's check the runtime dynamic shape inference: -def evaluate_dynamic_shape(lib: tvm.runtime.Module, m: int, n: int, k: int): - A = tvm.runtime.tensor(np.random.uniform(size=(m, k)).astype("float32")) - B = tvm.runtime.tensor(np.random.uniform(size=(k, n)).astype("float32")) - C = tvm.runtime.tensor(np.zeros((m, n), dtype="float32")) +def evaluate_dynamic_shape(lib: tvm_tirx.runtime.Module, m: int, n: int, k: int): + A = tvm_tirx.runtime.tensor(np.random.uniform(size=(m, k)).astype("float32")) + B = tvm_tirx.runtime.tensor(np.random.uniform(size=(k, n)).astype("float32")) + C = tvm_tirx.runtime.tensor(np.zeros((m, n), dtype="float32")) lib(A, B, C) return C.numpy() @@ -265,7 +266,7 @@ C = te.compute((128, 128), lambda i, j: te.max(Y[i, j], 0), name="C") # and one output parameter **C**. te_func = te.create_function([A, B, C]).with_attr({"global_symbol": "mm_relu"}) -TEModule = tvm.IRModule({"mm_relu": te_func}) +TEModule = tvm_tirx.IRModule({"mm_relu": te_func}) TEModule.show() ###################################################################### @@ -283,5 +284,5 @@ Y = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y") C = te.compute((M, N), lambda i, j: te.max(Y[i, j], 0), name="C") dyn_te_func = te.create_function([A, B, C]).with_attr({"global_symbol": "mm_relu"}) -DynamicTEModule = tvm.IRModule({"mm_relu": dyn_te_func}) +DynamicTEModule = tvm_tirx.IRModule({"mm_relu": dyn_te_func}) DynamicTEModule.show() diff --git a/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py b/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py index 6f53a1029c..7707419348 100644 --- a/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py +++ b/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py @@ -37,9 +37,10 @@ transformations of primitive tensor functions. # First, let's take a look at the implementation of ``mm_relu`` in the previous section: import tvm -from tvm.script import ir as I -from tvm.script import s_tir as Ts -from tvm.script import tirx as T +import tvm_tirx +from tvm.s_tir import script as Ts +from tvm_tirx.script import ir as I +from tvm_tirx.script import tirx as T @I.ir_module @@ -77,18 +78,18 @@ a_np = np.random.uniform(size=(128, 128)).astype("float32") b_np = np.random.uniform(size=(128, 128)).astype("float32") c_np = a_np @ b_np -a_nd = tvm.runtime.tensor(a_np) -b_nd = tvm.runtime.tensor(b_np) -c_nd = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32")) +a_nd = tvm_tirx.runtime.tensor(a_np) +b_nd = tvm_tirx.runtime.tensor(b_np) +c_nd = tvm_tirx.runtime.tensor(np.zeros((128, 128), dtype="float32")) -def evaluate(mod: tvm.IRModule): - lib = tvm.tirx.compile(mod, target="llvm").mod +def evaluate(mod: tvm_tirx.IRModule): + lib = tvm_tirx.tirx.compile(mod, target="llvm").mod # check correctness lib(a_nd, b_nd, c_nd) np.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-5) # evaluate performance - f_timer = lib.time_evaluator("main", tvm.cpu()) + f_timer = lib.time_evaluator("main", tvm_tirx.cpu()) print(f_timer(a_nd, b_nd, c_nd)) diff --git a/docs/get_started/tutorials/ir_module.py b/docs/get_started/tutorials/ir_module.py index beb5e7ed76..4219bb2ebb 100644 --- a/docs/get_started/tutorials/ir_module.py +++ b/docs/get_started/tutorials/ir_module.py @@ -42,6 +42,7 @@ from torch import nn from torch.export import export import tvm +import tvm_tirx from tvm import relax from tvm.relax.frontend.torch import from_exported_program @@ -118,8 +119,8 @@ mod_from_relax.show() # directly output the IRModule in the TVMScript syntax, or alternatively, # parse the TVMScript to obtain an IRModule. -from tvm.script import ir as I -from tvm.script import relax as R +from tvm.relax import script as R +from tvm_tirx.script import ir as I @I.ir_module @@ -234,11 +235,11 @@ mod.show() # We can deploy the IRModule on CPU by specifying the target as ``llvm``. exec = tvm.compile(mod, target="llvm") -dev = tvm.cpu() +dev = tvm_tirx.cpu() vm = relax.VirtualMachine(exec, dev) raw_data = np.random.rand(1, 784).astype("float32") -data = tvm.runtime.tensor(raw_data, dev) +data = tvm_tirx.runtime.tensor(raw_data, dev) cpu_out = vm["main"](data, *params_from_torch["main"]).numpy() print(cpu_out) @@ -255,7 +256,7 @@ print(cpu_out) from tvm.s_tir import dlight as dl -with tvm.target.Target("cuda"): +with tvm_tirx.target.Target("cuda"): gpu_mod = dl.ApplyDefaultSchedule( dl.gpu.Matmul(), dl.gpu.Fallback(), @@ -265,11 +266,11 @@ with tvm.target.Target("cuda"): # Now we can compile the IRModule on GPU, the similar way as we did on CPU. exec = tvm.compile(gpu_mod, target="cuda") -dev = tvm.cuda(0) +dev = tvm_tirx.cuda(0) vm = relax.VirtualMachine(exec, dev) # Need to allocate data and params on GPU device -data = tvm.runtime.tensor(raw_data, dev) -gpu_params = [tvm.runtime.tensor(p, dev) for p in params_from_torch["main"]] +data = tvm_tirx.runtime.tensor(raw_data, dev) +gpu_params = [tvm_tirx.runtime.tensor(p, dev) for p in params_from_torch["main"]] gpu_out = vm["main"](data, *gpu_params).numpy() print(gpu_out) diff --git a/docs/get_started/tutorials/quick_start.py b/docs/get_started/tutorials/quick_start.py index a1de5686a9..5f6ecd2919 100644 --- a/docs/get_started/tutorials/quick_start.py +++ b/docs/get_started/tutorials/quick_start.py @@ -80,6 +80,7 @@ to show how to use Apache TVM to compile a simple neural network. # import tvm +import tvm_tirx from tvm import relax from tvm.relax.frontend import nn @@ -137,14 +138,14 @@ mod = relax.get_pipeline("zero")(mod) import numpy as np -target = tvm.target.Target("llvm") +target = tvm_tirx.target.Target("llvm") ex = tvm.compile(mod, target) -device = tvm.cpu() +device = tvm_tirx.cpu() vm = relax.VirtualMachine(ex, device) data = np.random.rand(1, 784).astype("float32") -tvm_data = tvm.runtime.tensor(data, device=device) +tvm_data = tvm_tirx.runtime.tensor(data, device=device) params = [np.random.rand(*param.shape).astype("float32") for _, param in param_spec] -params = [tvm.runtime.tensor(param, device=device) for param in params] +params = [tvm_tirx.runtime.tensor(param, device=device) for param in params] print(vm["forward"](tvm_data, *params).numpy()) ################################################################################ diff --git a/docs/how_to/tutorials/bring_your_own_codegen.py b/docs/how_to/tutorials/bring_your_own_codegen.py index 96d94b3806..de1cebdbba 100644 --- a/docs/how_to/tutorials/bring_your_own_codegen.py +++ b/docs/how_to/tutorials/bring_your_own_codegen.py @@ -80,24 +80,24 @@ import tempfile import numpy as np -import tvm -import tvm.relax.backend.contrib.example_npu +import tvm.relax.backend.contrib.example_npu # noqa: F401 +import tvm_tirx from tvm import relax +from tvm.relax import script as R from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt from tvm.relax.backend.pattern_registry import get_patterns_with_prefix from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen -from tvm.script import relax as R -has_example_npu_codegen = tvm.get_global_func("relax.ext.example_npu", True) -has_example_npu_runtime = tvm.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True) +has_example_npu_codegen = tvm_tirx.get_global_func("relax.ext.example_npu", True) +has_example_npu_runtime = tvm_tirx.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True) has_example_npu = has_example_npu_codegen and has_example_npu_runtime -has_tensorrt_codegen = tvm.get_global_func("relax.ext.tensorrt", True) is not None -_is_trt_runtime_enabled = tvm.get_global_func("relax.is_tensorrt_runtime_enabled", True) +has_tensorrt_codegen = tvm_tirx.get_global_func("relax.ext.tensorrt", True) is not None +_is_trt_runtime_enabled = tvm_tirx.get_global_func("relax.is_tensorrt_runtime_enabled", True) has_tensorrt = ( has_tensorrt_codegen and _is_trt_runtime_enabled is not None and _is_trt_runtime_enabled() ) -has_cuda = tvm.cuda(0).exist +has_cuda = tvm_tirx.cuda(0).exist ###################################################################### # Step 2: Define the model @@ -107,7 +107,7 @@ has_cuda = tvm.cuda(0).exist # backends. [email protected]_module +@tvm_tirx.script.ir_module class ConvReLU: @R.function def main( @@ -166,12 +166,13 @@ weight_np = np.random.randn(16, 3, 3, 3).astype("float32") if has_example_npu: npu_mod = RunCodegen()(npu_mod) - with tvm.ir.transform.PassContext(opt_level=3): - npu_exec = relax.build(npu_mod, tvm.target.Target("llvm")) + with tvm_tirx.ir.transform.PassContext(opt_level=3): + npu_exec = relax.build(npu_mod, tvm_tirx.target.Target("llvm")) - npu_vm = relax.VirtualMachine(npu_exec, tvm.cpu()) + npu_vm = relax.VirtualMachine(npu_exec, tvm_tirx.cpu()) npu_out = npu_vm["main"]( - tvm.runtime.tensor(data_np, tvm.cpu()), tvm.runtime.tensor(weight_np, tvm.cpu()) + tvm_tirx.runtime.tensor(data_np, tvm_tirx.cpu()), + tvm_tirx.runtime.tensor(weight_np, tvm_tirx.cpu()), ) assert npu_out.numpy().shape == (1, 16, 30, 30) print("Example NPU run completed. Output shape:", npu_out.numpy().shape) @@ -206,17 +207,19 @@ print(trt_mod) # Build for CUDA, run on the GPU, and compare against the CPU reference. if has_tensorrt and has_cuda: - dev = tvm.cuda(0) - with tvm.ir.transform.PassContext(opt_level=3): + dev = tvm_tirx.cuda(0) + with tvm_tirx.ir.transform.PassContext(opt_level=3): trt_exec = relax.build(RunCodegen()(trt_mod), "cuda") - trt_out = relax.VirtualMachine(trt_exec, dev)["main"](tvm.runtime.tensor(data_np, dev)).numpy() + trt_out = relax.VirtualMachine(trt_exec, dev)["main"]( + tvm_tirx.runtime.tensor(data_np, dev) + ).numpy() cpu_mod = relax.transform.LegalizeOps()( relax.transform.BindParams("main", {"weight": weight_np})(ConvReLU) ) cpu_exec = relax.build(cpu_mod, "llvm") - cpu_out = relax.VirtualMachine(cpu_exec, tvm.cpu())["main"]( - tvm.runtime.tensor(data_np, tvm.cpu()) + cpu_out = relax.VirtualMachine(cpu_exec, tvm_tirx.cpu())["main"]( + tvm_tirx.runtime.tensor(data_np, tvm_tirx.cpu()) ).numpy() np.testing.assert_allclose(trt_out, cpu_out, rtol=1e-2, atol=1e-2) @@ -235,12 +238,12 @@ if has_tensorrt and has_cuda: fp16_mod = partition_for_tensorrt( relax.transform.BindParams("main", {"weight": weight_np})(ConvReLU) ) - with tvm.ir.transform.PassContext( + with tvm_tirx.ir.transform.PassContext( opt_level=3, config={"relax.ext.tensorrt.options": {"use_fp16": True}} ): fp16_exec = relax.build(RunCodegen()(fp16_mod), "cuda") - fp16_out = relax.VirtualMachine(fp16_exec, tvm.cuda(0))["main"]( - tvm.runtime.tensor(data_np, tvm.cuda(0)) + fp16_out = relax.VirtualMachine(fp16_exec, tvm_tirx.cuda(0))["main"]( + tvm_tirx.runtime.tensor(data_np, tvm_tirx.cuda(0)) ).numpy() np.testing.assert_allclose(fp16_out, cpu_out, rtol=5e-2, atol=5e-2) @@ -312,11 +315,11 @@ if has_torch and has_tensorrt and has_cuda: print("After importing and partitioning the PyTorch model:") print(torch_mod) - torch_dev = tvm.cuda(0) - with tvm.ir.transform.PassContext(opt_level=3): + torch_dev = tvm_tirx.cuda(0) + with tvm_tirx.ir.transform.PassContext(opt_level=3): torch_exec = relax.build(RunCodegen()(torch_mod), "cuda") deployed = relax.VirtualMachine(torch_exec, torch_dev)["main"]( - tvm.runtime.tensor(example_input.numpy(), torch_dev) + tvm_tirx.runtime.tensor(example_input.numpy(), torch_dev) )[0].numpy() np.testing.assert_allclose(deployed, torch_ref, rtol=1e-2, atol=1e-2) @@ -331,9 +334,9 @@ if has_torch and has_tensorrt and has_cuda: with tempfile.TemporaryDirectory() as tmpdir: lib_path = os.path.join(tmpdir, "deployed_trt.so") torch_exec.export_library(lib_path) - loaded = tvm.runtime.load_module(lib_path) + loaded = tvm_tirx.runtime.load_module(lib_path) reran = relax.VirtualMachine(loaded, torch_dev)["main"]( - tvm.runtime.tensor(example_input.numpy(), torch_dev) + tvm_tirx.runtime.tensor(example_input.numpy(), torch_dev) )[0].numpy() np.testing.assert_allclose(reran, torch_ref, rtol=1e-2, atol=1e-2) print("Reloaded the exported library and reran; output", reran.shape, "still matches.") diff --git a/docs/how_to/tutorials/cross_compilation_and_rpc.py b/docs/how_to/tutorials/cross_compilation_and_rpc.py index 9fc425c7a9..f218582b82 100644 --- a/docs/how_to/tutorials/cross_compilation_and_rpc.py +++ b/docs/how_to/tutorials/cross_compilation_and_rpc.py @@ -98,13 +98,14 @@ import numpy as np import tvm_ffi import tvm +import tvm_tirx from tvm import rpc, te -from tvm.support import utils +from tvm_tirx.support import utils -n = tvm.runtime.convert(1024) +n = tvm_tirx.runtime.convert(1024) A = te.placeholder((n,), name="A") B = te.compute((n,), lambda i: A[i] + 1.0, name="B") -mod = tvm.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "add_one")) +mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "add_one")) ###################################################################### # Then we cross compile the kernel. @@ -416,10 +417,10 @@ def run_opencl(): # them according to your environment. opencl_device_host = "10.77.1.145" opencl_device_port = 9090 - target = tvm.target.Target("opencl", host={"kind": "llvm", "mtriple": "aarch64-linux-gnu"}) + target = tvm_tirx.target.Target("opencl", host={"kind": "llvm", "mtriple": "aarch64-linux-gnu"}) # create schedule for the above "add one" compute declaration - mod = tvm.IRModule.from_expr(te.create_function([A, B])) + mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B])) sch = tvm.s_tir.Schedule(mod) (x,) = sch.get_loops(block=sch.get_sblock("B")) xo, xi = sch.split(x, [None, 32]) @@ -538,7 +539,7 @@ def run_pytorch_model_via_rpc(): if local_demo: # For demonstration on local machine, use local target - target = tvm.target.Target("llvm") + target = tvm_tirx.target.Target("llvm") print("Using local target for demonstration") else: # Choose the appropriate target for your device: @@ -560,7 +561,7 @@ def run_pytorch_model_via_rpc(): # - OpenCL: tvm.target.Target("opencl", host={"kind": "llvm", "mtriple": "aarch64-linux-gnu"}) # # For this example, we use ARM 64-bit - target = tvm.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu"}) + target = tvm_tirx.target.Target({"kind": "llvm", "mtriple": "aarch64-linux-gnu"}) print(f"Cross-compiling for target: {target}") # Apply optimization pipeline diff --git a/docs/how_to/tutorials/customize_opt.py b/docs/how_to/tutorials/customize_opt.py index e6e83c33c4..fc52128e6a 100644 --- a/docs/how_to/tutorials/customize_opt.py +++ b/docs/how_to/tutorials/customize_opt.py @@ -55,8 +55,10 @@ import tempfile import numpy as np import tvm -from tvm import IRModule, relax +import tvm_tirx +from tvm import relax from tvm.relax.frontend import nn +from tvm_tirx import IRModule ###################################################################### # Composable IRModule Optimization @@ -119,11 +121,11 @@ except ImportError as e: # Define a new pass for CUBLAS dispatch [email protected]_pass(opt_level=0, name="CublasDispatch") +@tvm_tirx.ir.transform.module_pass(opt_level=0, name="CublasDispatch") class CublasDispatch: - def transform_module(self, mod: IRModule, _ctx: tvm.ir.transform.PassContext) -> IRModule: + def transform_module(self, mod: IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> IRModule: # Check if CUBLAS is enabled - if not tvm.get_global_func("relax.ext.cublas", True): + if not tvm_tirx.get_global_func("relax.ext.cublas", True): raise Exception("CUBLAS is not enabled.") # Get interested patterns @@ -163,12 +165,12 @@ mod.show() # To save CI time and avoid flakiness, we skip the tuning process in CI environment. # -device = tvm.cuda(0) -target = tvm.target.Target.from_device(device) +device = tvm_tirx.cuda(0) +target = tvm_tirx.target.Target.from_device(device) if os.getenv("CI", "") != "true": trials = 2000 with target, tempfile.TemporaryDirectory() as tmp_dir: - mod = tvm.ir.transform.Sequential( + mod = tvm_tirx.ir.transform.Sequential( [ relax.get_pipeline("zero"), relax.transform.MetaScheduleTuneTIR(work_dir=tmp_dir, max_trials_global=trials), @@ -190,7 +192,7 @@ from tvm.s_tir import dlight as dl # Apply DLight rules with target: - mod = tvm.ir.transform.Sequential( + mod = tvm_tirx.ir.transform.Sequential( [ relax.get_pipeline("zero"), dl.ApplyDefaultSchedule( # pylint: disable=not-callable @@ -218,11 +220,13 @@ mod.show() # We can build and deploy the optimized model to the TVM runtime. ex = tvm.compile(mod, target="cuda") -dev = tvm.cuda(0) +dev = tvm_tirx.cuda(0) vm = relax.VirtualMachine(ex, dev) # Need to allocate data and params on GPU device -data = tvm.runtime.tensor(np.random.rand(*input_shape).astype("float32"), dev) -gpu_params = [tvm.runtime.tensor(np.random.rand(*p.shape).astype(p.dtype), dev) for _, p in params] +data = tvm_tirx.runtime.tensor(np.random.rand(*input_shape).astype("float32"), dev) +gpu_params = [ + tvm_tirx.runtime.tensor(np.random.rand(*p.shape).astype(p.dtype), dev) for _, p in params +] gpu_out = vm["forward"](data, *gpu_params).numpy() print(gpu_out) diff --git a/docs/how_to/tutorials/e2e_opt_model.py b/docs/how_to/tutorials/e2e_opt_model.py index 5de3d50485..8c35d1882e 100644 --- a/docs/how_to/tutorials/e2e_opt_model.py +++ b/docs/how_to/tutorials/e2e_opt_model.py @@ -68,6 +68,7 @@ torch_model = resnet18(weights=ResNet18_Weights.DEFAULT).eval() # optimization. import tvm +import tvm_tirx from tvm import relax from tvm.relax.frontend.torch import from_exported_program @@ -118,7 +119,7 @@ if not IS_IN_CI: TOTAL_TRIALS = 512 # Change to 20000 for better performance if needed MAX_TRIALS_PER_TASK = 16 # Change to more trials per task for better performance if needed -target = tvm.target.Target("nvidia/geforce-rtx-3090-ti") # Change to your target device +target = tvm_tirx.target.Target("nvidia/geforce-rtx-3090-ti") # Change to your target device work_dir = "tuning_logs" if not IS_IN_CI: @@ -143,11 +144,11 @@ if not IS_IN_CI: with target: mod = tvm.s_tir.transform.DefaultGPUSchedule()(mod) ex = tvm.compile(mod, target=target) - dev = tvm.cuda(0) + dev = tvm_tirx.cuda(0) vm = relax.VirtualMachine(ex, dev) # Need to allocate data and params on GPU device - gpu_data = tvm.runtime.tensor(np.random.rand(1, 3, 224, 224).astype("float32"), dev) - gpu_params = [tvm.runtime.tensor(p, dev) for p in params["main"]] + gpu_data = tvm_tirx.runtime.tensor(np.random.rand(1, 3, 224, 224).astype("float32"), dev) + gpu_params = [tvm_tirx.runtime.tensor(p, dev) for p in params["main"]] gpu_out = vm["main"](gpu_data, *gpu_params)[0].numpy() print(gpu_out.shape) diff --git a/docs/how_to/tutorials/export_and_load_executable.py b/docs/how_to/tutorials/export_and_load_executable.py index 9972861c12..542266803a 100644 --- a/docs/how_to/tutorials/export_and_load_executable.py +++ b/docs/how_to/tutorials/export_and_load_executable.py @@ -65,6 +65,7 @@ except ImportError: # pragma: no cover import tvm_ffi import tvm +import tvm_tirx from tvm import relax from tvm.relax.frontend.torch import from_exported_program @@ -115,7 +116,7 @@ if RUN_EXAMPLE: # executable. Passing ``workspace_dir`` keeps the intermediate packaging files, # which is useful to inspect what was produced. -TARGET = tvm.target.Target("llvm") +TARGET = tvm_tirx.target.Target("llvm") ARTIFACT_DIR = Path("relax_export_artifacts") ARTIFACT_DIR.mkdir(exist_ok=True) @@ -159,16 +160,16 @@ if RUN_EXAMPLE: # runtime module directly. if RUN_EXAMPLE: - loaded_rt_mod = tvm.runtime.load_module(str(library_path)) - dev = tvm.cpu(0) + loaded_rt_mod = tvm_tirx.runtime.load_module(str(library_path)) + dev = tvm_tirx.cpu(0) vm = relax.VirtualMachine(loaded_rt_mod, dev) # Prepare input data input_tensor = torch.randn(1, 1, 28, 28, dtype=torch.float32) - vm_input = tvm.runtime.tensor(input_tensor.numpy(), dev) + vm_input = tvm_tirx.runtime.tensor(input_tensor.numpy(), dev) # Prepare parameters (allocate on target device) - vm_params = [tvm.runtime.tensor(p, dev) for p in params["main"]] + vm_params = [tvm_tirx.runtime.tensor(p, dev) for p in params["main"]] # Run inference: pass input data followed by all parameters tvm_output = vm["main"](vm_input, *vm_params) diff --git a/docs/how_to/tutorials/import_model.py b/docs/how_to/tutorials/import_model.py index 888235a859..c8e31558fd 100644 --- a/docs/how_to/tutorials/import_model.py +++ b/docs/how_to/tutorials/import_model.py @@ -55,6 +55,7 @@ from torch import nn from torch.export import export import tvm +import tvm_tirx from tvm import relax from tvm.relax.frontend.torch import from_exported_program @@ -191,13 +192,13 @@ def convert_sigmoid(node: torch.fx.Node, importer: ExportedProgramImporter) -> r mod_compiled = relax.get_pipeline("zero")(mod) exec_module = tvm.compile(mod_compiled, target="llvm") -dev = tvm.cpu() +dev = tvm_tirx.cpu() vm = relax.VirtualMachine(exec_module, dev) # Run inference input_data = np.random.rand(1, 3, 32, 32).astype("float32") -tvm_input = tvm.runtime.tensor(input_data, dev) -tvm_params = [tvm.runtime.tensor(p, dev) for p in params["main"]] +tvm_input = tvm_tirx.runtime.tensor(input_data, dev) +tvm_params = [tvm_tirx.runtime.tensor(p, dev) for p in params["main"]] tvm_out = vm["main"](tvm_input, *tvm_params).numpy() # Compare with PyTorch diff --git a/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py b/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py index 35e4110f9e..e195f431ce 100644 --- a/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py +++ b/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py @@ -56,13 +56,13 @@ try: except ImportError: torch = None -import tvm +import tvm_tirx from tvm import relax +from tvm.relax import script as R from tvm.relax.base_py_module import BasePyModule -from tvm.script import ir as I -from tvm.script import relax as R -from tvm.script import s_tir as Ts -from tvm.script import tirx as T +from tvm.s_tir import script as Ts +from tvm_tirx.script import ir as I +from tvm_tirx.script import tirx as T IS_IN_CI = os.getenv("CI", "").lower() == "true" HAS_TORCH = torch is not None @@ -103,7 +103,7 @@ if RUN_EXAMPLE: return self._convert_tvm_to_pytorch(result) # TIR functions are JIT-compiled at instantiation - mod = MyFirstModule(device=tvm.cpu(0)) + mod = MyFirstModule(device=tvm_tirx.cpu(0)) x = torch.tensor([1.0, 2.0, 3.0, 4.0]) y = torch.tensor([10.0, 20.0, 30.0, 40.0]) @@ -171,7 +171,7 @@ if RUN_EXAMPLE: print(f" [DEBUG] probs sum: {result.sum(dim=-1)}") return result - mod = DebugModule(device=tvm.cpu(0)) + mod = DebugModule(device=tvm_tirx.cpu(0)) print("Running with debug prints:") probs = mod.forward(torch.randn(2, 4), torch.randn(4, 3)) @@ -197,7 +197,7 @@ if RUN_EXAMPLE: if RUN_EXAMPLE: # Register a packed function (simulating an external library binding) - @tvm.register_global_func("my_bias_add", override=True) + @tvm_tirx.register_global_func("my_bias_add", override=True) def my_bias_add(x, bias, out): """Packed function: adds bias to each row of x.""" @@ -243,7 +243,7 @@ if RUN_EXAMPLE: # 3. Python/PyTorch activation return F.relu(h_biased) - mod = PipelineModule(device=tvm.cpu(0)) + mod = PipelineModule(device=tvm_tirx.cpu(0)) x = torch.randn(2, 4) w = torch.randn(4, 3) @@ -365,7 +365,7 @@ if RUN_EXAMPLE: out = R.call_py_func("silu", (h,), ty_args=[R.Tensor((4, 8), "float32")]) return out - mod = HybridVMModule(device=tvm.cpu(0)) + mod = HybridVMModule(device=tvm_tirx.cpu(0)) x = torch.randn(4, 8) # call_py_func is also callable from Python directly @@ -412,7 +412,7 @@ if RUN_EXAMPLE: ) -> R.Tensor((n,), "float32"): return R.add(x, y) - mod = DynamicModule(device=tvm.cpu(0), target="llvm") + mod = DynamicModule(device=tvm_tirx.cpu(0), target="llvm") # Inspect what the module contains print("Functions:", mod.list_functions()) diff --git a/docs/how_to/tutorials/optimize_llm.py b/docs/how_to/tutorials/optimize_llm.py index 7faa4f43bc..2c01922452 100644 --- a/docs/how_to/tutorials/optimize_llm.py +++ b/docs/how_to/tutorials/optimize_llm.py @@ -64,12 +64,14 @@ from pprint import pprint from tvm_ffi import Shape import tvm -from tvm import relax, te, tirx +import tvm_tirx +from tvm import relax, te from tvm.relax import register_pipeline from tvm.relax.frontend import nn from tvm.relax.frontend.nn import Tensor, op from tvm.relax.frontend.nn.llm.kv_cache import PagedKVCache, TIRPagedKVCache from tvm.s_tir import dlight +from tvm_tirx import tirx ###################################################################### # First, we need to define the model configuration. The configuration includes the key parameters @@ -92,8 +94,8 @@ class LlamaConfig: head_dim: int = 64 # hidden_size // num_attention_heads -dev = tvm.cuda(0) -target = tvm.target.Target.from_device(dev) +dev = tvm_tirx.cuda(0) +target = tvm_tirx.target.Target.from_device(dev) ###################################################################### @@ -382,9 +384,11 @@ def _pipeline( # pylint: disable=too-many-arguments ): ext_mods = ext_mods or [] - @tvm.ir.transform.module_pass(opt_level=0) - def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule: - seq = tvm.ir.transform.Sequential( + @tvm_tirx.ir.transform.module_pass(opt_level=0) + def _pipeline( + mod: tvm_tirx.ir.IRModule, _ctx: tvm_tirx.ir.transform.PassContext + ) -> tvm_tirx.ir.IRModule: + seq = tvm_tirx.ir.transform.Sequential( [ # Phase 1. Passes on high-level operator graph # We can enable cublas for further optimization @@ -490,7 +494,8 @@ if not IS_IN_CI: # Convert params into ndarray params = [ - tvm.runtime.tensor(param_dict[k].astype("float16"), device=dev) for k in named_params.keys() + tvm_tirx.runtime.tensor(param_dict[k].astype("float16"), device=dev) + for k in named_params.keys() ] @@ -524,7 +529,7 @@ if not IS_IN_CI: input_len = len(prompt) # Load prompt tokens into TVM ndarray on the target device - tokens = tvm.runtime.tensor(np.array(prompt).astype("int32"), device=dev) + tokens = tvm_tirx.runtime.tensor(np.array(prompt).astype("int32"), device=dev) ###################################################################### # Create the KVCache @@ -548,7 +553,7 @@ if not IS_IN_CI: # The next step is to embed the tokens into the hidden states. We use the `embed` function # compiled in the Relax IRModule to embed the tokens into the hidden states. -nd_view_func = tvm.get_global_func("vm.builtin.reshape") +nd_view_func = tvm_tirx.get_global_func("vm.builtin.reshape") def embed(tokens, params): @@ -563,9 +568,9 @@ def embed(tokens, params): # ~~~~~~~ # Before running the forward pass, we first get some help functions for preparation. -add_sequence_func = tvm.get_global_func("vm.builtin.kv_state_add_sequence") -begin_forward_func = tvm.get_global_func("vm.builtin.kv_state_begin_forward") -end_forward_func = tvm.get_global_func("vm.builtin.kv_state_end_forward") +add_sequence_func = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence") +begin_forward_func = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward") +end_forward_func = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward") ###################################################################### # As we are creating a new sequence, we need to call `add_sequence_func` to initialize @@ -610,7 +615,7 @@ if not IS_IN_CI: print("The generated token:") while last_token != tokenizer.eos_token_id: - tokens = tvm.runtime.tensor(np.array([last_token]).astype("int32"), device=dev) + tokens = tvm_tirx.runtime.tensor(np.array([last_token]).astype("int32"), device=dev) hidden_states = embed(tokens, params) begin_forward_func(kv_cache, Shape([seq_id]), Shape([1])) logits, kv_cache = vm["decode"](hidden_states, kv_cache, params) diff --git a/docs/reference/api/python/codegen.rst b/docs/reference/api/python/codegen.rst index 0fca6ed79f..c1ee15cc0f 100644 --- a/docs/reference/api/python/codegen.rst +++ b/docs/reference/api/python/codegen.rst @@ -15,8 +15,8 @@ specific language governing permissions and limitations under the License. -tvm.codegen -=========== -.. automodule:: tvm.codegen +tvm_tirx.codegen +================ +.. automodule:: tvm_tirx.codegen :members: :show-inheritance: diff --git a/docs/reference/api/python/error.rst b/docs/reference/api/python/error.rst index a228f7b8bd..297e430093 100644 --- a/docs/reference/api/python/error.rst +++ b/docs/reference/api/python/error.rst @@ -15,9 +15,9 @@ specific language governing permissions and limitations under the License. -tvm.error ---------- -.. automodule:: tvm.error +tvm_tirx.error +-------------- +.. automodule:: tvm_tirx.error :members: :imported-members: :autosummary: diff --git a/docs/reference/api/python/instrument.rst b/docs/reference/api/python/instrument.rst index 46d193ce15..10b4bee978 100644 --- a/docs/reference/api/python/instrument.rst +++ b/docs/reference/api/python/instrument.rst @@ -15,7 +15,7 @@ specific language governing permissions and limitations under the License. -tvm.ir.transform.instrument -=========================== -.. automodule:: tvm.ir.transform.instrument +tvm_tirx.ir.transform.instrument +================================ +.. automodule:: tvm_tirx.ir.transform.instrument :members: diff --git a/docs/reference/api/python/ir.rst b/docs/reference/api/python/ir.rst index 1747dee7be..60c4289b77 100644 --- a/docs/reference/api/python/ir.rst +++ b/docs/reference/api/python/ir.rst @@ -15,9 +15,9 @@ specific language governing permissions and limitations under the License. -tvm.ir ------- -.. automodule:: tvm.ir +tvm_tirx.ir +----------- +.. automodule:: tvm_tirx.ir :members: :imported-members: :exclude-members: Expr diff --git a/docs/reference/api/python/runtime/runtime.rst b/docs/reference/api/python/runtime/runtime.rst index ae373080ae..04bcd64f2d 100644 --- a/docs/reference/api/python/runtime/runtime.rst +++ b/docs/reference/api/python/runtime/runtime.rst @@ -15,7 +15,7 @@ specific language governing permissions and limitations under the License. -tvm.runtime ------------ -.. automodule:: tvm.runtime +tvm_tirx.runtime +---------------- +.. automodule:: tvm_tirx.runtime :members: diff --git a/docs/reference/api/python/script/ir_builder.rst b/docs/reference/api/python/script/ir_builder.rst index f26f70e1c1..564e9c1db6 100644 --- a/docs/reference/api/python/script/ir_builder.rst +++ b/docs/reference/api/python/script/ir_builder.rst @@ -15,12 +15,12 @@ specific language governing permissions and limitations under the License. -tvm.script.ir_builder ---------------------- +tvm_tirx.script.ir_builder +-------------------------- -tvm.script.ir_builder -********************* -.. automodule:: tvm.script.ir_builder +tvm_tirx.script.ir_builder +************************** +.. automodule:: tvm_tirx.script.ir_builder :members: :imported-members: :exclude-members: Call, DataTypeImm, FuncType, GenericConst, MissingType, Op, PrimType, Range, StringImm, StringType, TensorRegion, TensorRegionType, Tuple, Type, Var diff --git a/docs/reference/api/python/script/parser.rst b/docs/reference/api/python/script/parser.rst index 31beb76afd..d4028219b8 100644 --- a/docs/reference/api/python/script/parser.rst +++ b/docs/reference/api/python/script/parser.rst @@ -15,44 +15,44 @@ specific language governing permissions and limitations under the License. -tvm.script.parser ------------------ +tvm_tirx.script.parser +---------------------- -tvm.script.parser -***************** -.. automodule:: tvm.script.parser +tvm_tirx.script.parser +********************** +.. automodule:: tvm_tirx.script.parser :members: :imported-members: -tvm.script.parser.jit_support -***************************** +tvm_tirx.script.parser.jit_support +********************************** These helpers carry ``const_args``, the selected fixed argument map, into generated builder execution. Parameter validation and caching remain with the JIT entry point. -.. automodule:: tvm.script.parser.jit_support +.. automodule:: tvm_tirx.script.parser.jit_support :members: use_specialization, read_specialization_bindings -tvm.script.parser.protocol_registry -*********************************** +tvm_tirx.script.parser.protocol_registry +**************************************** Language variants register source syntax policies independently of their construction hooks. Namespace objects and lazy initialization callbacks are -registered through :func:`tvm.script.parser.register_namespace` and -:func:`tvm.script.parser.register_namespace_initializer`. +registered through :func:`tvm_tirx.script.parser.register_namespace` and +:func:`tvm_tirx.script.parser.register_namespace_initializer`. -.. automodule:: tvm.script.parser.protocol_registry +.. automodule:: tvm_tirx.script.parser.protocol_registry :members: register_scalar_annotation, register_mutable_decl The language variant aliases below share the public construction namespaces documented in :doc:`script`. Parser entry points above use the canonical frontend. -tvm.script.parser.ir -******************** -.. automodule:: tvm.script.parser.ir +tvm_tirx.script.parser.ir +************************* +.. automodule:: tvm_tirx.script.parser.ir -tvm.script.parser.relax -*********************** -.. automodule:: tvm.script.parser.relax +tvm_tirx.script.parser.relax +**************************** +.. automodule:: tvm_tirx.script.parser.relax -tvm.script.parser.tirx -********************** -.. automodule:: tvm.script.parser.tirx +tvm_tirx.script.parser.tirx +*************************** +.. automodule:: tvm_tirx.script.parser.tirx diff --git a/docs/reference/api/python/script/printer.rst b/docs/reference/api/python/script/printer.rst index 098333de07..503ea719ad 100644 --- a/docs/reference/api/python/script/printer.rst +++ b/docs/reference/api/python/script/printer.rst @@ -15,12 +15,12 @@ specific language governing permissions and limitations under the License. -tvm.script.printer +tvm_tirx.script.printer ------------------------- -tvm.script.printer +tvm_tirx.script.printer ************************* -.. automodule:: tvm.script.printer +.. automodule:: tvm_tirx.script.printer :members: :imported-members: :exclude-members: Scriptable diff --git a/docs/reference/api/python/script/script.rst b/docs/reference/api/python/script/script.rst index dd287b4899..612b634c7a 100644 --- a/docs/reference/api/python/script/script.rst +++ b/docs/reference/api/python/script/script.rst @@ -15,22 +15,22 @@ specific language governing permissions and limitations under the License. -tvm.script ----------- +tvm_tirx.script +--------------- -tvm.script -********** -.. automodule:: tvm.script +tvm_tirx.script +*************** +.. automodule:: tvm_tirx.script :members: :imported-members: :exclude-members: Any -tvm.script.relax +tvm.relax.script **************** Operator attributes re-exported by this namespace are documented in :mod:`tvm.relax.op`. -.. automodule:: tvm.script.relax +.. automodule:: tvm.relax.script :members: :exclude-members: ExternFunc, ShapeExpr, TupleGetItem, Range, function, macro, ArgmaxArgminAttrs, ArgsortAttrs, AstypeAttrs, CallInplacePackedAttrs, CallTIRInplaceAttrs, @@ -42,28 +42,28 @@ Operator attributes re-exported by this namespace are documented in :mod:`tvm.re SplitAttrs, SqueezeAttrs, StackAttrs, StatisticalAttrs, StridedSliceAttrs, TakeAttrs, TileAttrs, TopKAttrs, ToVDeviceAttrs, TriluAttrs, WrapParamAttrs -.. autofunction:: tvm.script.relax.function +.. autofunction:: tvm.relax.script.function -.. autofunction:: tvm.script.relax.macro +.. autofunction:: tvm.relax.script.macro -tvm.script.tirx -*************** +tvm_tirx.script.tirx +******************** -Operator attributes re-exported by this namespace are documented in :mod:`tvm.tirx`. -Shared IR classes and builder infrastructure are documented in :mod:`tvm.ir` and -:mod:`tvm.script.ir_builder`. -The canonical ``MemoryScope`` class is documented in :mod:`tvm.runtime`. +Operator attributes re-exported by this namespace are documented in :mod:`tvm_tirx.tirx`. +Shared IR classes and builder infrastructure are documented in :mod:`tvm_tirx.ir` and +:mod:`tvm_tirx.script.ir_builder`. +The canonical ``MemoryScope`` class is documented in :mod:`tvm_tirx.runtime`. -.. automodule:: tvm.script.tirx +.. automodule:: tvm_tirx.script.tirx :members: :exclude-members: AlreadyEmitted, DataTypeImm, FuncType, GenericConst, IRBuilder, IRModuleFrame, MemoryScope, MissingType, Op, PrimType, StringImm, StringType, TensorRegion, TensorRegionType, Type, Range, meta_var, Var, Call, CommReducer, Reduce, FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod, LShift, RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min, Max, EQ, NE, LT, LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle, CallEffectKind, IterVar, ComposeLayout, DtypeConstructor, ThreadScope, Iter, Layout, LetAnnotation, LocalVectorAnnotation, TileLayout, CallFFIKernelAttr, TensorMapEncodeTiledAttr, function, jit, inli [...] -.. autofunction:: tvm.script.tirx.function +.. autofunction:: tvm_tirx.script.tirx.function -.. autofunction:: tvm.script.tirx.jit +.. autofunction:: tvm_tirx.script.tirx.jit -.. autofunction:: tvm.script.tirx.inline +.. autofunction:: tvm_tirx.script.tirx.inline -.. autofunction:: tvm.script.tirx.macro +.. autofunction:: tvm_tirx.script.tirx.macro diff --git a/docs/reference/api/python/support.rst b/docs/reference/api/python/support.rst index 8042ce34c1..2c5991081c 100644 --- a/docs/reference/api/python/support.rst +++ b/docs/reference/api/python/support.rst @@ -15,41 +15,41 @@ specific language governing permissions and limitations under the License. -tvm.support ------------ -.. automodule:: tvm.support +tvm_tirx.support +---------------- +.. automodule:: tvm_tirx.support :members: :imported-members: :autosummary: -tvm.support.cc -~~~~~~~~~~~~~~ -.. automodule:: tvm.support.cc +tvm_tirx.support.cc +~~~~~~~~~~~~~~~~~~~ +.. automodule:: tvm_tirx.support.cc :members: -tvm.support.nvcc -~~~~~~~~~~~~~~~~ -.. automodule:: tvm.support.nvcc +tvm_tirx.support.nvcc +~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: tvm_tirx.support.nvcc :members: -tvm.support.rocm -~~~~~~~~~~~~~~~~ -.. automodule:: tvm.support.rocm +tvm_tirx.support.rocm +~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: tvm_tirx.support.rocm :members: -tvm.support.ndk -~~~~~~~~~~~~~~~ -.. automodule:: tvm.support.ndk +tvm_tirx.support.ndk +~~~~~~~~~~~~~~~~~~~~ +.. automodule:: tvm_tirx.support.ndk :members: -tvm.support.xcode -~~~~~~~~~~~~~~~~~ -.. automodule:: tvm.support.xcode +tvm_tirx.support.xcode +~~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: tvm_tirx.support.xcode :members: -tvm.support.clang -~~~~~~~~~~~~~~~~~ -.. automodule:: tvm.support.clang +tvm_tirx.support.clang +~~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: tvm_tirx.support.clang :members: tvm.contrib.emcc @@ -57,17 +57,17 @@ tvm.contrib.emcc .. automodule:: tvm.contrib.emcc :members: -tvm.support.popen_pool -~~~~~~~~~~~~~~~~~~~~~~ -.. automodule:: tvm.support.popen_pool +tvm_tirx.support.popen_pool +~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: tvm_tirx.support.popen_pool :members: -tvm.support.utils -~~~~~~~~~~~~~~~~~ -.. automodule:: tvm.support.utils +tvm_tirx.support.utils +~~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: tvm_tirx.support.utils :members: -tvm.support.tar -~~~~~~~~~~~~~~~ -.. automodule:: tvm.support.tar +tvm_tirx.support.tar +~~~~~~~~~~~~~~~~~~~~ +.. automodule:: tvm_tirx.support.tar :members: diff --git a/docs/reference/api/python/sym.rst b/docs/reference/api/python/sym.rst index 0dbb6ee72b..86d2fe971b 100644 --- a/docs/reference/api/python/sym.rst +++ b/docs/reference/api/python/sym.rst @@ -15,8 +15,8 @@ specific language governing permissions and limitations under the License. -tvm.sym -------- -.. automodule:: tvm.sym +tvm_tirx.sym +------------ +.. automodule:: tvm_tirx.sym :members: :imported-members: diff --git a/docs/reference/api/python/target.rst b/docs/reference/api/python/target.rst index 9ccfb248d8..1e0b574914 100644 --- a/docs/reference/api/python/target.rst +++ b/docs/reference/api/python/target.rst @@ -15,15 +15,15 @@ specific language governing permissions and limitations under the License. -tvm.target ----------- -.. automodule:: tvm.target +tvm_tirx.target +--------------- +.. automodule:: tvm_tirx.target :members: :imported-members: :autosummary: -tvm.target.tag -~~~~~~~~~~~~~~ -.. automodule:: tvm.target.tag +tvm_tirx.target.tag +~~~~~~~~~~~~~~~~~~~ +.. automodule:: tvm_tirx.target.tag :members: :imported-members: diff --git a/docs/reference/api/python/testing.rst b/docs/reference/api/python/testing.rst index 5ef7181fd7..f16af323c9 100644 --- a/docs/reference/api/python/testing.rst +++ b/docs/reference/api/python/testing.rst @@ -15,8 +15,8 @@ specific language governing permissions and limitations under the License. -tvm.testing ------------ -.. automodule:: tvm.testing +tvm_tirx.testing +---------------- +.. automodule:: tvm_tirx.testing :members: :imported-members: diff --git a/docs/reference/api/python/transform.rst b/docs/reference/api/python/transform.rst index aebdbdbad5..3908c2c331 100644 --- a/docs/reference/api/python/transform.rst +++ b/docs/reference/api/python/transform.rst @@ -15,8 +15,8 @@ specific language governing permissions and limitations under the License. -tvm.ir.transform -================ -.. automodule:: tvm.ir.transform +tvm_tirx.ir.transform +===================== +.. automodule:: tvm_tirx.ir.transform :members: :imported-members: diff --git a/docs/tirx/api/analysis.rst b/docs/tirx/api/analysis.rst index 8efdf42349..82843eca8c 100644 --- a/docs/tirx/api/analysis.rst +++ b/docs/tirx/api/analysis.rst @@ -18,5 +18,5 @@ Analysis ======== -.. automodule:: tvm.tirx.analysis.analysis +.. automodule:: tvm_tirx.tirx.analysis.analysis :members: diff --git a/docs/tirx/api/backend.rst b/docs/tirx/api/backend.rst index 58f819ff2f..6730e79c1f 100644 --- a/docs/tirx/api/backend.rst +++ b/docs/tirx/api/backend.rst @@ -21,37 +21,37 @@ Backend Extension API Backend loading --------------- -``tvm.backend`` discovers and loads target-owned Python semantics. Depending +``tvm_tirx.backend`` discovers and loads target-owned Python semantics. Depending on the target, loading a backend registers its TVMScript namespaces, tile-dispatch implementations, target tags, compilation-pipeline entry points, and code-generation support. -.. automodule:: tvm.backend +.. automodule:: tvm_tirx.backend :members: :no-index: CUDA registration ----------------- -.. autofunction:: tvm.backend.cuda.register_backend +.. autofunction:: tvm_tirx.backend.cuda.register_backend :no-index: -.. autofunction:: tvm.backend.cuda.script_namespace +.. autofunction:: tvm_tirx.backend.cuda.script_namespace :no-index: -.. autofunction:: tvm.backend.cuda.script_namespaces +.. autofunction:: tvm_tirx.backend.cuda.script_namespaces :no-index: Trainium registration --------------------- -.. autofunction:: tvm.backend.trn.register_backend +.. autofunction:: tvm_tirx.backend.trn.register_backend :no-index: -.. autofunction:: tvm.backend.trn.script_namespace +.. autofunction:: tvm_tirx.backend.trn.script_namespace :no-index: -.. autofunction:: tvm.backend.trn.script_namespaces +.. autofunction:: tvm_tirx.backend.trn.script_namespaces :no-index: These functions are backend integration points. Kernel-facing CUDA APIs are diff --git a/docs/tirx/api/compilation.rst b/docs/tirx/api/compilation.rst index 95e180a1c9..e7e04c6d37 100644 --- a/docs/tirx/api/compilation.rst +++ b/docs/tirx/api/compilation.rst @@ -18,15 +18,15 @@ Compilation =========== -tvm.tirx.compile ----------------- -.. currentmodule:: tvm.tirx +tvm_tirx.tirx.compile +--------------------- +.. currentmodule:: tvm_tirx.tirx .. autofunction:: compile :no-index: -tvm.tirx.compilation_pipeline ------------------------------ -.. automodule:: tvm.tirx.compilation_pipeline +tvm_tirx.tirx.compilation_pipeline +---------------------------------- +.. automodule:: tvm_tirx.tirx.compilation_pipeline :members: :no-index: diff --git a/docs/tirx/api/cuda.rst b/docs/tirx/api/cuda.rst index 3762be0f6b..701263a222 100644 --- a/docs/tirx/api/cuda.rst +++ b/docs/tirx/api/cuda.rst @@ -23,7 +23,7 @@ CUDA helpers The CUDA backend installs the ``Tx.cuda`` namespace when it is loaded:: - from tvm.script import tirx as Tx + from tvm_tirx.script import tirx as Tx Tx.cuda.cta_sync() leader = Tx.cuda.elect_sync() @@ -91,56 +91,56 @@ provide ``.warp`` and ``.block`` variants. CUDA language utilities ----------------------- -``tvm.backend.cuda.lang`` provides reusable objects for hand-written CUDA TIRx +``tvm_tirx.backend.cuda.lang`` provides reusable objects for hand-written CUDA TIRx kernels, including schedulers, pipelines, barriers, and shared-memory pools. -.. autoclass:: tvm.backend.cuda.lang.BaseTileScheduler +.. autoclass:: tvm_tirx.backend.cuda.lang.BaseTileScheduler :members: -.. py:class:: tvm.backend.cuda.lang.ClusterPersistentScheduler2D(prefix, num_m_tiles, num_n_tiles, num_clusters, l2_group_size=8, cluster_m=1, cluster_n=1, serpentine=False) +.. py:class:: tvm_tirx.backend.cuda.lang.ClusterPersistentScheduler2D(prefix, num_m_tiles, num_n_tiles, num_clusters, l2_group_size=8, cluster_m=1, cluster_n=1, serpentine=False) Persistent two-dimensional tile scheduler with group-major or serpentine traversal. ``init`` selects the first tile for a cluster, ``next_tile`` or ``next_tile_stride`` advances it, and ``valid`` reports whether the current ``m_idx`` / ``n_idx`` is in range. -.. autoclass:: tvm.backend.cuda.lang.FlashAttentionLPTScheduler +.. autoclass:: tvm_tirx.backend.cuda.lang.FlashAttentionLPTScheduler :members: -.. py:class:: tvm.backend.cuda.lang.FlashAttentionLinearScheduler(prefix, num_batches, num_heads, num_m_blocks, num_ctas) +.. py:class:: tvm_tirx.backend.cuda.lang.FlashAttentionLinearScheduler(prefix, num_batches, num_heads, num_m_blocks, num_ctas) Linear persistent scheduler over ``(batch, head, m_block)`` tasks. ``init`` starts from a CTA index, ``next_tile`` advances by ``num_ctas``, and ``valid`` reports whether work remains. -.. autoclass:: tvm.backend.cuda.lang.MBarrier +.. autoclass:: tvm_tirx.backend.cuda.lang.MBarrier :members: -.. autoclass:: tvm.backend.cuda.lang.Pipeline +.. autoclass:: tvm_tirx.backend.cuda.lang.Pipeline :members: -.. autoclass:: tvm.backend.cuda.lang.PipelineState +.. autoclass:: tvm_tirx.backend.cuda.lang.PipelineState :members: -.. autoclass:: tvm.backend.cuda.lang.SMEMPool +.. autoclass:: tvm_tirx.backend.cuda.lang.SMEMPool :members: -.. autoclass:: tvm.backend.cuda.lang.TCGen05Bar +.. autoclass:: tvm_tirx.backend.cuda.lang.TCGen05Bar :members: -.. autoclass:: tvm.backend.cuda.lang.TMABar +.. autoclass:: tvm_tirx.backend.cuda.lang.TMABar :members: IKET profiling -------------- -.. automodule:: tvm.backend.cuda.iket +.. automodule:: tvm_tirx.backend.cuda.iket :members: :no-index: CUDA-specific transforms ------------------------ -.. automodule:: tvm.backend.cuda.transforms +.. automodule:: tvm_tirx.backend.cuda.transforms :members: :no-index: diff --git a/docs/tirx/api/cuda_compile.rst b/docs/tirx/api/cuda_compile.rst index 2f5c1d8e47..c93558bc5a 100644 --- a/docs/tirx/api/cuda_compile.rst +++ b/docs/tirx/api/cuda_compile.rst @@ -56,26 +56,26 @@ compiler arguments are not modeled as separate Python fields. .. code-block:: python - from tvm.backend.cuda import BackendConfig + from tvm_tirx.backend.cuda import BackendConfig cuda_config = BackendConfig( compiler="nvrtc", nvrtc=["--use_fast_math", "--ftz=false"], ) - target = tvm.target.Target({ + target = tvm_tirx.target.Target({ "kind": "cuda", "arch": "sm_100a", "backend_config": {"cuda": cuda_config}, }) executable = tvm.compile(func, target=target) -.. autoclass:: tvm.backend.cuda.BackendConfig +.. autoclass:: tvm_tirx.backend.cuda.BackendConfig :members: :undoc-members: -.. autoclass:: tvm.backend.cuda.compiler.CompilationResult +.. autoclass:: tvm_tirx.backend.cuda.compiler.CompilationResult :members: -.. autofunction:: tvm.backend.cuda.compiler.compile_source +.. autofunction:: tvm_tirx.backend.cuda.compiler.compile_source See :doc:`../native_basics/cuda/compiling` for target-tag defaults, inheritance, per-entry compilation, and artifact behavior. See :doc:`cuda_compile_options` diff --git a/docs/tirx/api/execution.rst b/docs/tirx/api/execution.rst index 6c22eebb30..d3acb99681 100644 --- a/docs/tirx/api/execution.rst +++ b/docs/tirx/api/execution.rst @@ -18,9 +18,9 @@ Thread Scope Enum Tags ====================== -tvm.tirx.thread_scope ---------------------- -.. automodule:: tvm.tirx.thread_scope +tvm_tirx.tirx.thread_scope +-------------------------- +.. automodule:: tvm_tirx.tirx.thread_scope :members: :no-index: @@ -43,7 +43,7 @@ Same-backend hierarchy comparisons use explicit ranks; cross-backend comparisons are rejected. For example, CUDA TMA instructions accept ``thread`` and ``tcgen05.ld/st`` accept ``warpgroup``. CUDA index bindings and predicate analysis live in -``tvm.backend.cuda.tile.exec_context``. +``tvm_tirx.backend.cuda.tile.exec_context``. Instruction APIs continue accepting ``scope="warp"`` and also accept qualified names or ``ThreadScope`` values. Construction resolves the value before @@ -52,10 +52,10 @@ Omitting scope uses the instruction backend's ``thread`` tag. Direct attribute construction must provide the tag or its qualified name; shared attribute schemas have no backend-independent default. -.. automodule:: tvm.backend.cuda.thread_scope +.. automodule:: tvm_tirx.backend.cuda.thread_scope :members: :no-index: -.. automodule:: tvm.backend.trn.thread_scope +.. automodule:: tvm_tirx.backend.trn.thread_scope :members: :no-index: diff --git a/docs/tirx/api/layout.rst b/docs/tirx/api/layout.rst index 158a0a2637..2b3b37172e 100644 --- a/docs/tirx/api/layout.rst +++ b/docs/tirx/api/layout.rst @@ -21,41 +21,41 @@ Layout IR Layouts describe how logical tensor coordinates map to storage and execution axes. See :doc:`../layout` for the programming model and worked examples. -The primary classes are re-exported from ``tvm.tirx``, but are defined in this +The primary classes are re-exported from ``tvm_tirx.tirx``, but are defined in this module: -.. autoclass:: tvm.tirx.layout.Layout +.. autoclass:: tvm_tirx.tirx.layout.Layout :members: :no-index: -.. autoclass:: tvm.tirx.layout.Axis +.. autoclass:: tvm_tirx.tirx.layout.Axis :members: :no-index: -.. autoclass:: tvm.tirx.layout.Iter +.. autoclass:: tvm_tirx.tirx.layout.Iter :members: :no-index: -.. autoclass:: tvm.tirx.layout.TileLayout +.. autoclass:: tvm_tirx.tirx.layout.TileLayout :members: :no-index: -.. autoclass:: tvm.tirx.layout.ComposeLayout +.. autoclass:: tvm_tirx.tirx.layout.ComposeLayout :members: :no-index: ``S[...]`` and ``R[...]`` build shard and replica layout specifications, respectively. Named axes such as ``laneid``, ``warpid``, ``tid_in_wg``, ``TLane``, and ``TCol`` belong to CUDA. Use ``T.cuda.Axis.laneid`` or -import axes from ``tvm.backend.cuda.axis``. Trainium's ``P``, ``F``, and +import axes from ``tvm_tirx.backend.cuda.axis``. Trainium's ``P``, ``F``, and ``Bank`` belong to ``T.trn.Axis``. The generic memory axis remains ``T.Axis.m``. CUDA memory/instruction layout factories are defined in -``tvm.backend.cuda.layout`` and exposed as ``T.cuda.layout``. Trainium +``tvm_tirx.backend.cuda.layout`` and exposed as ``T.cuda.layout``. Trainium layouts use explicit partition, free and bank axes. For example:: - from tvm.backend.trn.axis import P, F, Bank - from tvm.tirx.layout import TileLayout, S + from tvm_tirx.backend.trn.axis import P, F, Bank + from tvm_tirx.tirx.layout import TileLayout, S # A (128, 512) SBUF tile: partition p, free element f. sbuf_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)]) @@ -74,7 +74,7 @@ the full tensor shape before slicing an operand region. The query does not change the tensor type or its ``layout=None`` field. Noncompact global tensors and other storage scopes require an explicit layout when the instruction uses one. -.. automodule:: tvm.tirx.layout +.. automodule:: tvm_tirx.tirx.layout :members: :no-index: :exclude-members: Axis, ComposeLayout, Iter, Layout, TileLayout diff --git a/docs/tirx/api/ptx.rst b/docs/tirx/api/ptx.rst index c244389c77..f903f62097 100644 --- a/docs/tirx/api/ptx.rst +++ b/docs/tirx/api/ptx.rst @@ -22,11 +22,11 @@ The CUDA backend makes a table-driven PTX instruction namespace available when authoring TIRx with TVMScript. This is an authoring namespace, not a separate IR dialect:: - from tvm.script import tirx as Tx + from tvm_tirx.script import tirx as Tx Tx.ptx.add.rn.f32(dst, lhs, rhs) Each ``Tx.ptx`` call must match a form registered in -:mod:`tvm.backend.cuda.ptx.table`. The table validates the instruction family, +:mod:`tvm_tirx.backend.cuda.ptx.table`. The table validates the instruction family, modifier combination, operand count, and operand types while constructing one ``tirx.ptx.*`` IR call. A new instruction is added by extending the table and, when it needs a new representation, its code-generation support. @@ -187,11 +187,11 @@ TVM CUDA codegen does not emit these operations; use ``Tx.ptx`` in new kernels. Namespace API ------------- -.. autoclass:: tvm.backend.cuda.ptx.PTXNamespace +.. autoclass:: tvm_tirx.backend.cuda.ptx.PTXNamespace :members: :exclude-members: SINK :special-members: __getitem__ -.. automodule:: tvm.backend.cuda.ptx +.. automodule:: tvm_tirx.backend.cuda.ptx :members: :exclude-members: PTXNamespace diff --git a/docs/tirx/api/script.rst b/docs/tirx/api/script.rst index 196d407a9c..e4a85b764d 100644 --- a/docs/tirx/api/script.rst +++ b/docs/tirx/api/script.rst @@ -18,9 +18,9 @@ Core TVMScript ============== -TIRx kernels use ``tvm.script.tirx`` for the parser and core IR builders:: +TIRx kernels use ``tvm_tirx.script.tirx`` for the parser and core IR builders:: - from tvm.script import tirx as Tx + from tvm_tirx.script import tirx as Tx Tx.alloc_tensor(...) @@ -30,7 +30,7 @@ authoring layers and TIRx IR, see :ref:`tirx-programming-model`. Parser entry points ------------------- -.. automodule:: tvm.script.tirx +.. automodule:: tvm_tirx.script.tirx :members: function, jit, inline, macro :undoc-members: :no-index: @@ -39,18 +39,18 @@ Core IR builder --------------- Common statement frames and construction helpers live in -``tvm.script.ir_builder.frame`` and ``tvm.script.ir_builder.stmt``. The TIRx +``tvm_tirx.script.ir_builder.frame`` and ``tvm_tirx.script.ir_builder.stmt``. The TIRx namespace re-exports them alongside its function, tensor/layout, and execution scope builders. Native shared statement builders live under ``include/tvm/script/ir_builder`` and ``src/script/ir_builder``; dialect-specific thread-placement and tensor alias policies remain in the TIRx/S-TIR builders. -.. automodule:: tvm.tirx.script.ir_builder.ir +.. automodule:: tvm_tirx.tirx.script.ir_builder.ir :members: :no-index: :exclude-members: LetAnnotation -.. currentmodule:: tvm.tirx.script.ir_builder.ir +.. currentmodule:: tvm_tirx.tirx.script.ir_builder.ir .. py:class:: LetAnnotation(type_spec=None) @@ -65,9 +65,9 @@ construct IR. ``IRBuilder.line_trace_scope()`` records file-backed Python source lines. Script-facing void calls automatically emit into an active IRBuilder, with or without line tracing:: - from tvm.script import tirx as Tx - from tvm.script.ir_builder import IRBuilder - from tvm.tirx.script import ir_builder as I + from tvm_tirx.script import tirx as Tx + from tvm_tirx.script.ir_builder import IRBuilder + from tvm_tirx.tirx.script import ir_builder as I with IRBuilder() as builder: with I.function_(): @@ -83,7 +83,7 @@ Auto emission applies to script-facing calls through ``Tx``, including PTX and backend namespaces. A void call returns an ``AlreadyEmitted`` receipt; passing that receipt to ``evaluate`` or ``emit_`` does not emit it again. Calls producing values remain expressions. Raw IR constructors such as -``tvm.ir.Call``, ``tvm.tirx.op`` and raw backend APIs remain pure even inside +``tvm_tirx.ir.Call``, ``tvm_tirx.tirx.op`` and raw backend APIs remain pure even inside an IRBuilder. Line tracing only supplies source locations; it neither enables nor disables emission. The trace scope restores its previous trace and location state on exit, including exceptional exits, and applies only to its owning builder. @@ -121,9 +121,9 @@ A register fragment is an ordinary local tensor with an explicit layout. For a two-CTA Layout-B accumulator, define the TMEM and register layouts outside the kernel:: - from tvm.backend.cuda.axis import tid_in_wg - from tvm.backend.cuda.layout import tmem_mma_operand_layout - from tvm.tirx.layout import TileLayout, S + from tvm_tirx.backend.cuda.axis import tid_in_wg + from tvm_tirx.backend.cuda.layout import tmem_mma_operand_layout + from tvm_tirx.tirx.layout import TileLayout, S layout_b = tmem_mma_operand_layout( "D", (64, 128), "float32", M=128, cta_group=2) diff --git a/docs/tirx/api/tirx.rst b/docs/tirx/api/tirx.rst index f02bfe7720..c35e206a70 100644 --- a/docs/tirx/api/tirx.rst +++ b/docs/tirx/api/tirx.rst @@ -29,7 +29,7 @@ loads, and tensor-region constructors. Include ``tvm/tirx/type.h`` for ``TensorRegion`` expression from ``tvm/ir/expr.h`` and ``TensorRegionType`` from ``tvm/ir/type.h``. -.. automodule:: tvm.tirx +.. automodule:: tvm_tirx.tirx :members: :imported-members: :exclude-members: Call, CommReducer, ComposeLayout, DispatchContext, diff --git a/docs/tirx/api/trainium.rst b/docs/tirx/api/trainium.rst index 6794934a3e..c909e9ba53 100644 --- a/docs/tirx/api/trainium.rst +++ b/docs/tirx/api/trainium.rst @@ -27,35 +27,35 @@ The Trainium backend installs ``Tx.nki``. Its current operations are ``tensorscalar``, ``tensorscalar_reduce``, ``scalar_tensor_tensor``, ``scalar_tensor_scalar``, ``memset``, ``identity``, and ``affine_select``. -.. autoclass:: tvm.backend.trn.script.NKINamespace +.. autoclass:: tvm_tirx.backend.trn.script.NKINamespace :members: :no-index: Layout helpers -------------- -.. automodule:: tvm.backend.trn.layout +.. automodule:: tvm_tirx.backend.trn.layout :members: :no-index: Compilation pipeline -------------------- -.. automodule:: tvm.backend.trn.pipeline +.. automodule:: tvm_tirx.backend.trn.pipeline :members: :no-index: Transforms ---------- -.. automodule:: tvm.backend.trn.transform +.. automodule:: tvm_tirx.backend.trn.transform :members: :no-index: -.. autoclass:: tvm.backend.trn.transform.TrnNaiveAllocator +.. autoclass:: tvm_tirx.backend.trn.transform.TrnNaiveAllocator :members: :no-index: -.. autoclass:: tvm.backend.trn.transform.TrnPrivateTensorAlloc +.. autoclass:: tvm_tirx.backend.trn.transform.TrnPrivateTensorAlloc :members: :no-index: diff --git a/docs/tirx/api/transform.rst b/docs/tirx/api/transform.rst index 258e225d31..1d7a924b06 100644 --- a/docs/tirx/api/transform.rst +++ b/docs/tirx/api/transform.rst @@ -19,7 +19,7 @@ Compiler Transforms =================== -.. automodule:: tvm.tirx.transform +.. automodule:: tvm_tirx.tirx.transform :members: :exclude-members: Attrs, PrimExpr, Op, Call :imported-members:
