This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit c703cc213e373b43e401622a4217185baea5abf8 Author: Tianqi Chen <[email protected]> AuthorDate: Mon Oct 19 12:39:30 2026 -0400 [REFACTOR][TIRx] Use explicit loop bounds and steps (#20796) Represent For with var, begin, exclusive end, required typed step, kind, and body. Unit-step loops carry an explicit typed one in the core IR. Store Range as begin/extent for regions, retaining FromMinExtent construction. The static Range::GetExtentFromBeginEnd helper recovers widths from endpoints through shallow matching without analyzer work. Update C++ and Python construction, reflection, traversal, scripting, symbolic analysis, scheduling, and code generation while preserving supported index typing and stride-aware iteration counts. --- tirx/include/tvm/ir/expr.h | 21 ++-- tirx/include/tvm/ir/stmt.h | 48 +++---- .../tvm_tirx/backend/cuda/tile/patterns/common.py | 12 +- .../tvm_tirx/backend/cuda/tile/patterns/tma.py | 2 +- .../backend/trn/tile/instruction_generator.py | 6 +- .../backend/trn/transform/private_tensor_alloc.py | 2 +- tirx/python/tvm_tirx/ir/expr.py | 7 +- tirx/python/tvm_tirx/ir/stmt.py | 38 +++--- tirx/src/backend/cuda/codegen/codegen_cuda.cc | 11 +- tirx/src/backend/trn/codegen/codegen_trn.cc | 20 +-- tirx/src/backend/vulkan/codegen/codegen_spirv.cc | 21 +--- tirx/src/backend/webgpu/codegen/codegen_webgpu.cc | 11 +- tirx/src/codegen/c/codegen_c.cc | 11 +- tirx/src/codegen/llvm/codegen_cpu.cc | 22 ++-- tirx/src/codegen/llvm/codegen_llvm.cc | 10 +- tirx/src/ir/expr.cc | 36 ++++-- tirx/src/ir/stmt.cc | 138 ++++++++++----------- tirx/src/script/ir_builder/ir.cc | 24 ++-- tirx/src/script/printer/ir/ir.cc | 4 +- tirx/src/script/printer/ir/prim_expr.cc | 14 +-- tirx/src/sym/analyzer.cc | 2 +- tirx/src/sym/const_int_bound.cc | 2 +- tirx/src/sym/int_set.cc | 26 ++-- tirx/src/sym/iter_affine_map.cc | 14 +-- tirx/src/sym/known_comparison_analyzer.cc | 8 +- tirx/src/sym/z3_prover.cc | 22 ++-- tirx/src/tirx/analysis/check_well_formed.cc | 8 +- tirx/src/tirx/ir/data_type_rewriter.cc | 44 ++++--- tirx/src/tirx/ir/index_map.cc | 8 +- tirx/src/tirx/ir/ir_mutator_with_analyzer.cc | 35 +++--- tirx/src/tirx/ir/ir_visitor_with_analyzer.cc | 14 +-- tirx/src/tirx/ir/stmt_functor.cc | 24 ++-- tirx/src/tirx/ir/tensor.cc | 14 ++- tirx/src/tirx/script/printer/for_loop.cc | 27 ++-- tirx/src/tirx/tile/lower.cc | 11 +- tirx/src/tirx/transform/common_subexpr_elim.cc | 5 +- tirx/src/tirx/transform/hoist_if.cc | 8 +- tirx/src/tirx/transform/ir_utils.cc | 4 +- .../tirx/transform/lower_tirx_dedup_tensormap.cc | 30 +++-- tirx/src/tirx/transform/lower_tvm_builtin.cc | 25 ++-- tirx/src/tirx/transform/make_packed_api.cc | 6 +- tirx/src/tirx/transform/narrow_datatype.cc | 4 +- tirx/src/tirx/transform/remove_no_op.cc | 11 +- .../tirx/transform/renormalize_split_pattern.cc | 7 +- tirx/src/tirx/transform/stmt_simplify.cc | 7 +- tirx/src/tirx/transform/storage_rewrite.cc | 4 +- tirx/src/tirx/transform/unroll_loop.cc | 6 +- tirx/src/tirx/transform/vectorize_loop.cc | 49 ++++---- tirx/tests/python/ir/test_nonnullable.py | 4 +- tirx/tests/python/script/test_printer_entry.py | 2 +- .../python/script/test_special_parser_protocol.py | 4 +- .../tirx/analysis/test_check_well_defined.py | 1 + tirx/tests/python/tirx/base/test_nodes.py | 4 +- tirx/tests/python/tirx/codegen/test_blackwell.py | 12 +- tirx/tests/python/tirx/script/test_basic_usage.py | 40 +++--- tirx/tests/python/tirx/script/test_ir_builder.py | 42 ++++--- tirx/tests/python/tirx/script/test_tensors.py | 12 +- tirx/tests/python/tirx/test_constructor.py | 6 +- tirx/tests/python/tirx/test_tile_namespace.py | 2 +- tirx/tests/python/tirx/tile/cuda/test_cp_async.py | 4 +- .../python/tirx/tile/cuda/test_cp_async_bulk.py | 10 +- .../tirx/tile/cuda/test_cp_async_bulk_tensor.py | 2 +- .../python/tirx/tile/cuda/test_tcgen05_ld_st.py | 2 +- tirx/tests/python/tirx/transform/test_hoist_if.py | 6 +- .../tirx/transform/test_split_host_device.py | 2 +- tirx/tests/python/tirx/transform/test_vectorize.py | 2 +- 66 files changed, 527 insertions(+), 513 deletions(-) diff --git a/tirx/include/tvm/ir/expr.h b/tirx/include/tvm/ir/expr.h index 313c1409d7..add9df21c8 100644 --- a/tirx/include/tvm/ir/expr.h +++ b/tirx/include/tvm/ir/expr.h @@ -1011,21 +1011,21 @@ class Call : public Expr { /*! \brief range over one dimension */ class RangeNode : public ffi::Object { public: - /*! \brief beginning of the node */ - PrimExpr min; - /*! \brief the extend of range */ + /*! \brief Inclusive beginning of the range. */ + PrimExpr begin; + /*! \brief Extent of the range. */ PrimExpr extent; /*! \brief the location of this range in the source */ mutable Location loc = UnknownLoc(); /*! \brief constructor */ - explicit RangeNode(ffi::UnsafeInit tag) : min(tag), extent(tag) {} - RangeNode(PrimExpr min, PrimExpr extent, Location loc = UnknownLoc()) - : min(min), extent(extent), loc(loc) {} + explicit RangeNode(ffi::UnsafeInit tag) : begin(tag), extent(tag) {} + RangeNode(PrimExpr begin, PrimExpr extent, Location loc = UnknownLoc()) + : begin(begin), extent(extent), loc(loc) {} static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<RangeNode>() - .def_ro("min", &RangeNode::min) + .def_ro("begin", &RangeNode::begin) .def_ro("extent", &RangeNode::extent) .def_ro("loc", &RangeNode::loc, refl::AttachFieldFlag::SEqHashIgnore()); } @@ -1057,6 +1057,13 @@ class Range : public ffi::ObjectRef { */ TVM_TIRX_DLL static Range FromMinExtent(PrimExpr min, PrimExpr extent, Location loc = UnknownLoc()); + /*! + * \brief Extract the interval width from its endpoints without an analyzer. + * + * For begin + extent or extent + begin, return the original extent operand. + * Otherwise form end - begin using primitive arithmetic. + */ + TVM_TIRX_DLL static PrimExpr GetExtentFromBeginEnd(const PrimExpr& begin, const PrimExpr& end); // declare range. explicit Range(ffi::UnsafeInit tag, ffi::ObjectPtr<RangeNode> n) : ffi::ObjectRef(tag) { data_ = std::move(n); diff --git a/tirx/include/tvm/ir/stmt.h b/tirx/include/tvm/ir/stmt.h index 6813bd44d8..16768fced7 100644 --- a/tirx/include/tvm/ir/stmt.h +++ b/tirx/include/tvm/ir/stmt.h @@ -434,27 +434,31 @@ enum class ForKind : int { * * \code * - * for (loop_var = min; loop_var < min + extent; loop_var += step) { + * for (var = begin; var < end; var += step) { * // body * } * \endcode */ class ForNode : public StmtNode { public: - explicit ForNode(ffi::UnsafeInit tag) : loop_var(tag), min(tag), extent(tag), body(tag) {} - - ForNode(PrimVar loop_var, PrimExpr min, PrimExpr extent, SeqStmt body) - : loop_var(std::move(loop_var)), - min(std::move(min)), - extent(std::move(extent)), + explicit ForNode(ffi::UnsafeInit tag) : var(tag), begin(tag), end(tag), step(tag), body(tag) {} + + ForNode(PrimVar var, PrimExpr begin, PrimExpr end, PrimExpr step, ForKind kind, SeqStmt body) + : var(std::move(var)), + begin(std::move(begin)), + end(std::move(end)), + step(std::move(step)), + kind(kind), body(std::move(body)) {} /*! \brief The loop variable. */ - PrimVar loop_var; - /*! \brief The minimum value of iteration. */ - PrimExpr min; - /*! \brief The extent of the iteration. */ - PrimExpr extent; + PrimVar var; + /*! \brief The inclusive beginning of iteration. */ + PrimExpr begin; + /*! \brief The exclusive end of iteration. */ + PrimExpr end; + /*! \brief The required loop step, including an explicit typed one for unit stride. */ + PrimExpr step; /*! \brief The kind of the for loop. */ ForKind kind; /*! \brief The body of the for loop. */ @@ -466,21 +470,17 @@ class ForNode : public StmtNode { * annotations unless they explicitly handle the corresponding semantics. */ ffi::Map<ffi::String, ffi::Any> annotations; - /*! - * \brief The loop step. It is one if not specified. - */ - ffi::Optional<PrimExpr> step; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<ForNode>() - .def_ro("loop_var", &ForNode::loop_var, refl::AttachFieldFlag::SEqHashDefSimple()) - .def_ro("min", &ForNode::min) - .def_ro("extent", &ForNode::extent) + .def_ro("var", &ForNode::var, refl::AttachFieldFlag::SEqHashDefSimple()) + .def_ro("begin", &ForNode::begin) + .def_ro("end", &ForNode::end) + .def_ro("step", &ForNode::step) .def_ro("kind", &ForNode::kind) .def_ro("body", &ForNode::body) - .def_ro("annotations", &ForNode::annotations) - .def_ro("step", &ForNode::step); + .def_ro("annotations", &ForNode::annotations); } /*! \brief Check it is a loop without nontrivial loop step. */ @@ -495,9 +495,9 @@ class ForNode : public StmtNode { */ class For : public Stmt { public: - TVM_TIRX_DLL For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, SeqStmt body, - ffi::Map<ffi::String, ffi::Any> annotations = {}, - ffi::Optional<PrimExpr> step = std::nullopt, Location loc = UnknownLoc()); + TVM_TIRX_DLL For(PrimVar var, PrimExpr begin, PrimExpr end, PrimExpr step, ForKind kind, + SeqStmt body, ffi::Map<ffi::String, ffi::Any> annotations = {}, + Location loc = UnknownLoc()); explicit For(ffi::ObjectPtr<ForNode> node) : Stmt(std::move(node)) {} diff --git a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py index c89948448a..9d6daba5de 100644 --- a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py +++ b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py @@ -64,9 +64,9 @@ class Context: self.analyzer.bind(var, Range(0, extent)) for var, bound in sctx.var_range_map.items(): self.analyzer.bind( - var, Range.from_min_extent(self.index_coordinates(bound.min), bound.extent) + var, Range.from_min_extent(self.index_coordinates(bound.begin), bound.extent) ) - low, extent = (self.analyzer.simplify(x) for x in (bound.min, bound.extent)) + low, extent = (self.analyzer.simplify(x) for x in (bound.begin, bound.extent)) if isinstance(low, IntImm) and isinstance(extent, IntImm): self.static_ranges[var] = range(int(low), int(low) + int(extent)) self.bindings = [] @@ -218,7 +218,7 @@ class Context: ] ) for var, extent in reversed(list(zip(self.repeat or (), self.repeat_shape or ()))): - body = For(var, 0, extent, ForKind.UNROLLED, body) + body = For(var, 0, extent, IntImm(var.ty.dtype, 1), ForKind.UNROLLED, body) return SeqStmt([*self.bindings, body]) @@ -261,8 +261,8 @@ class Region: raise ValueError(f"Repeat shapes must agree: {name}={repeats}, expected {expected}") if bounds: for r, size in zip(region.region, self.ty.shape): - ctx.require(r.min >= 0, f"{name}: cannot prove region start is in bounds") - ctx.require(r.min + r.extent <= size, f"{name}: region exceeds tensor shape") + ctx.require(r.begin >= 0, f"{name}: cannot prove region start is in bounds") + ctx.require(r.begin + r.extent <= size, f"{name}: region exceeds tensor shape") self.raw = None def indices(self, atom, repeat=None): @@ -271,7 +271,7 @@ class Region: if len(coords) != len(self.region.region): raise ValueError(f"{self.name}: incorrect canonical coordinate rank") return tuple( - self.ctx.analyzer.simplify(r.min + i) for r, i in zip(self.region.region, coords) + self.ctx.analyzer.simplify(r.begin + i) for r, i in zip(self.region.region, coords) ) def physical(self, atom, repeat=None): diff --git a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py index 1601f2cb9f..272dd1209e 100644 --- a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py +++ b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py @@ -118,7 +118,7 @@ def _tensor_map(ctx, region, rank, box, encoding): _encode_requirement( ctx, dims[-1] % 128 == 0, "TMA packed FP4 inner dimension must be a multiple of 128" ) - starts = tuple(r.min for r in region.region) + starts = tuple(r.begin for r in region.region) if interleave: ctx.require(view_dims[-1] == channels, "TMA Ci extent must equal one interleave slice") # Tensor strides already specify a rectangular descriptor. A layout is diff --git a/tirx/python/tvm_tirx/backend/trn/tile/instruction_generator.py b/tirx/python/tvm_tirx/backend/trn/tile/instruction_generator.py index 964f1b1b99..11340c9b00 100644 --- a/tirx/python/tvm_tirx/backend/trn/tile/instruction_generator.py +++ b/tirx/python/tvm_tirx/backend/trn/tile/instruction_generator.py @@ -128,7 +128,7 @@ class InstructionGenerator: bound = self.analyzer.const_int_bound(r.extent) if not self.analyzer.can_prove_equal(bound.max_value, r.extent): changed = True - region.append(Range.from_min_extent(r.min, bound.max_value)) + region.append(Range.from_min_extent(r.begin, bound.max_value)) if changed: bound_region = TensorRegion(tensor_region.source, region, ty=TensorRegionType()) self.bound_regions[tensor_region] = bound_region @@ -143,7 +143,7 @@ class InstructionGenerator: new_seps = [0] for i in range(len(seps) - 1): r = tensor_region.region[i] - st = r.min + st = r.begin ext = r.extent reversed_shard = [] for j in reversed(range(seps[i], seps[i + 1])): @@ -466,7 +466,7 @@ class InstructionGenerator: def generate_indices(self, tensor_region: TensorRegion) -> list[Expr]: axes = self.generate_axes(tensor_region) - return [axes[i] + r.min for i, r in enumerate(tensor_region.region)] + return [axes[i] + r.begin for i, r in enumerate(tensor_region.region)] def _get_inst_logical_iter_list( self, diff --git a/tirx/python/tvm_tirx/backend/trn/transform/private_tensor_alloc.py b/tirx/python/tvm_tirx/backend/trn/transform/private_tensor_alloc.py index 54e55daa4f..81ef1ac0ce 100644 --- a/tirx/python/tvm_tirx/backend/trn/transform/private_tensor_alloc.py +++ b/tirx/python/tvm_tirx/backend/trn/transform/private_tensor_alloc.py @@ -54,7 +54,7 @@ def _collect_private_allocations(stmt: Stmt, target: Target): return tvm_ffi.WalkResult.SKIP def visit_for(op: For): - var_range_map[op.loop_var] = Range.from_min_extent(op.min, op.extent) + var_range_map[op.var] = Range(op.begin, op.end) def visit_op_call(op: Call): if getattr(op.op, "name", "") not in SPECS: diff --git a/tirx/python/tvm_tirx/ir/expr.py b/tirx/python/tvm_tirx/ir/expr.py index d026dda51c..8c847f9fd9 100644 --- a/tirx/python/tvm_tirx/ir/expr.py +++ b/tirx/python/tvm_tirx/ir/expr.py @@ -906,7 +906,7 @@ class Range(Node, Scriptable): if the end argument is not None. Otherwise, it creates `[0, begin)`. """ - min: Expr + begin: Expr extent: Expr loc: Location @@ -937,6 +937,11 @@ class Range(Node, Scriptable): """ return _ffi_api.RangeFromMinExtent(min_value, extent, loc) + @staticmethod + def get_extent_from_begin_end(begin: Expr, end: Expr) -> Expr: + """Extract the width from endpoints, reusing a shallow addition operand.""" + return _ffi_api.RangeGetExtentFromBeginEnd(begin, end) + def __eq__(self, other: tvm_ffi.Object) -> bool: return tvm_ffi.structural_equal(self, other) diff --git a/tirx/python/tvm_tirx/ir/stmt.py b/tirx/python/tvm_tirx/ir/stmt.py index 98ffef95d0..4b18b32698 100644 --- a/tirx/python/tvm_tirx/ir/stmt.py +++ b/tirx/python/tvm_tirx/ir/stmt.py @@ -130,14 +130,17 @@ class For(Stmt): Parameters ---------- - loop_var : Var + var : Var The loop variable. - min : Expr + begin : Expr The beginning value. - extent : Expr - The length of the loop. + end : Expr + The exclusive end of the loop. + + step : Expr + The required loop step, including a typed one for unit stride. kind : ForKind The type of the for. @@ -148,42 +151,39 @@ class For(Stmt): annotations: Optional[Mapping[str, tvm_ffi.Object]] Additional loop annotations interpreted by the consuming dialect. - step : Expr - The loop step. Defaults to None, which represents one. - loc : Location, optional The location of the stmt in the source code. """ - loop_var: Var - min: Expr - extent: Expr + var: Var + begin: Expr + end: Expr + step: Expr kind: ForKind body: "SeqStmt" annotations: Mapping[str, tvm_ffi.Object] - step: Expr | None loc: Location def __init__( self, - loop_var: Var, - min: Expr, # pylint: disable=redefined-builtin - extent: Expr, + var: Var, + begin: Expr, + end: Expr, + step: Expr, kind: ForKind, body: Stmt | Sequence[Stmt], annotations: Mapping[str, tvm_ffi.Object] | None = None, - step: Expr | None = None, loc: Location = UNKNOWN_LOC, ) -> None: self.__init_handle_by_constructor__( _ffi_api.For, # type: ignore - loop_var, - min, - extent, + var, + begin, + end, + step, kind, body, annotations, - step, loc, ) diff --git a/tirx/src/backend/cuda/codegen/codegen_cuda.cc b/tirx/src/backend/cuda/codegen/codegen_cuda.cc index 7ada9676e1..22cf57dfff 100644 --- a/tirx/src/backend/cuda/codegen/codegen_cuda.cc +++ b/tirx/src/backend/cuda/codegen/codegen_cuda.cc @@ -392,10 +392,9 @@ void CodeGenCUDA::Dispatch_(const ForNode* op) { // requires #pragma unroll to immediately precede the loop it controls; if // those declarations are printed between the pragma and the for statement, // nvcc is free to unroll the loop despite disable_unroll. - std::string begin_str = PrintExpr(op->min); - PrimExpr end = IsZero(op->min) ? op->extent : sym::Analyzer()->Simplify(op->min + op->extent); - std::string end_str = PrintExpr(end); - std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : ""; + std::string begin_str = PrintExpr(op->begin); + std::string end_str = PrintExpr(op->end); + std::string step_str = op->HasTrivialStep() ? "" : PrintExpr(op->step); if (op->annotations.count(tvm::tirx::attr::kDisableUnroll.name)) { PrintIndent(); stream << "#pragma unroll 1\n"; @@ -414,9 +413,9 @@ void CodeGenCUDA::Dispatch_(const ForNode* op) { stream << "\n"; } PrintIndent(); - std::string vid = AllocVarID(op->loop_var.get()); + std::string vid = AllocVarID(op->var.get()); stream << "for ("; - PrintType(op->loop_var.ty(), stream); + PrintType(op->var.ty(), stream); stream << ' ' << vid << " = " << begin_str << "; " << vid << " < " << end_str << "; "; if (step_str.empty()) { stream << "++" << vid; diff --git a/tirx/src/backend/trn/codegen/codegen_trn.cc b/tirx/src/backend/trn/codegen/codegen_trn.cc index ad64919928..260d9f3c82 100644 --- a/tirx/src/backend/trn/codegen/codegen_trn.cc +++ b/tirx/src/backend/trn/codegen/codegen_trn.cc @@ -297,28 +297,28 @@ void CodeGenTrainium::Dispatch_(const RegionStmtNode* op) { void CodeGenTrainium::Dispatch_(const ForNode* op) { bool is_outermost_loop = is_outermost_loop_; is_outermost_loop_ = false; - std::string extent = PrintExpr(op->extent); + std::string end = PrintExpr(op->end); PrintIndent(); - std::string vid = AllocVarID(op->loop_var.get()); - TVM_FFI_ICHECK(IsZero(op->min)); + std::string vid = AllocVarID(op->var.get()); + TVM_FFI_ICHECK(IsZero(op->begin)); + TVM_FFI_ICHECK(op->HasTrivialStep()) << "Trainium codegen requires unit-step loops"; if (ctx_.tensorizing) { - stream << vid << " = nl.arange(" << extent << ")\n"; + stream << vid << " = nl.arange(" << end << ")\n"; if (op->annotations.count(tvm::backend::trn::attr::kNkiDim.name)) { - ctx_.loopvar2dim[op->loop_var.get()] = + ctx_.loopvar2dim[op->var.get()] = op->annotations[tvm::backend::trn::attr::kNkiDim.name].as_or_throw<ffi::String>(); } - ctx_.tensorized_loop_vars.insert(op->loop_var.get()); + ctx_.tensorized_loop_vars.insert(op->var.get()); TVM_FFI_ICHECK(ctx_.loopvar2dim.empty() || ctx_.loopvar2dim.size() == ctx_.tensorized_loop_vars.size()) << "nki_dim attribute must be specified for all tensorized loop variables or none of them"; PrintStmt(op->body); - ctx_.tensorized_loop_vars.erase(op->loop_var.get()); + ctx_.tensorized_loop_vars.erase(op->var.get()); } else { if (is_outermost_loop) { - stream << "for " << vid << " in nl.sequential_range(" << extent - << ", body_no_reorder=True):\n"; + stream << "for " << vid << " in nl.sequential_range(" << end << ", body_no_reorder=True):\n"; } else { - stream << "for " << vid << " in nl.sequential_range(" << extent << "):\n"; + stream << "for " << vid << " in nl.sequential_range(" << end << "):\n"; } int for_scope = BeginScope(); PrintStmt(op->body); diff --git a/tirx/src/backend/vulkan/codegen/codegen_spirv.cc b/tirx/src/backend/vulkan/codegen/codegen_spirv.cc index 2b2bfa8fa9..f99a4d3f4a 100644 --- a/tirx/src/backend/vulkan/codegen/codegen_spirv.cc +++ b/tirx/src/backend/vulkan/codegen/codegen_spirv.cc @@ -801,20 +801,11 @@ void CodeGenSPIRV::Dispatch_(const TensorStoreNode* op) { } void CodeGenSPIRV::Dispatch_(const ForNode* op) { - analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); - spirv::Value init_value = MakeValue(op->min); - PrimExpr end = IsZero(op->min) ? op->extent : analyzer_->Simplify(op->min + op->extent); - spirv::Value end_value = MakeValue(end); - - // loop step - spirv::Value step; - if (op->HasTrivialStep()) { - step = op->loop_var.ty().MatchesCode(DLDataTypeCode::kDLInt) - ? builder_->IntImm(init_value.stype, 1) - : builder_->UIntImm(init_value.stype, 1); - } else { - step = MakeValue(tvm::prim::cast(*op->step, end.ty())); - } + analyzer_->Bind(op->var, Range(op->begin, op->end)); + spirv::Value init_value = MakeValue(op->begin); + spirv::Value end_value = MakeValue(op->end); + + spirv::Value step = MakeValue(op->step); // Must get init label after making value(to make sure they are correct) spirv::Label init_label = builder_->CurrentLabel(); @@ -841,7 +832,7 @@ void CodeGenSPIRV::Dispatch_(const ForNode* op) { // loop body builder_->StartLabel(body_label); - var_map_[op->loop_var.get()] = spirv::Value(loop_var); + var_map_[op->var.get()] = spirv::Value(loop_var); this->Dispatch(op->body); builder_->MakeInst(spv::OpBranch, continue_label); diff --git a/tirx/src/backend/webgpu/codegen/codegen_webgpu.cc b/tirx/src/backend/webgpu/codegen/codegen_webgpu.cc index 87043b378e..c4d6f3608d 100644 --- a/tirx/src/backend/webgpu/codegen/codegen_webgpu.cc +++ b/tirx/src/backend/webgpu/codegen/codegen_webgpu.cc @@ -850,14 +850,13 @@ void CodeGenWebGPU::DispatchAllocTensor(const BindNode* op, const CallNode* buff } void CodeGenWebGPU::Dispatch_(const ForNode* op) { - std::string begin_str = PrintExpr(op->min); - PrimExpr end = IsZero(op->min) ? op->extent : sym::Analyzer()->Simplify(op->min + op->extent); - std::string end_str = PrintExpr(end); - std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : ""; - std::string vid = AllocVarID(op->loop_var.get()); + std::string begin_str = PrintExpr(op->begin); + std::string end_str = PrintExpr(op->end); + std::string step_str = op->HasTrivialStep() ? "" : PrintExpr(op->step); + std::string vid = AllocVarID(op->var.get()); PrintIndent(); stream << "for (var " << vid << " : "; - PrintType(op->loop_var.ty(), stream); + PrintType(op->var.ty(), stream); stream << " = " << begin_str << "; " << vid << " < " << end_str << "; " << vid; if (step_str.empty()) { stream << "++"; diff --git a/tirx/src/codegen/c/codegen_c.cc b/tirx/src/codegen/c/codegen_c.cc index 23a74b5568..2784390185 100644 --- a/tirx/src/codegen/c/codegen_c.cc +++ b/tirx/src/codegen/c/codegen_c.cc @@ -1414,14 +1414,13 @@ void CodeGenC::Dispatch_(const AssertStmtNode* op) { } void CodeGenC::Dispatch_(const ForNode* op) { - std::string begin_str = PrintExpr(op->min); - PrimExpr end = IsZero(op->min) ? op->extent : sym::Analyzer()->Simplify(op->min + op->extent); - std::string end_str = PrintExpr(end); - std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : ""; + std::string begin_str = PrintExpr(op->begin); + std::string end_str = PrintExpr(op->end); + std::string step_str = op->HasTrivialStep() ? "" : PrintExpr(op->step); PrintIndent(); - std::string vid = AllocVarID(op->loop_var.get()); + std::string vid = AllocVarID(op->var.get()); stream << "for ("; - PrintType(op->loop_var.ty(), stream); + PrintType(op->var.ty(), stream); stream << ' ' << vid << " = " << begin_str << "; " << vid << " < " << end_str << "; "; if (step_str.empty()) { stream << "++" << vid; diff --git a/tirx/src/codegen/llvm/codegen_cpu.cc b/tirx/src/codegen/llvm/codegen_cpu.cc index 1a43d90c97..1b0bfb579a 100644 --- a/tirx/src/codegen/llvm/codegen_cpu.cc +++ b/tirx/src/codegen/llvm/codegen_cpu.cc @@ -1161,8 +1161,8 @@ void CodeGenCPU::Dispatch_(const ForNode* op) { EmitDebugLocation(op); // A single iteration needs no nested threadpool launch. Keep an explicit // parallel_launch region's first loop distributed across its existing tasks. - if (op->kind == ForKind::kParallel && IsOne(op->extent) && op->HasTrivialStep() && - op->annotations.empty() && + if (op->kind == ForKind::kParallel && IsOne(Range::GetExtentFromBeginEnd(op->begin, op->end)) && + op->HasTrivialStep() && op->annotations.empty() && (parallel_env_.penv == nullptr || parallel_env_.in_parallel_loop)) { auto serial = ffi::make_object<ForNode>(*op); serial->kind = ForKind::kDefault; @@ -1172,26 +1172,25 @@ void CodeGenCPU::Dispatch_(const ForNode* op) { if (op->kind == ForKind::kDefault || op->kind == ForKind::kUnrolled) { CodeGenLLVM::Dispatch_(op); } else if (op->kind == ForKind::kParallel) { - TVM_FFI_ICHECK(IsZero(op->min)) + TVM_FFI_ICHECK(IsZero(op->begin)) << "Parallel launch require canonical loop with zero start index"; TVM_FFI_ICHECK(op->HasTrivialStep()) << "Parallel launch require canonical loop with trivial loop step"; if (parallel_env_.penv == nullptr) { auto copy_node = For(ffi::make_object<ForNode>(*op)); - CreateParallelLaunch(copy_node, 0, - std::string("loop_parallel_") + op->loop_var->name.c_str()); + CreateParallelLaunch(copy_node, 0, std::string("loop_parallel_") + op->var->name.c_str()); } else { // already in parallel env. TVM_FFI_ICHECK(parallel_env_.task_id.has_value()); TVM_FFI_ICHECK(parallel_env_.num_task.has_value()); TVM_FFI_ICHECK(parallel_env_.penv != nullptr); - PrimType t(op->extent.ty()->dtype); + PrimType t = op->var.ty(); PrimExpr num_task = cast(parallel_env_.num_task.value().as_or_throw<PrimExpr>(), t); PrimExpr task_id = cast(parallel_env_.task_id.value().as_or_throw<PrimExpr>(), t); TVM_FFI_ICHECK(!parallel_env_.in_parallel_loop) << "Nested parallel loop is not supported by threadpool, try fuse them instead"; parallel_env_.in_parallel_loop = true; - PrimExpr end = IsZero(op->min) ? op->extent : analyzer_->Simplify(op->min + op->extent); + PrimExpr end = op->end; bool stride_pattern = false; if (auto it = op->annotations.find(tvm::codegen::attr::kParallelStridePattern.name); it != op->annotations.end()) { @@ -1206,14 +1205,13 @@ void CodeGenCPU::Dispatch_(const ForNode* op) { } } if (stride_pattern) { - CreateSerialFor(MakeValue(task_id), MakeValue(end), MakeValue(num_task), op->loop_var, - op->body); + CreateSerialFor(MakeValue(task_id), MakeValue(end), MakeValue(num_task), op->var, op->body); } else { - PrimExpr step = (op->extent + num_task - IntImm(t, 1)) / num_task; - PrimExpr begin = min(task_id * step, op->extent); + PrimExpr step = (op->end + num_task - IntImm(t, 1)) / num_task; + PrimExpr begin = min(task_id * step, op->end); end = min((task_id + IntImm(t, 1)) * step, end); CreateSerialFor(MakeValue(begin), MakeValue(end), - llvm::ConstantInt::getSigned(GetLLVMType(end), 1), op->loop_var, op->body); + llvm::ConstantInt::getSigned(GetLLVMType(end), 1), op->var, op->body); } parallel_env_.in_parallel_loop = false; ++parallel_env_.parallel_loop_count; diff --git a/tirx/src/codegen/llvm/codegen_llvm.cc b/tirx/src/codegen/llvm/codegen_llvm.cc index a602811097..855026f35e 100644 --- a/tirx/src/codegen/llvm/codegen_llvm.cc +++ b/tirx/src/codegen/llvm/codegen_llvm.cc @@ -2123,18 +2123,16 @@ void CodeGenLLVM::Dispatch_(const TensorStoreNode* op) { void CodeGenLLVM::Dispatch_(const ForNode* op) { EmitDebugLocation(op); - analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); + analyzer_->Bind(op->var, Range(op->begin, op->end)); if (op->kind == ForKind::kUnrolled) { LOG(WARNING) << "Unroll hint get ignore at CodeGenLLVM backend, " << " consider set unroll_explicit=True"; } else { TVM_FFI_ICHECK(op->kind == ForKind::kDefault); } - PrimExpr step = op->step.value_or(IntImm(op->extent.ty(), 1)); - PrimExpr end = IsZero(op->min) ? op->extent : analyzer_->Simplify(op->min + op->extent); - llvm::Value* begin_value = MakeValue(op->min); - llvm::Value* end_value = MakeValue(end); - CreateSerialFor(begin_value, end_value, MakeValue(step), op->loop_var, op->body); + llvm::Value* begin_value = MakeValue(op->begin); + llvm::Value* end_value = MakeValue(op->end); + CreateSerialFor(begin_value, end_value, MakeValue(op->step), op->var, op->body); } void CodeGenLLVM::Dispatch_(const WhileNode* op) { diff --git a/tirx/src/ir/expr.cc b/tirx/src/ir/expr.cc index 4cc801ae52..beebedaf04 100644 --- a/tirx/src/ir/expr.cc +++ b/tirx/src/ir/expr.cc @@ -393,7 +393,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> RangeVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { const RangeNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RangeNode>(value); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->min)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->begin)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->extent)); return std::nullopt; } @@ -402,16 +402,16 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> RangeMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { const RangeNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RangeNode>(value); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_min_u, - mutator->MutateExpected(self->min)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_begin_u, + mutator->MutateExpected(self->begin)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_extent_u, mutator->MutateExpected(self->extent)); - if (mapped_min_u.UnchangedOrSameAs(self->min) && + if (mapped_begin_u.UnchangedOrSameAs(self->begin) && mapped_extent_u.UnchangedOrSameAs(self->extent)) { return ffi::Unchanged(); } ffi::ObjectPtr<RangeNode> copy = ffi::make_object<RangeNode>(*self); - if (!mapped_min_u.IsUnchanged()) copy->min = std::move(mapped_min_u).ValueUnchecked(); + if (!mapped_begin_u.IsUnchanged()) copy->begin = std::move(mapped_begin_u).ValueUnchecked(); if (!mapped_extent_u.IsUnchanged()) copy->extent = std::move(mapped_extent_u).ValueUnchecked(); return ffi::Any(std::move(copy)); } @@ -420,16 +420,16 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> RangeMaybeInplaceMutate ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { RangeNode* self = const_cast<RangeNode*>( ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RangeNode>(value)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_min_u, - mutator->MutateExpected(self->min, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_begin_u, + mutator->MutateExpected(self->begin, ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr<PrimExpr>, mapped_extent_u, mutator->MutateExpected(self->extent, ffi::InplaceMode::kAllow)); - if (mapped_min_u.UnchangedOrSameAs(self->min) && + if (mapped_begin_u.UnchangedOrSameAs(self->begin) && mapped_extent_u.UnchangedOrSameAs(self->extent)) { return ffi::Unchanged(); } - if (!mapped_min_u.IsUnchanged()) self->min = std::move(mapped_min_u).ValueUnchecked(); + if (!mapped_begin_u.IsUnchanged()) self->begin = std::move(mapped_begin_u).ValueUnchecked(); if (!mapped_extent_u.IsUnchanged()) self->extent = std::move(mapped_extent_u).ValueUnchecked(); return ffi::Unchanged(); } @@ -1140,8 +1140,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { // Range Range::Range(PrimExpr begin, PrimExpr end, Location loc) - : Range(ffi::UnsafeInit{}, ffi::make_object<RangeNode>( - begin, tvm::prim::IsZero(begin) ? end : (end - begin), loc)) {} + : Range(ffi::UnsafeInit{}, + ffi::make_object<RangeNode>(begin, GetExtentFromBeginEnd(begin, end), loc)) {} TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; @@ -1166,9 +1166,21 @@ Range Range::FromMinExtent(PrimExpr min, PrimExpr extent, Location loc) { return Range(ffi::UnsafeInit{}, ffi::make_object<RangeNode>(min, extent, loc)); } +PrimExpr Range::GetExtentFromBeginEnd(const PrimExpr& begin, const PrimExpr& end) { + if (const auto* add = end.as<prim::AddNode>()) { + prim::ExprDeepEqual equal; + if (equal(add->a, begin)) return add->b; + if (equal(add->b, begin)) return add->a; + } + if (prim::IsZero(begin)) return end; + return end - begin; +} + TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("ir.RangeFromMinExtent", Range::FromMinExtent); + refl::GlobalDef() + .def("ir.RangeFromMinExtent", Range::FromMinExtent) + .def("ir.RangeGetExtentFromBeginEnd", Range::GetExtentFromBeginEnd); } // Var diff --git a/tirx/src/ir/stmt.cc b/tirx/src/ir/stmt.cc index 1396b7eb69..be6d5717a0 100644 --- a/tirx/src/ir/stmt.cc +++ b/tirx/src/ir/stmt.cc @@ -276,11 +276,11 @@ TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> ForVisit( const ForNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ForNode>(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( - kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->loop_var); })); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->min)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->extent)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body)); + kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->var); })); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->begin)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->end)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->step)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body)); return std::nullopt; } @@ -289,27 +289,27 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ForMutate( // skips: kind and constant annotations; unlike SBlock annotations, these carry no expressions. const ForNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ForNode>(value); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimVar>, mapped_loop_var, + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimVar>, mapped_var, mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { - return mutator->MutateExpected(self->loop_var); + return mutator->MutateExpected(self->var); })); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_min, - mutator->MutateExpected(self->min)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_extent, - mutator->MutateExpected(self->extent)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_begin, + mutator->MutateExpected(self->begin)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_end, + mutator->MutateExpected(self->end)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_step, + mutator->MutateExpected(self->step)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body, mutator->MutateExpected(self->body)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<PrimExpr>>, mapped_step, - mutator->MutateExpected(self->step)); - if (mapped_loop_var.UnchangedOrSameAs(self->loop_var) && - mapped_min.UnchangedOrSameAs(self->min) && mapped_extent.UnchangedOrSameAs(self->extent) && - mapped_body.UnchangedOrSameAs(self->body) && mapped_step.UnchangedOrSameAs(self->step)) { + if (mapped_var.UnchangedOrSameAs(self->var) && mapped_begin.UnchangedOrSameAs(self->begin) && + mapped_end.UnchangedOrSameAs(self->end) && mapped_body.UnchangedOrSameAs(self->body) && + mapped_step.UnchangedOrSameAs(self->step)) { return ffi::Unchanged(); } ffi::ObjectPtr<ForNode> copy = ffi::make_object<ForNode>(*self); - copy->loop_var = std::move(mapped_loop_var).ValueOrUnchanged(std::move(copy->loop_var)); - copy->min = std::move(mapped_min).ValueOrUnchanged(std::move(copy->min)); - copy->extent = std::move(mapped_extent).ValueOrUnchanged(std::move(copy->extent)); + copy->var = std::move(mapped_var).ValueOrUnchanged(std::move(copy->var)); + copy->begin = std::move(mapped_begin).ValueOrUnchanged(std::move(copy->begin)); + copy->end = std::move(mapped_end).ValueOrUnchanged(std::move(copy->end)); copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body)); copy->step = std::move(mapped_step).ValueOrUnchanged(std::move(copy->step)); return ffi::Any(std::move(copy)); @@ -320,28 +320,27 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> ForMaybeInplaceMutate( // skips: kind and constant annotations; unlike SBlock annotations, these carry no expressions. ForNode* self = const_cast<ForNode*>( ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ForNode>(value)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimVar>, mapped_loop_var, + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimVar>, mapped_var, mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { - return mutator->MutateExpected(self->loop_var, + return mutator->MutateExpected(self->var, ffi::InplaceMode::kAllow); })); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_min, - mutator->MutateExpected(self->min, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr<PrimExpr>, mapped_extent, - mutator->MutateExpected(self->extent, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_begin, + mutator->MutateExpected(self->begin, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_end, + mutator->MutateExpected(self->end, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_step, + mutator->MutateExpected(self->step, ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body, mutator->MutateExpected(self->body, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<PrimExpr>>, mapped_step, - mutator->MutateExpected(self->step, ffi::InplaceMode::kAllow)); - if (mapped_loop_var.UnchangedOrSameAs(self->loop_var) && - mapped_min.UnchangedOrSameAs(self->min) && mapped_extent.UnchangedOrSameAs(self->extent) && - mapped_body.UnchangedOrSameAs(self->body) && mapped_step.UnchangedOrSameAs(self->step)) { + if (mapped_var.UnchangedOrSameAs(self->var) && mapped_begin.UnchangedOrSameAs(self->begin) && + mapped_end.UnchangedOrSameAs(self->end) && mapped_body.UnchangedOrSameAs(self->body) && + mapped_step.UnchangedOrSameAs(self->step)) { return ffi::Unchanged(); } - if (!mapped_loop_var.IsUnchanged()) self->loop_var = std::move(mapped_loop_var).ValueUnchecked(); - if (!mapped_min.IsUnchanged()) self->min = std::move(mapped_min).ValueUnchecked(); - if (!mapped_extent.IsUnchanged()) self->extent = std::move(mapped_extent).ValueUnchecked(); + if (!mapped_var.IsUnchanged()) self->var = std::move(mapped_var).ValueUnchecked(); + if (!mapped_begin.IsUnchanged()) self->begin = std::move(mapped_begin).ValueUnchecked(); + if (!mapped_end.IsUnchanged()) self->end = std::move(mapped_end).ValueUnchecked(); if (!mapped_body.IsUnchanged()) self->body = SeqStmt(std::move(mapped_body).ValueUnchecked()); if (!mapped_step.IsUnchanged()) self->step = std::move(mapped_step).ValueUnchecked(); return ffi::Unchanged(); @@ -787,12 +786,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // For -For::For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, SeqStmt body, - ffi::Map<ffi::String, Any> annotations, ffi::Optional<PrimExpr> step, Location loc) +For::For(PrimVar var, PrimExpr begin, PrimExpr end, PrimExpr step, ForKind kind, SeqStmt body, + ffi::Map<ffi::String, Any> annotations, Location loc) : Stmt(ffi::UnsafeInit{}) { - TVM_FFI_ICHECK(loop_var.defined()); - TVM_FFI_ICHECK(min.defined()); - TVM_FFI_ICHECK(extent.defined()); + TVM_FFI_ICHECK(var.defined()); + TVM_FFI_ICHECK(begin.defined()); + TVM_FFI_ICHECK(end.defined()); + TVM_FFI_ICHECK(step.defined()); TVM_FFI_ICHECK(body.defined()); auto require_scalar_int_dtype = [&](PrimExpr expr, const char* field_name) { @@ -802,47 +802,42 @@ For::For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, SeqStmt << "TIR For nodes require a scalar integer as the " << field_name << ", but received " << expr << " with dtype " << dtype; }; - require_scalar_int_dtype(loop_var, "loop_var"); - require_scalar_int_dtype(min, "min"); - require_scalar_int_dtype(extent, "extent"); + require_scalar_int_dtype(var, "var"); + require_scalar_int_dtype(begin, "begin"); + require_scalar_int_dtype(end, "end"); + require_scalar_int_dtype(step, "step"); - // When extent, min or step is an IntImm whose dtype differs from loop_var's + // When end, begin or step is an IntImm whose dtype differs from var's // (narrower bits and/or a different signedness code), we directly promote it // to the loop var's dtype as long as the value stays representable. auto try_promote_imm_dtype = [&](const PrimExpr& e) -> PrimExpr { PrimType e_ty = e.ty(); - PrimType loop_var_ty = loop_var.ty(); - if (e_ty == loop_var_ty) return e; + PrimType var_ty = var.ty(); + if (e_ty == var_ty) return e; if (const IntImmNode* a = e.as<IntImmNode>()) { - TVM_FFI_ICHECK(IntImmValueFits(a->value, loop_var_ty)) + TVM_FFI_ICHECK(IntImmValueFits(a->value, var_ty)) << "Literal value " << a->value << " is not representable in the loop variable's dtype (" - << loop_var_ty << ")"; - return IntImm(loop_var_ty, a->value); + << var_ty << ")"; + return IntImm(var_ty, a->value); } - TVM_FFI_ICHECK(e_ty.bits() <= loop_var_ty.bits()) - << " Loop variable's dtype (" << loop_var_ty - << ") is narrower than that of `min` or `extent` (" << e_ty << ")"; + TVM_FFI_ICHECK(e_ty.bits() <= var_ty.bits()) + << " Loop variable's dtype (" << var_ty << ") is narrower than that of `begin` or `end` (" + << e_ty << ")"; return e; }; - min = try_promote_imm_dtype(min); - extent = try_promote_imm_dtype(extent); + begin = try_promote_imm_dtype(begin); + end = try_promote_imm_dtype(end); - TVM_FFI_ICHECK(loop_var.ty() == min.ty()) << loop_var.ty() << " vs " << min.ty(); - TVM_FFI_ICHECK(loop_var.ty() == extent.ty()) << loop_var.ty() << " vs " << extent.ty(); + TVM_FFI_ICHECK(var.ty() == begin.ty()) << var.ty() << " vs " << begin.ty(); + TVM_FFI_ICHECK(var.ty() == end.ty()) << var.ty() << " vs " << end.ty(); - if (step.has_value()) { - require_scalar_int_dtype(*step, "step"); - step = try_promote_imm_dtype(*step); - TVM_FFI_ICHECK(loop_var.ty() == step.value().ty()) - << loop_var.ty() << " vs " << step.value().ty(); - } + step = try_promote_imm_dtype(step); + TVM_FFI_ICHECK(var.ty() == step.ty()) << var.ty() << " vs " << step.ty(); - ffi::ObjectPtr<ForNode> node = ffi::make_object<ForNode>(std::move(loop_var), std::move(min), - std::move(extent), std::move(body)); - node->kind = kind; + ffi::ObjectPtr<ForNode> node = ffi::make_object<ForNode>( + std::move(var), std::move(begin), std::move(end), std::move(step), kind, std::move(body)); node->annotations = std::move(annotations); - node->step = std::move(step); node->loc = loc; data_ = std::move(node); } @@ -856,16 +851,15 @@ TVM_FFI_STATIC_INIT_BLOCK() { .attr(ffi::type_attr::kStructuralMaybeInplaceMutate, ffi::FStructuralMutate::FromNative<&ForMaybeInplaceMutate>()); - refl::GlobalDef().def("ir.For", - [](PrimVar loop_var, PrimExpr min, PrimExpr extent, int kind, SeqStmt body, - ffi::Optional<ffi::Map<ffi::String, Any>> annotations, - ffi::Optional<PrimExpr> step, Location loc) { - return For(loop_var, min, extent, static_cast<ForKind>(kind), body, - annotations.value_or(ffi::Map<ffi::String, Any>()), step, loc); - }); + refl::GlobalDef().def( + "ir.For", [](PrimVar var, PrimExpr begin, PrimExpr end, PrimExpr step, int kind, SeqStmt body, + ffi::Optional<ffi::Map<ffi::String, Any>> annotations, Location loc) { + return For(var, begin, end, step, static_cast<ForKind>(kind), body, + annotations.value_or(ffi::Map<ffi::String, Any>()), loc); + }); } -bool ForNode::HasTrivialStep() const { return !step.has_value() || IsOne(*step); } +bool ForNode::HasTrivialStep() const { return IsOne(step); } std::ostream& operator<<(std::ostream& out, ForKind type) { // NOLINT(*) switch (type) { diff --git a/tirx/src/script/ir_builder/ir.cc b/tirx/src/script/ir_builder/ir.cc index c29c352d58..9bb4e99059 100644 --- a/tirx/src/script/ir_builder/ir.cc +++ b/tirx/src/script/ir_builder/ir.cc @@ -211,23 +211,24 @@ PrimExpr ConvertLoopBound(const PrimExpr& e, const PrimType& var_ty) { ForFrame For(ForKind kind, PrimExpr start, PrimExpr stop, ffi::Optional<PrimExpr> step, ffi::Optional<PrimType> dtype, ffi::Optional<ffi::Map<ffi::String, Any>> annotations) { PrimType var_ty = InferLoopVarDtype(start, stop, dtype); - PrimExpr min = ConvertLoopBound(start, var_ty); - PrimExpr extent = sym::Analyzer()->Simplify(ConvertLoopBound(stop, var_ty) - min); + PrimExpr begin = ConvertLoopBound(start, var_ty); + PrimExpr end = ConvertLoopBound(stop, var_ty); if (step.has_value()) { step = ConvertLoopBound(step.value(), var_ty); } ffi::ObjectPtr<ForFrameObj> n = ffi::make_object<ForFrameObj>(); n->vars = {Var("v", var_ty)}; - n->doms = {Range::FromMinExtent(min, extent)}; + n->doms = {Range(begin, end)}; n->steps = {step}; - n->f_make_for_loop = [kind, annotations](ffi::Array<Var> vars, ffi::Array<Range> doms, - ffi::Array<ffi::Optional<PrimExpr>> steps, - tvm::SeqStmt body, Location loc) { + n->f_make_for_loop = [kind, annotations, end](ffi::Array<Var> vars, ffi::Array<Range> doms, + ffi::Array<ffi::Optional<PrimExpr>> steps, + tvm::SeqStmt body, Location loc) { TVM_FFI_ICHECK_EQ(vars.size(), 1); TVM_FFI_ICHECK_EQ(doms.size(), 1); TVM_FFI_ICHECK_EQ(steps.size(), 1); - return tvm::For(vars[0].as_or_throw<tvm::PrimVar>(), doms[0]->min, doms[0]->extent, kind, body, - annotations.value_or(ffi::Map<ffi::String, Any>()), steps[0], loc); + PrimVar var = vars[0].as_or_throw<PrimVar>(); + return tvm::For(var, doms[0]->begin, end, steps[0].value_or(IntImm(var.ty(), 1)), kind, body, + annotations.value_or(ffi::Map<ffi::String, Any>()), loc); }; return ForFrame(n); } @@ -267,9 +268,10 @@ ForFrame Grid(ffi::Array<ffi::Variant<PrimExpr, ffi::Tuple<PrimExpr, PrimExpr>>> for (int i = n - 1; i >= 0; --i) { Range dom = doms[i]; Var var = vars[i]; - result = tvm::For(var.as_or_throw<tvm::PrimVar>(), dom->min, dom->extent, ForKind::kDefault, - SeqStmt(std::move(result)), - /*annotations=*/{}, /*step=*/steps[i], loc); + PrimVar prim_var = var.as_or_throw<PrimVar>(); + result = tvm::For(prim_var, dom->begin, dom->begin + dom->extent, + steps[i].value_or(IntImm(prim_var.ty(), 1)), ForKind::kDefault, + SeqStmt(std::move(result)), /*annotations=*/{}, loc); } return result; }; diff --git a/tirx/src/script/printer/ir/ir.cc b/tirx/src/script/printer/ir/ir.cc index da5ab2551b..4c0aff5c36 100644 --- a/tirx/src/script/printer/ir/ir.cc +++ b/tirx/src/script/printer/ir/ir.cc @@ -155,8 +155,8 @@ ffi::Optional<ExprDoc> RangeDocTranslate(DocTranslatorObj* d, ffi::AnyView input const ffi::Object*) { const auto* range = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RangeNode>(input); - ExprDoc min = d->Translate(range->min).value(); - ExprDoc end = d->Translate(range->min + range->extent).value(); + ExprDoc min = d->Translate(range->begin).value(); + ExprDoc end = d->Translate(range->begin + range->extent).value(); return NamespaceDoc("ir")->Attr("Range")->Call({min, end}); } diff --git a/tirx/src/script/printer/ir/prim_expr.cc b/tirx/src/script/printer/ir/prim_expr.cc index 4536993a20..7c2270987c 100644 --- a/tirx/src/script/printer/ir/prim_expr.cc +++ b/tirx/src/script/printer/ir/prim_expr.cc @@ -109,20 +109,20 @@ ffi::Optional<ExprDoc> TensorRegionDocTranslate(DocTranslatorObj* d, ffi::AnyVie if (indexable) { for (size_t i = 0; i < region->region.size(); ++i) { const Range& range = region->region[i]; - if (range->min.ty()->dtype.lanes != 1 || range->extent.ty()->dtype.lanes != 1) { + if (range->begin.ty()->dtype.lanes != 1 || range->extent.ty()->dtype.lanes != 1) { indexable = false; break; } - ExprDoc start = d->Translate(range->min).value(); + ExprDoc start = d->Translate(range->begin).value(); const auto* extent = range->extent.as<IntImmNode>(); bool point = - extent && extent->value == 1 && ffi::StructuralEqual()(extent->ty, range->min.ty()); + extent && extent->value == 1 && ffi::StructuralEqual()(extent->ty, range->begin.ty()); if (point && (has_slice || i + 1 < region->region.size())) { slices.push_back(start); - indices.push_back(range->min); + indices.push_back(range->begin); } else { try { - PrimExpr stop = range->min + range->extent; + PrimExpr stop = range->begin + range->extent; // Scalar constructors accept signed 64-bit Python literals. A computed // endpoint outside that range cannot reconstruct through indexing. if (const auto* imm = stop.as<IntImmNode>(); @@ -132,7 +132,7 @@ ffi::Optional<ExprDoc> TensorRegionDocTranslate(DocTranslatorObj* d, ffi::AnyVie break; } slices.push_back(SliceDoc(start, d->Translate(stop).value(), std::nullopt)); - indices.push_back(Slice(range->min, stop, std::nullopt)); + indices.push_back(Slice(range->begin, stop, std::nullopt)); has_slice = true; } catch (const ffi::Error&) { indexable = false; @@ -161,7 +161,7 @@ ffi::Optional<ExprDoc> TensorRegionDocTranslate(DocTranslatorObj* d, ffi::AnyVie NamespaceDoc("ir") ->Attr("Range") ->Attr("from_min_extent") - ->Call({d->Translate(range->min).value(), d->Translate(range->extent).value()})); + ->Call({d->Translate(range->begin).value(), d->Translate(range->extent).value()})); } return NamespaceDoc("ir") ->Attr("TensorRegion") diff --git a/tirx/src/sym/analyzer.cc b/tirx/src/sym/analyzer.cc index 4aa8144e63..3ca4bcdcc2 100644 --- a/tirx/src/sym/analyzer.cc +++ b/tirx/src/sym/analyzer.cc @@ -70,7 +70,7 @@ void AnalyzerObj::Bind(const Var& var, const PrimExpr& expr, bool allow_override void AnalyzerObj::Bind(const Var& var, const Range& range, bool allow_override) { TVM_FFI_ICHECK(range.defined()); if (prim::IsOne(range->extent)) { - this->Bind(var, range->min, allow_override); + this->Bind(var, range->begin, allow_override); } else { this->const_int_bound.Bind(var, range, allow_override); this->int_set.Bind(var, range, allow_override); diff --git a/tirx/src/sym/const_int_bound.cc b/tirx/src/sym/const_int_bound.cc index b5c9c1e55b..ca876a12ea 100644 --- a/tirx/src/sym/const_int_bound.cc +++ b/tirx/src/sym/const_int_bound.cc @@ -109,7 +109,7 @@ class ConstIntBoundAnalyzer::Impl final bool IsBound(const Var& var) const { return var_map_.find(var) != var_map_.end(); } void Bind(const Var& var, const Range& range, bool allow_override) { - Entry a = Dispatch(range->min); + Entry a = Dispatch(range->begin); Entry b = Dispatch(range->extent); Entry ret; ret.min_value = a.min_value; diff --git a/tirx/src/sym/int_set.cc b/tirx/src/sym/int_set.cc index 2d5a7e5167..50c90fd876 100644 --- a/tirx/src/sym/int_set.cc +++ b/tirx/src/sym/int_set.cc @@ -916,9 +916,9 @@ IntSet IntSet::FromMinExtent(PrimExpr min, PrimExpr extent) { IntSet IntSet::FromRange(Range r) { // must make sure it can be matched back by MatchRange. if (IsOne(r->extent)) { - return IntSet::SinglePoint(r->min); + return IntSet::SinglePoint(r->begin); } - return IntervalSet(r->min, r->extent + r->min - 1); + return IntervalSet(r->begin, r->extent + r->begin - 1); } bool IntSet::MatchRange(const Range& b) const { @@ -927,8 +927,8 @@ bool IntSet::MatchRange(const Range& b) const { if (!a_int) return false; if (!a_int->HasUpperBound() || !a_int->HasLowerBound()) return false; Analyzer ana; - return ProveEqual(ana.get(), a_int->min_value, b->min) && - ProveEqual(ana.get(), a_int->max_value, b->extent + b->min - 1); + return ProveEqual(ana.get(), a_int->min_value, b->begin) && + ProveEqual(ana.get(), a_int->max_value, b->extent + b->begin - 1); } IntSet Union(const ffi::Array<IntSet>& sets) { @@ -1046,15 +1046,15 @@ IntSet IntSet::Vector(PrimExpr x) { IntSet EvalSet(Range r, const ffi::Map<Var, IntSet>& dom_map) { Analyzer ana; - PrimType min_ty = r->min.ty(); + PrimType min_ty = r->begin.ty(); if (min_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt) && ana->CanProveEqual(r->extent, 1)) { - return EvalSet(r->min, dom_map); + return EvalSet(r->begin, dom_map); } IntervalSetEvaluator m(ana.get(), dom_map); - // Simplifying first can give tighter bounds if r->min and r->extent share variables - PrimExpr sum = r->min + r->extent - 1; - auto res = m.Eval(IntervalSet(r->min, ana->Simplify(sum))); + // Simplifying first can give tighter bounds if r->begin and r->extent share variables + PrimExpr sum = r->begin + r->extent - 1; + auto res = m.Eval(IntervalSet(r->begin, ana->Simplify(sum))); return res; } @@ -1064,8 +1064,8 @@ ffi::Array<IntSet> EvalSet(const ffi::Array<Range>& region, const ffi::Map<Var, ffi::Array<IntSet> result; result.reserve(region.size()); for (const Range& r : region) { - PrimExpr sum = r->min + (r->extent - 1); - result.push_back(m.Eval(IntervalSet(r->min, ana->Simplify(sum)))); + PrimExpr sum = r->begin + (r->extent - 1); + result.push_back(m.Eval(IntervalSet(r->begin, ana->Simplify(sum)))); } return result; } @@ -1163,7 +1163,7 @@ ffi::Optional<ffi::Array<IntSet>> EstimateRegionStrictBound(const ffi::Array<Ran // dynamic extent is not supported yet. return std::nullopt; } - affine_indices.push_back(range->min); + affine_indices.push_back(range->begin); } auto res = DetectIterMap( /*indices=*/affine_indices, /*input_iters=*/input_iters, @@ -1215,7 +1215,7 @@ ffi::Array<IntSet> EstimateRegionUpperBound(const ffi::Array<Range>& region, // try estimate each dimension independently for (const Range& range : region) { auto res = DetectIterMap( - /*indices=*/{range->min}, /*input_iters=*/input_iters, + /*indices=*/{range->begin}, /*input_iters=*/input_iters, /*predicate=*/predicate, /*check_level=*/IterMapLevel::Surjective, analyzer); if (!res->indices.empty()) { TVM_FFI_ICHECK_EQ(res->indices.size(), 1U); diff --git a/tirx/src/sym/iter_affine_map.cc b/tirx/src/sym/iter_affine_map.cc index def26c2903..813bf66f75 100644 --- a/tirx/src/sym/iter_affine_map.cc +++ b/tirx/src/sym/iter_affine_map.cc @@ -176,15 +176,15 @@ class IterMapRewriter : public tvm::ExprMutator { const PrimVar& var = kv.first; const Range& vrng = kv.second; if (simplify_trivial_iterators && IsOne(vrng->extent)) { - var_map_.insert_or_assign(var, IterSumExpr({}, vrng->min)); - } else if (IsZero(vrng->min)) { + var_map_.insert_or_assign(var, IterSumExpr({}, vrng->begin)); + } else if (IsZero(vrng->begin)) { IterMark mark(var.as_or_throw<PrimExpr>(), vrng->extent); var_map_.insert_or_assign(var, IterSplitExpr(mark)); input_marks_.push_back(mark); } else { - IterMark mark(var.as_or_throw<PrimExpr>() - vrng->min, vrng->extent); + IterMark mark(var.as_or_throw<PrimExpr>() - vrng->begin, vrng->extent); IterSumExpr sum_expr = ToIterSumExpr(IterSplitExpr(mark)); - sum_expr.CopyOnWrite()->base = vrng->min; + sum_expr.CopyOnWrite()->base = vrng->begin; var_map_.insert_or_assign(var, sum_expr); input_marks_.push_back(mark); } @@ -1487,8 +1487,8 @@ bool MatchBoundConstraints(PrimExpr pred, ffi::Map<PrimVar, Range>* input_iters, auto var = opt.value(); auto it = input_iters->find(var); if (it != input_iters->end()) { - PrimExpr iter_min = (*it).second->min; - PrimExpr iter_max = (*it).second->min + (*it).second->extent; + PrimExpr iter_min = (*it).second->begin; + PrimExpr iter_max = (*it).second->begin + (*it).second->extent; if (lower_bound.has_value()) iter_min = max(iter_min, lower_bound.value()); if (upper_bound.has_value()) iter_max = min(iter_max, upper_bound.value()); input_iters->Set(var, Range(iter_min, iter_max)); @@ -1514,7 +1514,7 @@ bool IterRangeSanityCheck(const ffi::Map<PrimVar, Range>& iter_ranges) { : ffi::WalkResult::Advance(); }; for (const auto& it : iter_ranges) { - if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(it.second->min, walkfn).has_value() || + if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(it.second->begin, walkfn).has_value() || ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(it.second->extent, walkfn).has_value()) { return false; } diff --git a/tirx/src/sym/known_comparison_analyzer.cc b/tirx/src/sym/known_comparison_analyzer.cc index 16ae5f69b8..ab9850c2dc 100644 --- a/tirx/src/sym/known_comparison_analyzer.cc +++ b/tirx/src/sym/known_comparison_analyzer.cc @@ -418,7 +418,7 @@ void KnownComparisonAnalyzer::Impl::Bind(const Var& var, const Range& range, boo auto it = prev_bindings_.find(var); if (it != prev_bindings_.end()) { prim::ExprDeepEqual expr_equal; - bool differs_from_previous = !expr_equal(range->min, (*it).second->min) || + bool differs_from_previous = !expr_equal(range->begin, (*it).second->begin) || !expr_equal(range->extent, (*it).second->extent); if (differs_from_previous) { TVM_FFI_ICHECK(allow_override) << "Binding of variable " << var << " as " << range @@ -435,10 +435,10 @@ void KnownComparisonAnalyzer::Impl::Bind(const Var& var, const Range& range, boo PrimExpr prim_var = var.as_or_throw<PrimExpr>(); if (IsConstInt(range->extent, 1)) { - AddKnown(prim_var == range->min, &knowns_); + AddKnown(prim_var == range->begin, &knowns_); } else { - AddKnown(prim_var >= range->min, &knowns_); - AddKnown(prim_var < range->min + range->extent, &knowns_); + AddKnown(prim_var >= range->begin, &knowns_); + AddKnown(prim_var < range->begin + range->extent, &knowns_); } } diff --git a/tirx/src/sym/z3_prover.cc b/tirx/src/sym/z3_prover.cc index ef3be241a0..300efe1a5a 100644 --- a/tirx/src/sym/z3_prover.cc +++ b/tirx/src/sym/z3_prover.cc @@ -222,7 +222,7 @@ class Z3Prover::Impl final : tvm::ExprFunctor<z3::expr(const Expr&)> { ApplyBindValue(entry.var.value(), entry.value.value()); break; case Scope::BindRange: - ApplyBindRange(entry.var.value(), entry.min.value(), entry.extent.value()); + ApplyBindRange(entry.var.value(), entry.begin.value(), entry.extent.value()); break; case Scope::Constraint: ApplyConstraint(entry.constraint.value(), entry.is_assume, scope_index); @@ -264,7 +264,7 @@ class Z3Prover::Impl final : tvm::ExprFunctor<z3::expr(const Expr&)> { } kind; ffi::Optional<Var> var; ffi::Optional<PrimExpr> value; - ffi::Optional<PrimExpr> min; + ffi::Optional<PrimExpr> begin; ffi::Optional<PrimExpr> extent; ffi::Optional<PrimExpr> constraint; /// Assume-mode constraints keep side-effectful expressions memoized for @@ -407,12 +407,12 @@ class Z3Prover::Impl final : tvm::ExprFunctor<z3::expr(const Expr&)> { void Bind(const Var& var, const Range& range, bool allow_override = false) { if (!IsZ3SupportedExpr(var.get())) return; scope_stack_.back().push_back( - Scope{Scope::BindRange, var, std::nullopt, range->min, range->extent}); + Scope{Scope::BindRange, var, std::nullopt, range->begin, range->extent}); if (!solver) return; // journaled; translated when the solver materializes - ApplyBindRange(var, range->min, range->extent); + ApplyBindRange(var, range->begin, range->extent); } - void ApplyBindRange(const Var& var, const PrimExpr& min, const PrimExpr& extent) { + void ApplyBindRange(const Var& var, const PrimExpr& begin, const PrimExpr& extent) { // 1. Create a placeholder for the var, and save it in the memo // if the var is overrided later, we can just update the memo, and the old placeholder will // be ignored @@ -424,11 +424,11 @@ class Z3Prover::Impl final : tvm::ExprFunctor<z3::expr(const Expr&)> { // instead of adding an unsat constraint, we just skip the range constraint to leave it a // free var // - // NOTE: min + extent builds a fresh AddNode that is not folded, so we must - // test IsConstInt on min and extent individually and add the two constants + // NOTE: begin + extent builds a fresh AddNode that is not folded, so we must + // test IsConstInt on begin and extent individually and add the two constants // in C++. Otherwise this fast path is never taken and we always emit the more expensive // symbolic constraint below. - if (auto min_imm = min.as<IntImm>(), extent_imm = extent.as<IntImm>(); + if (auto min_imm = begin.as<IntImm>(), extent_imm = extent.as<IntImm>(); min_imm.has_value() && extent_imm.has_value()) { const ffi::BigInt& min_value = (*min_imm)->value; ffi::BigInt max_value = min_value + (*extent_imm)->value; @@ -438,7 +438,7 @@ class Z3Prover::Impl final : tvm::ExprFunctor<z3::expr(const Expr&)> { } } else { PrimExpr prim_var = var.as_or_throw<PrimExpr>(); - solver->add(ConvertBool(extent <= 0 || (min <= prim_var && prim_var < min + extent))); + solver->add(ConvertBool(extent <= 0 || (begin <= prim_var && prim_var < begin + extent))); } } @@ -529,8 +529,8 @@ class Z3Prover::Impl final : tvm::ExprFunctor<z3::expr(const Expr&)> { ss << "; bind value: " << s.var.value() << " = " << s.value.value() << "\n"; break; case Scope::BindRange: - ss << "; bind range: " << s.var.value() << " in [" << s.min.value() << ", " - << s.min.value() + s.extent.value() << ")\n"; + ss << "; bind range: " << s.var.value() << " in [" << s.begin.value() << ", " + << s.begin.value() + s.extent.value() << ")\n"; break; } } diff --git a/tirx/src/tirx/analysis/check_well_formed.cc b/tirx/src/tirx/analysis/check_well_formed.cc index 7be7d44887..ef9818b33c 100644 --- a/tirx/src/tirx/analysis/check_well_formed.cc +++ b/tirx/src/tirx/analysis/check_well_formed.cc @@ -366,17 +366,17 @@ class WellFormednessVerifier : public StmtExprVisitor { } Result Visit_(const ForNode* op) override { - this->Visit(op->min); - this->Visit(op->extent); + this->Visit(op->begin); + this->Visit(op->end); + this->Visit(op->step); if (check_ == Check::kLoopControl) { - this->Visit(op->step); ++loop_depth_; this->Visit(op->body); --loop_depth_; return std::nullopt; } return WithScope([&]() -> Result { - Define(op->loop_var, AccessPath::Root()->Attr("loop_var")); + Define(op->var, AccessPath::Root()->Attr("var")); return WithScope([&]() { return this->Visit(op->body); }); }); } diff --git a/tirx/src/tirx/ir/data_type_rewriter.cc b/tirx/src/tirx/ir/data_type_rewriter.cc index a7ca22f26b..b0126cc0b4 100644 --- a/tirx/src/tirx/ir/data_type_rewriter.cc +++ b/tirx/src/tirx/ir/data_type_rewriter.cc @@ -48,24 +48,20 @@ UnchangedOr<Stmt> DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode inpl << ffi::AnyView(result).GetTypeKey(); if (!op->unique()) inplace_mode = InplaceMode::kDisallow; } - PrimExpr e = Mutate(op->loop_var, inplace_mode).ValueOrUnchanged(op->loop_var); + PrimExpr e = Mutate(op->var, inplace_mode).ValueOrUnchanged(op->var); Var var = e.as_or_throw<Var>(); PrimType var_ty = var->ty.as_or_throw<PrimType>(); if (inplace_mode == InplaceMode::kAllow) { auto* n = const_cast<ForNode*>(op); - n->min = prim::cast(op->min, var_ty); - n->extent = prim::cast(op->extent, var_ty); - if (op->step.has_value()) { - n->step = prim::cast(*op->step, var_ty); - } + n->begin = prim::cast(op->begin, var_ty); + n->end = prim::cast(op->end, var_ty); + n->step = prim::cast(op->step, var_ty); return result; } auto n = ffi::make_object<ForNode>(*op); - n->min = prim::cast(op->min, var_ty); - n->extent = prim::cast(op->extent, var_ty); - if (op->step.has_value()) { - n->step = prim::cast(*op->step, var_ty); - } + n->begin = prim::cast(op->begin, var_ty); + n->end = prim::cast(op->end, var_ty); + n->step = prim::cast(op->step, var_ty); return For(n); } @@ -427,25 +423,27 @@ UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const ForNode* op, InplaceMode bool is_enabled = is_enabled_; is_enabled_ = true; PrimVar new_loop_var = - Mutate(op->loop_var, inplace_mode).ValueOrUnchanged(op->loop_var).as_or_throw<PrimVar>(); - auto min_result = Mutate(op->min, inplace_mode); - bool min_unchanged = min_result.UnchangedOrSameAs(op->min); - PrimExpr min = std::move(min_result).ValueOrUnchanged(op->min); - auto extent_result = Mutate(op->extent, inplace_mode); - bool extent_unchanged = extent_result.UnchangedOrSameAs(op->extent); - PrimExpr extent = std::move(extent_result).ValueOrUnchanged(op->extent); + Mutate(op->var, inplace_mode).ValueOrUnchanged(op->var).as_or_throw<PrimVar>(); + auto begin_result = Mutate(op->begin, inplace_mode); + bool begin_unchanged = begin_result.UnchangedOrSameAs(op->begin); + PrimExpr begin = std::move(begin_result).ValueOrUnchanged(op->begin); + auto end_result = Mutate(op->end, inplace_mode); + bool end_unchanged = end_result.UnchangedOrSameAs(op->end); + PrimExpr end = std::move(end_result).ValueOrUnchanged(op->end); + PrimExpr step = Mutate(op->step, inplace_mode).ValueOrUnchanged(op->step); is_enabled_ = is_enabled; auto new_body_result = Mutate(op->body, inplace_mode); bool new_body_unchanged = new_body_result.UnchangedOrSameAs(op->body); Stmt new_body = std::move(new_body_result).ValueOrUnchanged(op->body); - if (!new_loop_var.same_as(op->loop_var) || !min_unchanged || !extent_unchanged || - !new_body_unchanged) { + if (!new_loop_var.same_as(op->var) || !begin_unchanged || !end_unchanged || + !step.same_as(op->step) || !new_body_unchanged) { For new_for = ffi::GetRef<For>(op); auto* n = new_for.CopyOnWrite(); - n->loop_var = new_loop_var; - n->min = prim::cast(min, new_loop_var.ty()); - n->extent = prim::cast(extent, new_loop_var.ty()); + n->var = new_loop_var; + n->begin = prim::cast(begin, new_loop_var.ty()); + n->end = prim::cast(end, new_loop_var.ty()); + n->step = prim::cast(step, new_loop_var.ty()); n->body = new_body; return new_for; diff --git a/tirx/src/tirx/ir/index_map.cc b/tirx/src/tirx/ir/index_map.cc index 892d063bb9..fcbeac6ebf 100644 --- a/tirx/src/tirx/ir/index_map.cc +++ b/tirx/src/tirx/ir/index_map.cc @@ -127,7 +127,7 @@ std::pair<IndexMap, PrimExpr> IndexMapInverseImpl(const IndexMap& self, PrimVar index = self->initial_indices[i]; PrimExpr expr{ffi::UnsafeInit{}}; if (IsOne(initial_ranges[i]->extent) && !inverse_exprs_map.count(index)) { - expr = initial_ranges[i]->min; + expr = initial_ranges[i]->begin; } else { expr = inverse_exprs_map.at(index); } @@ -272,8 +272,8 @@ ffi::Array<Range> IndexMapNode::MapRanges(const ffi::Array<Range>& ranges, return PrimType::Int(max_bits); }(); output.MutateByApply([&](const Range& range) { - if (range->min.ty() != output_dtype || range->extent.ty() != output_dtype) { - return Range::FromMinExtent(prim::cast(range->min, output_dtype), + if (range->begin.ty() != output_dtype || range->extent.ty() != output_dtype) { + return Range::FromMinExtent(prim::cast(range->begin, output_dtype), prim::cast(range->extent, output_dtype)); } else { return range; @@ -299,7 +299,7 @@ ffi::Array<PrimExpr> IndexMapNode::MapShape(const ffi::Array<PrimExpr>& shape, ffi::Array<PrimExpr> output; for (auto& range : mapped) { - TVM_FFI_ICHECK(IsZero(range->min)); + TVM_FFI_ICHECK(IsZero(range->begin)); output.push_back(range->extent); } diff --git a/tirx/src/tirx/ir/ir_mutator_with_analyzer.cc b/tirx/src/tirx/ir/ir_mutator_with_analyzer.cc index ffb8c7d708..4101094cfa 100644 --- a/tirx/src/tirx/ir/ir_mutator_with_analyzer.cc +++ b/tirx/src/tirx/ir/ir_mutator_with_analyzer.cc @@ -89,38 +89,37 @@ UnchangedOr<Stmt> StmtExprMutatorWithAnalyzer::Mutate_(const ForNode* op, InplaceMode inplace_mode) { return constraint_scope_.WithNewScope([&]() -> UnchangedOr<Stmt> { // record the loop variable as iterators - Range dom = Range::FromMinExtent(op->min, op->extent); - analyzer_->Bind(op->loop_var, dom); - iter_vars_.Set(op->loop_var, dom); - auto min_result = this->Mutate(op->min, inplace_mode); - bool min_unchanged = min_result.UnchangedOrSameAs(op->min); - PrimExpr min = std::move(min_result).ValueOrUnchanged(op->min); - auto extent_result = this->Mutate(op->extent, inplace_mode); - bool extent_unchanged = extent_result.UnchangedOrSameAs(op->extent); - PrimExpr extent = std::move(extent_result).ValueOrUnchanged(op->extent); - ffi::Optional<PrimExpr> step{std::nullopt}; - if (op->step.has_value()) { - step = this->Mutate(*op->step, inplace_mode).ValueOrUnchanged(*op->step); - } + Range dom = Range(op->begin, op->end); + analyzer_->Bind(op->var, dom); + iter_vars_.Set(op->var, dom); + auto begin_result = this->Mutate(op->begin, inplace_mode); + bool begin_unchanged = begin_result.UnchangedOrSameAs(op->begin); + PrimExpr begin = std::move(begin_result).ValueOrUnchanged(op->begin); + auto end_result = this->Mutate(op->end, inplace_mode); + bool end_unchanged = end_result.UnchangedOrSameAs(op->end); + PrimExpr end = std::move(end_result).ValueOrUnchanged(op->end); + PrimExpr step = this->Mutate(op->step, inplace_mode).ValueOrUnchanged(op->step); Stmt body = constraint_scope_.WithNewScope([&]() -> Stmt { + EnterConstraintFacts(&constraint_scope_.Current(), analyzer_, end > begin); + PrimExpr extent = Range::GetExtentFromBeginEnd(begin, end); EnterConstraintFacts(&constraint_scope_.Current(), analyzer_, extent > IntImm(extent.ty(), 0)); return this->Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body); }); - if (min_unchanged && extent_unchanged && body.same_as(op->body) && step.same_as(op->step)) { + if (begin_unchanged && end_unchanged && body.same_as(op->body) && step.same_as(op->step)) { return ffi::Unchanged(); } else { if (inplace_mode == InplaceMode::kAllow) { auto* n = const_cast<ForNode*>(op); - n->min = std::move(min); - n->extent = std::move(extent); + n->begin = std::move(begin); + n->end = std::move(end); n->step = std::move(step); n->body = std::move(body); return ffi::Unchanged(); } auto n = ffi::make_object<ForNode>(*op); - n->min = std::move(min); - n->extent = std::move(extent); + n->begin = std::move(begin); + n->end = std::move(end); n->step = std::move(step); n->body = std::move(body); return Stmt(n); diff --git a/tirx/src/tirx/ir/ir_visitor_with_analyzer.cc b/tirx/src/tirx/ir/ir_visitor_with_analyzer.cc index 4f13c2009a..456806d722 100644 --- a/tirx/src/tirx/ir/ir_visitor_with_analyzer.cc +++ b/tirx/src/tirx/ir/ir_visitor_with_analyzer.cc @@ -35,14 +35,14 @@ void StmtExprVisitorWithAnalyzer::InitVTable(VTable* vtable) { ffi::Optional<VisitInterrupt> StmtExprVisitorWithAnalyzer::Visit_(const ForNode* op) { return constraint_scope_.WithNewScope([&]() -> ffi::Optional<VisitInterrupt> { - analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->min)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->extent)); - if (op->step.has_value()) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(*op->step)); - } + analyzer_->Bind(op->var, Range(op->begin, op->end)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->begin)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->end)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->step)); return constraint_scope_.WithNewScope([&]() -> ffi::Optional<VisitInterrupt> { - constraint_scope_.Current().Emplace(analyzer_, op->extent > IntImm(op->extent.ty(), 0)); + constraint_scope_.Current().Emplace(analyzer_, op->end > op->begin); + PrimExpr extent = Range::GetExtentFromBeginEnd(op->begin, op->end); + constraint_scope_.Current().Emplace(analyzer_, extent > IntImm(extent.ty(), 0)); return this->Visit(op->body); }); }); diff --git a/tirx/src/tirx/ir/stmt_functor.cc b/tirx/src/tirx/ir/stmt_functor.cc index 5f9e6e01b8..8034176cc9 100644 --- a/tirx/src/tirx/ir/stmt_functor.cc +++ b/tirx/src/tirx/ir/stmt_functor.cc @@ -146,11 +146,9 @@ ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const RegionStmtNode* op) } ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const ForNode* op) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->min)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->extent)); - if (op->step.has_value()) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(*op->step)); - } + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->begin)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->end)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->step)); return this->Visit(op->body); } @@ -277,24 +275,24 @@ UnchangedOr<Stmt> StmtExprMutator::Mutate_(const RegionStmtNode* op, InplaceMode } UnchangedOr<Stmt> StmtExprMutator::Mutate_(const ForNode* op, InplaceMode inplace_mode) { - auto min = Mutate(op->min, inplace_mode); - auto extent = Mutate(op->extent, inplace_mode); - auto step = Mutate(op->step, inplace_mode).as_or_throw<UnchangedOr<ffi::Optional<PrimExpr>>>(); + auto begin = Mutate(op->begin, inplace_mode); + auto end = Mutate(op->end, inplace_mode); + auto step = Mutate(op->step, inplace_mode); auto body = Mutate(op->body, inplace_mode); - if (min.UnchangedOrSameAs(op->min) && extent.UnchangedOrSameAs(op->extent) && + if (begin.UnchangedOrSameAs(op->begin) && end.UnchangedOrSameAs(op->end) && step.UnchangedOrSameAs(op->step) && body.UnchangedOrSameAs(op->body)) return ffi::Unchanged(); if (inplace_mode == InplaceMode::kAllow) { auto* writable = const_cast<ForNode*>(op); - if (!min.IsUnchanged()) writable->min = std::move(min).ValueUnchecked(); - if (!extent.IsUnchanged()) writable->extent = std::move(extent).ValueUnchecked(); + if (!begin.IsUnchanged()) writable->begin = std::move(begin).ValueUnchecked(); + if (!end.IsUnchanged()) writable->end = std::move(end).ValueUnchecked(); if (!step.IsUnchanged()) writable->step = std::move(step).ValueUnchecked(); if (!body.IsUnchanged()) writable->body = std::move(body).ValueUnchecked(); return ffi::Unchanged(); } auto copy = ffi::make_object<ForNode>(*op); - if (!min.IsUnchanged()) copy->min = std::move(min).ValueUnchecked(); - if (!extent.IsUnchanged()) copy->extent = std::move(extent).ValueUnchecked(); + if (!begin.IsUnchanged()) copy->begin = std::move(begin).ValueUnchecked(); + if (!end.IsUnchanged()) copy->end = std::move(end).ValueUnchecked(); if (!step.IsUnchanged()) copy->step = std::move(step).ValueUnchecked(); if (!body.IsUnchanged()) copy->body = std::move(body).ValueUnchecked(); return Stmt(std::move(copy)); diff --git a/tirx/src/tirx/ir/tensor.cc b/tirx/src/tirx/ir/tensor.cc index 9cb8d2f734..5890b18681 100644 --- a/tirx/src/tirx/ir/tensor.cc +++ b/tirx/src/tirx/ir/tensor.cc @@ -63,10 +63,12 @@ using SubscriptSlice = ffi::Array<ffi::Variant< PrimExpr>>; PrimExpr SliceExtent(const PrimExpr& start, const PrimExpr& stop, const sym::Analyzer& analyzer) { - if (const auto* add = stop.as<AddNode>(); add && ffi::StructuralEqual()(add->a, start)) { - return add->b; + PrimExpr extent = Range::GetExtentFromBeginEnd(start, stop); + if (const auto* add = stop.as<AddNode>(); + add && (extent.same_as(add->a) || extent.same_as(add->b))) { + return extent; } - return analyzer->Simplify(stop - start); + return analyzer->Simplify(extent); } ffi::ObjectRef RealizeTensorSubscript( @@ -150,7 +152,7 @@ ffi::ObjectRef RealizeTensorRegionSubscript(Expr value, SubscriptSlice slice, Lo ffi::Array<PrimExpr> indices; indices.reserve(slice.size()); for (size_t i = 0; i < slice.size(); ++i) { - indices.push_back(source->region[i]->min + slice[i].as<PrimExpr>().value()); + indices.push_back(source->region[i]->begin + slice[i].as<PrimExpr>().value()); } return TensorLoad(source->source.as_or_throw<TensorVar>(), indices, loc); } @@ -161,7 +163,7 @@ ffi::ObjectRef RealizeTensorRegionSubscript(Expr value, SubscriptSlice slice, Lo for (size_t i = 0; i < slice.size(); ++i) { const Range& old_range = source->region[i]; if (auto point = slice[i].as<PrimExpr>()) { - PrimExpr new_min = old_range->min + point.value(); + PrimExpr new_min = old_range->begin + point.value(); region.push_back(Range::FromMinExtent(new_min, IntImm(point.value().ty(), 1))); } else { auto descriptor = slice[i] @@ -171,7 +173,7 @@ ffi::ObjectRef RealizeTensorRegionSubscript(Expr value, SubscriptSlice slice, Lo PrimExpr start = descriptor.get<0>().value_or(IntImm(old_range->extent.ty(), 0)); PrimExpr stop = descriptor.get<1>().value_or(old_range->extent); region.push_back( - Range::FromMinExtent(old_range->min + start, SliceExtent(start, stop, analyzer))); + Range::FromMinExtent(old_range->begin + start, SliceExtent(start, stop, analyzer))); } } for (size_t i = slice.size(); i < source->region.size(); ++i) { diff --git a/tirx/src/tirx/script/printer/for_loop.cc b/tirx/src/tirx/script/printer/for_loop.cc index 609ff36b6e..a601011944 100644 --- a/tirx/src/tirx/script/printer/for_loop.cc +++ b/tirx/src/tirx/script/printer/for_loop.cc @@ -59,16 +59,8 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d, ffi::AnyView input, default: TVM_FFI_THROW(TypeError) << "printer unknown loop kind"; } - ExprDoc min = d->Translate(loop->min).value(); - ExprDoc extent = d->Translate(loop->extent).value(); - ExprDoc end = OperationDoc(OperationDocNode::Kind::kAdd, {min, extent}); - if (prim::IsZero(loop->min)) { - end = extent; - } else if (loop->min.as<IntImmNode>() && loop->extent.as<IntImmNode>()) { - // Python integer addition can widen the endpoint before the builder sees - // it. Keep the IR operation and its dtype for literal bounds as well. - end = NamespaceDoc("tirx")->Attr("Add")->Call({min, extent}); - } + ExprDoc begin = d->Translate(loop->begin).value(); + ExprDoc end = d->Translate(loop->end).value(); ffi::Array<ffi::String> keys; ffi::Array<ExprDoc> values; bool unroll_option = false; @@ -103,26 +95,25 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d, ffi::AnyView input, keys.push_back("annotations"); values.push_back(DictDoc(annotation_keys, annotation_values)); } - if (loop->step.has_value()) { + if (!prim::IsOne(loop->step)) { keys.push_back("step"); - values.push_back(d->Translate(loop->step.value()).value()); + values.push_back(d->Translate(loop->step).value()); } // For requires the bounds and index to have the same type. Their typed // expressions therefore carry the index dtype without a separate keyword. ForDoc doc(ffi::UnsafeInit{}); { - IdDoc var = VarDoc(d, loop->loop_var); - ffi::Array<ExprDoc> bounds = {min, end}; + IdDoc var = VarDoc(d, loop->var); + ffi::Array<ExprDoc> bounds = {begin, end}; ExprDoc callee = NamespaceDoc("tirx")->Attr(method); if (loop->kind == ForKind::kDefault && loop->annotations.empty()) { callee = IdDoc("range"); - // range's step is positional. Retain even an explicit unit step because - // it is part of the source For node, unlike an absent step. - if (loop->step.has_value()) { + // range's step is positional. The builder restores an omitted unit step. + if (!prim::IsOne(loop->step)) { bounds.push_back(values.back()); keys.pop_back(); values.pop_back(); - } else if (prim::IsZero(loop->min) && loop->min.ty() == PrimType::Int(32)) { + } else if (prim::IsZero(loop->begin) && loop->begin.ty() == PrimType::Int(32)) { bounds.erase(bounds.begin()); } } diff --git a/tirx/src/tirx/tile/lower.cc b/tirx/src/tirx/tile/lower.cc index 3a44490607..0639525105 100644 --- a/tirx/src/tirx/tile/lower.cc +++ b/tirx/src/tirx/tile/lower.cc @@ -107,9 +107,10 @@ class NoOpCallVerifier : public StmtExprVisitor { } Result Visit_(const ForNode* op) override { - this->Visit(op->min); - this->Visit(op->extent); - this->Visit(op->loop_var->ty); + this->Visit(op->begin); + this->Visit(op->end); + this->Visit(op->step); + this->Visit(op->var->ty); return this->Visit(op->body); } @@ -362,9 +363,9 @@ class TileDispatcher : public StmtExprMutator { UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final { // Collect the loop variables - auto loop_var = op->loop_var.as_or_throw<Var>(); + auto loop_var = op->var.as_or_throw<Var>(); TVM_FFI_ICHECK(!var_range_map_.count(loop_var)) << "Internal Error: Duplicate loop variable"; - var_range_map_.Set(loop_var, Range::FromMinExtent(op->min, op->extent)); + var_range_map_.Set(loop_var, Range(op->begin, op->end)); return StmtExprMutator::Mutate_(op, inplace_mode); } diff --git a/tirx/src/tirx/transform/common_subexpr_elim.cc b/tirx/src/tirx/transform/common_subexpr_elim.cc index 0f8969a875..dd7b0133fa 100644 --- a/tirx/src/tirx/transform/common_subexpr_elim.cc +++ b/tirx/src/tirx/transform/common_subexpr_elim.cc @@ -570,8 +570,9 @@ class CSEPlanner : public StmtExprVisitor { /*! \brief For loops: bounds in parent scope, body in child scope. */ ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) override { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->min)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->extent)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->begin)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->end)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->step)); int saved = current_scope_; current_scope_ = AllocScope(saved, ffi::GetRef<Stmt>(op)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->body)); diff --git a/tirx/src/tirx/transform/hoist_if.cc b/tirx/src/tirx/transform/hoist_if.cc index 093d064ee7..7a8f6639c5 100644 --- a/tirx/src/tirx/transform/hoist_if.cc +++ b/tirx/src/tirx/transform/hoist_if.cc @@ -155,9 +155,9 @@ class IfHoister : public StmtExprMutator { // Preserve the same synchronization boundary before thread-binding lowering. // Only ordinary serial loops permit code motion across their execution scope. // Hoisting must not suppress evaluation of an effectful loop header either. - if (op->kind != ForKind::kDefault || SideEffect(op->min) > CallEffectKind::kReadState || - SideEffect(op->extent) > CallEffectKind::kReadState || - (op->step && SideEffect(op->step.value()) > CallEffectKind::kReadState)) { + if (op->kind != ForKind::kDefault || SideEffect(op->begin) > CallEffectKind::kReadState || + SideEffect(op->end) > CallEffectKind::kReadState || + SideEffect(op->step) > CallEffectKind::kReadState) { return scopes_.WithNewScope([&] { return StmtExprMutator::Mutate_(op, inplace_mode); }); } auto& scope = scopes_.Current(); @@ -166,7 +166,7 @@ class IfHoister : public StmtExprMutator { if (active_loops_++ == 0) split_in_nest_ = false; // Snapshot ancestors outside this loop so they do not require a false-path copy. LoopState loop{else_depth_, {}}; - const VarNode* var = op->loop_var.get(); + const VarNode* var = op->var.get(); // Register this destination and its variable together. An index in loops // identifies both where to wrap a lifted If and which variables it must avoid. scope.loop_depths[var] = scope.loops.size(); diff --git a/tirx/src/tirx/transform/ir_utils.cc b/tirx/src/tirx/transform/ir_utils.cc index 6aa37aa957..689d5f8f81 100644 --- a/tirx/src/tirx/transform/ir_utils.cc +++ b/tirx/src/tirx/transform/ir_utils.cc @@ -189,7 +189,7 @@ UnchangedOr<Stmt> IRConvertSSA::Mutate_(const IfNode* op, InplaceMode inplace_mo } UnchangedOr<Stmt> IRConvertSSA::Mutate_(const ForNode* op, InplaceMode inplace_mode) { - const Var& v = op->loop_var; + const Var& v = op->var; if (defined_.count(v.get())) { return scope_.WithNewScope([&]() -> Stmt { Var new_var = MakeNewVar(v); @@ -197,7 +197,7 @@ UnchangedOr<Stmt> IRConvertSSA::Mutate_(const ForNode* op, InplaceMode inplace_m Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op)); auto n = ffi::make_object<ForNode>(*stmt.as<ForNode>()); - n->loop_var = new_var.as_or_throw<PrimVar>(); + n->var = new_var.as_or_throw<PrimVar>(); return For(n); }); } else { diff --git a/tirx/src/tirx/transform/lower_tirx_dedup_tensormap.cc b/tirx/src/tirx/transform/lower_tirx_dedup_tensormap.cc index ea237f6cce..342e5d3cf2 100644 --- a/tirx/src/tirx/transform/lower_tirx_dedup_tensormap.cc +++ b/tirx/src/tirx/transform/lower_tirx_dedup_tensormap.cc @@ -89,8 +89,9 @@ class CuTensorMapDedupAnalyzer : public StmtExprVisitor { CuTensorMapDedupAnalyzer() { canonical_list_.emplace_back(std::vector<std::pair<Call, Var>>()); } ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->min)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->extent)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->begin)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->end)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->step)); canonical_list_.emplace_back(std::vector<std::pair<Call, Var>>()); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->body)); canonical_list_.pop_back(); @@ -166,30 +167,33 @@ class CuTensorMapDedupRewriter : public StmtExprMutator { private: UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final { - auto min_result = Mutate(op->min, inplace_mode); - bool min_unchanged = min_result.UnchangedOrSameAs(op->min); - PrimExpr min = std::move(min_result).ValueOrUnchanged(op->min); - auto extent_result = Mutate(op->extent, inplace_mode); - bool extent_unchanged = extent_result.UnchangedOrSameAs(op->extent); - PrimExpr extent = std::move(extent_result).ValueOrUnchanged(op->extent); + auto begin_result = Mutate(op->begin, inplace_mode); + bool begin_unchanged = begin_result.UnchangedOrSameAs(op->begin); + PrimExpr begin = std::move(begin_result).ValueOrUnchanged(op->begin); + auto end_result = Mutate(op->end, inplace_mode); + bool end_unchanged = end_result.UnchangedOrSameAs(op->end); + PrimExpr end = std::move(end_result).ValueOrUnchanged(op->end); + PrimExpr step = Mutate(op->step, inplace_mode).ValueOrUnchanged(op->step); emitted_keys_.emplace_back(); auto body_result = Mutate(op->body, inplace_mode); emitted_keys_.pop_back(); bool body_unchanged = body_result.UnchangedOrSameAs(op->body); Stmt body = std::move(body_result).ValueOrUnchanged(op->body); - if (min_unchanged && extent_unchanged && body_unchanged) { + if (begin_unchanged && end_unchanged && step.same_as(op->step) && body_unchanged) { return ffi::Unchanged(); } else { if (inplace_mode == InplaceMode::kAllow) { auto* n = const_cast<ForNode*>(op); - n->min = std::move(min); - n->extent = std::move(extent); + n->begin = std::move(begin); + n->end = std::move(end); + n->step = step; n->body = std::move(body); return ffi::Unchanged(); } auto n = ffi::make_object<ForNode>(*op); - n->min = std::move(min); - n->extent = std::move(extent); + n->begin = std::move(begin); + n->end = std::move(end); + n->step = step; n->body = std::move(body); return Stmt(n); } diff --git a/tirx/src/tirx/transform/lower_tvm_builtin.cc b/tirx/src/tirx/transform/lower_tvm_builtin.cc index 0993422bbb..2b6e21f739 100644 --- a/tirx/src/tirx/transform/lower_tvm_builtin.cc +++ b/tirx/src/tirx/transform/lower_tvm_builtin.cc @@ -365,12 +365,13 @@ class BuiltinLower : public StmtExprMutator { } UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final { - auto min_result = this->Mutate(op->min, inplace_mode); - bool min_unchanged = min_result.UnchangedOrSameAs(op->min); - PrimExpr min = std::move(min_result).ValueOrUnchanged(op->min); - auto extent_result = this->Mutate(op->extent, inplace_mode); - bool extent_unchanged = extent_result.UnchangedOrSameAs(op->extent); - PrimExpr extent = std::move(extent_result).ValueOrUnchanged(op->extent); + auto begin_result = this->Mutate(op->begin, inplace_mode); + bool begin_unchanged = begin_result.UnchangedOrSameAs(op->begin); + PrimExpr begin = std::move(begin_result).ValueOrUnchanged(op->begin); + auto end_result = this->Mutate(op->end, inplace_mode); + bool end_unchanged = end_result.UnchangedOrSameAs(op->end); + PrimExpr end = std::move(end_result).ValueOrUnchanged(op->end); + PrimExpr step = Mutate(op->step, inplace_mode).ValueOrUnchanged(op->step); Stmt body = op->body; if (op->kind == ForKind::kParallel) { @@ -382,19 +383,21 @@ class BuiltinLower : public StmtExprMutator { }); } - if (min_unchanged && extent_unchanged && body.same_as(op->body)) { + if (begin_unchanged && end_unchanged && step.same_as(op->step) && body.same_as(op->body)) { return ffi::Unchanged(); } else { if (inplace_mode == InplaceMode::kAllow) { auto* n = const_cast<ForNode*>(op); - n->min = std::move(min); - n->extent = std::move(extent); + n->begin = std::move(begin); + n->end = std::move(end); + n->step = step; n->body = std::move(body); return ffi::Unchanged(); } auto n = ffi::make_object<ForNode>(*op); - n->min = std::move(min); - n->extent = std::move(extent); + n->begin = std::move(begin); + n->end = std::move(end); + n->step = step; n->body = std::move(body); return Stmt(n); } diff --git a/tirx/src/tirx/transform/make_packed_api.cc b/tirx/src/tirx/transform/make_packed_api.cc index 038a4d8ec4..58954d549b 100644 --- a/tirx/src/tirx/transform/make_packed_api.cc +++ b/tirx/src/tirx/transform/make_packed_api.cc @@ -62,9 +62,9 @@ class ReturnRewriter : public StmtExprMutator { UnchangedOr<Stmt> Mutate_(const ForNode* node, InplaceMode inplace_mode) override { // Unannotated single-iteration loops execute without a threadpool launch. - bool is_parallel = - node->kind == ForKind::kParallel && - !(prim::IsOne(node->extent) && node->HasTrivialStep() && node->annotations.empty()); + bool is_parallel = node->kind == ForKind::kParallel && + !(prim::IsOne(Range::GetExtentFromBeginEnd(node->begin, node->end)) && + node->HasTrivialStep() && node->annotations.empty()); if (is_parallel) in_parallel_ += 1; Stmt ret = StmtExprMutator::Mutate_(node, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(node)); diff --git a/tirx/src/tirx/transform/narrow_datatype.cc b/tirx/src/tirx/transform/narrow_datatype.cc index 2b5f1ab404..c1c548bc37 100644 --- a/tirx/src/tirx/transform/narrow_datatype.cc +++ b/tirx/src/tirx/transform/narrow_datatype.cc @@ -120,8 +120,8 @@ class DataTypeVisitor final : public StmtExprVisitor { } ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) { - analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); - vextent_.insert_or_assign(op->loop_var.as<VarNode>(), op->extent.ty()); + analyzer_->Bind(op->var, Range(op->begin, op->end)); + vextent_.insert_or_assign(op->var.as<VarNode>(), op->end.ty()); return StmtExprVisitor::Visit_(op); } diff --git a/tirx/src/tirx/transform/remove_no_op.cc b/tirx/src/tirx/transform/remove_no_op.cc index ed40b90b68..3fb1d69f03 100644 --- a/tirx/src/tirx/transform/remove_no_op.cc +++ b/tirx/src/tirx/transform/remove_no_op.cc @@ -118,19 +118,20 @@ class NoOpRemover : public StmtExprMutatorWithAnalyzer { } } UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final { - auto extent_range = sym::EvalSet(op->extent, var_range_map_); + auto extent_range = + sym::EvalSet(Range::GetExtentFromBeginEnd(op->begin, op->end), var_range_map_); if (!sym::is_neg_inf(extent_range.max()) && !sym::is_pos_inf(extent_range.max()) && analyzer_->CanProve(extent_range.max() <= 0)) { return Evaluate(0); } - var_range_map_.Set(op->loop_var, sym::IntSet::FromMinExtent(op->min, op->extent)); + var_range_map_.Set(op->var, sym::IntSet::FromRange(Range(op->begin, op->end))); Stmt stmt = Parent::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op)); - var_range_map_.erase(op->loop_var); + var_range_map_.erase(op->var); op = stmt.as<ForNode>(); - if (IsZero(op->extent)) { + if (IsZero(Range::GetExtentFromBeginEnd(op->begin, op->end))) { return Evaluate(0); } - return tirx::IsNoOp(op->body) ? MakeEvaluate({op->min, op->extent}) : stmt; + return tirx::IsNoOp(op->body) ? MakeEvaluate({op->begin, op->end, op->step}) : stmt; } UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode inplace_mode) final { diff --git a/tirx/src/tirx/transform/renormalize_split_pattern.cc b/tirx/src/tirx/transform/renormalize_split_pattern.cc index 9e3804751c..0493c7836c 100644 --- a/tirx/src/tirx/transform/renormalize_split_pattern.cc +++ b/tirx/src/tirx/transform/renormalize_split_pattern.cc @@ -186,10 +186,9 @@ class SplitPatternReNormalizer : public StmtExprMutatorWithAnalyzer { } UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final { - analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); - With<ConstraintContext> ctx1(analyzer_, op->loop_var >= op->min); - With<ConstraintContext> ctx2(analyzer_, - static_cast<PrimExpr>(op->loop_var) < op->min + op->extent); + analyzer_->Bind(op->var, Range(op->begin, op->end)); + With<ConstraintContext> ctx1(analyzer_, op->var >= op->begin); + With<ConstraintContext> ctx2(analyzer_, static_cast<PrimExpr>(op->var) < op->end); return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode); } diff --git a/tirx/src/tirx/transform/stmt_simplify.cc b/tirx/src/tirx/transform/stmt_simplify.cc index 8de2d373ad..d8c2d51b0e 100644 --- a/tirx/src/tirx/transform/stmt_simplify.cc +++ b/tirx/src/tirx/transform/stmt_simplify.cc @@ -117,10 +117,9 @@ UnchangedOr<ffi::Any> StmtSimplifier::Mutate(ffi::AnyView input, InplaceMode inp } UnchangedOr<Stmt> StmtSimplifier::Mutate_(const ForNode* op, InplaceMode inplace_mode) { - analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); - With<sym::ConstraintContext> ctx1(analyzer_, op->loop_var >= op->min); - With<sym::ConstraintContext> ctx2(analyzer_, - static_cast<PrimExpr>(op->loop_var) < op->min + op->extent); + analyzer_->Bind(op->var, Range(op->begin, op->end)); + With<sym::ConstraintContext> ctx1(analyzer_, op->var >= op->begin); + With<sym::ConstraintContext> ctx2(analyzer_, static_cast<PrimExpr>(op->var) < op->end); return Parent::Mutate_(op, inplace_mode); } diff --git a/tirx/src/tirx/transform/storage_rewrite.cc b/tirx/src/tirx/transform/storage_rewrite.cc index 9f128401b1..69406cda7d 100644 --- a/tirx/src/tirx/transform/storage_rewrite.cc +++ b/tirx/src/tirx/transform/storage_rewrite.cc @@ -710,8 +710,8 @@ class StoragePlanRewriter : public StmtExprMutator { Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op)); op = stmt.as<ForNode>(); - return For(op->loop_var, op->min, op->extent, op->kind, MakeAttach(svec, op->body), - op->annotations, op->step); + return For(op->var, op->begin, op->end, op->step, op->kind, MakeAttach(svec, op->body), + op->annotations); } else { return StmtExprMutator::Mutate_(op, inplace_mode); } diff --git a/tirx/src/tirx/transform/unroll_loop.cc b/tirx/src/tirx/transform/unroll_loop.cc index f5c8c2ac4c..f5259c9aeb 100644 --- a/tirx/src/tirx/transform/unroll_loop.cc +++ b/tirx/src/tirx/transform/unroll_loop.cc @@ -182,7 +182,7 @@ class LoopUnroller : public StmtExprMutator { // If a loop var is used as indices to a local memory, it must be unrolled so // the local memory access can be turned into register access. - if (this->var_touched_local_.count(op->loop_var) && value > 0 && unroll_local_access_) { + if (this->var_touched_local_.count(op->var) && value > 0 && unroll_local_access_) { auto_unroll = true; } @@ -298,7 +298,7 @@ class LoopUnroller : public StmtExprMutator { return ffi::Unchanged(); }; for (int i = 0; i < value; ++i) { - vmap.Set(op->loop_var, op->min + IntImm(op->loop_var.ty(), i)); + vmap.Set(op->var, op->begin + IntImm(op->var.ty(), i)); Stmt step = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(body, f_substitute).as_or_throw<Stmt>(); unrolled.push_back(step); @@ -310,7 +310,7 @@ class LoopUnroller : public StmtExprMutator { // returns the extent of the loop if it's a constant integer, otherwise return -1 int GetExtent(const ForNode* op) { // constant folding. - PrimExpr extent = analyzer_->Simplify(op->extent); + PrimExpr extent = analyzer_->Simplify(Range::GetExtentFromBeginEnd(op->begin, op->end)); const IntImmNode* v1 = extent.as<IntImmNode>(); int value = -1; // integers that do not fit in int32_t are treated as symbolic, diff --git a/tirx/src/tirx/transform/vectorize_loop.cc b/tirx/src/tirx/transform/vectorize_loop.cc index 34991aff73..d21ffde003 100644 --- a/tirx/src/tirx/transform/vectorize_loop.cc +++ b/tirx/src/tirx/transform/vectorize_loop.cc @@ -1074,22 +1074,25 @@ class Vectorizer : public StmtExprMutator { if (op->kind == ForKind::kVectorized) { LOG(WARNING) << "Detect vectorize inside vectorized loop, ignoring..."; } - TVM_FFI_ICHECK(IsZero(op->min)); - TVM_FFI_ICHECK(!op->extent.ty().IsScalableVector() && !op->extent.ty().IsFixedLengthVector()); - auto extent_update = this->Mutate(op->extent, inplace_mode); - bool extent_unchanged = extent_update.UnchangedOrSameAs(op->extent); - PrimExpr extent = std::move(extent_update).ValueOrUnchanged(op->extent); - if (extent.ty().IsScalableVector() || extent.ty().IsFixedLengthVector()) { + TVM_FFI_ICHECK(IsZero(op->begin)); + TVM_FFI_ICHECK(!op->end.ty().IsScalableVector() && !op->end.ty().IsFixedLengthVector()); + auto end_update = this->Mutate(op->end, inplace_mode); + bool end_unchanged = end_update.UnchangedOrSameAs(op->end); + PrimExpr end = std::move(end_update).ValueOrUnchanged(op->end); + PrimExpr step = this->Mutate(op->step, inplace_mode).ValueOrUnchanged(op->step); + if (end.ty().IsScalableVector() || end.ty().IsFixedLengthVector() || + step.ty().IsScalableVector() || step.ty().IsFixedLengthVector()) { return Scalarize(ffi::GetRef<Stmt>(op)); } auto body_update = this->Mutate(op->body, inplace_mode); bool body_unchanged = body_update.UnchangedOrSameAs(op->body); Stmt body = std::move(body_update).ValueOrUnchanged(op->body); - if (extent_unchanged && body_unchanged) { + if (end_unchanged && step.same_as(op->step) && body_unchanged) { return ffi::Unchanged(); } else { auto n = ffi::make_object<ForNode>(*op); - n->extent = extent; + n->end = end; + n->step = step; n->body = body; return For(n); } @@ -1178,7 +1181,8 @@ class Vectorizer : public StmtExprMutator { auto substituter = ffi::make_object<StmtExprMutator>(); substituter->VarRemapSet(var_, idx); stmt = substituter->Mutate(stmt).ValueOrUnchanged(stmt); - return For(idx.as_or_throw<PrimVar>(), IntImm(var_ty, 0), var_lanes_, ForKind::kDefault, stmt); + return For(idx.as_or_throw<PrimVar>(), IntImm(var_ty, 0), var_lanes_, IntImm(var_ty, 1), + ForKind::kDefault, stmt); } private: @@ -1361,9 +1365,10 @@ class LoopVectorizer : public StmtExprMutator { UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final { if (op->kind == ForKind::kVectorized) { - auto* extent_as_int = op->extent.as<IntImmNode>(); + PrimExpr extent = Range::GetExtentFromBeginEnd(op->begin, op->end); + auto* extent_as_int = extent.as<IntImmNode>(); - TVM_FFI_ICHECK(IsZero(op->min)); + TVM_FFI_ICHECK(IsZero(op->begin)); // General calls still have vectorization paths that query a compile-time // lane count, so keep them on the existing fixed-width path for now. if (extent_as_int && extent_as_int->value > 1 && TargetHasRVV(target_) && @@ -1380,12 +1385,12 @@ class LoopVectorizer : public StmtExprMutator { return StmtExprVisitor::Visit(value); } }; - bool is_scalable_expr = ffi::make_object<VScaleFinder>()->Visit(op->extent).has_value(); + bool is_scalable_expr = ffi::make_object<VScaleFinder>()->Visit(extent).has_value(); TVM_FFI_ICHECK(is_scalable_expr && TargetHasVLA(target_)) - << "Failed to vectorize loop with extent " << op->extent << " for target " << target_; + << "Failed to vectorize loop with extent " << extent << " for target " << target_; } Stmt body = PreserveUnrollPolicy(op->body, op->annotations); - return ffi::make_object<Vectorizer>(op->loop_var, op->extent, target_) + return ffi::make_object<Vectorizer>(op->var, extent, target_) ->Mutate(body, inplace_mode) .ValueOrUnchanged(body); } else { @@ -1398,7 +1403,7 @@ class LoopVectorizer : public StmtExprMutator { // Match the existing TIRx scalable-vector convention. LLVM/RVV still // selects the runtime vector length with vsetvli. static constexpr int kDefaultVScaleFactor = 4; - PrimType index_dtype = op->loop_var.ty(); + PrimType index_dtype = op->var.ty(); PrimExpr zero = IntImm(index_dtype, 0); PrimExpr fixed_extent = IntImm(index_dtype, extent); PrimExpr scalable_lanes = CreateNewLanes(/*is_scalable=*/true, kDefaultVScaleFactor); @@ -1409,21 +1414,21 @@ class LoopVectorizer : public StmtExprMutator { } PrimExpr num_chunks = ceildiv(fixed_extent, scalable_lanes_index); - PrimVar outer(op->loop_var->name + ".vla.o", index_dtype); - PrimVar inner(op->loop_var->name + ".vla.i", lane_dtype); + PrimVar outer(op->var->name + ".vla.o", index_dtype); + PrimVar inner(op->var->name + ".vla.i", lane_dtype); PrimExpr inner_index = inner; if (inner_index.ty() != index_dtype) { inner_index = prim::Cast(index_dtype, inner_index); } PrimExpr index = outer * scalable_lanes_index + inner_index; auto substituter = ffi::make_object<StmtExprMutator>(); - substituter->VarRemapSet(op->loop_var, index); + substituter->VarRemapSet(op->var, index); Stmt body = substituter->Mutate(op->body).ValueOrUnchanged(op->body); Stmt guarded_body = If(index < fixed_extent, body, std::nullopt, op->loc); - Stmt vector_loop = For(inner, IntImm(lane_dtype, 0), scalable_lanes, ForKind::kVectorized, - guarded_body, op->annotations, std::nullopt, op->loc); - Stmt loop = - For(outer, zero, num_chunks, ForKind::kDefault, vector_loop, {}, std::nullopt, op->loc); + Stmt vector_loop = For(inner, IntImm(lane_dtype, 0), scalable_lanes, IntImm(lane_dtype, 1), + ForKind::kVectorized, guarded_body, op->annotations, op->loc); + Stmt loop = For(outer, zero, num_chunks, IntImm(index_dtype, 1), ForKind::kDefault, vector_loop, + {}, op->loc); return this->Mutate(loop, InplaceMode::kDisallow).ValueOrUnchanged(loop); } diff --git a/tirx/tests/python/ir/test_nonnullable.py b/tirx/tests/python/ir/test_nonnullable.py index e8820d1e42..e2fbba4a57 100644 --- a/tirx/tests/python/ir/test_nonnullable.py +++ b/tirx/tests/python/ir/test_nonnullable.py @@ -49,11 +49,11 @@ def test_optional_statement_fields_and_roundtrip(): tvm_tirx.ir.Var("i", tvm_tirx.ir.PrimType("int32")), 0, 4, + 1, tvm_tirx.ir.ForKind.DEFAULT, body, - step=None, ) - assert loop.step is None + assert int(loop.step) == 1 for value in [conditional, loop]: restored = tvm_tirx.ir.load_json(tvm_tirx.ir.save_json(value)) tvm_tirx.ir.assert_structural_equal(value, restored, map_free_vars=True) diff --git a/tirx/tests/python/script/test_printer_entry.py b/tirx/tests/python/script/test_printer_entry.py index 924af38551..b61b938751 100644 --- a/tirx/tests/python/script/test_printer_entry.py +++ b/tirx/tests/python/script/test_printer_entry.py @@ -72,7 +72,7 @@ def test_plain_render_without_path_mapping(): def test_redirected_repr_translation_failure(): - value = tvm_tirx.ir.For(tvm_tirx.tirx.Var("i", "int32"), 0, 1, 99, tvm_tirx.ir.Evaluate(0)) + value = tvm_tirx.ir.For(tvm_tirx.tirx.Var("i", "int32"), 0, 1, 1, 99, tvm_tirx.ir.Evaluate(0)) with pytest.raises(TypeError, match="unknown loop kind"): value.script() assert re.fullmatch(r"ir\.For\((?:0x)?[0-9a-fA-F]+\)", repr(value)) diff --git a/tirx/tests/python/script/test_special_parser_protocol.py b/tirx/tests/python/script/test_special_parser_protocol.py index 96976dcdab..68a11d7075 100644 --- a/tirx/tests/python/script/test_special_parser_protocol.py +++ b/tirx/tests/python/script/test_special_parser_protocol.py @@ -114,5 +114,5 @@ def main(n: Script.int32): Script.evaluate(i) """ function = entry.parse(source) - assert function.body[0].extent.same_as(function.params[0]) - assert function.body[0].body[0].value.same_as(function.body[0].loop_var) + assert function.body[0].end.same_as(function.params[0]) + assert function.body[0].body[0].value.same_as(function.body[0].var) diff --git a/tirx/tests/python/tirx/analysis/test_check_well_defined.py b/tirx/tests/python/tirx/analysis/test_check_well_defined.py index bcf8e58f0f..c781d95e04 100644 --- a/tirx/tests/python/tirx/analysis/test_check_well_defined.py +++ b/tirx/tests/python/tirx/analysis/test_check_well_defined.py @@ -58,6 +58,7 @@ def test_error_for_out_of_scope_usage(): tvm_tirx.tirx.Var("j", "int32"), 0, 1, + 1, tvm_tirx.ir.ForKind.DEFAULT, tvm_tirx.ir.SeqStmt([tvm_tirx.ir.Bind(i, 42), tvm_tirx.ir.Evaluate(i)]), ) diff --git a/tirx/tests/python/tirx/base/test_nodes.py b/tirx/tests/python/tirx/base/test_nodes.py index 44b683f586..9a6e04dae7 100644 --- a/tirx/tests/python/tirx/base/test_nodes.py +++ b/tirx/tests/python/tirx/base/test_nodes.py @@ -132,8 +132,8 @@ def test_basic(): def test_stmt(): x = tvm_tirx.ir.Evaluate(0) - tvm_tirx.ir.For(tvm_tirx.tirx.Var("i", "int32"), 0, 1, tvm_tirx.ir.ForKind.DEFAULT, x) - tvm_tirx.ir.For(tvm_tirx.tirx.Var("i", "int32"), 0, 1, tvm_tirx.ir.ForKind.UNROLLED, x, step=2) + tvm_tirx.ir.For(tvm_tirx.tirx.Var("i", "int32"), 0, 1, 1, tvm_tirx.ir.ForKind.DEFAULT, x) + tvm_tirx.ir.For(tvm_tirx.tirx.Var("i", "int32"), 0, 1, 2, tvm_tirx.ir.ForKind.UNROLLED, x) def test_dir(): diff --git a/tirx/tests/python/tirx/codegen/test_blackwell.py b/tirx/tests/python/tirx/codegen/test_blackwell.py index 46f2f3d3f4..9f4c302c85 100644 --- a/tirx/tests/python/tirx/codegen/test_blackwell.py +++ b/tirx/tests/python/tirx/codegen/test_blackwell.py @@ -450,7 +450,7 @@ def test_tcgen05_cp_ld_roundtrip(): copy_value_1 = copy_src_tensor_1[ tuple( [ - copy_src_1.region[k].min + copy_src_1.region[k].begin + copy_index_1 // math.prod([int(s.extent) for s in copy_src_1.region[k + 1 :]]) % copy_src_1.region[k].extent @@ -461,7 +461,7 @@ def test_tcgen05_cp_ld_roundtrip(): copy_dst_index_1 = T.meta_var( tuple( [ - copy_dst_1.region[k].min + copy_dst_1.region[k].begin + copy_index_1 // math.prod([int(s.extent) for s in copy_dst_1.region[k + 1 :]]) % copy_dst_1.region[k].extent @@ -611,7 +611,7 @@ def test_tcgen05_mma_ss_no_tma(swizzle): copy_value_2 = copy_src_tensor_2[ tuple( [ - copy_src_2.region[k].min + copy_src_2.region[k].begin + copy_index_2 // math.prod([int(s.extent) for s in copy_src_2.region[k + 1 :]]) % copy_src_2.region[k].extent @@ -622,7 +622,7 @@ def test_tcgen05_mma_ss_no_tma(swizzle): copy_dst_index_2 = T.meta_var( tuple( [ - copy_dst_2.region[k].min + copy_dst_2.region[k].begin + copy_index_2 // math.prod([int(s.extent) for s in copy_dst_2.region[k + 1 :]]) % copy_dst_2.region[k].extent @@ -643,7 +643,7 @@ def test_tcgen05_mma_ss_no_tma(swizzle): copy_value_3 = copy_src_tensor_3[ tuple( [ - copy_src_3.region[k].min + copy_src_3.region[k].begin + copy_index_3 // math.prod([int(s.extent) for s in copy_src_3.region[k + 1 :]]) % copy_src_3.region[k].extent @@ -654,7 +654,7 @@ def test_tcgen05_mma_ss_no_tma(swizzle): copy_dst_index_3 = T.meta_var( tuple( [ - copy_dst_3.region[k].min + copy_dst_3.region[k].begin + copy_index_3 // math.prod([int(s.extent) for s in copy_dst_3.region[k + 1 :]]) % copy_dst_3.region[k].extent diff --git a/tirx/tests/python/tirx/script/test_basic_usage.py b/tirx/tests/python/tirx/script/test_basic_usage.py index 58b590c7c3..67199b35f9 100644 --- a/tirx/tests/python/tirx/script/test_basic_usage.py +++ b/tirx/tests/python/tirx/script/test_basic_usage.py @@ -863,9 +863,9 @@ def test_loop_var_dtype_uint32(): # fmt: on loop = func.body[0] - assert loop.loop_var.ty == PrimType("uint32") - assert loop.min.ty == PrimType("uint32") - assert loop.extent.ty == PrimType("uint32") + assert loop.var.ty == PrimType("uint32") + assert loop.begin.ty == PrimType("uint32") + assert loop.end.ty == PrimType("uint32") _assert_roundtrip(func) @@ -885,9 +885,9 @@ def test_loop_var_dtype_uint32_with_step(): # fmt: on loop = func.body[0] - assert loop.loop_var.ty == PrimType("uint32") - assert loop.min.ty == PrimType("uint32") - assert loop.extent.ty == PrimType("uint32") + assert loop.var.ty == PrimType("uint32") + assert loop.begin.ty == PrimType("uint32") + assert loop.end.ty == PrimType("uint32") assert loop.step.ty == PrimType("uint32") _assert_roundtrip(func) @@ -902,7 +902,7 @@ def test_loop_var_dtype_uint32_all_for_kinds(for_kind): A[i] = T.float32(1) # fmt: on - assert func.body[0].loop_var.ty == PrimType("uint32") + assert func.body[0].var.ty == PrimType("uint32") _assert_roundtrip(func) @@ -916,8 +916,8 @@ def test_grid_loop_var_dtype_uint32(): # fmt: on outer = func.body[0] - assert outer.loop_var.ty == PrimType("uint32") - assert outer.body[0].loop_var.ty == PrimType("uint32") + assert outer.var.ty == PrimType("uint32") + assert outer.body[0].var.ty == PrimType("uint32") _assert_roundtrip(func) @@ -930,7 +930,7 @@ def test_loop_var_dtype_defaults_to_int32(): A[i] = T.float32(1) # fmt: on - assert func.body[0].loop_var.ty == PrimType("int32") + assert func.body[0].var.ty == PrimType("int32") _assert_roundtrip(func) @@ -945,7 +945,7 @@ def test_loop_var_dtype_inferred_from_unsigned_extent(): A[i] = T.float32(1) # fmt: on - assert func.body[0].loop_var.ty == PrimType("uint32") + assert func.body[0].var.ty == PrimType("uint32") _assert_roundtrip(func) @@ -961,8 +961,8 @@ def test_loop_var_dtype_casts_mismatched_bound(): # fmt: on loop = func.body[0] - assert loop.loop_var.ty == PrimType("uint32") - assert loop.extent.ty == PrimType("uint32") + assert loop.var.ty == PrimType("uint32") + assert loop.end.ty == PrimType("uint32") _assert_roundtrip(func) @@ -975,15 +975,17 @@ def test_loop_var_dtype_rejects_unsupported(dtype): def test_hand_built_for_promotes_int_literal_bounds_to_uint32(): """The For constructor retypes literal bounds to the loop var's dtype.""" loop_var = tvm_tirx.tirx.Var("i", "uint32") - loop = tvm_tirx.ir.For(loop_var, 0, 128, tvm_tirx.ir.ForKind.DEFAULT, tvm_tirx.ir.Evaluate(0)) - assert loop.min.ty == PrimType("uint32") - assert loop.extent.ty == PrimType("uint32") + loop = tvm_tirx.ir.For( + loop_var, 0, 128, 1, tvm_tirx.ir.ForKind.DEFAULT, tvm_tirx.ir.Evaluate(0) + ) + assert loop.begin.ty == PrimType("uint32") + assert loop.end.ty == PrimType("uint32") def test_hand_built_for_rejects_negative_literal_for_uint32(): loop_var = tvm_tirx.tirx.Var("i", "uint32") with pytest.raises(Exception, match="not representable"): - tvm_tirx.ir.For(loop_var, -1, 128, tvm_tirx.ir.ForKind.DEFAULT, tvm_tirx.ir.Evaluate(0)) + tvm_tirx.ir.For(loop_var, -1, 128, 1, tvm_tirx.ir.ForKind.DEFAULT, tvm_tirx.ir.Evaluate(0)) def _assert_equal_functions(lhs, rhs, *args, **kwargs): @@ -1065,8 +1067,8 @@ def test_thread_binding_dtype(): loop_i = func.body[0] loop_j = loop_i.body[0] - assert loop_i.loop_var.ty.dtype == "int64" - assert loop_j.loop_var.ty.dtype == "int32" + assert loop_i.var.ty.dtype == "int64" + assert loop_j.var.ty.dtype == "int32" def test_inferred_ty_with_prim_args(): diff --git a/tirx/tests/python/tirx/script/test_ir_builder.py b/tirx/tests/python/tirx/script/test_ir_builder.py index 67230d0e5d..3b414f091e 100644 --- a/tirx/tests/python/tirx/script/test_ir_builder.py +++ b/tirx/tests/python/tirx/script/test_ir_builder.py @@ -53,38 +53,43 @@ def test_ir_builder_tir_for(): # the expected for thread_binding_expected = tvm_tirx.ir.For( - loop_var=tirx.Var("", "int32"), - min=0, - extent=8, + var=tirx.Var("", "int32"), + begin=0, + end=8, + step=1, kind=tvm_tirx.ir.ForKind.PARALLEL, body=tvm_tirx.ir.Evaluate(0), annotations={"thread_binding": "threadIdx.x"}, ) unroll_expected = tvm_tirx.ir.For( - loop_var=tirx.Var("", "int32"), - min=0, - extent=16, + var=tirx.Var("", "int32"), + begin=0, + end=16, + step=1, kind=tvm_tirx.ir.ForKind.UNROLLED, body=thread_binding_expected, ) vectorized_expected = tvm_tirx.ir.For( - loop_var=tirx.Var("", "int32"), - min=0, - extent=32, + var=tirx.Var("", "int32"), + begin=0, + end=32, + step=1, kind=tvm_tirx.ir.ForKind.VECTORIZED, body=unroll_expected, ) parallel_expected = tvm_tirx.ir.For( - loop_var=tirx.Var("", "int32"), - min=0, - extent=64, + var=tirx.Var("", "int32"), + begin=0, + end=64, + step=1, kind=tvm_tirx.ir.ForKind.PARALLEL, body=vectorized_expected, ) for_expected = tvm_tirx.ir.For( - loop_var=tirx.Var("", "int32"), - min=0, - extent=128, + var=tirx.Var("", "int32"), + begin=0, + end=128, + step=1, kind=tvm_tirx.ir.ForKind.DEFAULT, body=parallel_expected, ) @@ -102,9 +107,10 @@ def test_ir_builder_tir_for_uint(): for_actual = ib.get() for_expected = tvm_tirx.ir.For( - loop_var=tirx.Var("", "uint32"), - min=tirx.const(0, "uint32"), - extent=tirx.const(128, "uint32"), + var=tirx.Var("", "uint32"), + begin=tirx.const(0, "uint32"), + end=tirx.const(128, "uint32"), + step=1, kind=tvm_tirx.ir.ForKind.DEFAULT, body=tvm_tirx.ir.Evaluate(0), ) diff --git a/tirx/tests/python/tirx/script/test_tensors.py b/tirx/tests/python/tirx/script/test_tensors.py index 1816895209..c383f22add 100644 --- a/tirx/tests/python/tirx/script/test_tensors.py +++ b/tirx/tests/python/tirx/script/test_tensors.py @@ -1269,14 +1269,14 @@ def test_tensor_contiguous_regions(): reg = A[:, :, 8:16] assert isinstance(reg, TensorRegion) assert len(reg.region) == 3 - assert (int(reg.region[2].min), int(reg.region[2].extent)) == (8, 8) + assert (int(reg.region[2].begin), int(reg.region[2].extent)) == (8, 8) reg2 = A[3, :, 0:8] - assert (int(reg2.region[0].min), int(reg2.region[0].extent)) == (3, 1) + assert (int(reg2.region[0].begin), int(reg2.region[0].extent)) == (3, 1) reg_c = C[:, 256:384] - assert (int(reg_c.region[1].min), int(reg_c.region[1].extent)) == (256, 128) + assert (int(reg_c.region[1].begin), int(reg_c.region[1].extent)) == (256, 128) c = T.Var(name="c", ty="int32") dynamic = A[:, :, c * 8 : c * 8 + 8] - assert_structural_equal(dynamic.region[2].min, c * 8) + assert_structural_equal(dynamic.region[2].begin, c * 8) assert tvm_tirx.sym.Analyzer().can_prove_equal(dynamic.region[2].extent, 8) @@ -1325,7 +1325,7 @@ def test_tensor_slice_region(): narrowed = br[4:12, 2:10] assert isinstance(narrowed, TensorRegion) assert narrowed.source.same_as(buf) - assert [(int(dim.min), int(dim.extent)) for dim in narrowed.region] == [ + assert [(int(dim.begin), int(dim.extent)) for dim in narrowed.region] == [ (36, 8), (2, 8), ] @@ -1337,7 +1337,7 @@ def test_tensor_slice_region(): point_then_region = br[3] assert isinstance(point_then_region, TensorRegion) - assert [(int(dim.min), int(dim.extent)) for dim in point_then_region.region] == [ + assert [(int(dim.begin), int(dim.extent)) for dim in point_then_region.region] == [ (35, 1), (0, 32), ] diff --git a/tirx/tests/python/tirx/test_constructor.py b/tirx/tests/python/tirx/test_constructor.py index 47eb03ae0c..04c5077007 100644 --- a/tirx/tests/python/tirx/test_constructor.py +++ b/tirx/tests/python/tirx/test_constructor.py @@ -292,10 +292,10 @@ def test_stmt_constructor(): assert len(x.message_parts) == 1 assert x.message_parts[0].value == "hellow" - x = tvm_tirx.ir.For(tvm_tirx.tirx.Var("x", "int32"), 0, 10, tvm_tirx.ir.ForKind.DEFAULT, nop) + x = tvm_tirx.ir.For(tvm_tirx.tirx.Var("x", "int32"), 0, 10, 1, tvm_tirx.ir.ForKind.DEFAULT, nop) assert isinstance(x, tvm_tirx.ir.For) - assert x.min.value == 0 - assert x.extent.value == 10 + assert x.begin.value == 0 + assert x.end.value == 10 assert len(x.body) == 1 and x.body[0] == nop tensor_var = tvm_tirx.tirx.Var("buf", tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType("bool"))) diff --git a/tirx/tests/python/tirx/test_tile_namespace.py b/tirx/tests/python/tirx/test_tile_namespace.py index 1f8a564f3f..fbffa563c0 100644 --- a/tirx/tests/python/tirx/test_tile_namespace.py +++ b/tirx/tests/python/tirx/test_tile_namespace.py @@ -131,7 +131,7 @@ def test_effective_global_layout_preserves_source_and_region(): region = tensor[2:4, 3:7] layout = region.source.ty.effective_layout() assert int(layout.apply(2, 3, shape=tensor.shape)["m"]) == 48 - starts = [r.min for r in region.region] + starts = [r.begin for r in region.region] assert int(layout.apply(*starts, shape=tensor.shape)["m"]) == 48 assert int(layout.apply(starts[0] + 1, starts[1] + 3, shape=tensor.shape)["m"]) == 67 assert save_json(tensor) == before diff --git a/tirx/tests/python/tirx/tile/cuda/test_cp_async.py b/tirx/tests/python/tirx/tile/cuda/test_cp_async.py index 29a7d16528..b1e0f4ba46 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_cp_async.py +++ b/tirx/tests/python/tirx/tile/cuda/test_cp_async.py @@ -117,7 +117,7 @@ def test_copy_g2s_s2g_cta_vec_load(task, dtype): copy_value_1 = copy_src_tensor_1[ tuple( [ - copy_src_1.region[k].min + copy_src_1.region[k].begin + copy_index_1 // math.prod([int(s.extent) for s in copy_src_1.region[k + 1 :]]) % copy_src_1.region[k].extent @@ -128,7 +128,7 @@ def test_copy_g2s_s2g_cta_vec_load(task, dtype): copy_dst_index_1 = T.meta_var( tuple( [ - copy_dst_1.region[k].min + copy_dst_1.region[k].begin + copy_index_1 // math.prod([int(s.extent) for s in copy_dst_1.region[k + 1 :]]) % copy_dst_1.region[k].extent diff --git a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py index 19e86c38e2..1a954c9a29 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py +++ b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py @@ -67,7 +67,7 @@ def _count_s2c_ops(impl): total = 0 def visit_for(op, visitor): - loop_extents.append(op.extent) + loop_extents.append(tvm_tirx.ir.Range.get_extent_from_begin_end(op.begin, op.end)) visitor.default_visit(op) loop_extents.pop() @@ -205,7 +205,7 @@ def test_dsmem(shape, dtype, src_spec, dst_spec, expected): copy_value_1 = copy_src_tensor_1[ tuple( [ - copy_src_1.region[k].min + copy_src_1.region[k].begin + copy_index_1 // math.prod( [int(s.extent) for s in copy_src_1.region[k + 1 :]] @@ -218,7 +218,7 @@ def test_dsmem(shape, dtype, src_spec, dst_spec, expected): copy_dst_index_1 = T.meta_var( tuple( [ - copy_dst_1.region[k].min + copy_dst_1.region[k].begin + copy_index_1 // math.prod( [int(s.extent) for s in copy_dst_1.region[k + 1 :]] @@ -256,7 +256,7 @@ def test_dsmem(shape, dtype, src_spec, dst_spec, expected): copy_value_2 = copy_src_tensor_2[ tuple( [ - copy_src_2.region[k].min + copy_src_2.region[k].begin + copy_index_2 // math.prod( [int(s.extent) for s in copy_src_2.region[k + 1 :]] @@ -269,7 +269,7 @@ def test_dsmem(shape, dtype, src_spec, dst_spec, expected): copy_dst_index_2 = T.meta_var( tuple( [ - copy_dst_2.region[k].min + copy_dst_2.region[k].begin + copy_index_2 // math.prod( [int(s.extent) for s in copy_dst_2.region[k + 1 :]] diff --git a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py index a5f483754d..4ff6c24643 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py +++ b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py @@ -81,7 +81,7 @@ class _TMACounter: self.weighted_calls = [] def _visit_for(self, op, visitor): - self.loop_extents.append(op.extent) + self.loop_extents.append(tvm_tirx.ir.Range.get_extent_from_begin_end(op.begin, op.end)) visitor.default_visit(op) self.loop_extents.pop() diff --git a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py index 9430c1ce75..e672b0dd27 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py +++ b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py @@ -56,7 +56,7 @@ from tvm_tirx.tirx.layout import S, TileLayout def _window(region): """Keep the original region's storage and offsets in a zero-copy view.""" if isinstance(region, tvm_tirx.ir.TensorRegion): - return region.source.sub[tuple(slice(r.min, r.min + r.extent) for r in region.region)] + return region.source.sub[tuple(slice(r.begin, r.begin + r.extent) for r in region.region)] return region diff --git a/tirx/tests/python/tirx/transform/test_hoist_if.py b/tirx/tests/python/tirx/transform/test_hoist_if.py index 3820f31300..88ccc06b97 100644 --- a/tirx/tests/python/tirx/transform/test_hoist_if.py +++ b/tirx/tests/python/tirx/transform/test_hoist_if.py @@ -290,11 +290,11 @@ def test_effectful_enclosing_condition(): tvm_tirx.ir.assert_structural_equal(After, Expected) [email protected]("field", ["min", "extent", "step"]) [email protected]("field", ["begin", "end", "step"]) def test_effectful_loop_header(field): value = tvm_tirx.tirx.call_extern("loop_header", ty="int32") - start = value if field == "min" else 0 - stop = value if field == "extent" else 4 + start = value if field == "begin" else 0 + stop = value if field == "end" else 4 step = value if field == "step" else 1 @I.ir_module diff --git a/tirx/tests/python/tirx/transform/test_split_host_device.py b/tirx/tests/python/tirx/transform/test_split_host_device.py index 7623b50ebd..42c3b3a305 100644 --- a/tirx/tests/python/tirx/transform/test_split_host_device.py +++ b/tirx/tests/python/tirx/transform/test_split_host_device.py @@ -44,7 +44,7 @@ def test_ssa_across_entire_module(): T.evaluate(i) after = tvm_tirx.tirx.transform.SplitHostDevice()(before) - loop_var = after["main"].body[0].loop_var + loop_var = after["main"].body[0].var param_var = after["main_kernel"].params[0] assert not loop_var.same_as(param_var) diff --git a/tirx/tests/python/tirx/transform/test_vectorize.py b/tirx/tests/python/tirx/transform/test_vectorize.py index e7fcd4d9fd..6ee3d996b1 100644 --- a/tirx/tests/python/tirx/transform/test_vectorize.py +++ b/tirx/tests/python/tirx/transform/test_vectorize.py @@ -522,7 +522,7 @@ def test_illegal_vscale_in_non_sve_compilation(): for j in T.vectorized(0, 4 * T.vscale()): A[j] = 13 - msg = "Failed to vectorize loop with extent T.vscale\\(\\) \\* 4 for target" + msg = "Failed to vectorize loop with extent 4 \\* T.vscale\\(\\) for target" with pytest.raises(RuntimeError, match=msg): tvm_tirx.tirx.transform.VectorizeLoop()(Mod)
