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 7de260a4ec [REFACTOR][TIRx] Keep unroll policy on loops and retire
source imports (#20556)
7de260a4ec is described below
commit 7de260a4ec97351361dec3e6817cbf10e3134839
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 07:45:04 2026 -0400
[REFACTOR][TIRx] Keep unroll policy on loops and retire source imports
(#20556)
---
include/tvm/tirx/stmt.h | 22 +++--
python/tvm/s_tir/dlight/cpu/gemv.py | 4 +-
python/tvm/s_tir/dlight/cpu/reduction.py | 4 +-
python/tvm/s_tir/dlight/gpu/gemv.py | 24 ++---
python/tvm/s_tir/dlight/gpu/general_reduction.py | 12 +--
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 14 +--
python/tvm/s_tir/dlight/gpu/matmul.py | 4 +-
python/tvm/s_tir/dlight/gpu/reduction.py | 4 +-
python/tvm/s_tir/dlight/gpu/transpose.py | 4 +-
.../s_tir/meta_schedule/post_optimization/space.py | 2 +-
.../feature_extractor/per_store_feature.cc | 15 ++-
.../postproc/rewrite_parallel_vectorize_unroll.cc | 4 +-
src/s_tir/transform/annotate_irregular_loop.cc | 2 +-
src/s_tir/transform/lower_opaque_block.cc | 64 +++++++++++--
src/target/llvm/codegen_cpu.cc | 5 -
src/target/llvm/codegen_llvm.cc | 30 ------
src/target/llvm/codegen_llvm.h | 2 -
src/target/source/codegen_c.cc | 5 -
src/tirx/transform/lower_tirx_opaque.cc | 52 +++++++++--
src/tirx/transform/unroll_loop.cc | 74 +++++++++------
tests/python/codegen/test_target_codegen_llvm.py | 41 +-------
tests/python/s_tir/dlight/test_benchmark.py | 2 +-
tests/python/s_tir/dlight/test_cpu_gemv.py | 12 +--
tests/python/s_tir/dlight/test_gpu_conv.py | 2 +-
tests/python/s_tir/dlight/test_gpu_gemv.py | 36 +++----
.../s_tir/dlight/test_gpu_general_reduction.py | 42 ++++-----
.../python/s_tir/dlight/test_gpu_low_batch_gemv.py | 12 +--
tests/python/s_tir/dlight/test_gpu_matmul.py | 10 +-
.../s_tir/dlight/test_gpu_matmul_tensorize.py | 2 +-
tests/python/s_tir/dlight/test_gpu_reduction.py | 14 +--
tests/python/s_tir/dlight/test_gpu_transpose.py | 6 +-
...schedule_feature_extractor_per_store_feature.py | 6 +-
.../test_meta_schedule_mma_tensorize.py | 4 +-
.../test_meta_schedule_postproc_rewrite_layout.py | 6 +-
...e_postproc_rewrite_parallel_vectorize_unroll.py | 2 +-
...est_meta_schedule_postproc_verify_vtcm_limit.py | 2 +-
.../test_meta_schedule_trace_apply.py | 104 ++++++++++-----------
.../schedule/test_tir_schedule_compute_inline.py | 4 +-
.../test_s_tir_transform_lower_opaque_block.py | 2 +-
.../test_tir_transform_flatten_buffer.py | 4 +-
.../test_tir_transform_unroll_loop.py | 11 +--
.../tirx/transform/test_transform_lower_tirx.py | 10 +-
42 files changed, 345 insertions(+), 336 deletions(-)
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 1ef75452cd..86ea304051 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -850,7 +850,7 @@ class ScopeIdDefStmt : public Stmt {
TVM_DEFINE_OBJECT_REF_COW_METHOD(ScopeIdDefStmtNode);
};
-/*! \brief namespace of possible attributes in AttrStmt.attr_key */
+/*! \brief Statement attribute and loop annotation keys. */
namespace attr {
/*!
* \brief Mark the scope as when computation start to happen.
@@ -863,14 +863,18 @@ constexpr const char* device_id = "device_id";
constexpr const char* device_scope = "device_scope";
/*! \brief The device type. */
constexpr const char* device_type = "device_type";
-/*! \brief Pragma: auto-unroll, max_step */
-constexpr const char* pragma_auto_unroll_max_step =
"pragma_auto_unroll_max_step";
-/*! \brief Import C source or file into the final code gen module */
-constexpr const char* pragma_import_c = "pragma_import_c";
-/*! \brief Import llvm source or file into the final code gen module */
-constexpr const char* pragma_import_llvm = "pragma_import_llvm";
-/*! \brief Pragma: unroll explicit */
-constexpr const char* pragma_unroll_explicit = "pragma_unroll_explicit";
+/*!
+ * \brief For annotation: maximum work for automatic unrolling.
+ *
+ * Integer policy inherited by nested loops unless they override it. Consumed
by UnrollLoop.
+ */
+constexpr const char* auto_unroll_max_step = "auto_unroll_max_step";
+/*!
+ * \brief For annotation: expand unrolled bodies instead of preserving
unrolled loops.
+ *
+ * Integer policy inherited by nested loops unless they override it. Consumed
by UnrollLoop.
+ */
+constexpr const char* unroll_explicit = "unroll_explicit";
/*! \brief Mark storage alignment requirement of buffers */
constexpr const char* storage_alignment = "storage_alignment";
/*! \brief Shared execution attributes consumed before and after block
lowering. */
diff --git a/python/tvm/s_tir/dlight/cpu/gemv.py
b/python/tvm/s_tir/dlight/cpu/gemv.py
index 836c6de09a..e29d163103 100644
--- a/python/tvm/s_tir/dlight/cpu/gemv.py
+++ b/python/tvm/s_tir/dlight/cpu/gemv.py
@@ -121,8 +121,8 @@ class GEMV(CPUScheduleRule):
sch.parallel(s_outer) # Parallelize along the s-axis (major
computation loop)
# Apply loop unrolling for better CPU performance
- sch.annotate(r_outer, "pragma_auto_unroll_max_step", unroll_factor)
- sch.annotate(r_outer, "pragma_unroll_explicit", 1)
+ sch.annotate(r_outer, "auto_unroll_max_step", unroll_factor)
+ sch.annotate(r_outer, "unroll_explicit", 1)
return sch
return apply(
diff --git a/python/tvm/s_tir/dlight/cpu/reduction.py
b/python/tvm/s_tir/dlight/cpu/reduction.py
index 050e0c6a1b..a1747a4cb7 100644
--- a/python/tvm/s_tir/dlight/cpu/reduction.py
+++ b/python/tvm/s_tir/dlight/cpu/reduction.py
@@ -147,5 +147,5 @@ class Reduction(CPUScheduleRule):
if isinstance(extent, int) and extent <= vec_lanes:
return
_, inner_loop = sch.split(inner, factors=[None, vec_lanes])
- sch.annotate(inner_loop, ann_key="pragma_auto_unroll_max_step",
ann_val=vec_lanes)
- sch.annotate(inner_loop, ann_key="pragma_unroll_explicit", ann_val=1)
+ sch.annotate(inner_loop, ann_key="auto_unroll_max_step",
ann_val=vec_lanes)
+ sch.annotate(inner_loop, ann_key="unroll_explicit", ann_val=1)
diff --git a/python/tvm/s_tir/dlight/gpu/gemv.py
b/python/tvm/s_tir/dlight/gpu/gemv.py
index 6d6c9ed6b2..50c0635ed6 100644
--- a/python/tvm/s_tir/dlight/gpu/gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/gemv.py
@@ -256,26 +256,22 @@ class GEMV(GPUScheduleRule):
sch.annotate(
block_or_loop=sch.get_loops(rf)[3],
- ann_key="pragma_auto_unroll_max_step",
+ ann_key="auto_unroll_max_step",
ann_val=unroll_factor,
)
- sch.annotate(
- block_or_loop=sch.get_loops(rf)[3],
ann_key="pragma_unroll_explicit", ann_val=1
- )
+ sch.annotate(block_or_loop=sch.get_loops(rf)[3],
ann_key="unroll_explicit", ann_val=1)
sch.annotate(
block_or_loop=sch.get_loops(rf2)[3],
- ann_key="pragma_auto_unroll_max_step",
+ ann_key="auto_unroll_max_step",
ann_val=unroll_factor,
)
- sch.annotate(
- block_or_loop=sch.get_loops(rf2)[3],
ann_key="pragma_unroll_explicit", ann_val=1
- )
+ sch.annotate(block_or_loop=sch.get_loops(rf2)[3],
ann_key="unroll_explicit", ann_val=1)
if LOAD_V_SHARED:
sch.annotate(
block_or_loop=sch.get_loops(V_shared)[-4],
- ann_key="pragma_unroll_explicit",
+ ann_key="unroll_explicit",
ann_val=unroll_factor,
)
sch.annotate(
@@ -533,12 +529,10 @@ class GEMV(GPUScheduleRule):
sch.annotate(
block_or_loop=sch.get_loops(rf2)[3],
- ann_key="pragma_auto_unroll_max_step",
+ ann_key="auto_unroll_max_step",
ann_val=UNROLL,
)
- sch.annotate(
- block_or_loop=sch.get_loops(rf2)[3],
ann_key="pragma_unroll_explicit", ann_val=1
- )
+ sch.annotate(block_or_loop=sch.get_loops(rf2)[3],
ann_key="unroll_explicit", ann_val=1)
# Schedule epilogue
if epilogue_info is not None:
@@ -663,8 +657,8 @@ class GEMV(GPUScheduleRule):
sch.bind(tx, "threadIdx.x")
sch.reorder(bx, tx, r0, r1, c, vec)
- sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=8)
- sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1)
+ sch.annotate(tx, ann_key="auto_unroll_max_step", ann_val=8)
+ sch.annotate(tx, ann_key="unroll_explicit", ann_val=1)
if LOAD_V_SHARED:
V_shared = sch.cache_read(block, vector_input_buffers[0],
storage_scope="shared")
diff --git a/python/tvm/s_tir/dlight/gpu/general_reduction.py
b/python/tvm/s_tir/dlight/gpu/general_reduction.py
index 1bc14609cd..ea1acc990d 100644
--- a/python/tvm/s_tir/dlight/gpu/general_reduction.py
+++ b/python/tvm/s_tir/dlight/gpu/general_reduction.py
@@ -192,8 +192,8 @@ class GeneralReduction(GPUScheduleRule):
sch.reorder(tx, r_loop)
sch.bind(bx, "blockIdx.x")
sch.bind(tx, "threadIdx.x")
- sch.annotate(r_loop, ann_key="pragma_auto_unroll_max_step",
ann_val=unroll_depth)
- sch.annotate(r_loop, ann_key="pragma_unroll_explicit", ann_val=1)
+ sch.annotate(r_loop, ann_key="auto_unroll_max_step",
ann_val=unroll_depth)
+ sch.annotate(r_loop, ann_key="unroll_explicit", ann_val=1)
for block in reversed(block_infos[:-1]):
block = block.block_rv
@@ -204,11 +204,11 @@ class GeneralReduction(GPUScheduleRule):
r_loop, tx = sch.split(r_loop, [None, len_tx])
sch.reorder(tx, r_loop)
sch.bind(tx, "threadIdx.x")
- sch.annotate(r_loop, ann_key="pragma_auto_unroll_max_step",
ann_val=unroll_depth)
- sch.annotate(r_loop, ann_key="pragma_unroll_explicit", ann_val=1)
+ sch.annotate(r_loop, ann_key="auto_unroll_max_step",
ann_val=unroll_depth)
+ sch.annotate(r_loop, ann_key="unroll_explicit", ann_val=1)
# TODO: It's just a workaround to avoid unroll spatial loops, because
of the bug of
# the pass lower-thread-allreduce. We should fix it in the future.
- # sch.annotate(bx, ann_key="pragma_auto_unroll_max_step",
ann_val=unroll_depth)
- # sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
+ # sch.annotate(bx, ann_key="auto_unroll_max_step",
ann_val=unroll_depth)
+ # sch.annotate(bx, ann_key="unroll_explicit", ann_val=1)
return sch
diff --git a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
index e562d848b4..0c4c29a27a 100644
--- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
@@ -470,26 +470,22 @@ class LowBatchGEMV(GPUScheduleRule):
sch.annotate(
block_or_loop=sch.get_loops(rf)[4],
- ann_key="pragma_auto_unroll_max_step",
+ ann_key="auto_unroll_max_step",
ann_val=unroll_factor,
)
- sch.annotate(
- block_or_loop=sch.get_loops(rf)[4],
ann_key="pragma_unroll_explicit", ann_val=1
- )
+ sch.annotate(block_or_loop=sch.get_loops(rf)[4],
ann_key="unroll_explicit", ann_val=1)
sch.annotate(
block_or_loop=sch.get_loops(rf2)[4],
- ann_key="pragma_auto_unroll_max_step",
+ ann_key="auto_unroll_max_step",
ann_val=unroll_factor,
)
- sch.annotate(
- block_or_loop=sch.get_loops(rf2)[4],
ann_key="pragma_unroll_explicit", ann_val=1
- )
+ sch.annotate(block_or_loop=sch.get_loops(rf2)[4],
ann_key="unroll_explicit", ann_val=1)
if LOAD_V_SHARED:
sch.annotate(
block_or_loop=sch.get_loops(V_shared)[-4],
- ann_key="pragma_unroll_explicit",
+ ann_key="unroll_explicit",
ann_val=unroll_factor,
)
sch.annotate(
diff --git a/python/tvm/s_tir/dlight/gpu/matmul.py
b/python/tvm/s_tir/dlight/gpu/matmul.py
index b4bac8a2c7..e08fb6c2ae 100644
--- a/python/tvm/s_tir/dlight/gpu/matmul.py
+++ b/python/tvm/s_tir/dlight/gpu/matmul.py
@@ -1084,8 +1084,8 @@ class Matmul(GPUScheduleRule):
sch.vectorize(v)
if config.unroll > 0:
- sch.annotate(tx, ann_key="pragma_auto_unroll_max_step",
ann_val=config.unroll)
- sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1)
+ sch.annotate(tx, ann_key="auto_unroll_max_step",
ann_val=config.unroll)
+ sch.annotate(tx, ann_key="unroll_explicit", ann_val=1)
l2g = sch.cache_write(main_block, 0, "local")
sch.reverse_compute_at(l2g, tx, preserve_unit_loops=True)
diff --git a/python/tvm/s_tir/dlight/gpu/reduction.py
b/python/tvm/s_tir/dlight/gpu/reduction.py
index e16a87b52d..1f27e8a174 100644
--- a/python/tvm/s_tir/dlight/gpu/reduction.py
+++ b/python/tvm/s_tir/dlight/gpu/reduction.py
@@ -269,8 +269,8 @@ class Reduction(GPUScheduleRule):
sch.reorder(bx, tx, r)
sch.bind(bx, "blockIdx.x")
sch.bind(tx, "threadIdx.x")
- sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=256)
- sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1)
+ sch.annotate(tx, ann_key="auto_unroll_max_step", ann_val=256)
+ sch.annotate(tx, ann_key="unroll_explicit", ann_val=1)
sch.set_scope(rf, 0, "local")
sch.decompose_reduction(rf, r)
# Schedule the write back block
diff --git a/python/tvm/s_tir/dlight/gpu/transpose.py
b/python/tvm/s_tir/dlight/gpu/transpose.py
index d57ec7df21..412b31f1f0 100644
--- a/python/tvm/s_tir/dlight/gpu/transpose.py
+++ b/python/tvm/s_tir/dlight/gpu/transpose.py
@@ -122,8 +122,8 @@ class Transpose(GPUScheduleRule):
sch.unroll(v)
sch.storage_align(block=cache_read, buffer_index=0, axis=0, factor=32,
offset=1)
- sch.annotate(bi, ann_key="pragma_auto_unroll_max_step",
ann_val=unroll_depth)
- sch.annotate(bi, ann_key="pragma_unroll_explicit", ann_val=1)
+ sch.annotate(bi, ann_key="auto_unroll_max_step", ann_val=unroll_depth)
+ sch.annotate(bi, ann_key="unroll_explicit", ann_val=1)
if prologue is not None:
sch.compute_inline(prologue)
diff --git a/python/tvm/s_tir/meta_schedule/post_optimization/space.py
b/python/tvm/s_tir/meta_schedule/post_optimization/space.py
index 5cae34ef92..702a8e76d8 100644
--- a/python/tvm/s_tir/meta_schedule/post_optimization/space.py
+++ b/python/tvm/s_tir/meta_schedule/post_optimization/space.py
@@ -123,7 +123,7 @@ class Space:
interval = self.power_of_two(5, 9)
elif ann_key == ["meta_schedule.vectorize"]:
interval = self.power_of_two(4, 8)
- elif ann_key == ["pragma_auto_unroll_max_step"]:
+ elif ann_key == ["auto_unroll_max_step"]:
interval = self.power_of_two(7, 11)
elif ann_key == ["meta_schedule.thread_extent_low_inclusive"]:
interval = self.power_of_two(5, 6)
diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
index 34814c0872..4f6bbe0351 100644
--- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
+++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
@@ -87,13 +87,12 @@ std::vector<int64_t> GetBufferShape(const TensorVar&
buffer, sym::AnalyzerObj* a
}
/*!
- * \brief Given a loop, return its `pragma_auto_unroll_max_step` annotation if
it exists
+ * \brief Given a loop, return its `auto_unroll_max_step` annotation if it
exists
* \param loop The loop to be checked
- * \return The value of `pragma_auto_unroll_max_step` if it exists, or -1 if
it does not exist
+ * \return The value of `auto_unroll_max_step` if it exists, or -1 if it does
not exist
*/
-int64_t GetPragmaAutoUnroll(const ForNode* loop) {
- if (ffi::Optional<IntImm> auto_unroll =
- GetAnn<IntImm>(loop, tirx::attr::pragma_auto_unroll_max_step)) {
+int64_t GetAutoUnrollMaxStep(const ForNode* loop) {
+ if (ffi::Optional<IntImm> auto_unroll = GetAnn<IntImm>(loop,
tirx::attr::auto_unroll_max_step)) {
return static_cast<int64_t>(auto_unroll.value()->value);
}
return -1;
@@ -431,7 +430,7 @@ tvm::transform::Sequential PassListForPerStoreFeature() {
struct LoopNest {
int64_t prod = 1; // The product of the extents of all the loops
ForVec loops; // All the loops
- IntVec auto_unroll; // The loops with auto unroll pragma
+ IntVec auto_unroll; // The loops with auto unroll policy
ForVec parallel; // The loops whose ForKind are kParallel
ForVec vectorize; // The loops whose ForKind are kVectorized
ForVec unroll; // The loops whose ForKind are kUnrolled
@@ -456,7 +455,7 @@ struct LoopNest {
this->prod *= *extent;
}
this->loops.push_back(loop);
- if ((*auto_unroll_attr = utils::GetPragmaAutoUnroll(loop)) > 0) {
+ if ((*auto_unroll_attr = utils::GetAutoUnrollMaxStep(loop)) > 0) {
this->auto_unroll.push_back(*auto_unroll_attr);
}
ForVec* ref_loops = nullptr;
@@ -1305,7 +1304,7 @@ namespace group5 {
struct Feature {
int64_t outer_prod; // The product of lengths of outer loops
int num_loops; // The number of outer loops
- int auto_unroll_max_step; // The value of pragma "auto_unroll_max_step"
+ int auto_unroll_max_step; // The value of annotation "auto_unroll_max_step"
static constexpr int64_t kCount = 3;
diff --git
a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
index dd47377dc2..1121eedf1f 100644
--- a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
+++ b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
@@ -440,8 +440,8 @@ void RewriteUnroll(const Schedule& sch, int
unroll_explicit, int max_step, const
return;
}
- sch->Annotate(loop, tirx::attr::pragma_auto_unroll_max_step,
IntImm::Int32(max_step));
- sch->Annotate(loop, tirx::attr::pragma_unroll_explicit,
IntImm::Int32(unroll_explicit));
+ sch->Annotate(loop, tirx::attr::auto_unroll_max_step,
IntImm::Int32(max_step));
+ sch->Annotate(loop, tirx::attr::unroll_explicit,
IntImm::Int32(unroll_explicit));
}
} // namespace s_tir
diff --git a/src/s_tir/transform/annotate_irregular_loop.cc
b/src/s_tir/transform/annotate_irregular_loop.cc
index b56dd8a3d7..d957445b4e 100644
--- a/src/s_tir/transform/annotate_irregular_loop.cc
+++ b/src/s_tir/transform/annotate_irregular_loop.cc
@@ -58,7 +58,7 @@ class IrregularLoopAnnotator : public StmtExprMutator {
TVM_FFI_ICHECK(op->kind == ForKind::kSerial)
<< "Loop kind " << op->kind << " is invalid for irregular loop " <<
op->loop_var;
for (const char* key :
- {tirx::attr::pragma_auto_unroll_max_step,
tirx::attr::pragma_unroll_explicit,
+ {tirx::attr::auto_unroll_max_step, tirx::attr::unroll_explicit,
s_tir::attr::loop_partition_hint,
s_tir::attr::software_pipeline_stage}) {
TVM_FFI_ICHECK(!res->annotations.count(key))
<< "Annotation `" << key << "` is invalid for irregular loop " <<
op->loop_var;
diff --git a/src/s_tir/transform/lower_opaque_block.cc
b/src/s_tir/transform/lower_opaque_block.cc
index bdffbaae4c..908c3904db 100644
--- a/src/s_tir/transform/lower_opaque_block.cc
+++ b/src/s_tir/transform/lower_opaque_block.cc
@@ -23,10 +23,13 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/scope_stack.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
+#include <optional>
+
#include "ir_utils.h"
namespace tvm {
@@ -49,6 +52,11 @@ class OpaqueBlockLower : public StmtExprMutator {
}
private:
+ struct UnrollPolicy {
+ std::optional<ffi::Any> auto_unroll_max_step;
+ std::optional<ffi::Any> unroll_explicit;
+ };
+
UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
// Attribute metadata can reference a unit loop variable that is replaced.
// Rewrite it before descending into body-local definitions.
@@ -69,9 +77,15 @@ class OpaqueBlockLower : public StmtExprMutator {
TVM_FFI_ICHECK(op->iter_values.empty())
<< "Non-opaque blocks are not allowed in FlattenBuffer. Please "
"call pass ConvertBlocksToOpaque before.";
- // Step 1. Visit the body
- SBlock new_block =
- this->Mutate(op->block,
inplace_mode).ValueOrUnchanged(op->block).as_or_throw<SBlock>();
+ // Block policy belongs to the loops inside the block after opaque
lowering.
+ auto parent_policy = unroll_policy_.Current();
+ SBlock new_block = unroll_policy_.WithNewScope([&]() {
+ unroll_policy_.Current() = parent_policy;
+ UpdateUnrollPolicy(op->block->annotations);
+ return this->Mutate(op->block, inplace_mode)
+ .ValueOrUnchanged(op->block)
+ .as_or_throw<SBlock>();
+ });
PrimExpr predicate = this->Mutate(op->predicate,
inplace_mode).ValueOrUnchanged(op->predicate);
// Step 2. Transform the `predicate` to if-then-else
Stmt body = new_block->body;
@@ -123,14 +137,25 @@ class OpaqueBlockLower : public StmtExprMutator {
VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), min));
}
- // Step 2. Annotations may refer to the loop's own variable. Rewrite them
- // before visiting body-local definitions.
+ // Keep policy on surviving descendants when this owner is lowered away.
+ auto parent_policy = unroll_policy_.Current();
std::vector<std::pair<std::string, Expr>> pragma_attrs;
- ffi::Map<ffi::String, ffi::Any> new_annotations =
- HandleAnnotations(op->annotations, &pragma_attrs, /*is_block=*/false);
-
- // Step 3. Visit recursively.
- Stmt body = this->Mutate(op->body,
inplace_mode).ValueOrUnchanged(op->body);
+ ffi::Map<ffi::String, ffi::Any> new_annotations;
+ auto annotations = op->annotations;
+ Stmt body = unroll_policy_.WithNewScope([&]() {
+ unroll_policy_.Current() = parent_policy;
+ UpdateUnrollPolicy(op->annotations);
+ const auto& policy = unroll_policy_.Current();
+ if (policy.auto_unroll_max_step.has_value()) {
+ annotations.Set(tirx::attr::auto_unroll_max_step,
policy.auto_unroll_max_step.value());
+ }
+ if (policy.unroll_explicit.has_value()) {
+ annotations.Set(tirx::attr::unroll_explicit,
policy.unroll_explicit.value());
+ }
+ // Rewrite annotations before visiting body-local definitions.
+ new_annotations = HandleAnnotations(annotations, &pragma_attrs,
/*is_block=*/false);
+ return this->Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
+ });
VarRemapSet(op->loop_var, previous_remap);
// Step 4. Keep thread-binding loops until LowerThreadBinding.
@@ -153,6 +178,21 @@ class OpaqueBlockLower : public StmtExprMutator {
return body;
}
+ void UpdateUnrollPolicy(const ffi::Map<ffi::String, ffi::Any>& annotations) {
+ auto& policy = unroll_policy_.Current();
+ if (auto value = annotations.Get(tirx::attr::auto_unroll_max_step);
+ value.has_value() && value.value() != nullptr) {
+ policy.auto_unroll_max_step = value.value();
+ }
+ if (auto value = annotations.Get(tirx::attr::unroll_explicit);
+ value.has_value() && value.value() != nullptr) {
+ policy.unroll_explicit = value.value();
+ }
+ }
+
+ // Effective policy is materialized on each surviving loop, preserving
nested overrides.
+ ScopeStack<UnrollPolicy> unroll_policy_;
+
/*! \brief Convert attr value from annotation map into Expr. */
Expr ConvertAttrValue(const ffi::String& key, const Any& obj) {
if (auto expr = obj.try_cast<Expr>()) {
@@ -180,6 +220,10 @@ class OpaqueBlockLower : public StmtExprMutator {
pragma_attrs->clear();
for (const auto& kv : annotations) {
const ffi::String& key = kv.first;
+ if ((key == tirx::attr::auto_unroll_max_step || key ==
tirx::attr::unroll_explicit) &&
+ kv.second == nullptr) {
+ continue;
+ }
if (tirx::attr::IsPragmaKey(key)) {
if (kv.second == nullptr) {
continue;
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index 7833a32215..c71a659b64 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -1147,11 +1147,6 @@ void CodeGenCPU::Dispatch_(const AttrStmtNode* op) {
TVM_FFI_THROW(ValueError)
<< "pragma_parallel_barrier_when_finish is retired; place
cpu_parallel_barrier() "
<< "after the former attribute body inside the parallel launch";
- } else if (op->attr_key == tirx::attr::pragma_import_llvm) {
- const StringImmNode* value = op->value.as<StringImmNode>();
- TVM_FFI_ICHECK(value != nullptr);
- this->HandleImport(value->value);
- this->Dispatch(op->body);
} else {
LOG(WARNING) << "Unknown pragma " << op->attr_key;
this->Dispatch(op->body);
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 8635bb6550..5b1119f2ef 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -390,36 +390,6 @@ std::unique_ptr<llvm::Module> CodeGenLLVM::Finish() {
return std::move(module_);
}
-void CodeGenLLVM::HandleImport(const std::string& code) {
- llvm::StringRef code_str(code);
- std::unique_ptr<llvm::Module> mlib;
-#if TVM_LLVM_VERSION >= 180
- if (code_str.ends_with(".ll") || code_str.ends_with(".bc")) {
-#else
- if (code_str.endswith(".ll") || code_str.endswith(".bc")) {
-#endif
- mlib = llvm_target_->GetInstance().LoadIR(code);
- } else {
- mlib = llvm_target_->GetInstance().ParseIR(code);
- }
-
-#if TVM_LLVM_VERSION >= 210
- mlib->setTargetTriple(llvm::Triple(llvm_target_->GetTargetTriple()));
-#else
- mlib->setTargetTriple(llvm_target_->GetTargetTriple());
-#endif
-
mlib->setDataLayout(llvm_target_->GetOrCreateTargetMachine()->createDataLayout());
- // mark all the functions as force inline
- for (llvm::Function& f : mlib->functions()) {
- f.removeFnAttr(llvm::Attribute::OptimizeNone);
- f.removeFnAttr(llvm::Attribute::NoInline);
- f.addFnAttr(llvm::Attribute::AlwaysInline);
- f.setLinkage(llvm::GlobalValue::AvailableExternallyLinkage);
- }
- // add to linker libraries.
- this->AddLinkModule(std::move(mlib));
-}
-
void CodeGenLLVM::AddLinkModule(std::unique_ptr<llvm::Module>&& mod) {
link_modules_.emplace_back(std::move(mod));
}
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index bf80ab9cf5..c077d21563 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -481,8 +481,6 @@ class CodeGenLLVM : public
tirx::ExprFunctor<llvm::Value*(const Expr&)>,
// do a scalarize call with f
llvm::Value* CreateScalarizedCall(const CallNode* op, llvm::Function* f,
const std::vector<llvm::Value*>& args);
- // handle module import
- void HandleImport(const std::string& code);
// cast operatpr
llvm::Value* CreateCast(PrimType from, PrimType to, llvm::Value* value);
// comparison op
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 07e97e48b6..a0fdc0efba 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -1311,11 +1311,6 @@ void CodeGenC::Dispatch_(const RegionStmtNode* op) {
void CodeGenC::Dispatch_(const AttrStmtNode* op) {
TVM_FFI_CHECK(op->attr_key != "thread_extent" && op->attr_key !=
"virtual_thread", ValueError)
<< "Launch attributes are retired; use tirx.launch_thread RegionStmt";
- if (op->attr_key == tirx::attr::pragma_import_c) {
- const StringImmNode* value = op->value.as<StringImmNode>();
- TVM_FFI_ICHECK(value != nullptr);
- decl_stream << value->value;
- }
this->PrintStmt(op->body);
}
diff --git a/src/tirx/transform/lower_tirx_opaque.cc
b/src/tirx/transform/lower_tirx_opaque.cc
index c3b05e93c2..f30beeaaa4 100644
--- a/src/tirx/transform/lower_tirx_opaque.cc
+++ b/src/tirx/transform/lower_tirx_opaque.cc
@@ -27,11 +27,14 @@
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/scope_stack.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
+#include <optional>
+
#include "ir_utils.h"
namespace tvm {
@@ -55,6 +58,11 @@ class TIRxOpaqueLower : public StmtExprMutator {
}
private:
+ struct UnrollPolicy {
+ std::optional<ffi::Any> auto_unroll_max_step;
+ std::optional<ffi::Any> unroll_explicit;
+ };
+
UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
// Attribute subjects may contain references to the loop binding being
// replaced, including expressions and IterVar fields. Rewrite metadata
@@ -88,12 +96,25 @@ class TIRxOpaqueLower : public StmtExprMutator {
VarRemapSet(op->loop_var, prim::cast(op->loop_var.ty(), min));
}
- // Annotations share the loop binding's scope. Mutate them before the body
- // so body-local definitions cannot escape into annotation expressions.
- auto annotations = this->Mutate(op->annotations, inplace_mode)
- .as_or_throw<UnchangedOr<ffi::Map<ffi::String,
ffi::Any>>>()
- .ValueOrUnchanged(op->annotations);
- Stmt body = this->Mutate(op->body,
inplace_mode).ValueOrUnchanged(op->body);
+ // Keep policy on surviving descendants when this owner is lowered away.
+ auto parent_policy = unroll_policy_.Current();
+ auto annotations = op->annotations;
+ Stmt body = unroll_policy_.WithNewScope([&]() {
+ unroll_policy_.Current() = parent_policy;
+ UpdateUnrollPolicy(op->annotations);
+ const auto& policy = unroll_policy_.Current();
+ if (policy.auto_unroll_max_step.has_value()) {
+ annotations.Set(tirx::attr::auto_unroll_max_step,
policy.auto_unroll_max_step.value());
+ }
+ if (policy.unroll_explicit.has_value()) {
+ annotations.Set(tirx::attr::unroll_explicit,
policy.unroll_explicit.value());
+ }
+ // Rewrite annotations before visiting body-local definitions.
+ annotations = this->Mutate(annotations, inplace_mode)
+ .as_or_throw<UnchangedOr<ffi::Map<ffi::String,
ffi::Any>>>()
+ .ValueOrUnchanged(annotations);
+ return this->Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
+ });
VarRemapSet(op->loop_var, previous_remap);
// Step 3. Handle annotations
@@ -146,6 +167,21 @@ class TIRxOpaqueLower : public StmtExprMutator {
return body;
}
+ void UpdateUnrollPolicy(const ffi::Map<ffi::String, ffi::Any>& annotations) {
+ auto& policy = unroll_policy_.Current();
+ if (auto value = annotations.Get(tirx::attr::auto_unroll_max_step);
+ value.has_value() && value.value() != nullptr) {
+ policy.auto_unroll_max_step = value.value();
+ }
+ if (auto value = annotations.Get(tirx::attr::unroll_explicit);
+ value.has_value() && value.value() != nullptr) {
+ policy.unroll_explicit = value.value();
+ }
+ }
+
+ // Effective policy is materialized on each surviving loop, preserving
nested overrides.
+ ScopeStack<UnrollPolicy> unroll_policy_;
+
/*! \brief Convert attr value from annotation map into Expr. */
Expr ConvertAttrValue(const ffi::String& key, const Any& obj) {
if (auto expr = obj.try_cast<Expr>()) {
@@ -173,6 +209,10 @@ class TIRxOpaqueLower : public StmtExprMutator {
pragma_attrs->clear();
for (const auto& kv : annotations) {
const ffi::String& key = kv.first;
+ if ((key == tirx::attr::auto_unroll_max_step || key ==
tirx::attr::unroll_explicit) &&
+ kv.second == nullptr) {
+ continue;
+ }
if (key == "pragma_unroll") {
if (kv.second != nullptr) {
preserved_annotations.Set(key, kv.second);
diff --git a/src/tirx/transform/unroll_loop.cc
b/src/tirx/transform/unroll_loop.cc
index c994e008ed..7f77a383b2 100644
--- a/src/tirx/transform/unroll_loop.cc
+++ b/src/tirx/transform/unroll_loop.cc
@@ -27,11 +27,13 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/scope_stack.h>
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
+#include <optional>
#include <unordered_set>
#include "../../runtime/thread_storage_scope.h"
@@ -101,31 +103,30 @@ class LoopUnroller : public StmtExprMutator {
using StmtExprMutator::Mutate_;
explicit LoopUnroller(int auto_max_step, int auto_max_depth, int
auto_max_extent,
bool explicit_unroll, bool unroll_local_access)
- : auto_max_step_(auto_max_step),
- auto_max_depth_(auto_max_depth),
+ : auto_max_depth_(auto_max_depth),
auto_max_extent_(auto_max_extent),
- explicit_unroll_(explicit_unroll),
- unroll_local_access_(unroll_local_access) {}
-
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- if (op->attr_key == "pragma_auto_unroll_max_step") {
- int value = op->value.as_or_throw<IntImm>()->value.as<int>().value();
- std::swap(value, auto_max_step_);
- Stmt ret = this->Mutate(op->body,
inplace_mode).ValueOrUnchanged(op->body);
- std::swap(value, auto_max_step_);
- return ret;
- } else if (op->attr_key == "pragma_unroll_explicit") {
- bool explicit_unroll =
static_cast<bool>(op->value.as_or_throw<IntImm>()->value);
- std::swap(explicit_unroll, explicit_unroll_);
- Stmt ret = this->Mutate(op->body,
inplace_mode).ValueOrUnchanged(op->body);
- std::swap(explicit_unroll, explicit_unroll_);
- return ret;
- } else {
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
+ unroll_local_access_(unroll_local_access) {
+ unroll_policy_.Current() = {auto_max_step, explicit_unroll};
}
- UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) {
+ UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
+ auto parent_policy = unroll_policy_.Current();
+ return unroll_policy_.WithNewScope([&]() {
+ auto& policy = unroll_policy_.Current();
+ policy = parent_policy;
+ if (auto value = op->annotations.Get(attr::auto_unroll_max_step);
+ value.has_value() && value.value() != nullptr) {
+ policy.auto_unroll_max_step =
value.value().cast<IntImm>()->value.as<int>().value();
+ }
+ if (auto value = op->annotations.Get(attr::unroll_explicit);
+ value.has_value() && value.value() != nullptr) {
+ policy.unroll_explicit =
static_cast<bool>(value.value().cast<IntImm>()->value);
+ }
+ return RewriteLoop(op, inplace_mode);
+ });
+ }
+
+ UnchangedOr<Stmt> RewriteLoop(const ForNode* op, InplaceMode inplace_mode) {
// Post order so we can collect more information
auto result = StmtExprMutator::Mutate_(op, inplace_mode);
if (!result.IsUnchanged()) {
@@ -134,13 +135,28 @@ class LoopUnroller : public StmtExprMutator {
inplace_mode = InplaceMode::kDisallow;
}
}
+ if (op->annotations.count(attr::auto_unroll_max_step) ||
+ op->annotations.count(attr::unroll_explicit)) {
+ ForNode* node;
+ if (inplace_mode == InplaceMode::kAllow) {
+ node = const_cast<ForNode*>(op);
+ } else {
+ auto copy = ffi::make_object<ForNode>(*op);
+ node = copy.get();
+ result = For(std::move(copy));
+ }
+ node->annotations.erase(attr::auto_unroll_max_step);
+ node->annotations.erase(attr::unroll_explicit);
+ op = node;
+ }
int value = GetExtent(op);
// condition for auto unroll
bool auto_unroll = (op->kind == ForKind::kSerial && value >= 0 &&
normal_loop_depth_ == 0 &&
unroll_depth_ <= auto_max_depth_);
- auto_unroll =
- auto_unroll && (value * step_count_ <= auto_max_step_ || value <=
auto_max_extent_);
+ const auto& policy = unroll_policy_.Current();
+ auto_unroll = auto_unroll && (value * step_count_ <=
policy.auto_unroll_max_step.value() ||
+ value <= auto_max_extent_);
if (op->kind == ForKind::kUnrolled) {
TVM_FFI_ICHECK_GE(value, 0) << "Cannot unroll non-constant loop";
@@ -160,7 +176,7 @@ class LoopUnroller : public StmtExprMutator {
normal_loop_depth_ += 1;
}
- if ((auto_unroll && explicit_unroll_) ||
+ if ((auto_unroll && policy.unroll_explicit.value()) ||
// unroll loops with extent = 1, no matter how many steps in body
(0 <= value && value <= auto_max_extent_ && auto_max_extent_ == 1)) {
return Unroll(op);
@@ -275,13 +291,15 @@ class LoopUnroller : public StmtExprMutator {
return value;
}
- // maximum number of step to perform auto unroll.
- int auto_max_step_;
+ struct UnrollPolicy {
+ std::optional<int> auto_unroll_max_step;
+ std::optional<bool> unroll_explicit;
+ };
+ ScopeStack<UnrollPolicy> unroll_policy_;
int auto_max_depth_;
// max extent of loop to auto unroll
// this does not count the total steps, only count the number of loops
int auto_max_extent_;
- bool explicit_unroll_;
// Wether to unroll loops to local access.
bool unroll_local_access_{false};
// Number of normal loops in scope
diff --git a/tests/python/codegen/test_target_codegen_llvm.py
b/tests/python/codegen/test_target_codegen_llvm.py
index 87a6110522..cf3548c775 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -27,7 +27,7 @@ import tvm
import tvm.testing
from tvm.script import ir as I
from tvm.script import tirx as T
-from tvm.support import clang, utils
+from tvm.support import utils
from tvm.target.codegen import llvm_get_intrinsic_name,
llvm_lookup_intrinsic_id
from tvm.testing import env
@@ -872,45 +872,6 @@ def
test_llvm_large_stack_allocation_uses_64bit_extent(extent):
assert re.search(rf"alloca float, i64 {extent}(?:,|$)", llvm_ir)
[email protected](not env.has_llvm(), reason="need llvm")
[email protected]_if_32bit
-def test_llvm_import():
- """all-platform-minimal-test: check shell dependent clang behavior."""
- # extern "C" is necessary to get the correct signature
- cc_code = """
- extern "C" float my_add(float x, float y) {
- return x + y;
- }
- """
-
- def check_llvm(use_file):
- if not clang.find_clang(required=False):
- print("skip because clang is not available")
- return
- temp = utils.tempdir()
- ll_path = temp.relpath("temp.ll")
- ll_code = clang.create_llvm(cc_code, output=ll_path)
- import_val = ll_path if use_file else ll_code
-
- @I.ir_module
- class Module:
- @T.prim_func
- def main(A: T.Tensor((10,), "float32"), B: T.Tensor((10,),
"float32")):
- T.func_attr({"tirx.noalias": True})
- for i in T.serial(10, annotations={"pragma_import_llvm":
import_val}):
- B[i] = T.call_pure_extern("float32", "my_add", A[i],
T.float32(1.0))
-
- f = tvm.compile(Module, target="llvm")
- dev = tvm.cpu(0)
- a = tvm.runtime.tensor(np.random.uniform(size=10).astype("float32"),
dev)
- b = tvm.runtime.tensor(np.random.uniform(size=10).astype("float32"),
dev)
- f(a, b)
- tvm.testing.assert_allclose(b.numpy(), a.numpy() + 1.0)
-
- check_llvm(use_file=True)
- check_llvm(use_file=False)
-
-
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_llvm_scalar_concat():
@I.ir_module
diff --git a/tests/python/s_tir/dlight/test_benchmark.py
b/tests/python/s_tir/dlight/test_benchmark.py
index f5f488da20..9d510100a9 100644
--- a/tests/python/s_tir/dlight/test_benchmark.py
+++ b/tests/python/s_tir/dlight/test_benchmark.py
@@ -118,7 +118,7 @@ def cuda_workload(inp0: T.Tensor((T.int64(1), m,
T.int64(4096))), inp1: T.Tensor
for ax2_1 in T.thread_binding(T.int64(1), thread="vthread.y"):
for ax1_1 in T.thread_binding(T.int64(1),
thread="vthread.x"):
for ax2_2 in T.thread_binding(T.int64(16),
thread="threadIdx.y"):
- for ax1_2 in T.thread_binding(T.int64(8),
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_2 in T.thread_binding(T.int64(8),
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
for ax2_3_init, ax1_3_init in
T.grid(T.int64(4), T.int64(4)):
with Ts.sblock("matmul_init"):
v0 = Ts.axis.spatial(T.int64(1), ax0)
diff --git a/tests/python/s_tir/dlight/test_cpu_gemv.py
b/tests/python/s_tir/dlight/test_cpu_gemv.py
index e0d8cb53e7..e29b92ac94 100644
--- a/tests/python/s_tir/dlight/test_cpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_cpu_gemv.py
@@ -84,7 +84,7 @@ def test_gemv_basic():
for ax0_fused in range(32):
for ax1_fused_0 in T.parallel((n + 63) // 64):
for ax1_fused_1 in T.vectorized(64):
- for ax2_fused_0 in T.serial(2,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_0 in T.serial(2,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax2_fused_1, u_0, u_1 in T.grid(64, 1, 1):
with Ts.sblock("NT_matmul"):
v0 = Ts.axis.spatial(32, ax0_fused)
@@ -140,7 +140,7 @@ def test_decode_gemv_256_threads():
for u_fused in range(1):
for ax0_fused_0 in T.parallel(172):
for ax0_fused_1 in T.vectorized(128):
- for ax1_0_fused_0 in T.serial(8,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0_fused_0 in T.serial(8,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax1_0_fused_1, ax1_1_0, ax1_1_1 in T.grid(64, 1,
8):
with Ts.sblock("NT_matmul"):
v0 = Ts.axis.spatial(22016, ax0_fused_0 * 128
+ ax0_fused_1)
@@ -188,7 +188,7 @@ def test_decode_gemv1():
for u_fused in range(1):
for ax0_fused_0 in T.parallel(172):
for ax0_fused_1 in T.vectorized(128):
- for ax1_0_fused_0 in T.serial(8,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0_fused_0 in T.serial(8,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax1_0_fused_1, ax1_1_0, ax1_1_1 in T.grid(64, 1,
8):
with Ts.sblock("NT_matmul"):
v0 = Ts.axis.spatial(22016, ax0_fused_0 * 128
+ ax0_fused_1)
@@ -244,7 +244,7 @@ def test_decode_gemv2():
for u_fused in range(1):
for ax0_fused_0 in T.parallel(250):
for ax0_fused_1 in T.vectorized(128):
- for ax1_0_fused_0 in T.serial(8,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0_fused_0 in T.serial(8,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax1_0_fused_1, ax1_1_0, ax1_1_1 in T.grid(64, 1,
8):
with Ts.sblock("NT_matmul"):
v0 = Ts.axis.spatial(32000, ax0_fused_0 * 128
+ ax0_fused_1)
@@ -306,7 +306,7 @@ def test_decode_gemv3():
for u_fused in range(1):
for ax0_fused_0 in T.parallel(T.int64(64)):
for ax0_fused_1 in T.vectorized(T.int64(64)):
- for ax1_0_fused_0 in T.serial(T.int64(11),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0_fused_0 in T.serial(T.int64(11),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax1_0_fused_1, ax1_1_0, ax1_1_1 in
T.grid(T.int64(128), T.int64(1), T.int64(8)):
with Ts.sblock("NT_matmul"):
v0 = Ts.axis.spatial(T.int64(4096),
ax0_fused_0 * T.int64(64) + ax0_fused_1)
@@ -535,7 +535,7 @@ def test_blockized_gemv():
for u_fused in range(1):
for ax0_fused_0 in T.parallel(128):
for ax0_fused_1 in T.vectorized(128):
- for ax1_fused_0 in T.serial(64,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_fused_0 in T.serial(64,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax1_fused_1, u_0, u_1 in T.grid(64, 1, 1):
with Ts.sblock("gemv"):
v0 = Ts.axis.spatial(16384,
ax0_fused_0 * 128 + ax0_fused_1)
diff --git a/tests/python/s_tir/dlight/test_gpu_conv.py
b/tests/python/s_tir/dlight/test_gpu_conv.py
index 7cfb1058ca..7b37462996 100644
--- a/tests/python/s_tir/dlight/test_gpu_conv.py
+++ b/tests/python/s_tir/dlight/test_gpu_conv.py
@@ -56,7 +56,7 @@ def test_conv3d():
for ax2_1 in T.thread_binding(1, thread="vthread.y"):
for ax1_1 in T.thread_binding(1, thread="vthread.x"):
for ax2_2 in T.thread_binding(16,
thread="threadIdx.y"):
- for ax1_2 in T.thread_binding(8,
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_2 in T.thread_binding(8,
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
for ax1_3_init, ax2_3_0_init in T.grid(4, 2):
for ax2_3_1_init in T.vectorized(2):
with Ts.sblock("C_init"):
diff --git a/tests/python/s_tir/dlight/test_gpu_gemv.py
b/tests/python/s_tir/dlight/test_gpu_gemv.py
index c17fe5be3b..3a6c415e9d 100644
--- a/tests/python/s_tir/dlight/test_gpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_gemv.py
@@ -114,7 +114,7 @@ def test_gemv_basic():
for ax0_fused_ax1_fused_fused_1 in T.thread_binding(1,
thread="threadIdx.y"):
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 in
T.thread_binding(64, thread="threadIdx.x"):
for ax0, ax1, ax2 in T.grid(1, 1, 1):
- for ax3_0 in T.serial(1,
annotations={"pragma_unroll_explicit": 256, "pragma_vectorize": 1}):
+ for ax3_0 in T.serial(1,
annotations={"unroll_explicit": 256, "pragma_vectorize": 1}):
for ax3_1 in T.thread_binding(1,
thread="threadIdx.y"):
for ax3_2 in T.thread_binding(64,
thread="threadIdx.x"):
for ax3_3 in T.vectorized(2):
@@ -135,7 +135,7 @@ def test_gemv_basic():
Ts.reads()
Ts.writes(var_NT_matmul_intermediate_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused,
0, v0, 0, v1])
var_NT_matmul_intermediate_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused,
0, v0, 0, v1] = T.float16(0)
- for ax2_fused_u_fused_0 in T.serial(1,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_u_fused_0 in T.serial(1,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax0, ax1, ax2_ax3_fused_0 in T.grid(1, 1, 1):
for ax2_ax3_fused_1 in T.vectorized(2):
with Ts.sblock("lv1638_local"):
@@ -158,7 +158,7 @@ def test_gemv_basic():
var_NT_matmul_intermediate_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused,
0, v0, 0, v1] =
var_NT_matmul_intermediate_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused,
0, v0, 0, v1] + lv1637_shared[0, v0, 0, vax2_fused_u_fused_0 * 128 +
vax2_fused_u_fused_2 * 2 + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused] *
lv1638_local[0, v0, v1, vax2_fused_u_fused_0 * 128 + vax2_fused_u_fused_2 * 2 +
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused]
for ax2_ax3_fused_0 in T.thread_binding(1, thread="threadIdx.y"):
for ax0 in T.thread_binding(64, thread="threadIdx.x"):
- for ax2_ax3_fused_1_0 in T.serial(1,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_ax3_fused_1_0 in T.serial(1,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax2_ax3_fused_1_1 in T.vectorized(1):
with Ts.sblock("NT_matmul_rf_init"):
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 = Ts.axis.spatial(64, ax0)
@@ -244,7 +244,7 @@ def test_decode_gemv_256_threads():
Ts.reads()
Ts.writes(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0])
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0] = T.float16(0)
- for ax1_0_fused_ax1_1_fused_0 in T.serial(32,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0_fused_ax1_1_fused_0 in T.serial(32,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax0_ax1_fused in T.serial(1):
for ax0_1 in T.vectorized(1):
with Ts.sblock("lv571_local"):
@@ -264,7 +264,7 @@ def test_decode_gemv_256_threads():
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0] =
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0] + lv1654[0, 0, vax1_0_fused_ax1_1_fused_0 * 128 +
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused * 8 +
vax1_0_fused_ax1_1_fused_2] * ((T.Cast("float16",
T.bitwise_and(T.shift_right(lv571_local[v0, vax1_0_fused_ax1_1_fused_0 * 16 +
[...]
for ax2_fused_0 in T.thread_binding(4, thread="threadIdx.x"):
for ax0 in T.thread_binding(16, thread="threadIdx.y"):
- for ax2_fused_1_0 in T.serial(1,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_1_0 in T.serial(1,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax2_fused_1_1 in T.vectorized(1):
with Ts.sblock("NT_matmul_rf_init"):
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 =
Ts.axis.spatial(16, ax0)
@@ -334,7 +334,7 @@ def test_decode_gemv1():
for u_fused_ax0_fused_fused_1 in T.thread_binding(16,
thread="threadIdx.y"):
for
ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 in
T.thread_binding(32, thread="threadIdx.x"):
for ax0, ax1 in T.grid(1, 1):
- for ax2_0 in T.serial(1,
annotations={"pragma_unroll_explicit": 256, "pragma_vectorize": 1}):
+ for ax2_0 in T.serial(1,
annotations={"unroll_explicit": 256, "pragma_vectorize": 1}):
for ax2_1 in T.thread_binding(16,
thread="threadIdx.y"):
for ax2_2 in T.thread_binding(32,
thread="threadIdx.x"):
for ax2_3 in T.vectorized(8):
@@ -352,7 +352,7 @@ def test_decode_gemv1():
Ts.reads()
Ts.writes(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0])
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0] = T.float16(0)
- for ax1_0_fused_ax1_1_fused_0 in T.serial(16,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0_fused_ax1_1_fused_0 in T.serial(16,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax0_ax1_fused_0 in range(1):
for ax0_ax1_fused_1 in T.vectorized(1):
with Ts.sblock("lv571_local"):
@@ -372,7 +372,7 @@ def test_decode_gemv1():
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0] =
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0] + lv1654_shared[0, 0, vax1_0_fused_ax1_1_fused_0 * 256 +
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 * 8 +
vax1_0_fused_ax1_1_fused_2 * 4 +
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused % 4] *
((T.Cast("float1 [...]
for ax2_fused_0_ax2_fused_1_fused in T.thread_binding(16,
thread="threadIdx.y"):
for ax0 in T.thread_binding(32, thread="threadIdx.x"):
- for ax2_fused_2_0 in T.serial(1,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_2_0 in T.serial(1,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax2_fused_2_1 in T.vectorized(1):
with Ts.sblock("NT_matmul_rf_init"):
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 =
Ts.axis.spatial(32, ax0)
@@ -450,7 +450,7 @@ def test_decode_gemv2():
for u_fused_ax0_fused_fused_1 in T.thread_binding(16,
thread="threadIdx.y"):
for
ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 in
T.thread_binding(32, thread="threadIdx.x"):
for ax0, ax1 in T.grid(1, 1):
- for ax2_0 in T.serial(1,
annotations={"pragma_unroll_explicit": 256, "pragma_vectorize": 1}):
+ for ax2_0 in T.serial(1,
annotations={"unroll_explicit": 256, "pragma_vectorize": 1}):
for ax2_1 in T.thread_binding(16,
thread="threadIdx.y"):
for ax2_2 in T.thread_binding(32,
thread="threadIdx.x"):
for ax2_3 in T.vectorized(8):
@@ -468,7 +468,7 @@ def test_decode_gemv2():
Ts.reads()
Ts.writes(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0])
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0] = T.float16(0)
- for ax1_0_fused_ax1_1_fused_0 in T.serial(16,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0_fused_ax1_1_fused_0 in T.serial(16,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax0_ax1_fused_0 in range(1):
for ax0_ax1_fused_1 in T.vectorized(1):
with Ts.sblock("lv771_local"):
@@ -488,7 +488,7 @@ def test_decode_gemv2():
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0] =
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
0, 0, v0] + lv3216_shared[0, 0, vax1_0_fused_ax1_1_fused_0 * 256 +
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // 4 * 8 +
vax1_0_fused_ax1_1_fused_2 * 4 +
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused % 4] *
((T.Cast("float1 [...]
for ax2_fused_0_ax2_fused_1_fused in T.thread_binding(16,
thread="threadIdx.y"):
for ax0 in T.thread_binding(32, thread="threadIdx.x"):
- for ax2_fused_2_0 in T.serial(1,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_2_0 in T.serial(1,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax2_fused_2_1 in T.vectorized(1):
with Ts.sblock("NT_matmul_rf_init"):
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 =
Ts.axis.spatial(32, ax0)
@@ -573,7 +573,7 @@ def test_decode_gemv3():
for u_fused_ax0_fused_fused_1 in T.thread_binding(T.int64(16),
thread="threadIdx.y"):
for
ax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 in
T.thread_binding(T.int64(32), thread="threadIdx.x"):
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
- for ax2_0 in T.serial(T.int64(22),
annotations={"pragma_unroll_explicit": 256, "pragma_vectorize": 1}):
+ for ax2_0 in T.serial(T.int64(22),
annotations={"unroll_explicit": 256, "pragma_vectorize": 1}):
for ax2_1 in T.thread_binding(T.int64(16),
thread="threadIdx.y"):
for ax2_2 in T.thread_binding(T.int64(32),
thread="threadIdx.x"):
for ax2_3 in T.vectorized(T.int64(1)):
@@ -592,7 +592,7 @@ def test_decode_gemv3():
Ts.reads()
Ts.writes(var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
T.int64(0), T.int64(0), v0])
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
T.int64(0), T.int64(0), v0] = T.float16(0)
- for ax1_0_fused_ax1_1_fused_0 in T.serial(T.int64(43),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0_fused_ax1_1_fused_0 in T.serial(T.int64(43),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax0_ax1_fused_0 in T.serial(T.int64(0),
T.int64(1)):
for ax0_ax1_fused_1 in T.vectorized(T.int64(1)):
with Ts.sblock("lv575_local"):
@@ -612,7 +612,7 @@ def test_decode_gemv3():
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
T.int64(0), T.int64(0), v0] =
var_NT_matmul_intermediate_rf_local[vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused,
T.int64(0), T.int64(0), v0] + lv574_shared[T.int64(0), T.int64(0),
vax1_0_fused_ax1_1_fused_0 * T.int64(256) +
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused // T.int64(4) *
T.int64(8) + vax1_0_fused_ax1_1_fused_2 * T.int6 [...]
for ax2_fused_0 in T.thread_binding(T.int64(16),
thread="threadIdx.y"):
for ax0 in T.thread_binding(T.int64(32), thread="threadIdx.x"):
- for ax2_fused_1_0 in T.serial(T.int64(1),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_1_0 in T.serial(T.int64(1),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax2_fused_1_1 in T.vectorized(T.int64(1)):
with Ts.sblock("NT_matmul_rf_init"):
vax1_0_fused_ax1_1_fused_1_ax1_0_fused_ax1_1_fused_3_fused_0 =
Ts.axis.spatial(T.int64(32), ax0)
@@ -773,7 +773,7 @@ def test_outer_reduction_adreno():
Ts.reads()
Ts.writes(var_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_2_ax1_0_fused_ax1_1_fused_4_fused_0,
0, 0, v0])
var_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_2_ax1_0_fused_ax1_1_fused_4_fused_0,
0, 0, v0] = T.float16(0)
- for ax1 in T.serial(8,
annotations={"pragma_auto_unroll_max_step": 8, "pragma_unroll_explicit": 1}):
+ for ax1 in T.serial(8,
annotations={"auto_unroll_max_step": 8, "unroll_explicit": 1}):
with Ts.sblock("matmul_rf_update"):
vax1_0_fused_ax1_1_fused_2_ax1_0_fused_ax1_1_fused_4_fused_0,
vax1_0_fused_ax1_1_fused_2_ax1_0_fused_ax1_1_fused_4_fused_1 =
Ts.axis.remap("SR", [ax0, ax1])
v0 = Ts.axis.spatial(4096,
u_fused_ax0_fused_fused_0 * 64 + ax2)
@@ -909,7 +909,7 @@ def test_outer_reduction_adreno_dynamic():
Ts.reads()
Ts.writes(var_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_2_ax1_0_fused_ax1_1_fused_4_fused_0,
T.int64(0), T.int64(0), v0])
var_matmul_intermediate_rf_local_1[vax1_0_fused_ax1_1_fused_2_ax1_0_fused_ax1_1_fused_4_fused_0,
T.int64(0), T.int64(0), v0] = T.float16(0)
- for ax1 in T.serial(T.int64(8),
annotations={"pragma_auto_unroll_max_step": 8, "pragma_unroll_explicit": 1}):
+ for ax1 in T.serial(T.int64(8),
annotations={"auto_unroll_max_step": 8, "unroll_explicit": 1}):
with Ts.sblock("matmul_rf_update"):
vax1_0_fused_ax1_1_fused_2_ax1_0_fused_ax1_1_fused_4_fused_0,
vax1_0_fused_ax1_1_fused_2_ax1_0_fused_ax1_1_fused_4_fused_1 =
Ts.axis.remap("SR", [ax0, ax1])
v0 = Ts.axis.spatial(v, u_fused_ax0_fused_fused_0
* T.int64(256) + ax2)
@@ -990,7 +990,7 @@ def test_blockized_gemv():
Ts.reads()
Ts.writes(o_rf_local[vax1_fused_u_fused_1_ax1_fused_u_fused_3_fused,
v_expert_id_o, v0])
o_rf_local[vax1_fused_u_fused_1_ax1_fused_u_fused_3_fused, v_expert_id_o, v0] =
T.float16(0)
- for ax1_fused_u_fused_0 in T.serial(32,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_fused_u_fused_0 in T.serial(32,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax0 in range(1):
for ax1_ax2_fused_0 in range(8):
for ax1_ax2_fused_1 in T.vectorized(1):
@@ -1012,7 +1012,7 @@ def test_blockized_gemv():
o_rf_local[vax1_fused_u_fused_1_ax1_fused_u_fused_3_fused, v_expert_id_o, v0] =
o_rf_local[vax1_fused_u_fused_1_ax1_fused_u_fused_3_fused, v_expert_id_o, v0] +
x[0, vax1_fused_u_fused_0 * 128 +
vax1_fused_u_fused_1_ax1_fused_u_fused_3_fused * 8 + vax1_fused_u_fused_2] *
w_local[0, v0, vax1_fused_u_fused_0 * 128 +
vax1_fused_u_fused_1_ax1_fused_u_fused_3_fused * 8 + vax1_fused_u_fused_2]
for ax2_fused_0 in T.thread_binding(4,
thread="threadIdx.x"):
for ax0 in T.thread_binding(16, thread="threadIdx.y"):
- for ax2_fused_1_0 in T.serial(1,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_1_0 in T.serial(1,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax2_fused_1_1 in T.vectorized(1):
with Ts.sblock("gemv_rf_init"):
vax1_fused_u_fused_1_ax1_fused_u_fused_3_fused_0 = Ts.axis.spatial(16, ax0)
diff --git a/tests/python/s_tir/dlight/test_gpu_general_reduction.py
b/tests/python/s_tir/dlight/test_gpu_general_reduction.py
index 76d0dde1fe..e4558f047d 100644
--- a/tests/python/s_tir/dlight/test_gpu_general_reduction.py
+++ b/tests/python/s_tir/dlight/test_gpu_general_reduction.py
@@ -156,7 +156,7 @@ def test_softmax_1():
for ax0_ax1_fused in T.thread_binding(n * T.int64(32),
thread="blockIdx.x"):
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
for ax2_fused_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax2_fused_0 in T.serial((m + T.int64(255)) //
T.int64(256), annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax2_fused_0 in T.serial((m + T.int64(255)) //
T.int64(256), annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("T_softmax_maxelem"):
v0 = Ts.axis.spatial(T.int64(32),
ax0_ax1_fused // n + ax0)
v1 = Ts.axis.spatial(n, ax0_ax1_fused % n +
ax1)
@@ -169,7 +169,7 @@ def test_softmax_1():
T_softmax_maxelem_shared[T.int64(0), v0, v1] =
T.max(T_softmax_maxelem_shared[T.int64(0), v0, v1], lv44[T.int64(0), v0, v1,
v2])
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
for ax2_fused_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax2_fused_0 in T.serial((m + T.int64(255)) //
T.int64(256), annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax2_fused_0 in T.serial((m + T.int64(255)) //
T.int64(256), annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("T_softmax_expsum"):
v0 = Ts.axis.spatial(T.int64(32),
ax0_ax1_fused // n + ax0)
v1 = Ts.axis.spatial(n, ax0_ax1_fused % n +
ax1)
@@ -181,7 +181,7 @@ def test_softmax_1():
T_softmax_expsum_shared[T.int64(0), v0,
v1] = T.float32(0)
T_softmax_expsum_shared[T.int64(0), v0, v1] =
T_softmax_expsum_shared[T.int64(0), v0, v1] + T.exp(lv44[T.int64(0), v0, v1,
v2] - T_softmax_maxelem_shared[T.int64(0), v0, v1])
for ax2_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax2_0 in T.serial((m + T.int64(255)) // T.int64(256),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_0 in T.serial((m + T.int64(255)) // T.int64(256),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("compute"):
v0 = Ts.axis.spatial(T.int64(32), ax0_ax1_fused //
n)
v1 = Ts.axis.spatial(n, ax0_ax1_fused % n)
@@ -245,7 +245,7 @@ def test_softmax_2():
for ax0_fused in T.thread_binding(T.int64(1), thread="blockIdx.x"):
for ax0 in T.serial(T.int64(0), T.int64(1)):
for ax1_fused_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax1_fused_0 in T.serial(T.int64(125),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_fused_0 in T.serial(T.int64(125),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("T_softmax_maxelem"):
v0 = Ts.axis.spatial(T.int64(1), ax0)
v1 = Ts.axis.reduce(T.int64(32000),
ax1_fused_0 * T.int64(256) + ax1_fused_1)
@@ -256,7 +256,7 @@ def test_softmax_2():
T_softmax_maxelem_shared[T.int64(0),
T.int64(0)] = T.max(T_softmax_maxelem_shared[T.int64(0), T.int64(0)],
A[T.int64(0), T.int64(0), v1])
for ax0 in T.serial(T.int64(0), T.int64(1)):
for ax1_fused_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax1_fused_0 in T.serial(T.int64(125),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_fused_0 in T.serial(T.int64(125),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("T_softmax_expsum"):
v0 = Ts.axis.spatial(T.int64(1), ax0)
v1 = Ts.axis.reduce(T.int64(32000),
ax1_fused_0 * T.int64(256) + ax1_fused_1)
@@ -266,7 +266,7 @@ def test_softmax_2():
T_softmax_expsum_shared[T.int64(0),
T.int64(0)] = T.float32(0)
T_softmax_expsum_shared[T.int64(0),
T.int64(0)] = T_softmax_expsum_shared[T.int64(0), T.int64(0)] +
T.exp(A[T.int64(0), T.int64(0), v1] - T_softmax_maxelem_shared[T.int64(0),
T.int64(0)])
for ax1_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax1_0 in T.serial(T.int64(125),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0 in T.serial(T.int64(125),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("T_softmax_norm"):
v0 = Ts.axis.spatial(T.int64(1), T.int64(0))
v1 = Ts.axis.spatial(T.int64(32000), ax1_0 *
T.int64(256) + ax1_1)
@@ -330,7 +330,7 @@ def test_softmax_3():
for ax0_ax2_fused in T.thread_binding(T.int64(32768),
thread="blockIdx.x"):
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
for ax2_fused_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax2_fused_0 in T.serial(T.int64(1),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_0 in T.serial(T.int64(1),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("T_softmax_maxelem"):
v0 = Ts.axis.spatial(T.int64(4), ax0_ax2_fused
// T.int64(8192) + ax0)
v1 = Ts.axis.spatial(T.int64(8192),
ax0_ax2_fused % T.int64(8192) + ax1)
@@ -343,7 +343,7 @@ def test_softmax_3():
T_softmax_maxelem_shared[T.int64(0), v0, v1] =
T.max(T_softmax_maxelem_shared[T.int64(0), v0, v1], input[T.int64(0), v0, v2,
v1])
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
for ax2_fused_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax2_fused_0 in T.serial(T.int64(1),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_0 in T.serial(T.int64(1),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("T_softmax_expsum"):
v0 = Ts.axis.spatial(T.int64(4), ax0_ax2_fused
// T.int64(8192) + ax0)
v1 = Ts.axis.spatial(T.int64(8192),
ax0_ax2_fused % T.int64(8192) + ax1)
@@ -355,7 +355,7 @@ def test_softmax_3():
T_softmax_expsum_shared[T.int64(0), v0,
v1] = T.float32(0.0)
T_softmax_expsum_shared[T.int64(0), v0, v1] =
T_softmax_expsum_shared[T.int64(0), v0, v1] + T.exp(input[T.int64(0), v0, v2,
v1] - T_softmax_maxelem_shared[T.int64(0), v0, v1])
for ax1_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax1_0 in T.serial(T.int64(1),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0 in T.serial(T.int64(1),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("T_softmax_norm"):
v0 = Ts.axis.spatial(T.int64(4), ax0_ax2_fused //
T.int64(8192))
v1 = Ts.axis.spatial(T.int64(32), ax1_0 *
T.int64(256) + ax1_1)
@@ -422,7 +422,7 @@ def test_layer_norm():
for ax0_fused in T.thread_binding(n, thread="blockIdx.x"):
for ax0 in T.serial(T.int64(0), T.int64(1)):
for ax1_fused_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax1_fused_0 in T.serial(T.int64(10),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_fused_0 in T.serial(T.int64(10),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("A_red_temp"):
v0 = Ts.axis.spatial(n, ax0_fused + ax0)
v1 = Ts.axis.reduce(T.int64(2560), ax1_fused_0
* T.int64(256) + ax1_fused_1)
@@ -436,7 +436,7 @@ def test_layer_norm():
A_red_temp_v0_shared[T.int64(0), v0] =
v_A_red_temp_v0
A_red_temp_v1_shared[T.int64(0), v0] =
v_A_red_temp_v1
for ax1_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax1_0 in T.serial(T.int64(10),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0 in T.serial(T.int64(10),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("compute"):
v0 = Ts.axis.spatial(n, ax0_fused)
v1 = Ts.axis.spatial(T.int64(2560), ax1_0 *
T.int64(256) + ax1_1)
@@ -487,7 +487,7 @@ def test_rms_norm():
for ax0_fused in T.thread_binding(n, thread="blockIdx.x"):
for ax0 in T.serial(T.int64(0), T.int64(1)):
for ax1_fused_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax1_fused_0 in T.serial(T.int64(16),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_fused_0 in T.serial(T.int64(16),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("Ared_temp"):
v0 = Ts.axis.spatial(n, ax0_fused + ax0)
v1 = Ts.axis.reduce(T.int64(4096), ax1_fused_0
* T.int64(256) + ax1_fused_1)
@@ -497,7 +497,7 @@ def test_rms_norm():
Ared_temp_shared[T.int64(0), v0] =
T.float32(0)
Ared_temp_shared[T.int64(0), v0] =
Ared_temp_shared[T.int64(0), v0] + T.Cast("float32", A[T.int64(0), v0, v1]) *
T.Cast("float32", A[T.int64(0), v0, v1])
for ax1_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
- for ax1_0 in T.serial(T.int64(16),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0 in T.serial(T.int64(16),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("rms_norm"):
v0 = Ts.axis.spatial(n, ax0_fused)
v1 = Ts.axis.spatial(T.int64(4096), ax1_0 *
T.int64(256) + ax1_1)
@@ -575,7 +575,7 @@ def test_group_norm():
for ax0_fused in T.thread_binding(T.int64(1), thread="blockIdx.x"):
for ax0 in range(32):
for ax1_fused_1 in T.thread_binding(256,
thread="threadIdx.x"):
- for ax1_fused_0 in T.serial(1,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_fused_0 in T.serial(1,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("A_red_temp"):
v0 = Ts.axis.spatial(32, ax0)
v1 = Ts.axis.reduce(64, ax1_fused_0 * 256 +
ax1_fused_1)
@@ -590,7 +590,7 @@ def test_group_norm():
A_red_temp_v0_shared[0, v0] = v_A_red_temp_v0
A_red_temp_v1_shared[0, v0] = v_A_red_temp_v1
for ax1_1 in T.thread_binding(256, thread="threadIdx.x"):
- for ax1_0 in T.serial(8,
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_0 in T.serial(8,
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("T_reshape_3"):
v0 = Ts.axis.spatial(T.int64(1), T.int64(0))
v1 = Ts.axis.spatial(2048, ax1_0 * 256 + ax1_1)
@@ -670,8 +670,8 @@ def test_logsumexp():
for ax2_fused_0 in T.serial(
T.int64(16),
annotations={
- "pragma_auto_unroll_max_step": 256,
- "pragma_unroll_explicit": 1,
+ "auto_unroll_max_step": 256,
+ "unroll_explicit": 1,
},
):
with Ts.sblock("max"):
@@ -700,8 +700,8 @@ def test_logsumexp():
for ax2_fused_0 in T.serial(
T.int64(16),
annotations={
- "pragma_auto_unroll_max_step": 256,
- "pragma_unroll_explicit": 1,
+ "auto_unroll_max_step": 256,
+ "unroll_explicit": 1,
},
):
with Ts.sblock("sum_exp"):
@@ -733,8 +733,8 @@ def test_logsumexp():
for ax2_0 in T.serial(
T.int64(1),
annotations={
- "pragma_auto_unroll_max_step": 256,
- "pragma_unroll_explicit": 1,
+ "auto_unroll_max_step": 256,
+ "unroll_explicit": 1,
},
):
with Ts.sblock("log"):
diff --git a/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
index fb05d0e450..21710d6e34 100644
--- a/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
@@ -84,7 +84,7 @@ def test_batch_decode_gemv():
Ts.reads()
Ts.writes(NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused,
v0, T.int64(0), v1])
NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused,
v0, T.int64(0), v1] = T.float16(0)
- for ax2_fused_u_fused_0 in T.serial(T.int64(112),
annotations={"pragma_auto_unroll_max_step": 8, "pragma_unroll_explicit": 1}):
+ for ax2_fused_u_fused_0 in T.serial(T.int64(112),
annotations={"auto_unroll_max_step": 8, "unroll_explicit": 1}):
for ax0_0_1, ax1 in T.grid(T.int64(2), T.int64(8)):
for ax0_1 in T.vectorized(T.int64(1)):
with Ts.sblock("dequantize"):
@@ -105,7 +105,7 @@ def test_batch_decode_gemv():
NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused,
v0, T.int64(0), v1] =
NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused,
v0, T.int64(0), v1] + T.if_then_else(v0 < batch_size, lv807[v0, T.int64(0),
vax2_fused_u_fused_0 * T.int64(256) +
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) +
vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fus [...]
for ax3_fused_0_ax3_fused_1_fused in
T.thread_binding(T.int64(8), thread="threadIdx.x"):
for ax0 in T.thread_binding(T.int64(32),
thread="threadIdx.y"):
- for ax3_fused_2_0 in T.serial(T.int64(1),
annotations={"pragma_auto_unroll_max_step": 8, "pragma_unroll_explicit": 1}):
+ for ax3_fused_2_0 in T.serial(T.int64(1),
annotations={"auto_unroll_max_step": 8, "unroll_explicit": 1}):
for ax2 in T.serial(T.int64(0), T.int64(4)):
for ax3_fused_2_1 in T.vectorized(T.int64(2)):
with Ts.sblock("NT_matmul_rf_init"):
@@ -197,7 +197,7 @@ def test_batch_gemv():
Ts.reads()
Ts.writes(NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused,
v0, T.int64(0), v1])
NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0,
T.int64(0), v1] = T.float16(0)
- for ax2_fused_u_fused_0 in T.serial(T.int64(16),
annotations={"pragma_auto_unroll_max_step": 8, "pragma_unroll_explicit": 1}):
+ for ax2_fused_u_fused_0 in T.serial(T.int64(16),
annotations={"auto_unroll_max_step": 8, "unroll_explicit": 1}):
for ax0_1, u_fused_ax1_fused_fused_2,
ax2_fused_u_fused_2 in T.grid(T.int64(4), T.int64(2), T.int64(2)):
for
ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1 in T.vectorized(T.int64(4)):
with Ts.sblock("NT_matmul_rf_update"):
@@ -210,7 +210,7 @@ def test_batch_gemv():
NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0,
T.int64(0), v1] =
NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0,
T.int64(0), v1] + T.if_then_else(v0 < batch_size, A[v0, T.int64(0),
vax2_fused_u_fused_0 * T.int64(256) +
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) +
vax2_fused_u_fused_2 * T.int64(4) +
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)], T.fl [...]
for ax3_fused_0_ax3_fused_1_fused in
T.thread_binding(T.int64(8), thread="threadIdx.x"):
for ax0 in T.thread_binding(T.int64(32),
thread="threadIdx.y"):
- for ax3_fused_2_0 in T.serial(T.int64(1),
annotations={"pragma_auto_unroll_max_step": 8, "pragma_unroll_explicit": 1}):
+ for ax3_fused_2_0 in T.serial(T.int64(1),
annotations={"auto_unroll_max_step": 8, "unroll_explicit": 1}):
for ax2 in T.serial(T.int64(0), T.int64(4)):
for ax3_fused_2_1 in T.vectorized(T.int64(2)):
with Ts.sblock("NT_matmul_rf_init"):
@@ -326,7 +326,7 @@ def test_small_spatial_axis():
Ts.reads()
Ts.writes(C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0,
v1])
C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, v1] =
T.float16(0)
- for ax2_fused_u_fused_0 in T.serial(T.int64(16),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax2_fused_u_fused_0 in T.serial(T.int64(16),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax0_1, u_fused_ax1_fused_fused_2,
ax2_fused_u_fused_2 in T.grid(T.int64(4), T.int64(2), T.int64(2)):
for
ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1 in T.vectorized(T.int64(4)):
with Ts.sblock("NT_matmul_rf_update"):
@@ -340,7 +340,7 @@ def test_small_spatial_axis():
C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, v1] =
C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, v1] +
T.if_then_else(v0 < batch_size, A[v0, vax2_fused_u_fused_0 * T.int64(256) +
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) +
vax2_fused_u_fused_2 * T.int64(4) +
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)], T.float16(0)) *
B[v1, vax2_fused_u_fused_0 * T.int64(256 [...]
for ax3_fused_0_ax3_fused_1_fused in
T.thread_binding(T.int64(16), thread="threadIdx.y"):
for ax0 in T.thread_binding(T.int64(32),
thread="threadIdx.x"):
- for ax3_fused_2_0 in T.serial(T.int64(1),
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax3_fused_2_0 in T.serial(T.int64(1),
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax2 in T.serial(T.int64(0), T.int64(4)):
for ax3_fused_2_1 in T.vectorized(T.int64(2)):
with Ts.sblock("NT_matmul_rf_init"):
diff --git a/tests/python/s_tir/dlight/test_gpu_matmul.py
b/tests/python/s_tir/dlight/test_gpu_matmul.py
index 6063ea3f85..e62dbd07f1 100644
--- a/tests/python/s_tir/dlight/test_gpu_matmul.py
+++ b/tests/python/s_tir/dlight/test_gpu_matmul.py
@@ -53,7 +53,7 @@ def test_matmul():
for ax2_1 in T.thread_binding(T.int64(1), thread="vthread.y"):
for ax1_1 in T.thread_binding(T.int64(1),
thread="vthread.x"):
for ax2_2 in T.thread_binding(T.int64(16),
thread="threadIdx.y"):
- for ax1_2 in T.thread_binding(T.int64(8),
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_2 in T.thread_binding(T.int64(8),
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
for ax1_3_init, ax2_3_0_init in
T.grid(T.int64(4), T.int64(2)):
for ax2_3_1_init in
T.vectorized(T.int64(2)):
with Ts.sblock("matmul_init"):
@@ -145,7 +145,7 @@ def test_matmul_int32():
for ax2_1 in T.thread_binding(1, thread="vthread.y"):
for ax1_1 in T.thread_binding(1, thread="vthread.x"):
for ax2_2 in T.thread_binding(16,
thread="threadIdx.y"):
- for ax1_2 in T.thread_binding(8,
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_2 in T.thread_binding(8,
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
for ax1_3_init, ax2_3_0_init in T.grid(4, 2):
for ax2_3_1_init in T.vectorized(2):
with Ts.sblock("matmul_init"):
@@ -247,7 +247,7 @@ def test_fused_matmul():
for ax2_1 in T.thread_binding(T.int64(1), thread="vthread.y"):
for ax1_1 in T.thread_binding(T.int64(1),
thread="vthread.x"):
for ax2_2 in T.thread_binding(T.int64(16),
thread="threadIdx.y"):
- for ax1_2 in T.thread_binding(T.int64(8),
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_2 in T.thread_binding(T.int64(8),
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
for ax1_3_init, ax2_3_0_init in
T.grid(T.int64(4), T.int64(2)):
for ax2_3_1_init in
T.vectorized(T.int64(2)):
with Ts.sblock("matmul_init"):
@@ -416,7 +416,7 @@ def test_output_fp32():
for ax2_1 in T.thread_binding(T.int64(1), thread="vthread.y"):
for ax1_1 in T.thread_binding(T.int64(1),
thread="vthread.x"):
for ax2_2 in T.thread_binding(T.int64(16),
thread="threadIdx.y"):
- for ax1_2 in T.thread_binding(T.int64(8),
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_2 in T.thread_binding(T.int64(8),
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
for ax1_3_init, ax2_3_0_init in
T.grid(T.int64(4), T.int64(2)):
for ax2_3_1_init in
T.vectorized(T.int64(2)):
with Ts.sblock("matmul_init"):
@@ -548,7 +548,7 @@ def test_inline_consumer_chain():
for ax2_1 in T.thread_binding(T.int64(1), thread="vthread.y"):
for ax1_1 in T.thread_binding(T.int64(1),
thread="vthread.x"):
for ax2_2 in T.thread_binding(T.int64(16),
thread="threadIdx.y"):
- for ax1_2 in T.thread_binding(T.int64(8),
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_2 in T.thread_binding(T.int64(8),
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
for ax1_3_init, ax2_3_0_init in
T.grid(T.int64(4), T.int64(2)):
for ax2_3_1_init in
T.vectorized(T.int64(2)):
with Ts.sblock("NT_matmul_init"):
diff --git a/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
index a59ec2c817..3eeab227a1 100644
--- a/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
+++ b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
@@ -218,7 +218,7 @@ def test_matmul_tensorize_too_small():
for ax2_1 in T.thread_binding(1, thread="vthread.y"):
for ax1_1 in T.thread_binding(1, thread="vthread.x"):
for ax2_2 in T.thread_binding(16,
thread="threadIdx.y"):
- for ax1_2 in T.thread_binding(8,
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_2 in T.thread_binding(8,
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
for ax1_3_init, ax2_3_0_init in T.grid(4, 2):
for ax2_3_1_init in T.vectorized(2):
with Ts.sblock("compute_init"):
diff --git a/tests/python/s_tir/dlight/test_gpu_reduction.py
b/tests/python/s_tir/dlight/test_gpu_reduction.py
index 41b134f664..c8e92982d0 100644
--- a/tests/python/s_tir/dlight/test_gpu_reduction.py
+++ b/tests/python/s_tir/dlight/test_gpu_reduction.py
@@ -62,7 +62,7 @@ def test_decode_gemv_1():
Ts.writes()
C_rf_local = Ts.sblock_alloc_buffer((512, 1, 1, 4096),
"float16", scope="local")
for ax0_fused in T.thread_binding(4096, thread="blockIdx.x"):
- for ax1_0_fused_1 in T.thread_binding(512,
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_0_fused_1 in T.thread_binding(512,
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
with Ts.sblock("matmul_rf_init"):
vax1_0_fused_1 = Ts.axis.spatial(512,
ax1_0_fused_1)
v0 = Ts.axis.spatial(4096, ax0_fused)
@@ -195,7 +195,7 @@ def test_decode_gemv_3():
Ts.writes()
C_rf_local = Ts.sblock_alloc_buffer((1024, 1, 1, 4096),
"float16", scope="local")
for ax0_0_fused in T.thread_binding(512, thread="blockIdx.x"):
- for ax1_fused_1 in T.thread_binding(1024,
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_fused_1 in T.thread_binding(1024,
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
for ax0_1_init in range(8):
with Ts.sblock("matmul_rf_init"):
vax1_fused_1 = Ts.axis.spatial(1024,
ax1_fused_1)
@@ -340,7 +340,7 @@ def test_decode_gemv_sigmoid():
C_local = Ts.sblock_alloc_buffer((1, 1, 4096), "float16",
scope="local")
C_rf_local = Ts.sblock_alloc_buffer((512, 1, 1, 4096),
"float16", scope="local")
for ax0_fused in T.thread_binding(4096, thread="blockIdx.x"):
- for ax1_0_fused_1 in T.thread_binding(512,
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_0_fused_1 in T.thread_binding(512,
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
with Ts.sblock("matmul_rf_init"):
vax1_0_fused_1 = Ts.axis.spatial(512,
ax1_0_fused_1)
v0 = Ts.axis.spatial(4096, ax0_fused)
@@ -426,7 +426,7 @@ def test_decode_gemv_1_fp32():
C_fp32_local = Ts.sblock_alloc_buffer((1, 1, 4096),
scope="local")
C_fp32_rf_local = Ts.sblock_alloc_buffer((512, 1, 1, 4096),
scope="local")
for ax0_fused in T.thread_binding(4096, thread="blockIdx.x"):
- for ax1_0_fused_1 in T.thread_binding(512,
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_0_fused_1 in T.thread_binding(512,
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
with Ts.sblock("matmul_rf_init"):
vax1_0_fused_1 = Ts.axis.spatial(512,
ax1_0_fused_1)
v0 = Ts.axis.spatial(4096, ax0_fused)
@@ -499,7 +499,7 @@ def test_reduction_no_spatial():
Ared_temp_shared = Ts.sblock_alloc_buffer((1, 1),
scope="shared")
Ared_temp_rf_local = Ts.sblock_alloc_buffer((1024, 1, 1),
scope="local")
for ax0_fused in T.thread_binding(T.int64(1),
thread="blockIdx.x"):
- for ax1_fused_1 in T.thread_binding(1024,
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_fused_1 in T.thread_binding(1024,
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
with Ts.sblock("Ared_temp_rf_init"):
vax1_fused_1 = Ts.axis.spatial(1024, ax1_fused_1)
v0 = Ts.axis.spatial(T.int64(1), T.int64(0))
@@ -726,7 +726,7 @@ def test_reduction_inner_no_broadcasting():
temp_local_local = Ts.sblock_alloc_buffer((256,), scope="local")
temp_local_rf_local = Ts.sblock_alloc_buffer((256, 256),
scope="local")
for ax0_fused in T.thread_binding(256, thread="blockIdx.x"):
- for ax1_fused_1 in T.thread_binding(256, thread="threadIdx.x",
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax1_fused_1 in T.thread_binding(256, thread="threadIdx.x",
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
with Ts.sblock("sum_rf_init"):
vax1_fused_1, v0 = Ts.axis.remap("SS", [ax1_fused_1,
ax0_fused])
Ts.reads()
@@ -1274,7 +1274,7 @@ def test_gemv_output_one_element():
NT_matmul_intermediate_shared =
Ts.sblock_alloc_buffer((T.int64(1), T.int64(1)), "float16", scope="shared")
NT_matmul_intermediate_rf_local =
Ts.sblock_alloc_buffer((T.int64(1024), T.int64(1), T.int64(1)), "float16",
scope="local")
for ax0_fused in T.thread_binding(T.int64(1), thread="blockIdx.x"):
- for ax1_fused_1 in T.thread_binding(T.int64(1024),
thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256,
"pragma_unroll_explicit": 1}):
+ for ax1_fused_1 in T.thread_binding(T.int64(1024),
thread="threadIdx.x", annotations={"auto_unroll_max_step": 256,
"unroll_explicit": 1}):
with Ts.sblock("NT_matmul_rf_init"):
vax1_fused_1 = Ts.axis.spatial(T.int64(1024),
ax1_fused_1)
v0 = Ts.axis.spatial(T.int64(1), T.int64(0))
diff --git a/tests/python/s_tir/dlight/test_gpu_transpose.py
b/tests/python/s_tir/dlight/test_gpu_transpose.py
index 7ce9a60f80..6ebe9437ea 100644
--- a/tests/python/s_tir/dlight/test_gpu_transpose.py
+++ b/tests/python/s_tir/dlight/test_gpu_transpose.py
@@ -53,7 +53,7 @@ def test_transpose():
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
# with Ts.sblock("root"):
rxplaceholder_shared = Ts.sblock_alloc_buffer((T.int64(512),
T.int64(4096)), scope="shared")
- for ax0_0_0 in T.thread_binding(T.int64(512), thread="blockIdx.y",
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax0_0_0 in T.thread_binding(T.int64(512), thread="blockIdx.y",
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax1_0 in T.thread_binding(T.int64(32),
thread="blockIdx.x"):
for ax0_ax1_fused_0 in range(T.int64(1)):
for ax0_ax1_fused_1 in T.thread_binding(T.int64(8),
thread="threadIdx.y"):
@@ -107,7 +107,7 @@ def test_decode_transpose():
def main(rxplaceholder: T.Tensor((T.int64(512), T.int64(4096)),
"uint32"), rxplaceholder_1: T.Tensor((T.int64(128), T.int64(4096)), "uint32"),
T_transpose: T.Tensor((T.int64(4096), T.int64(4096)), "float32")):
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
decode_shared = Ts.sblock_alloc_buffer((T.int64(4096),
T.int64(4096)), scope="shared")
- for ax0_0_0 in T.thread_binding(T.int64(64), thread="blockIdx.y",
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax0_0_0 in T.thread_binding(T.int64(64), thread="blockIdx.y",
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax1_0 in T.thread_binding(T.int64(256),
thread="blockIdx.x"):
for ax0_ax1_fused_0 in range(T.int64(1)):
for ax0_ax1_fused_1 in T.thread_binding(T.int64(8),
thread="threadIdx.y"):
@@ -162,7 +162,7 @@ def test_decode_int3_transpose():
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
# with Ts.sblock("root"):
decode_1_shared = Ts.sblock_alloc_buffer((T.int64(4096),
T.int64(4096)), "float16", scope="shared")
- for ax0_0_0 in T.thread_binding(T.int64(52), thread="blockIdx.y",
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax0_0_0 in T.thread_binding(T.int64(52), thread="blockIdx.y",
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax1_0 in T.thread_binding(T.int64(256),
thread="blockIdx.x"):
for ax0_ax1_fused_0 in range(T.int64(2)):
for ax0_ax1_fused_1 in T.thread_binding(T.int64(8),
thread="threadIdx.y"):
diff --git
a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
index 4348f8ba3d..da010baa8e 100644
---
a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
+++
b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
@@ -64,7 +64,7 @@ class LayoutTransform:
T.func_attr({"tirx.noalias": True, "global_symbol": "main"})
# body
# with Ts.sblock("root")
- for i0_i1_i2_i3_i4_fused in T.parallel(25088,
annotations={"pragma_auto_unroll_max_step":64, "pragma_unroll_explicit":1}):
+ for i0_i1_i2_i3_i4_fused in T.parallel(25088,
annotations={"auto_unroll_max_step":64, "unroll_explicit":1}):
with Ts.sblock("T_layout_trans_1"):
ax0 = Ts.axis.spatial(1, 0)
ax1 = Ts.axis.spatial(1, 0)
@@ -784,8 +784,8 @@ def test_gpu():
_, b_j = sch.split(b_ij, factors=[None, 16]) # outer: 8
sch.bind(b_j, "threadIdx.x")
# auto unroll
- sch.annotate(i0_j0, "pragma_auto_unroll_max_step",
tirx.IntImm("int32", 1024))
- sch.annotate(i0_j0, "pragma_unroll_explicit", tirx.IntImm("int32", 1))
+ sch.annotate(i0_j0, "auto_unroll_max_step", tirx.IntImm("int32", 1024))
+ sch.annotate(i0_j0, "unroll_explicit", tirx.IntImm("int32", 1))
return sch
extractor = ms.feature_extractor.PerStoreFeature()
diff --git
a/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
b/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
index 5ed900b0be..3ecb329093 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
@@ -312,8 +312,8 @@ def test_f16f16f32_mma_gemm():
l118, l119, l120, l121, l122, l123, l124 = sch.get_loops(block=b106)
l125, l126, l127, l128, l129, l130, l131 = sch.get_loops(block=b107)
l132, l133, l134, l135, l136, l137, l138, l139, l140, l141 =
sch.get_loops(block=b108)
- sch.annotate(block_or_loop=l132, ann_key="pragma_auto_unroll_max_step",
ann_val=0)
- sch.annotate(block_or_loop=l132, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l132, ann_key="auto_unroll_max_step", ann_val=0)
+ sch.annotate(block_or_loop=l132, ann_key="unroll_explicit", ann_val=1)
l142, l143, l144 = sch.get_loops(block=b109)
b145 = sch.get_sblock(name="C_o", func_name="main")
l146, l147, l148, l149, l150, l151, l152, l153, l154, l155 =
sch.get_loops(block=b145)
diff --git
a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
index 63ca669dbf..11698b8c9a 100644
---
a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
+++
b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
@@ -232,7 +232,7 @@ class Conv2dCacheRead:
conv2d_nhwc_global = Ts.sblock_alloc_buffer([1, 56, 56, 64],
dtype="float32")
pad_temp_global = Ts.sblock_alloc_buffer([1, 58, 58, 64],
dtype="float32")
p1_global = Ts.sblock_alloc_buffer([3, 3, 64, 64], dtype="float32")
- for i0_0_i1_0_i2_0_fused in T.parallel(4,
annotations={"pragma_auto_unroll_max_step":16, "pragma_unroll_explicit":1}):
+ for i0_0_i1_0_i2_0_fused in T.parallel(4,
annotations={"auto_unroll_max_step":16, "unroll_explicit":1}):
for ax0, ax1, ax2 in T.grid(1, 30, 30):
for ax3_fused in T.vectorized(64):
with Ts.sblock("pad_temp"):
@@ -317,7 +317,7 @@ class Conv2dCacheReadRewritten:
Ts.writes(p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1,
v2 % 32, v3 % 2])
Ts.sblock_attr({"meta_schedule.layout_rewrite_preproc":True})
p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32,
v3 % 2] = p1[v0, v1, v2, v3]
- for i0_0_i1_0_i2_0_fused in T.parallel(4,
annotations={"pragma_auto_unroll_max_step":16, "pragma_unroll_explicit":1}):
+ for i0_0_i1_0_i2_0_fused in T.parallel(4,
annotations={"auto_unroll_max_step":16, "unroll_explicit":1}):
for ax0, ax1, ax2 in T.grid(1, 30, 30):
for ax3_fused in T.vectorized(64):
with Ts.sblock("pad_temp"):
@@ -409,7 +409,7 @@ class Conv2dCacheReadMultipleRewritten:
Ts.reads(p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1,
v2 % 32, v3 % 2])
Ts.writes(p1_global2[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1,
v2 % 32, v3 % 2])
p1_global2[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3
% 2] = p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2]
- for i0_0_i1_0_i2_0_fused in T.parallel(4,
annotations={"pragma_auto_unroll_max_step":16, "pragma_unroll_explicit":1}):
+ for i0_0_i1_0_i2_0_fused in T.parallel(4,
annotations={"auto_unroll_max_step":16, "unroll_explicit":1}):
for ax0, ax1, ax2 in T.grid(1, 30, 30):
for ax3_fused in T.vectorized(64):
with Ts.sblock("pad_temp"):
diff --git
a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
index d5b2f0a0b8..fbca89c11e 100644
---
a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
+++
b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
@@ -250,7 +250,7 @@ def test_no_unroll_for_spatial_block():
with Ts.sblock("root"):
A_red_temp_v0 = Ts.sblock_alloc_buffer((1,))
A_red_temp_v1 = Ts.sblock_alloc_buffer((1,))
- for ax0 in T.serial(1, annotations={"pragma_auto_unroll_max_step":
512, "pragma_unroll_explicit": 1}):
+ for ax0 in T.serial(1, annotations={"auto_unroll_max_step": 512,
"unroll_explicit": 1}):
for k1, k2, k3 in T.grid(4, 4, 32):
with Ts.sblock("A_red_temp"):
v_ax0 = Ts.axis.spatial(1, 0)
diff --git
a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
index 5d14da347d..437c084e3b 100644
---
a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
+++
b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
@@ -48,7 +48,7 @@ class Conv2dNCHWcVTCM:
T.func_attr({"tirx.noalias": True, "global_symbol": "main"})
p0_global_vtcm = Ts.sblock_alloc_buffer([T.int64(1), T.int64(2),
T.int64(56), T.int64(56), T.int64(32)], dtype="uint8", scope="global.vtcm")
p1_global_vtcm = Ts.sblock_alloc_buffer([T.int64(2), T.int64(2),
T.int64(3), T.int64(3), T.int64(8), T.int64(32), T.int64(4)], dtype="uint8",
scope="global.vtcm")
- for n_0 in T.serial(T.int64(1),
annotations={"pragma_auto_unroll_max_step":16, "pragma_unroll_explicit":1}):
+ for n_0 in T.serial(T.int64(1),
annotations={"auto_unroll_max_step":16, "unroll_explicit":1}):
for oc_chunk_0, oh_0, ow_0, oc_block_0_0 in T.grid(T.int64(2),
T.int64(2), T.int64(2), T.int64(1)):
for oc_chunk_1_init, oh_1_init, ow_1_init, oc_chunk_2_init,
oh_2_init, ow_2_init in T.grid(T.int64(1), T.int64(27), T.int64(3), T.int64(1),
T.int64(1), T.int64(9)):
with Ts.sblock("conv2d_NCHWc_int8_o_init"):
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
index 05483bcbc1..a195bd7241 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
@@ -187,7 +187,7 @@ class DenseAdd_cpu_no_write_cache:
Ts.writes(p1_global[v1 // 16, v0 // 32, v1 % 16, v0 % 32])
Ts.sblock_attr({"meta_schedule.layout_rewrite_preproc": True})
p1_global[v1 // 16, v0 // 32, v1 % 16, v0 % 32] = p1[v0, v1]
- for i0_0_i1_0_i0_1_i1_1_fused in T.parallel(16,
annotations={"pragma_auto_unroll_max_step":16, "pragma_unroll_explicit":1}):
+ for i0_0_i1_0_i0_1_i1_1_fused in T.parallel(16,
annotations={"auto_unroll_max_step":16, "unroll_explicit":1}):
for i0_2_init, i1_2_init, i0_3_init in T.grid(4, 4, 2):
for i1_3_fused_init in T.vectorized(32):
with Ts.sblock("T_matmul_NT_init"):
@@ -233,7 +233,7 @@ class DenseAdd_scheduled_gpu:
for i0_0_i1_0_fused in T.thread_binding(
32,
thread="blockIdx.x",
- annotations={"pragma_auto_unroll_max_step": 64,
"pragma_unroll_explicit": 1},
+ annotations={"auto_unroll_max_step": 64, "unroll_explicit": 1},
):
for i0_1_i1_1_fused in T.thread_binding(1, thread="vthread.x"):
for i0_2_i1_2_fused in T.thread_binding(128,
thread="threadIdx.x"):
@@ -659,7 +659,7 @@ class Conv2dInt8_tensorcore_scheduled:
p1_reindex_shared = Ts.sblock_alloc_buffer((1, 1, 256, 64), "int8",
scope="shared")
pad_temp_reindex_shared_wmma_matrix_a = Ts.sblock_alloc_buffer((50176,
64), "int8", scope="wmma.matrix_a")
p1_reindex_shared_wmma_matrix_b = Ts.sblock_alloc_buffer((1, 1, 256,
64), "int8", scope="wmma.matrix_b")
- for ax2_0_0_ax3_0_0_fused in T.thread_binding(3136,
thread="blockIdx.x", annotations={"pragma_auto_unroll_max_step": 512,
"pragma_unroll_explicit": 1}):
+ for ax2_0_0_ax3_0_0_fused in T.thread_binding(3136,
thread="blockIdx.x", annotations={"auto_unroll_max_step": 512,
"unroll_explicit": 1}):
for ax2_0_1_ax3_0_1_fused in T.thread_binding(1,
thread="vthread.x"):
for ax2_0_2_ax3_0_2_fused in T.thread_binding(16,
thread="threadIdx.x"):
for ax0_0_init, ax1_0_init, ax0_1_init, ax1_1_init,
ax2_0_3_init, ax3_0_3_init, ax0_2_init, ax1_2_init, ax2_0_4_init, ax3_0_4_init
in T.grid(1, 1, 1, 1, 1, 1, 1, 1, 1, 1):
@@ -1136,7 +1136,7 @@ def get_conv2d_vnni_mod(intrin_id):
# body
# with Ts.sblock("root")
conv2d_NCHWc_int8 = Ts.sblock_alloc_buffer([1, 128, 7, 7, 16],
dtype="int32")
- for i0_0_i1_0_i2_0_i3_0_i4_0_0_i0_1_i1_1_fused in T.parallel(128,
annotations={"pragma_auto_unroll_max_step":64, "pragma_unroll_explicit":1}):
+ for i0_0_i1_0_i2_0_i3_0_i4_0_0_i0_1_i1_1_fused in T.parallel(128,
annotations={"auto_unroll_max_step":64, "unroll_explicit":1}):
for i2_1, i3_1, i4_0_1 in T.grid(7, 1, 1):
for i0_2_init, i1_2_init, i2_2_init, i3_2_init,
i4_0_2_init, i0_3_init, i1_3_init, i2_3_init, i3_3_init, i4_0_3_init in
T.grid(1, 1, 1, 1, 1, 1, 1, 1, 7, 1):
with Ts.sblock("conv2d_NCHWc_int8_o_init"):
@@ -1393,7 +1393,7 @@ class Conv2dWinogradAddResidualRelu_scheduled:
bgemm_local = Ts.sblock_alloc_buffer([6, 6, 196, 64], dtype="float32",
scope="local")
data_pack_shared = Ts.sblock_alloc_buffer([6, 6, 196, 64],
dtype="float32", scope="shared")
p1_shared = Ts.sblock_alloc_buffer([6, 6, 64, 64], dtype="float32",
scope="shared")
- for i2_0_i3_0_i2_1_i3_1_fused_0 in T.thread_binding(98,
thread="blockIdx.x", annotations={"pragma_auto_unroll_max_step":1024,
"pragma_unroll_explicit":1}):
+ for i2_0_i3_0_i2_1_i3_1_fused_0 in T.thread_binding(98,
thread="blockIdx.x", annotations={"auto_unroll_max_step":1024,
"unroll_explicit":1}):
for i2_0_i3_0_i2_1_i3_1_fused_1 in T.thread_binding(128,
thread="threadIdx.x"):
for ax0, ax1, ax2, ax3 in T.grid(6, 6, 1, 1):
with Ts.sblock("input_tile"):
@@ -1425,7 +1425,7 @@ class Conv2dWinogradAddResidualRelu_scheduled:
Ts.writes(data_pack[eps, nu, p, ci])
Ts.sblock_attr({"auto_scheduler_simplify_const_tensor_indices":["eps", "nu",
"r_a", "r_b"], "schedule_rule":"meta_schedule.winograd_data_pack.cuda"})
data_pack[eps, nu, p, ci] = data_pack[eps,
nu, p, ci] + input_tile_local[r_a, r_b, p, ci] * T.Select(r_a % 6 == 5 and eps
% 6 == 5, T.float32(1), T.Select(r_a % 6 == 5 and eps % 6 == 4, T.float32(0),
T.Select(r_a % 6 == 5 and eps % 6 == 3, T.float32(0), T.Select(r_a % 6 == 5 and
eps % 6 == 2, T.float32(0), T.Select(r_a % 6 == 5 and eps % 6 == 1,
T.float32(0), T.Select(r_a % 6 == 5 and eps % 6 == 0, T.float32(0),
T.Select(r_a % 6 == 4 and eps % 6 == 5, [...]
- for i0_0_i1_0_i2_0_i3_0_fused in T.thread_binding(168,
thread="blockIdx.x", annotations={"pragma_auto_unroll_max_step":1024,
"pragma_unroll_explicit":1}):
+ for i0_0_i1_0_i2_0_i3_0_fused in T.thread_binding(168,
thread="blockIdx.x", annotations={"auto_unroll_max_step":1024,
"unroll_explicit":1}):
for i0_1_i1_1_i2_1_i3_1_fused in T.thread_binding(4,
thread="vthread.x"):
for i0_2_i1_2_i2_2_i3_2_fused in T.thread_binding(48,
thread="threadIdx.x"):
for i0_3_init, i1_3_init, i2_3_init, i3_3_init, i0_4_init,
i1_4_init, i2_4_init, i3_4_init in T.grid(1, 1, 14, 1, 1, 1, 1, 1):
@@ -1481,7 +1481,7 @@ class Conv2dWinogradAddResidualRelu_scheduled:
Ts.reads(bgemm_local[v0, v1, v2, v3])
Ts.writes(bgemm[v0, v1, v2, v3])
bgemm[v0, v1, v2, v3] = bgemm_local[v0, v1, v2, v3]
- for i2_0_i3_0_i2_1_i3_1_fused_0 in T.thread_binding(25,
thread="blockIdx.x", annotations={"pragma_auto_unroll_max_step":1024,
"pragma_unroll_explicit":1}):
+ for i2_0_i3_0_i2_1_i3_1_fused_0 in T.thread_binding(25,
thread="blockIdx.x", annotations={"auto_unroll_max_step":1024,
"unroll_explicit":1}):
for i2_0_i3_0_i2_1_i3_1_fused_1 in T.thread_binding(512,
thread="threadIdx.x"):
for i0 in T.unroll(4):
for i1 in T.unroll(4):
@@ -1506,7 +1506,7 @@ class Conv2dWinogradAddResidualRelu_scheduled:
Ts.writes(inverse[vh, vw, p, co])
Ts.sblock_attr({"auto_scheduler_simplify_const_tensor_indices":["vh", "vw",
"r_a", "r_b"], "schedule_rule":"meta_schedule.winograd_inverse.cuda"})
inverse[vh, vw, p, co] = inverse[vh, vw,
p, co] + bgemm[r_a, r_b, p, co] * T.Select(r_a % 6 == 5 and vh % 4 == 3,
T.float32(1), T.Select(r_a % 6 == 5 and vh % 4 == 2, T.float32(0), T.Select(r_a
% 6 == 5 and vh % 4 == 1, T.float32(0), T.Select(r_a % 6 == 5 and vh % 4 == 0,
T.float32(0), T.Select(r_a % 6 == 4 and vh % 4 == 3, T.float32(-8),
T.Select(r_a % 6 == 4 and vh % 4 == 2, T.float32(4), T.Select(r_a % 6 == 4 and
vh % 4 == 1, T.float32(-2), T.Select [...]
- for i0_i1_i2_i3_fused_0 in T.thread_binding(1568, thread="blockIdx.x",
annotations={"pragma_auto_unroll_max_step":1024, "pragma_unroll_explicit":1}):
+ for i0_i1_i2_i3_fused_0 in T.thread_binding(1568, thread="blockIdx.x",
annotations={"auto_unroll_max_step":1024, "unroll_explicit":1}):
for i0_i1_i2_i3_fused_1 in T.thread_binding(128,
thread="threadIdx.x"):
with Ts.sblock("conv2d_winograd"):
n = Ts.axis.spatial(1, 0)
@@ -1916,8 +1916,8 @@ def test_dense_add_cpu_no_write_cache():
sch.parallel(loop=l38)
l39 = sch.fuse(l37, preserve_unit_iters=True)
sch.vectorize(loop=l39)
- sch.annotate(block_or_loop=l38, ann_key="pragma_auto_unroll_max_step",
ann_val=16)
- sch.annotate(block_or_loop=l38, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l38, ann_key="auto_unroll_max_step",
ann_val=16)
+ sch.annotate(block_or_loop=l38, ann_key="unroll_explicit", ann_val=1)
b40 = sch.get_sblock(name="T_matmul_NT", func_name="main")
l41, l42, l43, l44, l45, l46, l47 = sch.get_loops(block=b40)
b48 = sch.decompose_reduction(block=b40, loop=l42)
@@ -2032,17 +2032,17 @@ def test_dense_add_gpu():
sch.unannotate(block_or_loop=b70,
ann_key="meta_schedule.unroll_explicit")
b71, b72, b73, b74 = sch.get_child_blocks(b70)
l75, l76, l77, l78, l79, l80, l81 = sch.get_loops(block=b71)
- sch.annotate(block_or_loop=l75, ann_key="pragma_auto_unroll_max_step",
ann_val=64)
- sch.annotate(block_or_loop=l75, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l75, ann_key="auto_unroll_max_step",
ann_val=64)
+ sch.annotate(block_or_loop=l75, ann_key="unroll_explicit", ann_val=1)
l82, l83, l84, l85, l86, l87, l88 = sch.get_loops(block=b72)
- sch.annotate(block_or_loop=l82, ann_key="pragma_auto_unroll_max_step",
ann_val=64)
- sch.annotate(block_or_loop=l82, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l82, ann_key="auto_unroll_max_step",
ann_val=64)
+ sch.annotate(block_or_loop=l82, ann_key="unroll_explicit", ann_val=1)
l89, l90, l91, l92, l93, l94, l95, l96, l97, l98 =
sch.get_loops(block=b73)
- sch.annotate(block_or_loop=l89, ann_key="pragma_auto_unroll_max_step",
ann_val=64)
- sch.annotate(block_or_loop=l89, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l89, ann_key="auto_unroll_max_step",
ann_val=64)
+ sch.annotate(block_or_loop=l89, ann_key="unroll_explicit", ann_val=1)
l99, l100, l101, l102, l103 = sch.get_loops(block=b74)
- sch.annotate(block_or_loop=l99, ann_key="pragma_auto_unroll_max_step",
ann_val=64)
- sch.annotate(block_or_loop=l99, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l99, ann_key="auto_unroll_max_step",
ann_val=64)
+ sch.annotate(block_or_loop=l99, ann_key="unroll_explicit", ann_val=1)
b104 = sch.get_sblock(name="T_matmul_NT", func_name="main")
l105, l106, l107, l108, l109, l110, l111, l112, l113, l114 =
sch.get_loops(block=b104)
b115 = sch.decompose_reduction(block=b104, loop=l108)
@@ -2396,14 +2396,14 @@ def test_conv2d_int8_tensorcore():
sch.unannotate(block_or_loop=b220,
ann_key="meta_schedule.unroll_explicit")
b221, b222, b223, b224, b225, b226, b227 = sch.get_child_blocks(b220)
l228, l229, l230, l231, l232, l233, l234, l235, l236 =
sch.get_loops(block=b221)
- sch.annotate(block_or_loop=l228,
ann_key="pragma_auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l228, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l228, ann_key="auto_unroll_max_step",
ann_val=512)
+ sch.annotate(block_or_loop=l228, ann_key="unroll_explicit", ann_val=1)
l237, l238, l239, l240, l241, l242, l243, l244, l245 =
sch.get_loops(block=b222)
- sch.annotate(block_or_loop=l237,
ann_key="pragma_auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l237, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l237, ann_key="auto_unroll_max_step",
ann_val=512)
+ sch.annotate(block_or_loop=l237, ann_key="unroll_explicit", ann_val=1)
l246, l247, l248, l249, l250, l251, l252, l253, l254, l255, l256 =
sch.get_loops(block=b223)
- sch.annotate(block_or_loop=l246,
ann_key="pragma_auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l246, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l246, ann_key="auto_unroll_max_step",
ann_val=512)
+ sch.annotate(block_or_loop=l246, ann_key="unroll_explicit", ann_val=1)
(
l257,
l258,
@@ -2419,8 +2419,8 @@ def test_conv2d_int8_tensorcore():
l268,
l269,
) = sch.get_loops(block=b224)
- sch.annotate(block_or_loop=l257,
ann_key="pragma_auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l257, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l257, ann_key="auto_unroll_max_step",
ann_val=512)
+ sch.annotate(block_or_loop=l257, ann_key="unroll_explicit", ann_val=1)
(
l270,
l271,
@@ -2439,14 +2439,14 @@ def test_conv2d_int8_tensorcore():
l284,
l285,
) = sch.get_loops(block=b225)
- sch.annotate(block_or_loop=l270,
ann_key="pragma_auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l270, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l270, ann_key="auto_unroll_max_step",
ann_val=512)
+ sch.annotate(block_or_loop=l270, ann_key="unroll_explicit", ann_val=1)
l286, l287, l288, l289, l290 = sch.get_loops(block=b226)
- sch.annotate(block_or_loop=l286,
ann_key="pragma_auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l286, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l286, ann_key="auto_unroll_max_step",
ann_val=512)
+ sch.annotate(block_or_loop=l286, ann_key="unroll_explicit", ann_val=1)
l291, l292, l293, l294, l295 = sch.get_loops(block=b227)
- sch.annotate(block_or_loop=l291,
ann_key="pragma_auto_unroll_max_step", ann_val=512)
- sch.annotate(block_or_loop=l291, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l291, ann_key="auto_unroll_max_step",
ann_val=512)
+ sch.annotate(block_or_loop=l291, ann_key="unroll_explicit", ann_val=1)
b296 = sch.get_sblock(name="conv2d_nhwc_o", func_name="main")
(
l297,
@@ -2675,13 +2675,13 @@ def test_conv2d_int8_vnni():
) = sch.get_loops(block=b122)
l154 = sch.fuse(l124, l125, l126, l127, l128, l129, l130,
preserve_unit_iters=True)
sch.parallel(loop=l154)
- sch.annotate(block_or_loop=l154,
ann_key="pragma_auto_unroll_max_step", ann_val=64)
- sch.annotate(block_or_loop=l154, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l154, ann_key="auto_unroll_max_step",
ann_val=64)
+ sch.annotate(block_or_loop=l154, ann_key="unroll_explicit", ann_val=1)
l155, l156, l157, l158, l159, l160, l161, l162, l163 =
sch.get_loops(block=b123)
l164 = sch.fuse(l163, preserve_unit_iters=True)
sch.vectorize(loop=l164)
- sch.annotate(block_or_loop=l155,
ann_key="pragma_auto_unroll_max_step", ann_val=64)
- sch.annotate(block_or_loop=l155, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l155, ann_key="auto_unroll_max_step",
ann_val=64)
+ sch.annotate(block_or_loop=l155, ann_key="unroll_explicit", ann_val=1)
b165 = sch.get_sblock(name="conv2d_NCHWc_int8_o", func_name="main")
(
l166,
@@ -2940,17 +2940,17 @@ def test_winograd_gpu():
sch.unannotate(block_or_loop=b165,
ann_key="meta_schedule.unroll_explicit")
b166, b167, b168, b169, b170, b171, b172, b173 =
sch.get_child_blocks(b165)
l174, l175, l176, l177, l178, l179 = sch.get_loops(block=b166)
- sch.annotate(block_or_loop=l174,
ann_key="pragma_auto_unroll_max_step", ann_val=1024)
- sch.annotate(block_or_loop=l174, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l174, ann_key="auto_unroll_max_step",
ann_val=1024)
+ sch.annotate(block_or_loop=l174, ann_key="unroll_explicit", ann_val=1)
l180, l181, l182, l183, l184, l185 = sch.get_loops(block=b167)
- sch.annotate(block_or_loop=l180,
ann_key="pragma_auto_unroll_max_step", ann_val=1024)
- sch.annotate(block_or_loop=l180, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l180, ann_key="auto_unroll_max_step",
ann_val=1024)
+ sch.annotate(block_or_loop=l180, ann_key="unroll_explicit", ann_val=1)
l186, l187, l188, l189, l190, l191, l192 = sch.get_loops(block=b168)
- sch.annotate(block_or_loop=l186,
ann_key="pragma_auto_unroll_max_step", ann_val=1024)
- sch.annotate(block_or_loop=l186, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l186, ann_key="auto_unroll_max_step",
ann_val=1024)
+ sch.annotate(block_or_loop=l186, ann_key="unroll_explicit", ann_val=1)
l193, l194, l195, l196, l197, l198, l199 = sch.get_loops(block=b169)
- sch.annotate(block_or_loop=l193,
ann_key="pragma_auto_unroll_max_step", ann_val=1024)
- sch.annotate(block_or_loop=l193, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l193, ann_key="auto_unroll_max_step",
ann_val=1024)
+ sch.annotate(block_or_loop=l193, ann_key="unroll_explicit", ann_val=1)
(
l200,
l201,
@@ -2967,17 +2967,17 @@ def test_winograd_gpu():
l212,
l213,
) = sch.get_loops(block=b170)
- sch.annotate(block_or_loop=l200,
ann_key="pragma_auto_unroll_max_step", ann_val=1024)
- sch.annotate(block_or_loop=l200, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l200, ann_key="auto_unroll_max_step",
ann_val=1024)
+ sch.annotate(block_or_loop=l200, ann_key="unroll_explicit", ann_val=1)
l214, l215, l216, l217, l218, l219, l220 = sch.get_loops(block=b171)
- sch.annotate(block_or_loop=l214,
ann_key="pragma_auto_unroll_max_step", ann_val=1024)
- sch.annotate(block_or_loop=l214, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l214, ann_key="auto_unroll_max_step",
ann_val=1024)
+ sch.annotate(block_or_loop=l214, ann_key="unroll_explicit", ann_val=1)
l221, l222, l223, l224, l225, l226 = sch.get_loops(block=b172)
- sch.annotate(block_or_loop=l221,
ann_key="pragma_auto_unroll_max_step", ann_val=1024)
- sch.annotate(block_or_loop=l221, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l221, ann_key="auto_unroll_max_step",
ann_val=1024)
+ sch.annotate(block_or_loop=l221, ann_key="unroll_explicit", ann_val=1)
l227, l228 = sch.get_loops(block=b173)
- sch.annotate(block_or_loop=l227,
ann_key="pragma_auto_unroll_max_step", ann_val=1024)
- sch.annotate(block_or_loop=l227, ann_key="pragma_unroll_explicit",
ann_val=1)
+ sch.annotate(block_or_loop=l227, ann_key="auto_unroll_max_step",
ann_val=1024)
+ sch.annotate(block_or_loop=l227, ann_key="unroll_explicit", ann_val=1)
b229 = sch.get_sblock(name="data_pack", func_name="main")
l230, l231, l232, l233, l234, l235 = sch.get_loops(block=b229)
b236 = sch.decompose_reduction(block=b229, loop=l234)
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
b/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
index 3662f6ee50..66f383eae3 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
@@ -1390,7 +1390,7 @@ def test_reverse_compute_inline_layer_norm():
A_red_temp_v0_shared = Ts.sblock_alloc_buffer((T.int64(1), n),
scope="shared")
A_red_temp_v1_shared = Ts.sblock_alloc_buffer((T.int64(1), n),
scope="shared")
var_T_layer_norm_intermediate = Ts.sblock_alloc_buffer((T.int64(1), n,
T.int64(2560)))
- for ax0_ax1_fused in T.thread_binding(n, thread="blockIdx.x",
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax0_ax1_fused in T.thread_binding(n, thread="blockIdx.x",
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax0, ax1, ax2_0 in T.grid(T.int64(1), T.int64(1), T.int64(10)):
for ax2_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
with Ts.sblock("A_red_temp"):
@@ -1431,7 +1431,7 @@ def test_reverse_compute_inline_layer_norm():
# with Ts.sblock("root"):
A_red_temp_v0_shared = Ts.sblock_alloc_buffer((T.int64(1), n),
scope="shared")
A_red_temp_v1_shared = Ts.sblock_alloc_buffer((T.int64(1), n),
scope="shared")
- for ax0_ax1_fused in T.thread_binding(n, thread="blockIdx.x",
annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
+ for ax0_ax1_fused in T.thread_binding(n, thread="blockIdx.x",
annotations={"auto_unroll_max_step": 256, "unroll_explicit": 1}):
for ax0, ax1, ax2_0 in T.grid(T.int64(1), T.int64(1), T.int64(10)):
for ax2_1 in T.thread_binding(T.int64(256),
thread="threadIdx.x"):
with Ts.sblock("A_red_temp"):
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
index 0093acb378..caa0a5fa58 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
@@ -410,7 +410,7 @@ def test_preserved_annotations():
def test_none_pragma_annotation():
@Ts.prim_func
def before(A: T.Tensor(8, "float32"), B: T.Tensor(8, "float32")):
- for i in T.serial(8, annotations={"pragma_unroll_explicit": None}):
+ for i in T.serial(8, annotations={"unroll_explicit": None}):
with Ts.sblock("block"):
B[i] = A[i] + 1.0
diff --git a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py
b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py
index f64441aa87..0aca0c3b25 100644
--- a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py
+++ b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py
@@ -370,13 +370,13 @@ def test_flatten_inside_block():
tvm.ir.assert_structural_equal(After, Expected)
-def test_build_with_optional_pragma_unroll_explicit():
+def test_build_with_optional_unroll_explicit():
def check(value):
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Tensor((4, 5, 6), "int16"), B: T.Tensor((4, 5, 6),
"int16")):
- for ax0 in T.serial(4, annotations={"pragma_unroll_explicit":
value}):
+ for ax0 in T.serial(4, annotations={"unroll_explicit": value}):
for ax1, ax2 in T.grid(5, 6):
B[ax0, ax1, ax2] = A[ax0, ax1, ax2]
diff --git a/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
b/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
index 8fd535ac6e..835198276c 100644
--- a/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
+++ b/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
@@ -51,13 +51,12 @@ def test_unroll_loop():
assert ret.kind == tvm.tirx.ForKind.UNROLLED
@I.ir_module
- class ModuleWithPragma:
+ class ModuleWithPolicy:
@T.prim_func
def main(Ab: T.Tensor((n,), "int64"), n: T.int64): # noqa: F821
- with T.attr(T.int32(0), "pragma_auto_unroll_max_step", 16):
- for i in T.serial(n, n + 2):
- for j in T.unroll(8):
- Ab[j + 1] = Ab[i] + T.int64(1)
+ for i in T.serial(n, n + 2, annotations={"auto_unroll_max_step":
16}):
+ for j in T.unroll(8):
+ Ab[j + 1] = Ab[i] + T.int64(1)
for i in T.serial(n, n + 2):
for j in T.unroll(8):
Ab[j + 1] = Ab[i] + T.int64(1)
@@ -65,7 +64,7 @@ def test_unroll_loop():
with tvm.transform.PassContext(
config={"tirx.UnrollLoop": {"auto_max_depth": 8, "explicit_unroll":
False}}
):
- ret = tvm.tirx.transform.UnrollLoop()(ModuleWithPragma)["main"].body
+ ret = tvm.tirx.transform.UnrollLoop()(ModuleWithPolicy)["main"].body
assert isinstance(ret[0], tvm.tirx.For)
assert ret[0].kind == tvm.tirx.ForKind.UNROLLED
assert isinstance(ret[1], tvm.tirx.For)
diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py
b/tests/python/tirx/transform/test_transform_lower_tirx.py
index 0c9bd7f840..5199840c01 100644
--- a/tests/python/tirx/transform/test_transform_lower_tirx.py
+++ b/tests/python/tirx/transform/test_transform_lower_tirx.py
@@ -72,11 +72,11 @@ def test_lower_tirx_opaque_optional_pragma_annotations():
def before(A: T.Tensor(8, "float32"), B: T.Tensor(8, "float32")):
for i in T.serial(8, annotations={"pragma_unroll": None}):
B[i] = A[i] + 1.0
- for i in T.serial(8, annotations={"pragma_unroll_explicit": None}):
+ for i in T.serial(8, annotations={"unroll_explicit": None}):
B[i] = A[i] + 2.0
for i in T.serial(8, annotations={"pragma_unroll": False}):
B[i] = A[i] + 3.0
- for i in T.serial(8, annotations={"pragma_unroll_explicit": 0}):
+ for i in T.serial(8, annotations={"unroll_explicit": 0}):
B[i] = A[i] + 4.0
@T.prim_func(private=True)
@@ -87,13 +87,9 @@ def test_lower_tirx_opaque_optional_pragma_annotations():
B[i] = A[i] + 2.0
for i in T.serial(8, annotations={"pragma_unroll": False}):
B[i] = A[i] + 3.0
- for i in T.serial(8):
+ for i in T.serial(8, annotations={"unroll_explicit": 0}):
B[i] = A[i] + 4.0
- # The pragma refers to the variable bound by the loop inside its body.
- loop = after.body.seq[3]
- pragma = tvm.tirx.AttrStmt(loop.loop_var, "pragma_unroll_explicit", 0,
loop)
- after = after.with_body(tvm.tirx.SeqStmt([*after.body.seq[:3], pragma]))
lowered = tvm.tirx.transform.LowerTIRxOpaque()(tvm.IRModule({"main":
before}))
tvm.ir.assert_structural_equal(lowered["main"], after, map_free_vars=True)
