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 7e913b5826 [REFACTOR] Localize and simplify internal TIRx helpers
(#20683)
7e913b5826 is described below
commit 7e913b58268d88ee36428c9e22c58830b0eb9269
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 19:11:54 2026 -0400
[REFACTOR] Localize and simplify internal TIRx helpers (#20683)
Move expression replacement and tensor allocation-shape handling into
S-TIR, and keep WMMA fragment parsing with SPIR-V code generation.
Replace trivial ABI construction, pointer access, storage-scope and
statement-prefix helpers at their use sites, and remove unused
utilities. This reduces private TIRx helper dependencies while
preserving existing traversal, lowering and code-generation behavior.
---
src/backend/vulkan/codegen/codegen_spirv.cc | 33 +++-
src/s_tir/schedule/primitive/cache_index.cc | 2 +-
src/s_tir/transform/inject_ptx_async_copy.cc | 15 +-
.../transform/merge_shared_memory_allocations.cc | 23 ++-
.../transform/replace_selected_expr.cc | 8 +-
.../transform/replace_selected_expr.h | 18 +--
src/s_tir/transform/storage_access.cc | 6 +-
src/s_tir/transform/tensorcore_infer_fragment.cc | 8 +-
src/target/llvm/codegen_llvm.cc | 5 +-
src/target/source/codegen_c.cc | 9 +-
src/tirx/ir/pointer_type.h | 57 -------
src/tirx/transform/ir_utils.cc | 100 ------------
src/tirx/transform/ir_utils.h | 169 ---------------------
src/tirx/transform/lower_tvm_builtin.cc | 153 +++++++++++++------
src/tirx/transform/make_packed_api.cc | 13 +-
src/tirx/transform/storage_rewrite.cc | 38 +++--
src/tirx/transform/tvm_ffi_binder.cc | 55 ++++---
17 files changed, 265 insertions(+), 447 deletions(-)
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index 6b3823d544..70674d4292 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -32,10 +32,11 @@
#include <tvm/tirx/op/region.h>
#include <tvm/tirx/stmt.h>
+#include <stdexcept>
#include <string>
+#include <utility>
#include "../../../runtime/pack_args.h"
-#include "../../../tirx/transform/ir_utils.h"
#include "../runtime/vulkan_common.h"
namespace tvm {
@@ -44,6 +45,36 @@ using namespace tvm::prim;
namespace {
+int Stoi(const std::string& str) {
+ try {
+ return std::stoi(str);
+ } catch (std::invalid_argument& e) {
+ TVM_FFI_THROW(InternalError) << "Cannot convert \"" << str << "\" to int";
+ throw;
+ }
+}
+
+std::pair<int32_t, int32_t> GetWmmaFragmentDimSize(const std::string&
shape_str,
+ const std::string& scope) {
+ size_t m, n, k;
+ size_t last_pos = 0, pos = 0;
+ pos = shape_str.find(", ", last_pos);
+ m = Stoi(shape_str.substr(last_pos, pos - last_pos));
+ last_pos = pos + 2;
+ pos = shape_str.find(", ", last_pos);
+ n = Stoi(shape_str.substr(last_pos, pos - last_pos));
+ last_pos = pos + 2;
+ k = Stoi(shape_str.substr(last_pos, shape_str.length() - last_pos));
+ if (scope == "wmma.matrix_a") {
+ return std::pair<int32_t, int32_t>(m, k);
+ } else if (scope == "wmma.matrix_b") {
+ return std::pair<int32_t, int32_t>(k, n);
+ } else if (scope == "wmma.accumulator") {
+ return std::pair<int32_t, int32_t>(m, n);
+ }
+ return std::pair<int32_t, int32_t>(0, 0);
+}
+
const IntImmNode* AsIntImmNode(const Expr& expr) {
const IntImmNode* node = expr.as<IntImmNode>();
TVM_FFI_ICHECK(node);
diff --git a/src/s_tir/schedule/primitive/cache_index.cc
b/src/s_tir/schedule/primitive/cache_index.cc
index 3f0f7e308d..4e3db2e8ec 100644
--- a/src/s_tir/schedule/primitive/cache_index.cc
+++ b/src/s_tir/schedule/primitive/cache_index.cc
@@ -22,7 +22,7 @@
#include <tvm/s_tir/stmt.h>
#include <tvm/sym/int_set.h>
-#include "../../../tirx/transform/replace_selected_expr.h"
+#include "../../transform/replace_selected_expr.h"
#include "../utils.h"
#include "cache_index_helpers.h"
diff --git a/src/s_tir/transform/inject_ptx_async_copy.cc
b/src/s_tir/transform/inject_ptx_async_copy.cc
index 05171cb68a..99c904c8b5 100644
--- a/src/s_tir/transform/inject_ptx_async_copy.cc
+++ b/src/s_tir/transform/inject_ptx_async_copy.cc
@@ -31,7 +31,6 @@
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op/memory.h>
-#include "../../tirx/ir/pointer_type.h"
#include "storage_access.h"
#include "tvm/s_tir/stmt.h"
#include "tvm/tirx/stmt.h"
@@ -73,25 +72,21 @@ class PTXAsyncCopyInjector : public StmtExprMutator {
static_cast<int>(load->source.as_or_throw<tvm::tirx::TensorVar>()->dtype.StorageBytes());
if (bytes == 4 || bytes == 8 || bytes == 16) {
- auto dst_elem_type =
-
GetPtrElementType(store->dest.as_or_throw<TensorVar>().type()->DataPtrType());
- auto src_elem_type = GetPtrElementType(
-
load->source.as_or_throw<tvm::tirx::TensorVar>().type()->DataPtrType());
- TVM_FFI_ICHECK(dst_elem_type.has_value() && src_elem_type.has_value())
- << "Both store and load tensor should have a pointer type
annotation.";
+ PrimType dst_elem_type = store->dest.as_or_throw<TensorVar>()->dtype;
+ PrimType src_elem_type = load->source.as_or_throw<TensorVar>()->dtype;
int index_factor = 1;
- if (dst_elem_type.value() != src_elem_type.value()) {
+ if (dst_elem_type != src_elem_type) {
// The only case where src and dst have different dtypes is when the
dst shared memory
// is a byte tensor generated by merging dynamic shared memory.
TVM_FFI_ICHECK(store->dest.as_or_throw<TensorVar>().scope() ==
"shared.dyn");
- TVM_FFI_ICHECK((dst_elem_type.value() == PrimType::UInt(8)));
+ TVM_FFI_ICHECK((dst_elem_type == PrimType::UInt(8)));
// TensorStore/Load have the "pointer reinterpret" semantics
according to their
// "value" dtype. Their "indices" are supposed to be applied after
such pointer cast,
// for example: ((*float16)(byte_buffer))[tensor->indices] =
fp16_value;
// To replace TensorStore/Load with cp.async, we need to multiply
the store index by
// the byte size of the "value" dtype, to get the correct offset
into the byte tensor.
- index_factor =
static_cast<int>(src_elem_type.value().StorageBytes());
+ index_factor = static_cast<int>(src_elem_type.StorageBytes());
}
if (indices_lanes == 1) {
diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc
b/src/s_tir/transform/merge_shared_memory_allocations.cc
index 35282988ef..4777468db6 100644
--- a/src/s_tir/transform/merge_shared_memory_allocations.cc
+++ b/src/s_tir/transform/merge_shared_memory_allocations.cc
@@ -34,6 +34,7 @@
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
+#include <tvm/sym/analyzer.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/op/region.h>
#include <tvm/tirx/stmt.h>
@@ -45,7 +46,6 @@
#include "../../runtime/thread_storage_scope.h"
#include "../../support/arena.h"
-#include "../../tirx/transform/ir_utils.h"
namespace tvm {
namespace s_tir {
@@ -56,6 +56,19 @@ using runtime::StorageScope;
namespace {
+ffi::Array<PrimExpr> GetTensorAllocationShape(const TensorVar& tensor) {
+ ffi::Array<PrimExpr> alloc_shape = tensor->shape;
+ if (tensor->strides.size()) {
+ TVM_FFI_ICHECK_EQ(tensor->shape.size(), tensor->strides.size());
+ for (size_t i = tensor->strides.size() - 1; i > 0; --i) {
+ TVM_FFI_ICHECK(
+ sym::Analyzer()->CanProveEqual(floormod(tensor->strides[i - 1],
tensor->strides[i]), 0));
+ alloc_shape.Set(i, tensor->strides[i - 1] / tensor->strides[i]);
+ }
+ }
+ return alloc_shape;
+}
+
ffi::Optional<Var> GetBufferDataVar(const ffi::Any& data) {
if (auto var = data.as<Var>()) {
return var;
@@ -70,12 +83,16 @@ ffi::Optional<Var> GetBufferDataVar(const ffi::Any& data) {
} // namespace
bool IsDynamicSharedMemory(Var buffer_var) {
- StorageScope storage_scope =
runtime::StorageScope::Create(GetPtrStorageScope(buffer_var));
+ StorageScope storage_scope = runtime::StorageScope::Create(
+ (buffer_var->ty.as<TensorTypeNode>() ?
buffer_var->ty.as<TensorTypeNode>()->storage_scope
+ :
buffer_var->ty.as_or_throw<PtrType>()->storage_scope));
return storage_scope.rank == runtime::StorageRank::kShared &&
storage_scope.tag == ".dyn";
}
bool IsStaticSharedMemory(Var buffer_var) {
- StorageScope storage_scope =
runtime::StorageScope::Create(GetPtrStorageScope(buffer_var));
+ StorageScope storage_scope = runtime::StorageScope::Create(
+ (buffer_var->ty.as<TensorTypeNode>() ?
buffer_var->ty.as<TensorTypeNode>()->storage_scope
+ :
buffer_var->ty.as_or_throw<PtrType>()->storage_scope));
return storage_scope.rank == runtime::StorageRank::kShared &&
storage_scope.tag == "";
}
diff --git a/src/tirx/transform/replace_selected_expr.cc
b/src/s_tir/transform/replace_selected_expr.cc
similarity index 96%
rename from src/tirx/transform/replace_selected_expr.cc
rename to src/s_tir/transform/replace_selected_expr.cc
index 93e96481de..178ecb9b3d 100644
--- a/src/tirx/transform/replace_selected_expr.cc
+++ b/src/s_tir/transform/replace_selected_expr.cc
@@ -35,7 +35,7 @@
#include <tvm/tirx/transform.h> // For the declaration of the pass
namespace tvm {
-namespace tirx {
+namespace s_tir {
/*!
* \brief Toplevel (static) function that replace in an expression
@@ -87,7 +87,7 @@
ReplaceSelectedExpr::ReplaceSelectedExpr(std::function<bool(const PrimExpr&)> pr
can_replace_inside_(can_replace_inside) {}
/*!
- * \brief The method which overrides the generic dispatcher of StmtExprMutator
+ * \brief The method which overrides the generic dispatcher of
tirx::StmtExprMutator
* \param expr The expression to mutate
*/
UnchangedOr<ffi::Any> ReplaceSelectedExpr::Mutate(ffi::AnyView input,
InplaceMode inplace_mode) {
@@ -95,8 +95,8 @@ UnchangedOr<ffi::Any>
ReplaceSelectedExpr::Mutate(ffi::AnyView input, InplaceMod
if (predicate_selector_(prim_expr.value())) return new_expr_;
if (!can_replace_inside_(prim_expr.value())) return ffi::Unchanged();
}
- return StmtExprMutator::Mutate(input, inplace_mode);
+ return tirx::StmtExprMutator::Mutate(input, inplace_mode);
}
-} // namespace tirx
+} // namespace s_tir
} // namespace tvm
diff --git a/src/tirx/transform/replace_selected_expr.h
b/src/s_tir/transform/replace_selected_expr.h
similarity index 86%
rename from src/tirx/transform/replace_selected_expr.h
rename to src/s_tir/transform/replace_selected_expr.h
index efa95e6edb..0b79ad12d4 100644
--- a/src/tirx/transform/replace_selected_expr.h
+++ b/src/s_tir/transform/replace_selected_expr.h
@@ -24,26 +24,26 @@
with a predicate by another expression.
*/
-#ifndef TVM_TIR_TRANSFORM_REPLACE_SELECTED_EXPR_H_
-#define TVM_TIR_TRANSFORM_REPLACE_SELECTED_EXPR_H_
+#ifndef TVM_S_TIR_TRANSFORM_REPLACE_SELECTED_EXPR_H_
+#define TVM_S_TIR_TRANSFORM_REPLACE_SELECTED_EXPR_H_
#include <tvm/ir/expr_functor.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h> // For the class StmtExprMutator
+#include <tvm/tirx/stmt_functor.h> // For the class tirx::StmtExprMutator
namespace tvm {
-namespace tirx {
+namespace s_tir {
/*!
* \brief Mutator for replacing the expressions selected by a predicate in a
statement and/or
in an expression, which only replace inside of nodes in which it is
allowed to perform
replacecements (given by a second predicate)
*/
-class ReplaceSelectedExpr : public StmtExprMutator {
+class ReplaceSelectedExpr : public tirx::StmtExprMutator {
public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
+ using tirx::StmtExprMutator::Mutate;
+ using tirx::StmtExprMutator::Mutate_;
// Toplevel (static) functions
static PrimExpr ReplaceSelectedExprInExpr(
const PrimExpr& expr, std::function<bool(const PrimExpr&)>
predicate_selector,
@@ -69,7 +69,7 @@ class ReplaceSelectedExpr : public StmtExprMutator {
std::function<bool(const PrimExpr&)> can_replace_inside_;
};
-} // namespace tirx
+} // namespace s_tir
} // namespace tvm
-#endif // TVM_TIR_TRANSFORM_REPLACE_SELECTED_EXPR_H_
+#endif // TVM_S_TIR_TRANSFORM_REPLACE_SELECTED_EXPR_H_
diff --git a/src/s_tir/transform/storage_access.cc
b/src/s_tir/transform/storage_access.cc
index 2632b4f9b1..3d894b0dd5 100644
--- a/src/s_tir/transform/storage_access.cc
+++ b/src/s_tir/transform/storage_access.cc
@@ -32,8 +32,6 @@
#include <string>
#include <utility>
-#include "../../tirx/transform/ir_utils.h"
-
namespace tvm {
namespace s_tir {
using namespace tvm::tirx;
@@ -410,8 +408,8 @@ StorageScope StorageAccessVisitor::GetScope(Var buffer_var)
const {
if (auto tensor_type = buffer_var->ty.as<TensorType>()) {
return StorageScope::Create(tensor_type.value()->storage_scope);
}
- if (buffer_var->ty.as<PtrTypeNode>()) {
- return StorageScope::Create(GetPtrStorageScope(buffer_var));
+ if (const auto* pointer_type = buffer_var->ty.as<PtrTypeNode>()) {
+ return StorageScope::Create(pointer_type->storage_scope);
}
return StorageScope(); // global by default
}
diff --git a/src/s_tir/transform/tensorcore_infer_fragment.cc
b/src/s_tir/transform/tensorcore_infer_fragment.cc
index 8b6d24ba9f..477c793e16 100644
--- a/src/s_tir/transform/tensorcore_infer_fragment.cc
+++ b/src/s_tir/transform/tensorcore_infer_fragment.cc
@@ -79,7 +79,9 @@ class FragmentGetter : public s_tir::StmtExprVisitor {
TVM_FFI_ICHECK(k);
TVM_FFI_ICHECK(layout);
- std::string scope = GetPtrStorageScope(ffi::GetRef<Var>(buffer_var));
+ std::string scope = (buffer_var->ty.as<TensorTypeNode>()
+ ?
buffer_var->ty.as<TensorTypeNode>()->storage_scope
+ :
buffer_var->ty.as_or_throw<PtrType>()->storage_scope);
if (fragments.count(buffer_var)) {
// check if the fragment has met before
FragmentInfo info = fragments[buffer_var];
@@ -114,7 +116,9 @@ class FragmentGetter : public s_tir::StmtExprVisitor {
TVM_FFI_ICHECK(n);
TVM_FFI_ICHECK(k);
- std::string scope = GetPtrStorageScope(ffi::GetRef<Var>(buffer_var));
+ std::string scope = (buffer_var->ty.as<TensorTypeNode>()
+ ?
buffer_var->ty.as<TensorTypeNode>()->storage_scope
+ :
buffer_var->ty.as_or_throw<PtrType>()->storage_scope);
if (fragments.count(buffer_var)) {
FragmentInfo info = fragments[buffer_var];
TVM_FFI_ICHECK_EQ(m->value, info.m);
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 658c3aa178..c341747acb 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -712,8 +712,9 @@ void CodeGenLLVM::GetAlignment(PrimType t, const VarNode*
buf_var, const PrimExp
auto it = alloc_storage_info_.find(buf_var);
if (it != alloc_storage_info_.end()) {
const StorageInfo& info = it->second;
- *p_native_bits = NativeVectorBits(
-
runtime::StorageScope::Create(GetPtrStorageScope(ffi::GetRef<Var>(buf_var))));
+ *p_native_bits = NativeVectorBits(runtime::StorageScope::Create(
+ (buf_var->ty.as<TensorTypeNode>() ?
buf_var->ty.as<TensorTypeNode>()->storage_scope
+ :
buf_var->ty.as_or_throw<PtrType>()->storage_scope)));
max_align_bits = info.alignment * 8;
} else {
*p_native_bits = native_vector_bits_;
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 294735afc7..75639ffe5f 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -41,7 +41,6 @@
#include <limits>
#include "../../sym/pattern_match.h"
-#include "../../tirx/ir/pointer_type.h"
#include "codegen_params.h"
namespace tvm {
@@ -436,13 +435,15 @@ void CodeGenC::RegisterHandleType(const VarNode* buf_var,
const PrimType& t) {
void CodeGenC::RegisterHandleTypeFromPointer(const tvm::Var& var, const Expr*
value) {
if (value == nullptr) return;
- std::optional<PrimType> value_dtype = tirx::GetPtrElementType((*value)->ty);
- if (!value_dtype.has_value()) return;
+ const auto* pointer_type = (*value)->ty.as<PtrTypeNode>();
+ if (!pointer_type) return;
+ const auto* value_dtype = pointer_type->element_type.as<PrimTypeNode>();
+ if (!value_dtype) return;
auto* call = value->as<CallNode>();
if (call != nullptr && call->op.same_as(tirx::ptr_byte_offset_op())) {
pointer_offset_vars_.insert(var.get());
}
- RegisterHandleType(var.get(), value_dtype.value());
+ RegisterHandleType(var.get(), ffi::GetRef<PrimType>(value_dtype));
}
void CodeGenC::PrintVecElemLoad(const std::string& vec, const PrimType& t, int
i,
diff --git a/src/tirx/ir/pointer_type.h b/src/tirx/ir/pointer_type.h
deleted file mode 100644
index c78c74c3a0..0000000000
--- a/src/tirx/ir/pointer_type.h
+++ /dev/null
@@ -1,57 +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.
- */
-/*!
- * \file tirx/ir/pointer_type.h
- * \brief Primitive pointer type utility
- */
-#ifndef TVM_TIR_IR_POINTER_TYPE_H_
-#define TVM_TIR_IR_POINTER_TYPE_H_
-
-#include <tvm/ffi/dtype.h>
-#include <tvm/ir/type.h>
-
-#include <optional>
-
-namespace tvm {
-namespace tirx {
-
-/*!
- * \brief Returns the type of object pointed to.
- *
- * \param type The type to be checked.
- *
- * \return An std::optional<PrimType> object. If the type is a pointer
- * to a primitive type, the object has a value which is the pointed-to
- * type. Otherwise the object is nullopt.
- */
-inline std::optional<PrimType> GetPtrElementType(const Type& type) {
- if (!type.as<MissingType>().has_value()) {
- if (auto* ptr_type = type.as<PtrTypeNode>()) {
- if (auto* prim_type = ptr_type->element_type.as<PrimTypeNode>()) {
- return ffi::GetRef<PrimType>(prim_type);
- }
- }
- }
-
- return std::nullopt;
-}
-
-} // namespace tirx
-} // namespace tvm
-#endif // TVM_TIR_IR_POINTER_TYPE_H_
diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc
index 8e7b15201f..536eb6da8e 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -43,43 +43,6 @@ namespace tvm {
namespace tirx {
using namespace tvm::prim;
-Stmt MergeNest(const std::vector<Stmt>& nest, Stmt body) {
- // use reverse iteration
- for (auto ri = nest.rbegin(); ri != nest.rend(); ++ri) {
- Stmt s = *ri;
- if (const auto* for_ = s.as<ForNode>()) {
- auto n = ffi::make_object<ForNode>(*for_);
- TVM_FFI_ICHECK(tirx::IsNoOp(n->body));
- n->body = body;
- body = Stmt(n);
- } else if (const auto* bind = s.as<BindNode>()) {
- // Bind has no body -- prepend it before the accumulated body in a
SeqStmt.
- body = SeqStmt({ffi::GetRef<Stmt>(bind), body});
- } else if (const auto* ite = s.as<IfNode>()) {
- auto n = ffi::make_object<IfNode>(*ite);
- TVM_FFI_ICHECK(tirx::IsNoOp(n->then_case));
- TVM_FFI_ICHECK(!n->else_case);
- n->then_case = body;
- body = Stmt(n);
- } else if (const auto* seq = s.as<SeqStmtNode>()) {
- // A sequence nest is a prefix, followed by the accumulated body.
- body = SeqStmt({s, body}, seq->loc);
- } else if (s.as<AssertStmtNode>()) {
- body = SeqStmt({s, body});
- } else {
- TVM_FFI_THROW(InternalError) << "not supported nest type";
- }
- }
- return body;
-}
-
-Stmt MergeNest(const std::vector<std::vector<Stmt>>& nest, Stmt body) {
- for (auto ri = nest.rbegin(); ri != nest.rend(); ++ri) {
- body = MergeNest(*ri, body);
- }
- return body;
-}
-
Function IRConvertSSA::VisitFunction(Function func) {
std::unordered_set<const VarNode*> parameter_symbols;
// Define explicit parameters before the symbolic values in their types.
@@ -296,69 +259,6 @@ Stmt ConvertSSA(Stmt stmt) {
return ffi::make_object<IRConvertSSA>()->Mutate(stmt,
InplaceMode::kAllow).ValueOrUnchanged(stmt);
}
-ffi::String GetPtrStorageScope(Var buffer_var) {
- if (const auto* tensor_type = buffer_var->ty.as<TensorTypeNode>()) {
- return tensor_type->storage_scope;
- }
- const auto* ptr_type = buffer_var->ty.as<PtrTypeNode>();
- TVM_FFI_ICHECK(ptr_type)
- << "The provided variable is neither a pointer nor a tensor-typed
variable";
- return ptr_type->storage_scope;
-}
-
-ffi::Array<PrimExpr> GetTensorAllocationShape(const TensorVar& tensor) {
- ffi::Array<PrimExpr> alloc_shape = tensor->shape;
- if (tensor->strides.size()) {
- TVM_FFI_ICHECK_EQ(tensor->shape.size(), tensor->strides.size());
- for (size_t i = tensor->strides.size() - 1; i > 0; --i) {
- TVM_FFI_ICHECK(
- sym::Analyzer()->CanProveEqual(floormod(tensor->strides[i - 1],
tensor->strides[i]), 0));
- alloc_shape.Set(i, tensor->strides[i - 1] / tensor->strides[i]);
- }
- }
- return alloc_shape;
-}
-
-int Stoi(const std::string& str) {
- try {
- return std::stoi(str);
- } catch (std::invalid_argument& e) {
- TVM_FFI_THROW(InternalError) << "Cannot convert \"" << str << "\" to int";
- throw;
- }
-}
-
-std::pair<int32_t, int32_t> GetWmmaFragmentDimSize(const std::string&
shape_str,
- const std::string& scope) {
- size_t m, n, k;
- size_t last_pos = 0, pos = 0;
- pos = shape_str.find(", ", last_pos);
- m = Stoi(shape_str.substr(last_pos, pos - last_pos));
- last_pos = pos + 2;
- pos = shape_str.find(", ", last_pos);
- n = Stoi(shape_str.substr(last_pos, pos - last_pos));
- last_pos = pos + 2;
- k = Stoi(shape_str.substr(last_pos, shape_str.length() - last_pos));
- if (scope == "wmma.matrix_a") {
- return std::pair<int32_t, int32_t>(m, k);
- } else if (scope == "wmma.matrix_b") {
- return std::pair<int32_t, int32_t>(k, n);
- } else if (scope == "wmma.accumulator") {
- return std::pair<int32_t, int32_t>(m, n);
- }
- return std::pair<int32_t, int32_t>(0, 0);
-}
-
-std::optional<bool> IsHostFunc(const Function& func) {
- if (func->HasNonzeroAttr(tvm::tirx::attr::kIsHostFunc)) {
- return true;
- } else if (auto target = func->GetAttr<Target>(tvm::attr::kTarget)) {
- return target.value()->HasKey("cpu");
- } else {
- return std::nullopt;
- }
-}
-
IRModule IRConvertSSA::VisitIRModule(IRModule mod) {
ffi::Map<GlobalVar, BaseFunc> functions;
bool made_change = false;
diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h
index c9f5e45911..614fe5e7e9 100644
--- a/src/tirx/transform/ir_utils.h
+++ b/src/tirx/transform/ir_utils.h
@@ -38,7 +38,6 @@
#include <tvm/tirx/stmt_functor.h>
#include <functional>
-#include <limits>
#include <optional>
#include <string>
#include <unordered_map>
@@ -49,118 +48,6 @@
namespace tvm {
namespace tirx {
-/*!
- * \brief Fill empty loop/conditional bodies and prepend statement prefixes.
- * \param nest Loops, conditionals with empty bodies, or Bind/Assert/SeqStmt
prefixes.
- * \param body body
- * \return The combined Stmt
- */
-Stmt MergeNest(const std::vector<Stmt>& nest, Stmt body);
-
-/*!
- * \brief Fill empty loop/conditional bodies and prepend statement prefixes.
- * \param nest Groups of loops, conditionals with empty bodies, or statement
prefixes.
- * \param body body
- * \return The combined Stmt
- */
-Stmt MergeNest(const std::vector<std::vector<Stmt>>& nest, Stmt body);
-
-/*!
- * \brief update array with an unary function
- * \param arr array
- * \param fupdate an unary function
- * \tparam T type of array element
- * \tparam F type of the unary function
- * \return if update happens, return the new array, else return the
- * original array
- */
-template <typename T, typename F>
-inline ffi::Array<T> UpdateArray(ffi::Array<T> arr, F fupdate) {
- std::vector<T> new_arr(arr.size());
- bool changed = false;
- for (size_t i = 0; i < arr.size(); ++i) {
- T old_elem = arr[i];
- T new_elem = fupdate(old_elem);
- if (!new_elem.same_as(old_elem)) changed = true;
- new_arr[i] = new_elem;
- }
- if (!changed) {
- return arr;
- } else {
- return ffi::Array<T>(new_arr);
- }
-}
-
-/*!
- * \brief Get construct from struct
- * \param dtype The data type.
- * \param handle the struct handle.
- * \param index the offset index.
- * \param kind The data kind.
- * \return the get expression.
- */
-inline Expr TVMStructGet(Type type, Var handle, int index,
tirx::TVMStructFieldKind kind) {
- ffi::Array<Expr> args = {handle, IntImm::Int32(index),
IntImm::Int32(static_cast<int>(kind))};
- return Call(std::move(type), tirx::abi_field_get_op(), args);
-}
-
-inline PrimExpr TVMStructGet(PrimType type, Var handle, int index,
tirx::TVMStructFieldKind kind) {
- return TVMStructGet(Type(type), std::move(handle), index,
kind).as_or_throw<PrimExpr>();
-}
-
-/*!
- * \brief Address of handle + offset
- * \param handle the array handle.
- * \param dtype The data type.
- * \param offset the offset index.
- */
-inline Call AddressOffset(Var handle, PrimType dtype, int offset) {
- PrimExpr offset_expr = IntImm::Int32(offset * dtype.lanes());
- ffi::Array<PrimExpr> shape = {offset_expr + 1};
- auto pointer_type = handle->ty.as_or_throw<PtrType>();
- TensorVar dummy_buf(handle->name,
- TensorType(pointer_type->storage_scope, dtype, shape,
{}, 0, 0, 0));
- TensorLoad buf_load = MakeTensorLoad(dummy_buf, {offset_expr});
-
- return Call(handle->ty, tirx::address_of_op(), {buf_load});
-}
-
-/*!
- * \brief Address of handle + offset
- * \param handle the array handle.
- * \param dtype The data type.
- * \param offset the offset index.
- */
-inline Call AddressOffset(Var handle, PrimType dtype, PrimExpr offset) {
- if (dtype.lanes() != 1) {
- PrimType offset_ty = offset.ty();
- offset = offset * IntImm(offset_ty, dtype.lanes());
- offset = prim::Ramp(offset, IntImm(offset_ty, 1), dtype.lanes());
- }
-
- ffi::Array<PrimExpr> shape = {offset + 1};
- auto pointer_type = handle->ty.as_or_throw<PtrType>();
- TensorVar dummy_buf(handle->name, TensorType(pointer_type->storage_scope,
dtype.WithLanes(1),
- shape, {}, 0, 0, 0));
- TensorLoad buf_load = MakeTensorLoad(dummy_buf, {offset});
-
- return Call(handle->ty, tirx::address_of_op(), {buf_load});
-}
-
-/*!
- * \brief Set value into struct.
- * \param handle the struct handle.
- * \param index the offset index.
- * \param kind The data kind.
- * \param value The value to be set.
- * \return the set stmt.
- */
-inline Stmt TVMStructSet(Var handle, int index, tirx::TVMStructFieldKind kind,
Expr value) {
- ffi::Array<Expr> args = {handle, IntImm::Int32(index),
IntImm::Int32(static_cast<int>(kind)),
- value};
- return Evaluate(Call(PrimType::Int(32), tirx::abi_field_set_op(),
args).as_or_throw<PrimExpr>());
-}
-
/*!
* \brief Get the type that is passed around TVM ffi::Function API.
* \param t The original type.
@@ -198,28 +85,6 @@ inline int GetTempAllocaAlignment(const PrimType& type,
int64_t const_size) {
return align;
}
-/*!
- * \brief Create an int32 constant
- * \param index the value of the constant
- * \return the PrimExpr that represents the constant
- */
-inline PrimExpr ConstInt32(size_t index) {
- TVM_FFI_ICHECK_LE(index, std::numeric_limits<int>::max());
- return IntImm::Int32(static_cast<int>(index));
-}
-
-/*!
- * \brief Allocate TVMValues on the stack
- * \param ret_type exact pointer type returned by the allocation
- * \param type type of allocation
- * \param num number of TVMValues to allocate
- * \return Call representing the allocated pointer
- */
-inline Call StackAlloca(Type ret_type, std::string type, size_t num) {
- ffi::Array<Expr> args = {StringImm(type), ConstInt32(num)};
- return Call(std::move(ret_type), tirx::stack_alloca_op(), args);
-}
-
/*!
* \brief Convert a IR node to be SSA form.
* \param stmt The source statement to be converted.
@@ -309,40 +174,6 @@ class IRConvertSSA : public StmtExprMutator {
ScopeStack<ScopeLevel> scope_;
};
-/*!
- * \brief Return the storage scope associated with a buffer variable.
- * \param buffer_var The input buffer variable.
- * \return A string representing the storage scope of this buffer variable.
- */
-ffi::String GetPtrStorageScope(Var buffer_var);
-
-/*!
- * \brief Get stride aware tensor allocation shape from tensor.
- * \param tensor The tensor object.
- * \return shape The shape considering tensor strides.
- */
-ffi::Array<PrimExpr> GetTensorAllocationShape(const TensorVar& tensor);
-
-/*!
- * \brief Split string separated by "," to get wmma fragment dimension size.
- * \param shape_str The string to split.
- * \param scope The scope to match.
- * \return The result pair of fragment dimension size.
- */
-std::pair<int32_t, int32_t> GetWmmaFragmentDimSize(const std::string&
shape_str,
- const std::string& scope);
-
-/*! \brief Check if a Function is a host function
- *
- * \param func The function to be inspected
- *
- * \return True if the function is known to run on the host, false if
- * the function is known to run on the device. If it cannot be
- * determined (e.g. a function without a tvm::attr::kTarget
- * attribute), returns std::nullopt.
- */
-std::optional<bool> IsHostFunc(const Function& func);
-
} // namespace tirx
} // namespace tvm
#endif // TVM_TIR_TRANSFORM_IR_UTILS_H_
diff --git a/src/tirx/transform/lower_tvm_builtin.cc
b/src/tirx/transform/lower_tvm_builtin.cc
index 386437bf75..ecea4e5a77 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -184,23 +184,28 @@ class BuiltinLower : public StmtExprMutator {
if (scope.max_sizes.arg_stack != 0) {
alloca_stmts.push_back(
Bind(scope.stack_ffi_any,
- StackAlloca(scope.stack_ffi_any->ty, "tvm_ffi_any",
scope.max_sizes.arg_stack)));
+ Call(scope.stack_ffi_any->ty, tirx::stack_alloca_op(),
+ {StringImm("tvm_ffi_any"),
IntImm::Int32(scope.max_sizes.arg_stack)})));
}
if (scope.max_sizes.array_stack != 0) {
- alloca_stmts.push_back(Bind(scope.stack_array,
StackAlloca(scope.stack_array->ty, "array",
-
scope.max_sizes.array_stack)));
+ alloca_stmts.push_back(
+ Bind(scope.stack_array,
+ Call(scope.stack_array->ty, tirx::stack_alloca_op(),
+ {StringImm("array"),
IntImm::Int32(scope.max_sizes.array_stack)})));
}
if (scope.max_sizes.shape_stack != -1) {
scope.stack_shape =
decl_tensor({IntImm::Int64(scope.max_sizes.shape_stack)},
PrimType::Int(64), "stack_shape");
- stmt = SeqStmt({Bind(scope.stack_shape.value(),
- Call(scope.stack_shape.value().type(),
decl_tensor_op(),
-
{StackAlloca(scope.stack_shape.value().type()->DataPtrType(),
- "shape",
scope.max_sizes.shape_stack)},
- {}, {scope.stack_shape.value().type()})),
- stmt});
+ stmt = SeqStmt(
+ {Bind(scope.stack_shape.value(),
+ Call(scope.stack_shape.value().type(), decl_tensor_op(),
+ {Call(scope.stack_shape.value().type()->DataPtrType(),
+ tirx::stack_alloca_op(),
+ {StringImm("shape"),
IntImm::Int32(scope.max_sizes.shape_stack)})},
+ {}, {scope.stack_shape.value().type()})),
+ stmt});
}
if (!alloca_stmts.empty()) {
@@ -491,10 +496,10 @@ class BuiltinLower : public StmtExprMutator {
// no need to perform any store for a scalar shape
for (size_t i = 0; i < op->args.size(); ++i) {
prep_seq.emplace_back(
- TensorStore(scope.stack_shape.value(), {ConstInt32(stack_begin + i)},
+ TensorStore(scope.stack_shape.value(), {IntImm::Int32(stack_begin +
i)},
prim::cast(PrimType::Int(64),
op->args[i].as_or_throw<PrimExpr>())));
}
- PrimExpr offset = ConstInt32(stack_begin);
+ PrimExpr offset = IntImm::Int32(stack_begin);
TensorLoad load = MakeTensorLoad(scope.stack_shape.value(), {offset});
return Call(scope.stack_shape.value().type()->DataPtrType(),
address_of_op(), {load});
}
@@ -510,22 +515,43 @@ class BuiltinLower : public StmtExprMutator {
.ValueOrUnchanged(ffi::GetRef<Expr>(op));
op = expr.as<CallNode>();
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, kDLTensorData,
op->args[0]));
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, kDLTensorShape,
op->args[1]));
+ prep_seq.emplace_back(Evaluate(
+ Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorData), op->args[0]})
+ .as_or_throw<PrimExpr>()));
+ prep_seq.emplace_back(Evaluate(
+ Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorShape), op->args[1]})
+ .as_or_throw<PrimExpr>()));
Expr strides = op->args[2];
if (auto prim_strides = strides.as<PrimExpr>(); prim_strides &&
IsZero(prim_strides.value())) {
strides = tirx::ConstHandle(0);
}
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx,
kDLTensorStrides, strides));
+ prep_seq.emplace_back(Evaluate(
+ Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorStrides), strides})
+ .as_or_throw<PrimExpr>()));
prep_seq.emplace_back(
- TVMStructSet(scope.stack_array, idx, kDLTensorNDim,
op->args[3].as_or_throw<PrimExpr>()));
+ Evaluate(Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorNDim),
+ op->args[3].as_or_throw<PrimExpr>()})
+ .as_or_throw<PrimExpr>()));
PrimType dtype = op->args[4].as_or_throw<PrimExpr>().ty();
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx,
kDLTensorTypeCode,
- IntImm(PrimType::UInt(8),
static_cast<int>(dtype.code()))));
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx,
kDLTensorTypeBits,
- IntImm(PrimType::UInt(8),
dtype.bits())));
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx,
kDLTensorTypeLanes,
- IntImm(PrimType::UInt(16),
dtype.lanes())));
+ prep_seq.emplace_back(
+ Evaluate(Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorTypeCode),
+ IntImm(PrimType::UInt(8),
static_cast<int>(dtype.code()))})
+ .as_or_throw<PrimExpr>()));
+ prep_seq.emplace_back(
+ Evaluate(Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorTypeBits),
+ IntImm(PrimType::UInt(8), dtype.bits())})
+ .as_or_throw<PrimExpr>()));
+ prep_seq.emplace_back(
+ Evaluate(Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorTypeLanes),
+ IntImm(PrimType::UInt(16), dtype.lanes())})
+ .as_or_throw<PrimExpr>()));
// set byte offset
int data_bytes = GetVectorBytes(dtype);
PrimExpr elem_offset = op->args[5].as_or_throw<PrimExpr>();
@@ -535,15 +561,25 @@ class BuiltinLower : public StmtExprMutator {
} else {
byte_offset = elem_offset;
}
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx,
kDLTensorByteOffset,
- prim::cast(PrimType::UInt(64),
byte_offset)));
+ prep_seq.emplace_back(
+ Evaluate(Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorByteOffset),
+ prim::cast(PrimType::UInt(64), byte_offset)})
+ .as_or_throw<PrimExpr>()));
TVM_FFI_ICHECK(device_type_) << "Unknown device type in current IR";
TVM_FFI_ICHECK(device_id_) << "Unknown device id in current IR";
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx,
kDLTensorDeviceId,
- prim::cast(PrimType::Int(32),
device_id_.value())));
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx,
kDLTensorDeviceType,
- prim::cast(PrimType::Int(32),
device_type_.value())));
- return TVMStructGet(PtrType::VoidPtrType(), scope.stack_array, idx,
kDLTensorAddr);
+ prep_seq.emplace_back(
+ Evaluate(Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorDeviceId),
+ prim::cast(PrimType::Int(32), device_id_.value())})
+ .as_or_throw<PrimExpr>()));
+ prep_seq.emplace_back(
+ Evaluate(Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorDeviceType),
+ prim::cast(PrimType::Int(32), device_type_.value())})
+ .as_or_throw<PrimExpr>()));
+ return Call(PtrType::VoidPtrType(), tirx::abi_field_get_op(),
+ {scope.stack_array, IntImm::Int32(idx),
IntImm::Int32(kDLTensorAddr)});
}
void SetPackedArg(Expr arg, const Var& args_stack, size_t stack_offset,
@@ -577,21 +613,35 @@ class BuiltinLower : public StmtExprMutator {
if (arg_type_index == ffi::TypeIndex::kTVMFFIOpaquePtr) {
prep_seq->emplace_back(
If(Call(PrimType::Bool(), isnullptr_op(),
{arg}).as_or_throw<PrimExpr>(),
- TVMStructSet(args_stack, stack_offset, kTVMFFIAnyTypeIndex,
- ConstInt32(ffi::TypeIndex::kTVMFFINone)),
- SeqStmt(TVMStructSet(args_stack, stack_offset,
kTVMFFIAnyTypeIndex,
-
ConstInt32(ffi::TypeIndex::kTVMFFIOpaquePtr)))));
+ Evaluate(
+ Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {args_stack, IntImm::Int32(stack_offset),
IntImm::Int32(kTVMFFIAnyTypeIndex),
+ IntImm::Int32(ffi::TypeIndex::kTVMFFINone)})
+ .as_or_throw<PrimExpr>()),
+ SeqStmt(Evaluate(
+ Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {args_stack, IntImm::Int32(stack_offset),
IntImm::Int32(kTVMFFIAnyTypeIndex),
+ IntImm::Int32(ffi::TypeIndex::kTVMFFIOpaquePtr)})
+ .as_or_throw<PrimExpr>()))));
} else {
prep_seq->emplace_back(
- TVMStructSet(args_stack, stack_offset, kTVMFFIAnyTypeIndex,
ConstInt32(arg_type_index)));
+ Evaluate(Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {args_stack, IntImm::Int32(stack_offset),
+ IntImm::Int32(kTVMFFIAnyTypeIndex),
IntImm::Int32(arg_type_index)})
+ .as_or_throw<PrimExpr>()));
}
// set zero padding to ensure compatibility with FFI convention
- prep_seq->emplace_back(
- TVMStructSet(args_stack, stack_offset, kTVMFFIAnyZeroPadding,
ConstInt32(0)));
+ prep_seq->emplace_back(Evaluate(Call(PrimType::Int(32),
tirx::abi_field_set_op(),
+ {args_stack,
IntImm::Int32(stack_offset),
+
IntImm::Int32(kTVMFFIAnyZeroPadding), IntImm::Int32(0)})
+ .as_or_throw<PrimExpr>()));
// handle arg value
// NOTE: the intrinsic codegen will handle padding value clear for 32bit
// types or types that are smaller than 64 bits.
- prep_seq->emplace_back(TVMStructSet(args_stack, stack_offset,
kTVMFFIAnyUnionValue, arg));
+ prep_seq->emplace_back(Evaluate(
+ Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {args_stack, IntImm::Int32(stack_offset),
IntImm::Int32(kTVMFFIAnyUnionValue), arg})
+ .as_or_throw<PrimExpr>()));
}
/*!
@@ -631,12 +681,19 @@ class BuiltinLower : public StmtExprMutator {
this->SetPackedArg(arg, scope.stack_ffi_any, arg_stack_begin + i,
&prep_seq);
}
// explicitly set return value to None to avoid bad state interpretation
- prep_seq.emplace_back(TVMStructSet(scope.stack_ffi_any, num_args,
kTVMFFIAnyTypeIndex,
-
ConstInt32(ffi::TypeIndex::kTVMFFINone)));
- prep_seq.emplace_back(
- TVMStructSet(scope.stack_ffi_any, num_args, kTVMFFIAnyZeroPadding,
ConstInt32(0)));
- prep_seq.emplace_back(
- TVMStructSet(scope.stack_ffi_any, num_args, kTVMFFIAnyUnionValue,
IntImm::Int64(0)));
+ prep_seq.emplace_back(Evaluate(
+ Call(PrimType::Int(32), tirx::abi_field_set_op(),
+ {scope.stack_ffi_any, IntImm::Int32(num_args),
IntImm::Int32(kTVMFFIAnyTypeIndex),
+ IntImm::Int32(ffi::TypeIndex::kTVMFFINone)})
+ .as_or_throw<PrimExpr>()));
+ prep_seq.emplace_back(Evaluate(Call(PrimType::Int(32),
tirx::abi_field_set_op(),
+ {scope.stack_ffi_any,
IntImm::Int32(num_args),
+ IntImm::Int32(kTVMFFIAnyZeroPadding),
IntImm::Int32(0)})
+ .as_or_throw<PrimExpr>()));
+ prep_seq.emplace_back(Evaluate(Call(PrimType::Int(32),
tirx::abi_field_set_op(),
+ {scope.stack_ffi_any,
IntImm::Int32(num_args),
+ IntImm::Int32(kTVMFFIAnyUnionValue),
IntImm::Int64(0)})
+ .as_or_throw<PrimExpr>()));
// Verify stack size matches earlier value.
if (is_precheck_) {
scope.UpdateMax();
@@ -647,8 +704,8 @@ class BuiltinLower : public StmtExprMutator {
scope.run_sizes.array_stack = restore_array_stack;
scope.run_sizes.arg_stack = arg_stack_begin;
ffi::Array<Expr> packed_args = {op->args[name_offset], scope.stack_ffi_any,
- ConstInt32(arg_stack_begin),
- ConstInt32(arg_stack_begin + num_args)};
+ IntImm::Int32(arg_stack_begin),
+ IntImm::Int32(arg_stack_begin + num_args)};
return Call(op->ty, lowered_packed_op, packed_args);
}
@@ -753,7 +810,13 @@ namespace transform {
Pass LowerTVMBuiltin() {
auto pass_func = [](Function func, IRModule m, PassContext ctx) {
- if (IsHostFunc(func).value_or(false)) {
+ bool is_host = func->HasNonzeroAttr(tvm::tirx::attr::kIsHostFunc);
+ if (!is_host) {
+ if (auto target = func->GetAttr<Target>(tvm::attr::kTarget)) {
+ is_host = target.value()->HasKey("cpu");
+ }
+ }
+ if (is_host) {
func = BuiltinLower::Build(func);
VLOG(2) << "LowerTVMBuiltin: " << func;
}
diff --git a/src/tirx/transform/make_packed_api.cc
b/src/tirx/transform/make_packed_api.cc
index 0cca17b74d..3fbfd66e58 100644
--- a/src/tirx/transform/make_packed_api.cc
+++ b/src/tirx/transform/make_packed_api.cc
@@ -294,15 +294,20 @@ Function MakePackedAPI(Function func) {
// Return error code of zero on success
body = SeqStmt({body, Return(IntImm::Int32(0))});
- // Tensor declarations and alignment assumptions are ordinary statements,
- // not scopes with a body hole for MergeNest to fill.
+ // Prepend tensor declarations and alignment assumptions.
result.decl_tensors.push_back(body);
body = SeqStmt(result.decl_tensors);
- body = MergeNest(std::move(result.asserts), body);
+ if (!result.asserts.empty()) {
+ result.asserts.push_back(body);
+ body = SeqStmt(result.asserts);
+ }
if (need_set_device) {
body = RegionStmt(tirx::device_context_op(), {device_type, device_id}, {},
DictAttrs(), body);
}
- body = MergeNest(std::move(result.init_nest), body);
+ if (!result.init_nest.empty()) {
+ result.init_nest.push_back(body);
+ body = SeqStmt(result.init_nest);
+ }
func_ptr->body = body;
func_ptr->params = args;
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index 4094f7e84d..5b62aacd0d 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -47,7 +47,6 @@
#include "../../runtime/thread_storage_scope.h"
#include "../../sym/int_operator.h"
-#include "../ir/pointer_type.h"
#include "ir_utils.h"
namespace tvm {
@@ -825,7 +824,11 @@ class StoragePlanRewriter : public StmtExprMutator {
Stmt MakeAttach(const std::vector<StorageEntry*>& svec, Stmt body) {
for (auto it = svec.rbegin(); it != svec.rend(); it++) {
- body = MergeNest((*it)->alloc_nest, body);
+ if (!(*it)->alloc_nest.empty()) {
+ ffi::Array<Stmt> prefix((*it)->alloc_nest);
+ prefix.push_back(body);
+ body = SeqStmt(prefix);
+ }
}
return body;
}
@@ -1450,12 +1453,15 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
BufferVarInfo::kFunctionTensorParam);
continue;
}
- auto pointer_type = GetPtrElementType(buffer_var->ty);
- if (pointer_type.has_value() && !pointer_type.value().IsVoid()) {
- PrimType dtype = pointer_type.value();
+ const auto* pointer_type = buffer_var->ty.as<PtrTypeNode>();
+ if (!pointer_type) continue;
+ const auto* element_type = pointer_type->element_type.as<PrimTypeNode>();
+ if (!element_type) continue;
+ PrimType dtype = ffi::GetRef<PrimType>(element_type);
+ if (!dtype.IsVoid()) {
PrimExpr extent = 0;
OnArrayDeclaration(buffer_var, dtype, extent,
BufferVarInfo::kFunctionPointerParam);
- } else if (pointer_type.has_value() && allow_untyped_pointers_) {
+ } else if (allow_untyped_pointers_) {
OnArrayDeclaration(buffer_var, PrimType::Void(), 0,
BufferVarInfo::kFunctionPointerParam);
}
}
@@ -1536,14 +1542,15 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
}
void HandleLetNode(Var let_var) {
- auto pointer_type = GetPtrElementType(let_var->ty);
- if (pointer_type.has_value()) {
- if (!pointer_type.value().IsVoid()) {
- OnArrayDeclaration(let_var, pointer_type.value(), 0,
BufferVarInfo::kLetNode);
- } else if (allow_untyped_pointers_) {
- OnArrayDeclaration(let_var, PrimType::Void(), 0,
BufferVarInfo::kLetNode);
- }
- return;
+ const auto* pointer_type = let_var->ty.as<PtrTypeNode>();
+ if (!pointer_type) return;
+ const auto* element_type = pointer_type->element_type.as<PrimTypeNode>();
+ if (!element_type) return;
+ PrimType dtype = ffi::GetRef<PrimType>(element_type);
+ if (!dtype.IsVoid()) {
+ OnArrayDeclaration(let_var, dtype, 0, BufferVarInfo::kLetNode);
+ } else if (allow_untyped_pointers_) {
+ OnArrayDeclaration(let_var, PrimType::Void(), 0,
BufferVarInfo::kLetNode);
}
}
@@ -1792,7 +1799,8 @@ class VectorTypeRewriter : public StmtExprMutator {
type->layout = std::nullopt;
return RebuildTensorVar(old_tensor, std::move(type)).var();
}
- return Var(old_tensor_var->name, PtrType(preferred,
GetPtrStorageScope(old_tensor_var)),
+ return Var(old_tensor_var->name,
+ PtrType(preferred,
old_tensor_var->ty.as_or_throw<PtrType>()->storage_scope),
old_tensor_var->loc);
}();
diff --git a/src/tirx/transform/tvm_ffi_binder.cc
b/src/tirx/transform/tvm_ffi_binder.cc
index 017967b713..c978d0c748 100644
--- a/src/tirx/transform/tvm_ffi_binder.cc
+++ b/src/tirx/transform/tvm_ffi_binder.cc
@@ -650,9 +650,8 @@ Var TVMFFIABIBuilder::DLTensorGetFieldPtr(const Var&
handle, int field_kind,
const std::string& var_name) {
Type pointer_type = PtrType(DefaultIndexPrimType());
Var ptr(var_name, pointer_type);
- init_nest_.emplace_back(Bind(
- ptr,
- TVMStructGet(pointer_type, handle, 0,
static_cast<tirx::TVMStructFieldKind>(field_kind))));
+ init_nest_.emplace_back(Bind(ptr, Call(pointer_type,
tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(field_kind)})));
return ptr;
}
@@ -661,7 +660,9 @@ Var TVMFFIABIBuilder::DLTensorGetFieldPtr(const Var&
handle, int field_kind,
// ============================================================
PrimExpr TVMFFIABIBuilder::LoadInt64ArrayElem(const Var& ptr, int index) {
- return TVMStructGet(DefaultIndexPrimType(), ptr, index,
tirx::kInt64ArrayElem);
+ return Call(DefaultIndexPrimType(), tirx::abi_field_get_op(),
+ {ptr, IntImm::Int32(index),
IntImm::Int32(tirx::kInt64ArrayElem)})
+ .as_or_throw<PrimExpr>();
}
// ============================================================
@@ -730,7 +731,9 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const TensorVar&
tensor, const PrimEx
"Tensor");
// ── Section: ndim ────────────────────────────────────────────
- PrimExpr v_ndim = TVMStructGet(tvm_ndim_type, handle, 0,
tirx::kDLTensorNDim);
+ PrimExpr v_ndim = Call(tvm_ndim_type, tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(tirx::kDLTensorNDim)})
+ .as_or_throw<PrimExpr>();
PrimExpr a_ndim = IntImm(tvm_ndim_type,
static_cast<int64_t>(tensor->shape.size()));
EmitAssert(a_ndim == v_ndim, "ValueError", //
"Mismatched ", buf_name, ".ndim on argument #",
std::to_string(param_index),
@@ -738,13 +741,18 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const
TensorVar& tensor, const PrimEx
// ── Section: dtype ───────────────────────────────────────────
{
- PrimExpr code_matches = TVMStructGet(PrimType::UInt(8), handle, 0,
tirx::kDLTensorTypeCode) ==
- IntImm(PrimType::UInt(8), tensor->dtype.code());
- PrimExpr bits_matches = TVMStructGet(PrimType::UInt(8), handle, 0,
tirx::kDLTensorTypeBits) ==
- IntImm(PrimType::UInt(8), tensor->dtype.bits());
+ PrimExpr code_matches =
+ Call(PrimType::UInt(8), tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(tirx::kDLTensorTypeCode)})
+ .as_or_throw<PrimExpr>() == IntImm(PrimType::UInt(8),
tensor->dtype.code());
+ PrimExpr bits_matches =
+ Call(PrimType::UInt(8), tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(tirx::kDLTensorTypeBits)})
+ .as_or_throw<PrimExpr>() == IntImm(PrimType::UInt(8),
tensor->dtype.bits());
PrimExpr lanes_matches =
- TVMStructGet(PrimType::UInt(16), handle, 0, tirx::kDLTensorTypeLanes)
==
- IntImm(PrimType::UInt(16), tensor->dtype.lanes());
+ Call(PrimType::UInt(16), tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(tirx::kDLTensorTypeLanes)})
+ .as_or_throw<PrimExpr>() == IntImm(PrimType::UInt(16),
tensor->dtype.lanes());
PrimExpr cond = code_matches && bits_matches && lanes_matches;
if (!(tensor->dtype == PrimType::Int(1) || tensor->dtype ==
PrimType::Int(4) ||
tensor->dtype == PrimType::UInt(4))) {
@@ -783,16 +791,23 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const
TensorVar& tensor, const PrimEx
ffi::reflection::AccessPath byte_offset_path =
param_path->Attr(ffi::String("byte_offset"));
if (const auto* const_offset = tensor->elem_offset.as<IntImmNode>()) {
BindScalar(IntImm(PrimType::UInt(64), const_offset->value * data_bytes),
- TVMStructGet(PrimType::UInt(64), handle, 0,
tirx::kDLTensorByteOffset),
+ Call(PrimType::UInt(64), tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(tirx::kDLTensorByteOffset)})
+ .as_or_throw<PrimExpr>(),
byte_offset_path, true);
} else {
- PrimExpr byte_offset = TVMStructGet(PrimType::UInt(64), handle, 0,
tirx::kDLTensorByteOffset);
+ PrimExpr byte_offset =
+ Call(PrimType::UInt(64), tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(tirx::kDLTensorByteOffset)})
+ .as_or_throw<PrimExpr>();
EmitAssert(truncmod(byte_offset, MakeConst(PrimType::UInt(64),
data_bytes)) ==
MakeConst(PrimType::UInt(64), 0),
"ValueError", "Tensor byte offset must be divisible by its
element size");
if (BindScalar(tensor->elem_offset,
cast(tensor->elem_offset.ty(),
- (TVMStructGet(PrimType::UInt(64), handle, 0,
tirx::kDLTensorByteOffset) /
+ (Call(PrimType::UInt(64), tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(tirx::kDLTensorByteOffset)})
+ .as_or_throw<PrimExpr>() /
MakeConst(PrimType::UInt(64), data_bytes))),
byte_offset_path, true)) {
if (tensor->offset_factor > 1) {
@@ -817,7 +832,9 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const TensorVar&
tensor, const PrimEx
// ── Section: device ──────────────────────────────────────────
{
PrimExpr actual_device_type =
- TVMStructGet(PrimType::Int(32), handle, 0, tirx::kDLTensorDeviceType);
+ Call(PrimType::Int(32), tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(tirx::kDLTensorDeviceType)})
+ .as_or_throw<PrimExpr>();
// Use custom assertion for device_type to show human-readable device name
if (const auto* const_dt = device_type_.as<IntImmNode>()) {
PrimExpr cond = analyzer_->Simplify(IntImm::Int32(const_dt->value) ==
actual_device_type);
@@ -833,14 +850,18 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const
TensorVar& tensor, const PrimEx
BindScalar(device_type_, actual_device_type, device_type_path, true);
}
ffi::reflection::AccessPath device_id_path =
param_path->Attr(ffi::String("device_id"));
- BindScalar(device_id_, TVMStructGet(PrimType::Int(32), handle, 0,
tirx::kDLTensorDeviceId),
+ BindScalar(device_id_,
+ Call(PrimType::Int(32), tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(tirx::kDLTensorDeviceId)})
+ .as_or_throw<PrimExpr>(),
device_id_path, true);
}
// ── Section: data pointer ────────────────────────────────────
{
ffi::reflection::AccessPath data_path =
param_path->Attr(ffi::String("data"));
- Expr raw_data = TVMStructGet(PtrType::VoidPtrType(), handle, 0,
tirx::kDLTensorData);
+ Expr raw_data = Call(PtrType::VoidPtrType(), tirx::abi_field_get_op(),
+ {handle, IntImm::Int32(0),
IntImm::Int32(tirx::kDLTensorData)});
Expr typed_data = Call(tensor.type()->DataPtrType(),
tirx::reinterpret_op(), {raw_data});
{
Expr vptr = typed_data;
