This is an automated email from the ASF dual-hosted git repository.
tqchen 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 d942d73d16 [REFACTOR][S-TIR] Own printer hooks and retire legacy
predicated loads (#20517)
d942d73d16 is described below
commit d942d73d16d4aac497386a2e55e5da759d730a62
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 2 13:41:14 2026 -0400
[REFACTOR][S-TIR] Own printer hooks and retire legacy predicated loads
(#20517)
Move S-TIR block, match-buffer, representation and operation printer
hooks into the S-TIR dialect. Keep shared PrimFunc construction in TIRx
while S-TIR owns its decorator and buffer-layout policy.
Retire the obsolete InjectPTXLDG32 pass and s_tir.ldg32 intrinsic,
including their pipeline configuration, Python API, dedicated
printer/CUDA lowering and legacy examples. Predicated loads use the
existing PTX dialect with explicit destination initialization and
preservation.
---
include/tvm/s_tir/transform.h | 7 -
python/tvm/backend/cuda/script.py | 14 +-
python/tvm/s_tir/backend/adreno/pipeline.py | 4 -
python/tvm/s_tir/pipeline.py | 4 -
python/tvm/s_tir/transform/transform.py | 11 --
src/backend/cuda/codegen/codegen_cuda.cc | 33 ----
src/backend/cuda/op/target_builtin.cc | 9 -
src/s_tir/script/printer/expr.cc | 56 ++++++
.../utils.h => s_tir/script/printer/function.cc} | 25 +--
.../block.cc => s_tir/script/printer/stmt.cc} | 123 +++----------
src/{ => s_tir}/script/printer/utils.h | 24 +--
src/s_tir/transform/inject_ptx_ldg32.cc | 199 ---------------------
src/script/printer/script_printer.cc | 16 +-
src/script/printer/utils.h | 15 ++
src/tirx/ir/transform.cc | 1 -
src/tirx/script/printer/block.cc | 143 ---------------
src/tirx/script/printer/buffer.cc | 25 +--
src/tirx/script/printer/expr.cc | 44 +----
src/tirx/script/printer/function.cc | 35 ++--
src/tirx/script/printer/utils.h | 4 +
tests/python/codegen/test_inject_ptx_ldg32.py | 70 --------
.../test_s_tir_transform_inject_ptx_ldg32.py | 94 ----------
tests/python/tirx/test_op_namespace_cleanup.py | 6 -
23 files changed, 144 insertions(+), 818 deletions(-)
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 8750ea4b7d..54493013f1 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -288,13 +288,6 @@ TVM_DLL Pass RewriteUnsafeSelect();
*/
TVM_DLL Pass InstrumentBoundCheckers();
-/*!
- * \brief Rewrite global to local memory copy on CUDA with ldg32 instruction.
- * \param enable_inject Whether to enable injection.
- * \return The pass.
- */
-TVM_DLL Pass InjectPTXLDG32(bool enable_inject = true);
-
/*!
* \brief Lower VTCM allocations.
* \return The pass.
diff --git a/python/tvm/backend/cuda/script.py
b/python/tvm/backend/cuda/script.py
index ccfd258975..28f6109b99 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -29,15 +29,6 @@ from tvm.tirx.script.ir_builder.op import _dtype_forward,
_op_wrapper
# pylint: disable=protected-access
-def _s_tir_ldg32(reg, guard, addr, local_addr):
- if is_buffer_var(addr):
- addr = addr[0]
- return _tir_op.call_intrin(reg.ty, "tirx.s_tir.ldg32", reg, guard, addr,
local_addr)
-
-
-_s_tir_ldg32.__tir_op_name__ = "s_tir.ldg32"
-
-
class _CpAsyncRaw:
"""The raw cp.async node's printer surface."""
@@ -71,14 +62,13 @@ class _CpAsyncRaw:
class STIRNamespace:
"""Nodes the s_tir pipeline's own passes build.
- Nothing here is meant to be written by hand: ``InjectPTXLDG32`` and
- ``InjectPTXAsyncCopy`` construct these, later passes match on them, and
+ Nothing here is meant to be written by hand: ``InjectPTXAsyncCopy``
+ constructs these, later passes match on them, and
codegen turns them into asm. They have a script spelling only so printed
IR round-trips.
"""
def __init__(self):
- self.ldg32 = _s_tir_ldg32
self.cp_async_raw = _CpAsyncRaw()
diff --git a/python/tvm/s_tir/backend/adreno/pipeline.py
b/python/tvm/s_tir/backend/adreno/pipeline.py
index dbd7f28fc3..32f486d4de 100644
--- a/python/tvm/s_tir/backend/adreno/pipeline.py
+++ b/python/tvm/s_tir/backend/adreno/pipeline.py
@@ -74,8 +74,6 @@ def default_tir_pipeline(*, prepare_only=False):
# Additional passes based on configuration.
if bool(config.get("tirx.instrument_bound_checkers", False)):
passes.append(s_tir.transform.InstrumentBoundCheckers())
- if bool(config.get("tirx.s_tir.ldg32", False)):
- passes.append(s_tir.transform.InjectPTXLDG32(True))
if not bool(config.get("tirx.disable_cse_tir", False)):
passes.append(tirx.transform.CommonSubexprElim())
passes.extend(
@@ -99,8 +97,6 @@ def default_tir_pipeline(*, prepare_only=False):
)
if bool(config.get("tirx.use_async_copy", False)):
passes.append(s_tir.transform.InjectPTXAsyncCopy())
- if bool(config.get("tirx.s_tir.ldg32", False)):
- passes.append(s_tir.transform.InjectPTXLDG32())
passes.append(s_tir.transform.MergeSharedMemoryAllocations())
if not prepare_only:
passes.extend(
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index e91b6756e0..ee069b4fb8 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -74,8 +74,6 @@ def default_s_tir_pipeline(*, prepare_only=False):
# Additional passes based on configuration.
if bool(config.get("tirx.instrument_bound_checkers", False)):
passes.append(s_tir.transform.InstrumentBoundCheckers())
- if bool(config.get("tirx.s_tir.ldg32", False)):
- passes.append(s_tir.transform.InjectPTXLDG32(True))
if not bool(config.get("tirx.disable_cse_tir", False)):
passes.append(tirx.transform.CommonSubexprElim())
passes.extend(
@@ -99,8 +97,6 @@ def default_s_tir_pipeline(*, prepare_only=False):
)
if bool(config.get("tirx.use_async_copy", False)):
passes.append(s_tir.transform.InjectPTXAsyncCopy())
- if bool(config.get("tirx.s_tir.ldg32", False)):
- passes.append(s_tir.transform.InjectPTXLDG32())
passes.append(s_tir.transform.MergeSharedMemoryAllocations())
if not prepare_only:
passes.extend(
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index c6b2e94a19..e6095a656b 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -331,17 +331,6 @@ def InstrumentBoundCheckers():
return _ffi_api.InstrumentBoundCheckers() # type: ignore
-def InjectPTXLDG32(enable_inject_ptx_intrin=True):
- """Inject ptx.ldg.32 intrinsics.
-
- Parameters
- ----------
- enable_inject_ptx_intrin : bool
- If True, inject ptx.ldg.32 intrinsics.
- """
- return _ffi_api.InjectPTXLDG32(enable_inject_ptx_intrin) # type: ignore
-
-
def VerifyVTCMLimit(default_target=None):
"""Verify if the size of the allocated vtcm memory satisfies the limit.
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 8b365cdfa5..31ce6f3692 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1046,7 +1046,6 @@ void CodeGenCUDA::Dispatch_(const CallNode* op,
std::ostream& os) {
static const Op ptx_ldmatrix_legacy_op = Op::Get("tirx.ptx_legacy.ldmatrix");
static const Op mma_store_legacy_op = Op::Get("tirx.mma_store_legacy");
static const Op mma_fill_legacy_op = Op::Get("tirx.mma_fill_legacy");
- static const Op ptx_ldg32_op = Op::Get("tirx.s_tir.ldg32");
static const Op cuda_func_call_op = Op::Get("tirx.cuda.func_call");
if (op->op.same_as(tvm_fill_fragment_op)) {
@@ -1274,38 +1273,6 @@ void CodeGenCUDA::Dispatch_(const CallNode* op,
std::ostream& os) {
os << "for (int i = 0; i < " << num_elem << "; ++i) {\n";
os << dst << "[" << dst_offset << " + i] = 0.0;";
os << "}\n";
- } else if (IsOp(op, ptx_ldg32_op, "tirx.s_tir.ldg32")) {
- /*
- asm volatile (
- "{.reg .pred p;\n"
- " setp.ne.b32 p, %2, 0;\n"
- // " @p ld.global.nc.f32 %0, [%1];}\n"t
- " @p ld.global.nc.L2::128B.f32 %0, [%1];}\n"
- : "=f"(reg)
- : "l"(addr), "r"((int)guard)
- );
- */
-
- // get local
- std::string reg = this->PrintExpr(op->args[0]);
- // get guard
- std::string guard = this->PrintExpr(op->args[1]);
- const TensorLoadNode* addr_buffer = op->args[2].as<TensorLoadNode>();
- std::string global_addr = this->PrintExpr(addr_buffer->indices[0]);
- std::string global_buffer =
-
this->PrintExpr(addr_buffer->source.as_or_throw<tvm::tirx::BufferVar>().data());
- std::string local_addr = this->PrintExpr(op->args[3]);
- this->stream << "asm volatile (\n";
- this->stream << "\"{.reg .pred p;\\n\"\n";
- this->stream << "\" setp.ne.b32 p, %2, 0;\\n\"\n";
- this->stream << "\" @!p mov.b32 %0, 0;\\n\"\n";
- this->stream << "\" @p ld.global.nc.f32 %0, [%1];}\\n\"\n";
- // stream << "\" @p ld.global.nc.L2::128B.f32 %0, [%1];}\\n\"\n" ;
- stream << ": \"=f\"(" << reg << "[" << local_addr << "]"
- << ")\n";
- stream << ": \"l\"((void*)(" << global_buffer << "+" << global_addr <<
")), \"r\"((int)"
- << guard << ")\n";
- stream << ");\n";
} else if (op->op.same_as(tirx::builtin::reinterpret())) {
// Compile-time pointer reinterpret of a literal (e.g. a tcgen05 descriptor
// template encoded at address 0): emit C++-style reinterpret_cast<T*>(...)
diff --git a/src/backend/cuda/op/target_builtin.cc
b/src/backend/cuda/op/target_builtin.cc
index 660d113f94..73b511e96a 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -129,14 +129,6 @@ void RegisterCudaTargetBuiltins() {
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
- OpDef("tirx.s_tir.ldg32")
- .signature(sig::arg("reg", "The register."), sig::arg("guard", "The
guard predicate."),
- sig::arg("addr", "The address."), sig::arg("local_addr", "The
local address."))
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure))
- .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.s_tir.ldg32"))
- .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
- .set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("s_tir"));
-
// Raw legacy cp.async form emitted by InjectPTXAsyncCopy (and round-tripped
by
// the T.s_tir.cp_async_raw.legacy 6-arg surface). It carries the element
dtype in Call.dtype
// and prints it dtype-first; user-issued copies go through T.ptx instead.
@@ -463,7 +455,6 @@ void RegisterDeviceIntrinsicAliases() {
OpDef("tirx.cuda.elect_sync").set_attr<TFixedReturnType>("TFixedReturnType",
PrimType::UInt(32));
OpDef("tirx.cuda.ldg")
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeCudaLdg>());
- OpDef("tirx.s_tir.ldg32").set_attr<TFixedReturnType>("TFixedReturnType",
PrimType::Float(32));
OpDef("tirx.cuda.atomic_add")
.set_attr<FInferType>("FInferType",
FInferType::FromNative<&InferTypeReturnArgType<1>>());
OpDef("tirx.cuda.atomic_cas")
diff --git a/src/s_tir/script/printer/expr.cc b/src/s_tir/script/printer/expr.cc
new file mode 100644
index 0000000000..0d78e365f0
--- /dev/null
+++ b/src/s_tir/script/printer/expr.cc
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include <tvm/ir/op.h>
+
+#include "../../../tirx/script/printer/utils.h"
+
+namespace tvm {
+namespace script {
+namespace printer {
+namespace details {
+namespace {
+
+ffi::Optional<ExprDoc> CpAsyncRawDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* call =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
CallNode>(input);
+ // This constructor takes an element dtype followed by the five stored
operands.
+ // The dtype is carried by the Call itself, not derived from a pointer
argument.
+ if (!CanTranslateExplicitResultCall(call) || call->args.size() != 5) {
+ return RawCall(d, call, false);
+ }
+ ffi::Array<ExprDoc> args = {TypeValue(d, call->ty)};
+ for (const Expr& arg : call->args) {
+ args.push_back(MaterializeCallArgument(d, arg, d->Translate(arg).value()));
+ }
+ return NamespaceDoc("tirx")->Attr("s_tir")->Attr("cp_async_raw")->Call(args);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ OpDef("tirx.s_tir.cp_async_raw")
+ .set_attr<FDocTranslate>(kOpCallDocTranslate,
+
FDocTranslate::FromNative<&CpAsyncRawDocTranslate>());
+}
+
+} // namespace
+
+} // namespace details
+} // namespace printer
+} // namespace script
+} // namespace tvm
diff --git a/src/script/printer/utils.h b/src/s_tir/script/printer/function.cc
similarity index 63%
copy from src/script/printer/utils.h
copy to src/s_tir/script/printer/function.cc
index ebb9581b8f..a488007154 100644
--- a/src/script/printer/utils.h
+++ b/src/s_tir/script/printer/function.cc
@@ -16,30 +16,23 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_SCRIPT_PRINTER_UTILS_H_
-#define TVM_SCRIPT_PRINTER_UTILS_H_
+#include <tvm/ir/function.h>
-#include <tvm/ffi/dtype.h>
-
-#include <string>
+#include "../../../tirx/script/printer/utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-/*!
- * \brief Return the script spelling of a data type.
- * \param dtype The data type to format.
- * \return The canonical dtype name, including the void sentinel spelling.
- */
-inline std::string DType2Str(DLDataType dtype) {
- return (((dtype).code == kDLOpaqueHandle) && ((dtype).bits == 0) &&
((dtype).lanes == 0))
- ? "void"
- : ffi::DLDataTypeToString(dtype);
+void PrintSTirPrimFunc(DocTranslatorObj* d, const tirx::PrimFuncNode* func) {
+ d->SetExtraState("tirx.buffer_default_layout_none", true);
+ PrintPrimFunc(d, func, NamespaceDoc("s_tir")->Attr("prim_func"),
tvm::attr::kSTir);
+ d->SetExtraState("tirx.buffer_default_layout_none", std::nullopt);
}
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
-
-#endif // TVM_SCRIPT_PRINTER_UTILS_H_
diff --git a/src/tirx/script/printer/block.cc b/src/s_tir/script/printer/stmt.cc
similarity index 60%
copy from src/tirx/script/printer/block.cc
copy to src/s_tir/script/printer/stmt.cc
index bff72e63a2..a57326ac57 100644
--- a/src/tirx/script/printer/block.cc
+++ b/src/s_tir/script/printer/stmt.cc
@@ -18,18 +18,14 @@
*/
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/stmt.h>
-#include <tvm/tirx/exec_scope.h>
-#include <optional>
-
-#include "../../../script/printer/ir/utils.h"
-#include "utils.h"
+#include "../../../script/printer/utils.h"
+#include "../../../tirx/script/printer/utils.h"
namespace tvm {
namespace script {
namespace printer {
namespace details {
-
namespace {
ffi::Array<StmtDoc> SBlockBody(DocTranslatorObj* d, const s_tir::SBlockNode*
block,
@@ -114,14 +110,7 @@ ffi::Array<StmtDoc> SBlockBody(DocTranslatorObj* d, const
s_tir::SBlockNode* blo
d->Emit(AssignDoc(lhs, rhs, std::nullopt),
ffi::GetRef<ffi::ObjectRef>(buffer.get()));
}
for (const s_tir::MatchBufferRegion& match : block->match_buffers) {
- ExprDoc source = d->Translate(match->source).value();
- CallDoc rhs =
d->Translate(match->buffer.var()->ty).value().as_or_throw<CallDoc>();
- TVM_FFI_CHECK(rhs->callee.as_or_throw<AttrAccessDoc>()->name ==
"Buffer", TypeError)
- << "Ts.match_buffer cannot reconstruct this nonrepresentable
BufferType";
- rhs->callee = NamespaceDoc("s_tir")->Attr("match_buffer");
- rhs->args.insert(rhs->args.begin(), source);
- IdDoc lhs = VarDoc(d, match->buffer);
- d->Emit(AssignDoc(lhs, rhs, std::nullopt),
ffi::GetRef<ffi::ObjectRef>(match.get()));
+ d->Translate(match);
}
if (block->init.has_value()) {
d->Emit(ScopeDoc(std::nullopt,
NamespaceDoc("s_tir")->Attr("init")->Call({}),
@@ -174,104 +163,32 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kDocTranslate, FDocTranslate::FromNative<&SBlockDocTranslate>());
}
-ffi::String ScopeIdApi(tirx::ScopeBinding scope) {
- auto [parent, current] = tirx::ScopeBindingToStringPair(scope);
- if (parent == "kernel" && current == "cluster") return "cluster_id";
- if (parent == "kernel" && current == "cta") return "cta_id";
- if (parent == "cluster" && current == "cta") return "cta_id_in_cluster";
- if (parent == "cluster" && current == "cta_pair") return "cta_id_in_pair";
- if (parent == "cta" && current == "warpgroup") return "warpgroup_id";
- if (parent == "cta" && current == "warp") return "warp_id";
- if (parent == "warpgroup" && current == "warp") return "warp_id_in_wg";
- if (parent == "warp" && current == "thread") return "lane_id";
- if (parent == "cta" && current == "thread") return "thread_id";
- if (parent == "warpgroup" && current == "thread") return "thread_id_in_wg";
- TVM_FFI_THROW(ValueError) << "printer unknown scope-id binding " << parent
<< "/" << current;
-}
-
-ExprDoc ScopeExtents(DocTranslatorObj* d, const ffi::Array<PrimExpr>& values) {
- ffi::Array<ExprDoc> items;
- for (const PrimExpr& value : values)
items.push_back(d->Translate(value).value());
- return ListDoc(items);
-}
-
-ffi::Optional<ExprDoc> ScopeIdDefStmtDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
- const ffi::Object*
destination) {
- const auto* stmt =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
tirx::ScopeIdDefStmtNode>(
- input);
+ffi::Optional<ExprDoc> MatchBufferRegionDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
+ const ffi::Object*
destination) {
+ const auto* match =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<
+ const s_tir::MatchBufferRegionNode>(input);
TVM_FFI_CHECK(destination == nullptr, TypeError)
<< "printer statement-only node cannot fulfill a destination";
- const tirx::ScopeIdDef& def = stmt->def;
- ffi::Array<ExprDoc> lhs;
- for (const PrimVar& var : def->def_ids) lhs.push_back(VarDoc(d, var));
- ffi::Array<ExprDoc> args;
- if (def->scope != tirx::ScopeBinding::kClusterCtaPair &&
def->extents.has_value()) {
- args.push_back(ScopeExtents(d, def->extents.value()));
- }
- ffi::Array<ffi::String> keys;
- ffi::Array<ExprDoc> values;
- if (def->preferred_extents.has_value()) {
- keys.push_back("preferred");
- values.push_back(ScopeExtents(d, def->preferred_extents.value()));
- }
- if (!def->def_ids.empty()) {
- PrimType dtype = def->def_ids[0].ty();
- for (const PrimVar& var : def->def_ids) {
- TVM_FFI_CHECK(var.ty() == dtype, TypeError) << "printer mixed scope-id
dtypes";
- }
- if (dtype != PrimType::Int(32)) {
- keys.push_back("dtype");
- values.push_back(LiteralDoc::Str(ffi::DLDataTypeToString(dtype->dtype),
std::nullopt));
- }
- }
- d->Emit(AssignDoc(TupleDoc(lhs),
-
NamespaceDoc("tirx")->Attr(ScopeIdApi(def->scope))->Call(args, keys, values),
- std::nullopt),
- ffi::GetRef<ffi::ObjectRef>(stmt));
+ ExprDoc source = d->Translate(match->source).value();
+ CallDoc rhs =
d->Translate(match->buffer.var()->ty).value().as_or_throw<CallDoc>();
+ TVM_FFI_CHECK(rhs->callee.as_or_throw<AttrAccessDoc>()->name == "Buffer",
TypeError)
+ << "Ts.match_buffer cannot reconstruct this nonrepresentable BufferType";
+ rhs->callee = NamespaceDoc("s_tir")->Attr("match_buffer");
+ rhs->args.insert(rhs->args.begin(), source);
+ IdDoc lhs = VarDoc(d, match->buffer);
+ d->Emit(AssignDoc(lhs, rhs, std::nullopt),
ffi::GetRef<ffi::ObjectRef>(match));
return std::nullopt;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<tirx::ScopeIdDefStmtNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&ScopeIdDefStmtDocTranslate>());
-}
-
-ffi::Optional<ExprDoc> ExecScopeDocTranslate(DocTranslatorObj*, ffi::AnyView
input,
- const ffi::Object*) {
- const auto* scope =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
tirx::ExecScopeNode>(input);
- return NamespaceDoc("tirx")
- ->Attr("ExecScope")
- ->Call({LiteralDoc::Str(scope->name(), std::nullopt)});
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<tirx::ExecScopeNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&ExecScopeDocTranslate>());
-}
-
-ffi::Optional<ExprDoc> ScopeIdDefDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
- const ffi::Object*) {
- const auto* def =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
tirx::ScopeIdDefNode>(input);
- auto [parent, current] = tirx::ScopeBindingToStringPair(def->scope);
- ExprDoc ids = AnyValue(d, def->def_ids);
- ExprDoc extents = AnyValue(d, def->extents);
- ExprDoc preferred = AnyValue(d, def->preferred_extents);
- d->RecordOrigin(ids, def->def_ids);
- if (def->extents.has_value()) d->RecordOrigin(extents, def->extents.value());
- if (def->preferred_extents.has_value())
- d->RecordOrigin(preferred, def->preferred_extents.value());
- return NamespaceDoc("tirx")
- ->Attr("ScopeIdDef")
- ->Call({ids, extents, LiteralDoc::Str(parent, std::nullopt),
- LiteralDoc::Str(current, std::nullopt), preferred});
+ ffi::reflection::TypeAttrDef<s_tir::MatchBufferRegionNode>().attr(
+ kDocTranslate,
FDocTranslate::FromNative<&MatchBufferRegionDocTranslate>());
}
TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<tirx::ScopeIdDefNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&ScopeIdDefDocTranslate>());
+ RegisterScriptRepr<s_tir::MatchBufferRegionNode>();
+ RegisterScriptRepr<s_tir::SBlockNode>();
+ RegisterScriptRepr<s_tir::SBlockRealizeNode>();
}
} // namespace
diff --git a/src/script/printer/utils.h b/src/s_tir/script/printer/utils.h
similarity index 63%
copy from src/script/printer/utils.h
copy to src/s_tir/script/printer/utils.h
index ebb9581b8f..58a716435b 100644
--- a/src/script/printer/utils.h
+++ b/src/s_tir/script/printer/utils.h
@@ -16,30 +16,22 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_SCRIPT_PRINTER_UTILS_H_
-#define TVM_SCRIPT_PRINTER_UTILS_H_
+#ifndef SRC_S_TIR_SCRIPT_PRINTER_UTILS_H_
+#define SRC_S_TIR_SCRIPT_PRINTER_UTILS_H_
-#include <tvm/ffi/dtype.h>
-
-#include <string>
+#include <tvm/script/printer/doc_translator.h>
+#include <tvm/tirx/function.h>
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-/*!
- * \brief Return the script spelling of a data type.
- * \param dtype The data type to format.
- * \return The canonical dtype name, including the void sentinel spelling.
- */
-inline std::string DType2Str(DLDataType dtype) {
- return (((dtype).code == kDLOpaqueHandle) && ((dtype).bits == 0) &&
((dtype).lanes == 0))
- ? "void"
- : ffi::DLDataTypeToString(dtype);
-}
+void PrintSTirPrimFunc(DocTranslatorObj* d, const tirx::PrimFuncNode* func);
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
-#endif // TVM_SCRIPT_PRINTER_UTILS_H_
+#endif // SRC_S_TIR_SCRIPT_PRINTER_UTILS_H_
diff --git a/src/s_tir/transform/inject_ptx_ldg32.cc
b/src/s_tir/transform/inject_ptx_ldg32.cc
deleted file mode 100644
index 34e68068fb..0000000000
--- a/src/s_tir/transform/inject_ptx_ldg32.cc
+++ /dev/null
@@ -1,199 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/op.h>
-#include <tvm/s_tir/analysis.h>
-#include <tvm/s_tir/stmt_functor.h>
-#include <tvm/s_tir/transform.h>
-#include <tvm/sym/analyzer.h>
-#include <tvm/sym/iter_affine_map.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt.h>
-
-#include "../../sym/const_fold.h"
-#include "../../sym/pattern_match.h"
-
-namespace tvm {
-namespace s_tir {
-using namespace tvm::tirx;
-
-class PTXRewriter : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
- UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode)
override {
- if (value.as<ExprNode>()) return ffi::Unchanged();
- return StmtExprMutator::Mutate(value, inplace_mode);
- }
-
- Stmt AddAllocationsIfNeeded(Stmt body) {
- if (!needs_buffer || has_buffer_2) {
- return body;
- }
- EnsureBuffers();
- body = SeqStmt::Flatten(
- Bind(predicate_buffer.var(),
- Call(predicate_buffer.type(), tirx::builtin::alloc_buffer(),
- {tvm::Tuple(predicate_buffer->shape),
DataTypeImm(predicate_buffer->dtype->dtype),
- StringImm(predicate_buffer.scope())},
- DictAttrs())),
- Bind(addr_buffer.var(),
- Call(addr_buffer.type(), tirx::builtin::alloc_buffer(),
- {tvm::Tuple(addr_buffer->shape),
DataTypeImm(addr_buffer->dtype->dtype),
- StringImm(addr_buffer.scope())},
- DictAttrs())),
- body);
- has_buffer_2 = true;
- return body;
- }
-
- UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
- if (const auto* call = op->value.as<CallNode>();
- call && call->op.same_as(tirx::builtin::alloc_buffer())) {
- return Mutate_AllocBuffer(op, inplace_mode);
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- UnchangedOr<Stmt> Mutate_AllocBuffer(const BindNode* op, InplaceMode
inplace_mode) {
- Stmt result =
- StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- if (needs_buffer && !has_buffer_2) {
- EnsureBuffers();
- has_buffer_2 = true;
- result = SeqStmt::Flatten(
- Bind(predicate_buffer.var(), Call(predicate_buffer.type(),
tirx::builtin::alloc_buffer(),
-
{tvm::Tuple(predicate_buffer->shape),
-
DataTypeImm(predicate_buffer->dtype->dtype),
-
StringImm(predicate_buffer.scope())},
- DictAttrs())),
- Bind(addr_buffer.var(),
- Call(addr_buffer.type(), tirx::builtin::alloc_buffer(),
- {tvm::Tuple(addr_buffer->shape),
DataTypeImm(addr_buffer->dtype->dtype),
- StringImm(addr_buffer.scope())},
- DictAttrs())),
- result);
- }
- return result;
- }
-
- UnchangedOr<Stmt> Mutate_(const BufferStoreNode* store, InplaceMode
inplace_mode) final {
- Stmt result =
- StmtExprMutator::Mutate_(store,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(store));
- BufferVar load_buffer = store->buffer;
- PrimExpr load_value = store->value;
- // const TensorLoadNode* gload = load_value.as<TensorLoadNode>(); // take
- // the place of instance of
- const CallNode* call = load_value.as<CallNode>();
- if (call != nullptr) {
- const OpNode* op = call->op.as<OpNode>();
- if (op != nullptr && op->name == "prim.if_then_else") {
- PrimExpr predicate = call->args[0].as_or_throw<PrimExpr>();
- PrimExpr lhs = call->args[1].as_or_throw<PrimExpr>();
- PrimExpr rhs = call->args[2].as_or_throw<PrimExpr>();
- PrimExpr global_addr, local_addr;
- const TensorLoadNode* load = lhs.as<TensorLoadNode>();
- PrimExpr imm_value = rhs;
- if (load == nullptr) {
- load = rhs.as<TensorLoadNode>();
- imm_value = lhs;
- if (load == nullptr) {
- return result;
- }
- }
- global_addr = load->indices[0];
- const prim::RampNode* ramp = global_addr.as<prim::RampNode>();
- if (ramp != nullptr) {
- return result;
- }
- EnsureBuffers();
- needs_buffer = true;
- local_addr = store->indices[0];
- BufferStore addr_store(addr_buffer, global_addr, {IntImm::Int32(0)});
- BufferStore local_addr_store(addr_buffer, local_addr,
{IntImm::Int32(1)});
- BufferStore predicate_store(predicate_buffer, predicate,
{IntImm::Int32(0)});
- PrimExpr new_lhs, new_rhs, new_predicate, new_indice;
- new_lhs = BufferLoad(load->source.as_or_throw<tvm::tirx::BufferVar>(),
- {BufferLoad(addr_buffer, {IntImm::Int32(0)})});
- new_rhs = IntImm::Int32(0);
- new_predicate = BufferLoad(predicate_buffer, {IntImm::Int32(0)});
- new_indice = BufferLoad(addr_buffer, {IntImm::Int32(1)});
- BufferStore value_store(store->buffer, imm_value, {new_indice});
- static const Op ptx_ldg32_op = Op::Get("tirx.s_tir.ldg32");
- Evaluate ptx_load(Call(store->buffer->dtype, ptx_ldg32_op,
- {store->buffer.data(), new_predicate, new_lhs,
new_indice})
- .as_or_throw<PrimExpr>());
- ffi::Array<Stmt> tmp_seq = {addr_store, local_addr_store,
predicate_store, value_store,
- ptx_load};
- SeqStmt seq_stmt = SeqStmt(tmp_seq);
- return seq_stmt;
- }
- }
- return result;
- }
-
- void EnsureBuffers() {
- if (has_buffer_1) {
- return;
- }
- has_buffer_1 = true;
- // addr[0] -> global_addr / addr[1] -> local_addr
- addr_buffer = decl_buffer({IntImm::Int32(2)}, PrimType::Int(32), "addr",
"local");
- predicate_buffer = decl_buffer({IntImm::Int32(1)}, PrimType::Bool(),
"predicate", "local");
- }
-
- bool has_buffer_1 = false, has_buffer_2 = false;
- bool needs_buffer = false;
- BufferVar addr_buffer, predicate_buffer;
-};
-
-namespace transform {
-
-Pass InjectPTXLDG32(bool enable_inject_ptx_intrin) {
- auto pass_func = [enable_inject_ptx_intrin](PrimFunc f, IRModule m,
PassContext ctx) {
- if (enable_inject_ptx_intrin) {
- auto target = f->GetAttr<Target>("target");
- if (!target.has_value() || target.value()->kind->name != "cuda") {
- return f;
- }
- auto* n = f.CopyOnWrite();
- auto rewriter = ffi::make_object<PTXRewriter>();
- Stmt body = rewriter->Mutate(n->body).ValueOrUnchanged(n->body);
- n->body = rewriter->AddAllocationsIfNeeded(body);
- // inject ptx
- }
- return f;
- };
- return CreatePrimFuncPass(pass_func, 0, "s_tir.InjectPTXLDG32", {});
-}
-
-// The pass can now be invoked via the pass infrastructure, but we also add a
-// Python binding for it
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.InjectPTXLDG32", InjectPTXLDG32);
-}
-
-} // namespace transform
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/script/printer/script_printer.cc
b/src/script/printer/script_printer.cc
index 56d831ad55..a2fb3728e3 100644
--- a/src/script/printer/script_printer.cc
+++ b/src/script/printer/script_printer.cc
@@ -28,7 +28,6 @@
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
#include <tvm/runtime/logging.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/script/printer/doc_translator.h>
#include <tvm/script/printer/printer.h>
#include <tvm/te/operation.h>
@@ -42,6 +41,8 @@
#include <sstream>
#include <utility>
+#include "utils.h"
+
namespace tvm {
namespace {
@@ -61,15 +62,6 @@ std::string RenderFallbackWithInvisiblePathInfo(const
ffi::String& script,
return os.str();
}
-template <typename ObjectType>
-void RegisterScriptRepr() {
- namespace refl = ffi::reflection;
- refl::TypeAttrDef<ObjectType>().def(refl::type_attr::kRepr,
- [](ffi::ObjectRef obj, ffi::Function) ->
ffi::String {
- return
RedirectedReprPrinterMethod(obj);
- });
-}
-
} // namespace
std::string Script(const ffi::ObjectRef& node, const
ffi::Optional<PrinterConfig>& cfg) {
@@ -100,6 +92,7 @@ std::string RedirectedReprPrinterMethod(const
ffi::ObjectRef& obj) {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = ffi::reflection;
+ using script::printer::details::RegisterScriptRepr;
refl::GlobalDef()
.def("node.TVMScriptPrinterScript", tvm::Script)
.def("script.printer.ReprPrintRelax",
@@ -173,9 +166,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
RegisterScriptRepr<relax::distributed::DTensorTypeNode>();
RegisterScriptRepr<relax::distributed::DeviceMeshNode>();
RegisterScriptRepr<relax::distributed::PlacementNode>();
- RegisterScriptRepr<s_tir::MatchBufferRegionNode>();
- RegisterScriptRepr<s_tir::SBlockNode>();
- RegisterScriptRepr<s_tir::SBlockRealizeNode>();
RegisterScriptRepr<te::CommReducerNode>();
RegisterScriptRepr<te::ReduceNode>();
RegisterScriptRepr<tirx::AssertStmtNode>();
diff --git a/src/script/printer/utils.h b/src/script/printer/utils.h
index ebb9581b8f..b0e916eb76 100644
--- a/src/script/printer/utils.h
+++ b/src/script/printer/utils.h
@@ -20,6 +20,8 @@
#define TVM_SCRIPT_PRINTER_UTILS_H_
#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/script/printer/config.h>
#include <string>
@@ -38,6 +40,19 @@ inline std::string DType2Str(DLDataType dtype) {
: ffi::DLDataTypeToString(dtype);
}
+namespace details {
+
+template <typename ObjectType>
+void RegisterScriptRepr() {
+ namespace refl = ffi::reflection;
+ refl::TypeAttrDef<ObjectType>().def(refl::type_attr::kRepr,
+ [](ffi::ObjectRef obj, ffi::Function) ->
ffi::String {
+ return
RedirectedReprPrinterMethod(obj);
+ });
+}
+
+} // namespace details
+
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/tirx/ir/transform.cc b/src/tirx/ir/transform.cc
index eb512dd90b..f702004c45 100644
--- a/src/tirx/ir/transform.cc
+++ b/src/tirx/ir/transform.cc
@@ -46,7 +46,6 @@
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.debug_keep_trivial_loop", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.use_async_copy", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.merge_static_smem", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.vtcm_capacity", int64_t);
-TVM_REGISTER_PASS_CONFIG_OPTION("tirx.s_tir.ldg32", bool);
TVM_REGISTER_PASS_CONFIG_OPTION("tirx.enable_fast_math", bool);
/*!
diff --git a/src/tirx/script/printer/block.cc b/src/tirx/script/printer/block.cc
index bff72e63a2..80cfafe921 100644
--- a/src/tirx/script/printer/block.cc
+++ b/src/tirx/script/printer/block.cc
@@ -17,7 +17,6 @@
* under the License.
*/
#include <tvm/ir/prim/op.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/exec_scope.h>
#include <optional>
@@ -32,148 +31,6 @@ namespace details {
namespace {
-ffi::Array<StmtDoc> SBlockBody(DocTranslatorObj* d, const s_tir::SBlockNode*
block,
- const s_tir::SBlockRealizeNode* realize) {
- auto docs = d->WithDocScope([&]() {
- TVM_FFI_CHECK(!realize || realize->iter_values.size() ==
block->iter_vars.size(), TypeError)
- << "SBlockRealize binding count must match its iter vars";
- ffi::Array<AssignDoc> axes;
- for (size_t i = 0; i < block->iter_vars.size(); ++i) {
- const tirx::IterVar& iter = block->iter_vars[i];
- ffi::String kind;
- switch (iter->iter_type) {
- case tirx::IterVarType::kDataPar:
- kind = "spatial";
- break;
- case tirx::IterVarType::kCommReduce:
- kind = "reduce";
- break;
- case tirx::IterVarType::kOrdered:
- kind = "scan";
- break;
- case tirx::IterVarType::kOpaque:
- kind = "opaque";
- break;
- default:
- TVM_FFI_THROW(TypeError) << "printer unsupported SBlock iter var
kind";
- }
- ExprDoc domain =
- prim::is_zero(iter->dom->min) && iter->dom->min.ty() ==
iter->dom->extent.ty()
- ? d->Translate(iter->dom->extent).value()
- : NamespaceDoc("ir")
- ->Attr("Range")
- ->Attr("from_min_extent")
- ->Call({d->Translate(iter->dom->min).value(),
- d->Translate(iter->dom->extent).value()});
- ffi::Array<ExprDoc> args = {domain};
- if (realize)
args.push_back(d->Translate(realize->iter_values[i]).value());
- IdDoc lhs = VarDoc(d, iter->var);
- ffi::Array<ffi::String> keys;
- ffi::Array<ExprDoc> values;
- if (iter->var.ty() != PrimType::Int(32)) {
- keys.push_back("dtype");
- values.push_back(
- LiteralDoc::Str(ffi::DLDataTypeToString(iter->var.ty()->dtype),
std::nullopt));
- }
- axes.push_back(
- AssignDoc(lhs,
NamespaceDoc("s_tir")->Attr("axis")->Attr(kind)->Call(args, keys, values),
- std::nullopt));
- }
- for (size_t i = 0; i < axes.size(); ++i) d->Emit(axes[i],
block->iter_vars[i]);
- if (realize && !tvm::prim::is_one(realize->predicate)) {
- d->Emit(ExprStmtDoc(NamespaceDoc("s_tir")->Attr("where")->Call(
- {d->Translate(realize->predicate).value()})),
- realize->predicate);
- }
- ffi::Array<ExprDoc> reads;
- for (const TensorRegion& region : block->reads)
reads.push_back(d->Translate(region).value());
- d->Emit(ExprStmtDoc(NamespaceDoc("s_tir")->Attr("reads")->Call(reads)),
block->reads);
- ffi::Array<ExprDoc> writes;
- for (const TensorRegion& region : block->writes)
writes.push_back(d->Translate(region).value());
- d->Emit(ExprStmtDoc(NamespaceDoc("s_tir")->Attr("writes")->Call(writes)),
block->writes);
- if (!block->annotations.empty()) {
- d->Emit(
- ExprStmtDoc(
- NamespaceDoc("s_tir")->Attr("sblock_attr")->Call({AnyValue(d,
block->annotations)})),
- block->annotations);
- }
- for (const tirx::BufferVar& buffer : block->alloc_buffers) {
- CallDoc rhs =
d->Translate(buffer.var()->ty).value().as_or_throw<CallDoc>();
- TVM_FFI_CHECK(rhs->callee.as_or_throw<AttrAccessDoc>()->name ==
"Buffer", TypeError)
- << "Ts.sblock_alloc_buffer cannot reconstruct this nonrepresentable
BufferType";
- const auto* buffer_type = buffer.var()->ty.as<tirx::BufferTypeNode>();
- TVM_FFI_CHECK(
- buffer_type->allocated_addr.empty() ||
- (buffer_type->storage_scope != "global" &&
buffer_type->storage_scope != "shared" &&
- buffer_type->storage_scope != "shared.dyn" &&
buffer_type->storage_scope != "local"),
- TypeError)
- << "Ts.sblock_alloc_buffer does not accept allocated_addr in "
- << buffer_type->storage_scope;
- rhs->callee = NamespaceDoc("s_tir")->Attr("sblock_alloc_buffer");
- IdDoc lhs = VarDoc(d, buffer);
- d->Emit(AssignDoc(lhs, rhs, std::nullopt),
ffi::GetRef<ffi::ObjectRef>(buffer.get()));
- }
- for (const s_tir::MatchBufferRegion& match : block->match_buffers) {
- ExprDoc source = d->Translate(match->source).value();
- CallDoc rhs =
d->Translate(match->buffer.var()->ty).value().as_or_throw<CallDoc>();
- TVM_FFI_CHECK(rhs->callee.as_or_throw<AttrAccessDoc>()->name ==
"Buffer", TypeError)
- << "Ts.match_buffer cannot reconstruct this nonrepresentable
BufferType";
- rhs->callee = NamespaceDoc("s_tir")->Attr("match_buffer");
- rhs->args.insert(rhs->args.begin(), source);
- IdDoc lhs = VarDoc(d, match->buffer);
- d->Emit(AssignDoc(lhs, rhs, std::nullopt),
ffi::GetRef<ffi::ObjectRef>(match.get()));
- }
- if (block->init.has_value()) {
- d->Emit(ScopeDoc(std::nullopt,
NamespaceDoc("s_tir")->Attr("init")->Call({}),
- Body(block->init.value(), d)),
- block->init.value());
- }
- d->Translate(block->body);
- });
- return ToStmtDocArray(docs);
-}
-
-ffi::Optional<ExprDoc> SBlockRealizeDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
- const ffi::Object*
destination) {
- const auto* realize =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
s_tir::SBlockRealizeNode>(
- input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- const s_tir::SBlockNode* block = realize->block.get();
- d->Emit(ScopeDoc(std::nullopt,
- NamespaceDoc("s_tir")->Attr("sblock")->Call(
- {LiteralDoc::Str(block->name_hint, std::nullopt)}),
- SBlockBody(d, block, realize)),
- ffi::GetRef<ffi::ObjectRef>(realize));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<s_tir::SBlockRealizeNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&SBlockRealizeDocTranslate>());
-}
-
-ffi::Optional<ExprDoc> SBlockDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
- const ffi::Object* destination) {
- const auto* block =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
s_tir::SBlockNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- d->Emit(ScopeDoc(std::nullopt,
- NamespaceDoc("s_tir")->Attr("sblock")->Call(
- {LiteralDoc::Str(block->name_hint, std::nullopt)},
{"no_realize"},
- {LiteralDoc::Boolean(true, std::nullopt)}),
- SBlockBody(d, block, nullptr)),
- ffi::GetRef<ffi::ObjectRef>(block));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<s_tir::SBlockNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&SBlockDocTranslate>());
-}
-
ffi::String ScopeIdApi(tirx::ScopeBinding scope) {
auto [parent, current] = tirx::ScopeBindingToStringPair(scope);
if (parent == "kernel" && current == "cluster") return "cluster_id";
diff --git a/src/tirx/script/printer/buffer.cc
b/src/tirx/script/printer/buffer.cc
index ed5536ef2e..c6d69053fc 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -20,7 +20,6 @@
#include <tvm/ir/op.h>
#include <tvm/ir/prim/op.h>
#include <tvm/ir/prim/vector_expr.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/layout.h>
@@ -214,7 +213,7 @@ ffi::Optional<ExprDoc>
BufferTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
}
}
if (!buffer->layout.has_value()) {
- bool defaults_to_none =
d->GetOrCreateExtraState<bool>("s_tir.function_scope") ||
+ bool defaults_to_none =
d->GetOrCreateExtraState<bool>("tirx.buffer_default_layout_none") ||
buffer->storage_scope == "trn.sbuf" ||
buffer->storage_scope == "trn.psum";
if (!defaults_to_none) {
@@ -447,28 +446,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kDocTranslate, FDocTranslate::FromNative<&ComposeLayoutDocTranslate>());
}
-ffi::Optional<ExprDoc> MatchBufferRegionDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
- const ffi::Object*
destination) {
- const auto* match =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<
- const s_tir::MatchBufferRegionNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- ExprDoc source = d->Translate(match->source).value();
- CallDoc rhs =
d->Translate(match->buffer.var()->ty).value().as_or_throw<CallDoc>();
- TVM_FFI_CHECK(rhs->callee.as_or_throw<AttrAccessDoc>()->name == "Buffer",
TypeError)
- << "Ts.match_buffer cannot reconstruct this nonrepresentable BufferType";
- rhs->callee = NamespaceDoc("s_tir")->Attr("match_buffer");
- rhs->args.insert(rhs->args.begin(), source);
- IdDoc lhs = VarDoc(d, match->buffer);
- d->Emit(AssignDoc(lhs, rhs, std::nullopt),
ffi::GetRef<ffi::ObjectRef>(match));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<s_tir::MatchBufferRegionNode>().attr(
- kDocTranslate,
FDocTranslate::FromNative<&MatchBufferRegionDocTranslate>());
-}
-
} // namespace
} // namespace details
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index 895f4c657a..4b106b672b 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -88,6 +88,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FDocTranslate::FromNative<&VarDocTranslate>());
}
+bool CanTranslateExplicitResultCall(const CallNode* call) {
+ return !call->attrs.defined() && call->ty_args.empty() &&
call->ty.as<PrimType>() &&
+ std::all_of(call->args.begin(), call->args.end(), [](const Expr& arg)
{
+ return !arg->ty.IsMissing() && !arg.as<TensorRegionNode>();
+ });
+}
+
namespace {
ffi::Optional<ExprDoc> IterVarDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
@@ -176,13 +183,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kDocTranslate, FDocTranslate::FromNative<&LambdaExprDocTranslate>());
}
-bool CanTranslateExplicitResultCall(const CallNode* call) {
- return !call->attrs.defined() && call->ty_args.empty() &&
call->ty.as<PrimType>() &&
- std::all_of(call->args.begin(), call->args.end(), [](const Expr& arg)
{
- return !arg->ty.IsMissing() && !arg.as<TensorRegionNode>();
- });
-}
-
ffi::Optional<ExprDoc> StorageSyncDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
const ffi::Object*) {
const auto* call =
@@ -237,28 +237,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FDocTranslate::FromNative<&CallExternDocTranslate>());
}
-ffi::Optional<ExprDoc> CpAsyncRawDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
- const ffi::Object*) {
- const auto* call =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
CallNode>(input);
- // This constructor takes an element dtype followed by the five stored
operands.
- // The dtype is carried by the Call itself, not derived from a pointer
argument.
- if (!CanTranslateExplicitResultCall(call) || call->args.size() != 5) {
- return RawCall(d, call, false);
- }
- ffi::Array<ExprDoc> args = {TypeValue(d, call->ty)};
- for (const Expr& arg : call->args) {
- args.push_back(MaterializeCallArgument(d, arg, d->Translate(arg).value()));
- }
- return NamespaceDoc("tirx")->Attr("s_tir")->Attr("cp_async_raw")->Call(args);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- OpDef("tirx.s_tir.cp_async_raw")
- .set_attr<FDocTranslate>(kOpCallDocTranslate,
-
FDocTranslate::FromNative<&CpAsyncRawDocTranslate>());
-}
-
ffi::Optional<ExprDoc> CUDAFuncCallDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
const ffi::Object*) {
const auto* call =
@@ -533,14 +511,6 @@ ffi::Optional<ExprDoc>
TIRCallDocTranslate(DocTranslatorObj* d, const CallNode*
// dedicated translation covers each signature.
bool incompatible_signature = (op.value()->name == "tirx.cuda.ldg" &&
call->args.size() != 2) ||
op.value()->name == "tirx.cuda.wait_until";
- // The named ldg32 constructor uses the destination expression's type and
- // converts a source BufferVar to its first load. Preserve other native
forms.
- if (op.value()->name == "tirx.s_tir.ldg32") {
- incompatible_signature =
- call->args.size() != 4 || !call->args[0].as<TensorLoadNode>() ||
- !ffi::StructuralEqual()(call->args[0]->ty, result_type) ||
- (call->args[2].as<VarNode>() &&
call->args[2]->ty.as<tirx::BufferTypeNode>());
- }
// IKET lowering accepts a variadic IR signature, but its named
constructors
// expose just the event/token and an optional payload.
if (op.value()->name == "tirx.cuda.iket_mark" ||
diff --git a/src/tirx/script/printer/function.cc
b/src/tirx/script/printer/function.cc
index 233f71a9c5..2e7333594c 100644
--- a/src/tirx/script/printer/function.cc
+++ b/src/tirx/script/printer/function.cc
@@ -20,7 +20,6 @@
#include <tvm/ir/function.h>
#include <tvm/ir/global_info.h>
#include <tvm/ir/prim/op.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/script/printer/printer.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt_functor.h>
@@ -32,6 +31,7 @@
#include <utility>
#include <vector>
+#include "../../../s_tir/script/printer/utils.h"
#include "../../../script/printer/ir/utils.h"
#include "utils.h"
@@ -40,19 +40,10 @@ namespace script {
namespace printer {
namespace details {
-namespace {
-
-ffi::Optional<ExprDoc> PrimFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
- const ffi::Object* destination) {
- const auto* func =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
tirx::PrimFuncNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
+void PrintPrimFunc(DocTranslatorObj* d, const tirx::PrimFuncNode* func,
ExprDoc decorator,
+ const ffi::String& dialect_attr) {
VarScope vars(d);
- bool legacy_s_tir = func->attrs->dict.count(tvm::attr::kSTir);
- if (legacy_s_tir) d->SetExtraState("s_tir.function_scope", true);
-
ffi::String name =
func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).value_or("main");
FunctionDoc doc(ffi::UnsafeInit{});
{
@@ -109,8 +100,6 @@ ffi::Optional<ExprDoc>
PrimFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView in
std::nullopt));
}
}
- ExprDoc decorator = legacy_s_tir ? NamespaceDoc("s_tir")->Attr("prim_func")
- : NamespaceDoc("tirx")->Attr("prim_func");
ffi::Array<ffi::String> decorator_keys;
ffi::Array<ExprDoc> decorator_values;
if (!func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) {
@@ -127,7 +116,7 @@ ffi::Optional<ExprDoc>
PrimFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView in
body.insert(body.begin(), thread_declarations.begin(),
thread_declarations.end());
std::vector<std::pair<ffi::String, ffi::Any>> attrs;
for (const auto& [key, value] : func->attrs->dict) {
- if (key != tvm::attr::kGlobalSymbol && key != tvm::attr::kSTir &&
+ if (key != tvm::attr::kGlobalSymbol && (dialect_attr.empty() || key !=
dialect_attr) &&
key != tirx::attr::kPersistentKernel)
attrs.emplace_back(key, value);
}
@@ -151,9 +140,23 @@ ffi::Optional<ExprDoc>
PrimFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView in
doc = FunctionDoc(IdDoc(name), args, {decorator}, ret_type, body);
FinalizeFunctionDefinitions(d, signature_candidates, doc);
}
- if (legacy_s_tir) d->SetExtraState("s_tir.function_scope", std::nullopt);
vars.Close();
d->Emit(doc, ffi::GetRef<ffi::ObjectRef>(func));
+}
+
+namespace {
+
+ffi::Optional<ExprDoc> PrimFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
+ const ffi::Object* destination) {
+ const auto* func =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
tirx::PrimFuncNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ if (func->attrs->dict.count(tvm::attr::kSTir)) {
+ PrintSTirPrimFunc(d, func);
+ } else {
+ PrintPrimFunc(d, func, NamespaceDoc("tirx")->Attr("prim_func"), "");
+ }
return std::nullopt;
}
diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h
index 0b0213f514..acbb73c0b5 100644
--- a/src/tirx/script/printer/utils.h
+++ b/src/tirx/script/printer/utils.h
@@ -19,6 +19,7 @@
#ifndef SRC_TIRX_SCRIPT_PRINTER_UTILS_H_
#define SRC_TIRX_SCRIPT_PRINTER_UTILS_H_
+#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/var.h>
@@ -31,6 +32,9 @@ namespace script {
namespace printer {
namespace details {
+void PrintPrimFunc(DocTranslatorObj* d, const tirx::PrimFuncNode* func,
ExprDoc decorator,
+ const ffi::String& dialect_attr);
+bool CanTranslateExplicitResultCall(const CallNode* call);
bool IsScalarBuffer(DocTranslatorObj* d, const Expr& source);
ffi::Array<StmtDoc> Body(const tirx::Stmt& stmt, DocTranslatorObj* d);
ffi::Optional<ExprDoc> TIRCallPrefixDocTranslate(DocTranslatorObj* d, const
CallNode* call);
diff --git a/tests/python/codegen/test_inject_ptx_ldg32.py
b/tests/python/codegen/test_inject_ptx_ldg32.py
deleted file mode 100644
index 9ccf0a0a00..0000000000
--- a/tests/python/codegen/test_inject_ptx_ldg32.py
+++ /dev/null
@@ -1,70 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-import numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-from tvm.script import tirx as T
-from tvm.testing import env
-
-
[email protected]_func
-def vector_add(A: T.Buffer((16), "float32"), B: T.Buffer((32), "float32")) ->
None:
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
- bx = T.env_thread("blockIdx.x")
- tx = T.env_thread("threadIdx.x")
- T.launch_thread(bx, 1)
- T.launch_thread(tx, 32)
- A_local = T.alloc_buffer(32, "float32", scope="local")
- T.evaluate(
- T.call_intrin("float32", "tirx.s_tir.ldg32", A_local.data, tx % 2 ==
0, A[tx // 2], tx)
- )
- B[tx] = A_local[tx] + 1.0
-
-
[email protected]
[email protected](not env.has_cuda(), reason="need cuda")
-def test_inject_ptx_intrin():
- f = vector_add
- arch = tvm.support.nvcc.get_target_compute_version()
- major, _ = tvm.support.nvcc.parse_compute_version(arch)
- if major < 8:
- # Require at least SM80
- return
- mod = tvm.compile(f, target="cuda")
- A_np = np.random.rand(16).astype("float32")
- B_np = np.zeros(32).astype("float32")
- C_np = np.zeros(32).astype("float32")
-
- for i in range(32):
- if i % 2 == 0:
- C_np[i] = A_np[i // 2]
- C_np[i] += 1.0
-
- def run_and_check():
- dev = tvm.cuda(0)
- A_nd = tvm.runtime.tensor(A_np, device=dev)
- B_nd = tvm.runtime.tensor(B_np, device=dev)
- mod(A_nd, B_nd)
- tvm.testing.assert_allclose(B_nd.numpy(), C_np)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-if __name__ == "__main__":
- test_inject_ptx_intrin()
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_ldg32.py
b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_ldg32.py
deleted file mode 100644
index f99f70e13d..0000000000
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_ldg32.py
+++ /dev/null
@@ -1,94 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# ruff: noqa: F401
-
-import tvm_ffi
-
-import tvm
-import tvm.testing
-from tvm import s_tir
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-
-def _is_buffer_binding(node, *op_names):
- return (
- isinstance(node, tvm.tirx.Bind)
- and isinstance(node.value, tvm.ir.Call)
- and isinstance(node.value.op, tvm.ir.Op)
- and node.value.op.name in op_names
- )
-
-
-def _count_alloc(stmt):
- num_alloc = [0]
-
- def visit(n):
- if _is_buffer_binding(n, "tirx.alloc_buffer"):
- num_alloc[0] += 1
-
- tvm_ffi.structural_walk(stmt, visit)
- return num_alloc[0]
-
-
-def _count_ptx_ldg32(stmt):
- num_call = [0]
-
- def visit(n):
- if isinstance(n, tvm.ir.Call) and n.op.name == "tirx.s_tir.ldg32":
- num_call[0] += 1
-
- tvm_ffi.structural_walk(stmt, visit)
- return num_call[0]
-
-
[email protected]_func
-def where_no_alloc(A: T.Buffer((4,), "float32"), C: T.Buffer((4,), "float32"))
-> None:
- T.func_attr({"global_symbol": "main", "tirx.noalias": True, "target":
T.target("cuda")})
- for i in range(4):
- C[i] = T.if_then_else(A[i] > T.float32(0), A[i], T.float32(0))
-
-
[email protected]_func
-def where_no_alloc_cpu(A: T.Buffer((4,), "float32"), C: T.Buffer((4,),
"float32")) -> None:
- T.func_attr({"global_symbol": "main", "tirx.noalias": True, "target":
T.target("llvm")})
- for i in range(4):
- C[i] = T.if_then_else(A[i] > T.float32(0), A[i], T.float32(0))
-
-
-def test_inject_ptx_ldg32_inserts_alloc_for_no_alloc_func():
- mod = tvm.IRModule.from_expr(where_no_alloc)
- assert _count_alloc(mod["main"].body) == 0
-
- mod = tvm.s_tir.transform.InjectPTXLDG32()(mod)
- assert _count_alloc(mod["main"].body) > 0
- assert _count_ptx_ldg32(mod["main"].body) == 1
-
-
-def test_inject_ptx_ldg32_skip_non_cuda_target():
- mod = tvm.IRModule.from_expr(where_no_alloc_cpu)
- cpu_target = tvm.target.Target("llvm")
- mod = tvm.IRModule({"main": mod["main"].with_attr("target", cpu_target)})
- assert _count_alloc(mod["main"].body) == 0
-
- mod = tvm.s_tir.transform.InjectPTXLDG32()(mod)
- assert _count_alloc(mod["main"].body) == 0
- assert _count_ptx_ldg32(mod["main"].body) == 0
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/tirx/test_op_namespace_cleanup.py
b/tests/python/tirx/test_op_namespace_cleanup.py
index 8742546b81..c2337421b5 100644
--- a/tests/python/tirx/test_op_namespace_cleanup.py
+++ b/tests/python/tirx/test_op_namespace_cleanup.py
@@ -290,25 +290,20 @@ def
test_backend_load_updates_tirx_alias_and_script_facades(monkeypatch):
def test_device_intrinsic_printer_roundtrips_canonical_namespaces():
@T.prim_func
def device_namespaces(dst: T.handle, A: T.Buffer((1,), "float32")):
- Result = T.alloc_buffer((1,), "float32", scope="local")
T.cuda.cta_sync()
- T.s_tir.ldg32(Result[0], 1, A[0], 0)
T.metal.simd_shuffle(A[0], 0)
T.metal.simd_shuffle_up(A[0], 1)
T.metal.simd_shuffle_down(A[0], 1)
calls = _expr_calls(device_namespaces)
assert [call.op.name for call in calls] == [
- "tirx.alloc_buffer",
"tirx.cuda.cta_sync",
- "tirx.s_tir.ldg32",
"tirx.metal.simd_shuffle",
"tirx.metal.simd_shuffle_up",
"tirx.metal.simd_shuffle_down",
]
for op_name, namespace in [
("tirx.cuda.cta_sync", "cuda"),
- ("tirx.s_tir.ldg32", "s_tir"),
("tirx.metal.simd_shuffle", "metal"),
("tirx.metal.simd_shuffle_up", "metal"),
("tirx.metal.simd_shuffle_down", "metal"),
@@ -319,7 +314,6 @@ def
test_device_intrinsic_printer_roundtrips_canonical_namespaces():
code = device_namespaces.script()
assert "T.cuda.cta_sync(" in code
- assert "T.s_tir.ldg32(" in code
assert "T.metal.simd_shuffle(" in code
assert "T.metal.simd_shuffle_up(" in code
assert "T.metal.simd_shuffle_down(" in code
