src/tirx/transform/stmt_simplify.cc - tvm - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 /*!
  * \file stmt_simplify.cc
  * \brief Statement simplifier based on analyzer
  */

 #include "../../tirx/transform/stmt_simplify.h"

 #include <tvm/arith/analyzer.h>
 #include <tvm/ffi/cast.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/tirx/analysis.h>
 #include <tvm/tirx/builtin.h>
 #include <tvm/tirx/expr.h>
 #include <tvm/tirx/op.h>
 #include <tvm/tirx/transform.h>

 #include "../../arith/ir_mutator_with_analyzer.h"

 namespace tvm {
 namespace arith {

 using namespace tirx;

 struct StmtSimplifyConfigNode : public ffi::Object {
   bool transitively_prove_inequalities;
   bool convert_boolean_to_and_of_ors;
   bool apply_constraints_to_boolean_branches;

   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
     refl::ObjectDef<StmtSimplifyConfigNode>()
         .def_ro("transitively_prove_inequalities",
                 &StmtSimplifyConfigNode::transitively_prove_inequalities,
                 "If true, simplify conditionals with transitive combinations of scoped constraints",
                 refl::DefaultValue(false))
         .def_ro("convert_boolean_to_and_of_ors",
                 &StmtSimplifyConfigNode::convert_boolean_to_and_of_ors,
                 "If true, simplify conditionals into an AND of ORs", refl::DefaultValue(false))
         .def_ro("apply_constraints_to_boolean_branches",
                 &StmtSimplifyConfigNode::apply_constraints_to_boolean_branches,
                 "If true, simplify each branch of AND/OR under constraints provided by the other "
                 "branch",
                 refl::DefaultValue(false));
   }
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.transform.StmtSimplifyConfig", StmtSimplifyConfigNode,
                                     ffi::Object);

   RewriteSimplifier::Extension GetEnabledExtensions() const {
     RewriteSimplifier::Extension flags = RewriteSimplifier::kNone;
     if (transitively_prove_inequalities) {
       flags =
           RewriteSimplifier::Extension(flags | RewriteSimplifier::kTransitivelyProveInequalities);
     }
     if (convert_boolean_to_and_of_ors) {
       flags = RewriteSimplifier::Extension(flags | RewriteSimplifier::kConvertBooleanToAndOfOrs);
     }
     if (apply_constraints_to_boolean_branches) {
       flags = RewriteSimplifier::Extension(flags |
                                            RewriteSimplifier::kApplyConstraintsToBooleanBranches);
     }
     return flags;
   }
 };

 class StmtSimplifyConfig : public ffi::ObjectRef {
  public:
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StmtSimplifyConfig, ffi::ObjectRef,
                                                 StmtSimplifyConfigNode);
 };

 static StmtSimplifyConfig MakeDefaultStmtSimplifyConfig() {
   return tvm::transform::PassConfigWithDefaults<StmtSimplifyConfig>();
 }

 TVM_FFI_STATIC_INIT_BLOCK() { StmtSimplifyConfigNode::RegisterReflection(); }

 TVM_REGISTER_PASS_CONFIG_OPTION("tirx.StmtSimplify", StmtSimplifyConfig);

 class StmtSimplifier : public IRMutatorWithAnalyzer {
  public:
   static PrimFunc Apply(PrimFunc func, const Analyzer& analyzer,
                         ffi::Optional<StmtSimplifyConfig> config_opt = std::nullopt) {
     auto config = config_opt.value_or(MakeDefaultStmtSimplifyConfig());
     analyzer->rewrite_simplify.SetEnabledExtensions(config->GetEnabledExtensions());

     StmtSimplifier simplifier(analyzer, config);
     simplifier.MarkBufferParamShapes(func);
     func.CopyOnWrite()->body = simplifier(func->body);
     return func;
   }

  private:
   explicit StmtSimplifier(const Analyzer& analyzer, StmtSimplifyConfig config)
       : IRMutatorWithAnalyzer(analyzer), config_(config) {}

   using Parent = IRMutatorWithAnalyzer;
   using Parent::VisitExpr_;
   using Parent::VisitStmt;
   using Parent::VisitStmt_;

   // Do not simplify buffer definition fields (shape, strides, elem_offset).
   //
   // The simplifier's VisitExpr override calls analyzer_->Simplify() directly,
   // bypassing the normal ExprMutator dispatch. This means BufferLoad expressions
   // inside values (e.g., BufferStore value) skip VisitExpr_(BufferLoadNode*) and
   // thus skip VisitBufferUse. If VisitBufferDef remaps buffers at DeclBuffer sites,
   // the BufferLoad use sites won't pick up the remap, causing DeclBuffer/BufferLoad
   // buffer identity divergence and well-formedness violations.
   //
   // Instead, we keep buffer definitions unchanged and rely on used_in_buffer_def_
   // to prevent inlining LetStmt vars that appear in buffer definitions.
   BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) override { return buffer; }

   Expr VisitExpr(const Expr& expr) final {
     if (auto prim_expr = expr.as<PrimExpr>()) {
       return analyzer_->Simplify(prim_expr.value());
     }
     return Parent::VisitExpr(expr);
   }

   Stmt Simplify(Stmt stmt) { return operator()(std::move(stmt)); }

   Stmt VisitStmt_(const ForNode* op) 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);
     return Parent::VisitStmt_(op);
   }

   Stmt VisitStmt_(const BindNode* op) override {
     auto prim_value = op->value.as<PrimExpr>();
     if (!prim_value) {
       return Parent::VisitStmt_(op);
     }
     PrimExpr value = this->VisitPrimExpr(prim_value.value());
     // Bind in analyzer for constraint proving and simplification of
     // subsequent expressions.  Don't remove the Bind statement --
     // with flat Bind there's no body to inspect for usage patterns,
     // so we always keep the Bind.
     if (SideEffect(value) <= CallEffectKind::kPure) {
       analyzer_->Bind(op->var, value);
       // Record the binding so we can substitute it into assert conditions
       // (see VisitStmt_(const AssertStmtNode*)).  Under SSA each var is
       // bound exactly once, so the map grows monotonically without key
       // conflicts.  No scope-based cleanup is needed because vars bound
       // in inner scopes are only referenced within those scopes; stale
       // entries are harmless and never consulted again.
       non_inlined_bindings_.Set(op->var, value);
     }

     if (value.same_as(op->value)) {
       return ffi::GetRef<Stmt>(op);
     } else {
       auto n = this->CopyOnWrite(op);
       n->value = std::move(value);
       return Stmt(n);
     }
   }

   Stmt VisitStmt_(const IfThenElseNode* op) override {
     if (ffi::Optional<bool> cond = ProveCondition(op->condition)) {
       if (cond.value()) {
         return this->VisitStmt(op->then_case);
       } else if (op->else_case) {
         return this->VisitStmt(op->else_case.value());
       } else {
         return Evaluate(0);
       }
     } else {
       return Parent::VisitStmt_(op);
     }
   }

   Expr VisitExpr_(const CallNode* op) override {
     if (op->op.same_as(builtin::if_then_else())) {
       if (ffi::Optional<bool> cond = ProveCondition(op->args[0].as_or_throw<PrimExpr>())) {
         if (cond.value()) {
           return this->VisitExpr(op->args[1]);
         } else {
           return this->VisitExpr(op->args[2]);
         }
       }
     }
     return Parent::VisitExpr_(op);
   }

   Expr VisitExpr_(const BufferLoadNode* op) override { return Parent::VisitExpr_(op); }

   // eliminate useless stores
   Stmt VisitStmt_(const BufferStoreNode* op) override {
     BufferStore store = Parent::VisitStmt_(op).as_or_throw<BufferStore>();
     if (const BufferLoadNode* load = store->value.as<BufferLoadNode>()) {
       if (load->buffer.same_as(store->buffer) && ArrayDeepEqual(load->indices, store->indices) &&
           tirx::ExprDeepEqual()(load->buffer->elem_offset, store->buffer->elem_offset) &&
           ArrayDeepEqual(load->buffer->shape, store->buffer->shape) &&
           ArrayDeepEqual(load->buffer->strides, store->buffer->strides)) {
         return Evaluate(0);
       }
     }
     return store;
   }

  private:
   bool ArrayDeepEqual(const ffi::Array<PrimExpr>& lhs, const ffi::Array<PrimExpr>& rhs) {
     if (lhs.size() != rhs.size()) {
       return false;
     }
     for (size_t i = 0; i < lhs.size(); i++) {
       if (!tirx::ExprDeepEqual()(lhs[i], rhs[i])) {
         return false;
       }
     }
     return true;
   }

   /* \brief Internal utility for checking conditionals
    *
    * Substitutes any known Bind values and then simplifies with the analyzer.
    */
   ffi::Optional<bool> ProveCondition(PrimExpr condition) const {
     condition = Substitute(condition, non_inlined_bindings_);
     condition = analyzer_->Simplify(condition);
     if (const int64_t* as_int = as_const_int(condition)) {
       return *as_int != 0;
     } else {
       return std::nullopt;
     }
   }

   StmtSimplifyConfig config_;

   // Pure Bind values kept for substitution into assert conditions.
   // Grows monotonically under SSA — no scope-based cleanup required.
   ffi::Map<Var, PrimExpr> non_inlined_bindings_;
 };

 }  // namespace arith

 namespace tirx {

 PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer) {
   return arith::StmtSimplifier::Apply(std::move(func), analyzer);
 }

 namespace transform {

 Pass StmtSimplify() {
   auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
     arith::Analyzer analyzer;
     auto cfg = ctx->GetConfig<arith::StmtSimplifyConfig>("tirx.StmtSimplify");

     return arith::StmtSimplifier::Apply(f, analyzer, cfg);
   };
   return CreatePrimFuncPass(pass_func, 0, "tirx.StmtSimplify", {});
 }

 TVM_FFI_STATIC_INIT_BLOCK() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def("tirx.transform.StmtSimplify", StmtSimplify);
 }

 }  // namespace transform
 }  // namespace tirx
 }  // namespace tvm
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	/*!
	* \file stmt_simplify.cc
	* \brief Statement simplifier based on analyzer
	*/

	#include "../../tirx/transform/stmt_simplify.h"

	#include <tvm/arith/analyzer.h>
	#include <tvm/ffi/cast.h>
	#include <tvm/ffi/function.h>
	#include <tvm/ffi/reflection/registry.h>
	#include <tvm/tirx/analysis.h>
	#include <tvm/tirx/builtin.h>
	#include <tvm/tirx/expr.h>
	#include <tvm/tirx/op.h>
	#include <tvm/tirx/transform.h>

	#include "../../arith/ir_mutator_with_analyzer.h"

	namespace tvm {
	namespace arith {

	using namespace tirx;

	struct StmtSimplifyConfigNode : public ffi::Object {
	bool transitively_prove_inequalities;
	bool convert_boolean_to_and_of_ors;
	bool apply_constraints_to_boolean_branches;

	static void RegisterReflection() {
	namespace refl = tvm::ffi::reflection;
	refl::ObjectDef<StmtSimplifyConfigNode>()
	.def_ro("transitively_prove_inequalities",
	&StmtSimplifyConfigNode::transitively_prove_inequalities,
	"If true, simplify conditionals with transitive combinations of scoped constraints",
	refl::DefaultValue(false))
	.def_ro("convert_boolean_to_and_of_ors",
	&StmtSimplifyConfigNode::convert_boolean_to_and_of_ors,
	"If true, simplify conditionals into an AND of ORs", refl::DefaultValue(false))
	.def_ro("apply_constraints_to_boolean_branches",
	&StmtSimplifyConfigNode::apply_constraints_to_boolean_branches,
	"If true, simplify each branch of AND/OR under constraints provided by the other "
	"branch",
	refl::DefaultValue(false));
	}
	TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.transform.StmtSimplifyConfig", StmtSimplifyConfigNode,
	ffi::Object);

	RewriteSimplifier::Extension GetEnabledExtensions() const {
	RewriteSimplifier::Extension flags = RewriteSimplifier::kNone;
	if (transitively_prove_inequalities) {
	flags =
	RewriteSimplifier::Extension(flags \| RewriteSimplifier::kTransitivelyProveInequalities);
	}
	if (convert_boolean_to_and_of_ors) {
	flags = RewriteSimplifier::Extension(flags \| RewriteSimplifier::kConvertBooleanToAndOfOrs);
	}
	if (apply_constraints_to_boolean_branches) {
	flags = RewriteSimplifier::Extension(flags \|
	RewriteSimplifier::kApplyConstraintsToBooleanBranches);
	}
	return flags;
	}
	};

	class StmtSimplifyConfig : public ffi::ObjectRef {
	public:
	TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StmtSimplifyConfig, ffi::ObjectRef,
	StmtSimplifyConfigNode);
	};

	static StmtSimplifyConfig MakeDefaultStmtSimplifyConfig() {
	return tvm::transform::PassConfigWithDefaults<StmtSimplifyConfig>();
	}

	TVM_FFI_STATIC_INIT_BLOCK() { StmtSimplifyConfigNode::RegisterReflection(); }

	TVM_REGISTER_PASS_CONFIG_OPTION("tirx.StmtSimplify", StmtSimplifyConfig);

	class StmtSimplifier : public IRMutatorWithAnalyzer {
	public:
	static PrimFunc Apply(PrimFunc func, const Analyzer& analyzer,
	ffi::Optional<StmtSimplifyConfig> config_opt = std::nullopt) {
	auto config = config_opt.value_or(MakeDefaultStmtSimplifyConfig());
	analyzer->rewrite_simplify.SetEnabledExtensions(config->GetEnabledExtensions());

	StmtSimplifier simplifier(analyzer, config);
	simplifier.MarkBufferParamShapes(func);
	func.CopyOnWrite()->body = simplifier(func->body);
	return func;
	}

	private:
	explicit StmtSimplifier(const Analyzer& analyzer, StmtSimplifyConfig config)
	: IRMutatorWithAnalyzer(analyzer), config_(config) {}

	using Parent = IRMutatorWithAnalyzer;
	using Parent::VisitExpr_;
	using Parent::VisitStmt;
	using Parent::VisitStmt_;

	// Do not simplify buffer definition fields (shape, strides, elem_offset).
	//
	// The simplifier's VisitExpr override calls analyzer_->Simplify() directly,
	// bypassing the normal ExprMutator dispatch. This means BufferLoad expressions
	// inside values (e.g., BufferStore value) skip VisitExpr_(BufferLoadNode*) and
	// thus skip VisitBufferUse. If VisitBufferDef remaps buffers at DeclBuffer sites,
	// the BufferLoad use sites won't pick up the remap, causing DeclBuffer/BufferLoad
	// buffer identity divergence and well-formedness violations.
	//
	// Instead, we keep buffer definitions unchanged and rely on used_in_buffer_def_
	// to prevent inlining LetStmt vars that appear in buffer definitions.
	BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) override { return buffer; }

	Expr VisitExpr(const Expr& expr) final {
	if (auto prim_expr = expr.as<PrimExpr>()) {
	return analyzer_->Simplify(prim_expr.value());
	}
	return Parent::VisitExpr(expr);
	}

	Stmt Simplify(Stmt stmt) { return operator()(std::move(stmt)); }

	Stmt VisitStmt_(const ForNode* op) 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);
	return Parent::VisitStmt_(op);
	}

	Stmt VisitStmt_(const BindNode* op) override {
	auto prim_value = op->value.as<PrimExpr>();
	if (!prim_value) {
	return Parent::VisitStmt_(op);
	}
	PrimExpr value = this->VisitPrimExpr(prim_value.value());
	// Bind in analyzer for constraint proving and simplification of
	// subsequent expressions. Don't remove the Bind statement --
	// with flat Bind there's no body to inspect for usage patterns,
	// so we always keep the Bind.
	if (SideEffect(value) <= CallEffectKind::kPure) {
	analyzer_->Bind(op->var, value);
	// Record the binding so we can substitute it into assert conditions
	// (see VisitStmt_(const AssertStmtNode*)). Under SSA each var is
	// bound exactly once, so the map grows monotonically without key
	// conflicts. No scope-based cleanup is needed because vars bound
	// in inner scopes are only referenced within those scopes; stale
	// entries are harmless and never consulted again.
	non_inlined_bindings_.Set(op->var, value);
	}

	if (value.same_as(op->value)) {
	return ffi::GetRef<Stmt>(op);
	} else {
	auto n = this->CopyOnWrite(op);
	n->value = std::move(value);
	return Stmt(n);
	}
	}

	Stmt VisitStmt_(const IfThenElseNode* op) override {
	if (ffi::Optional<bool> cond = ProveCondition(op->condition)) {
	if (cond.value()) {
	return this->VisitStmt(op->then_case);
	} else if (op->else_case) {
	return this->VisitStmt(op->else_case.value());
	} else {
	return Evaluate(0);
	}
	} else {
	return Parent::VisitStmt_(op);
	}
	}

	Expr VisitExpr_(const CallNode* op) override {
	if (op->op.same_as(builtin::if_then_else())) {
	if (ffi::Optional<bool> cond = ProveCondition(op->args[0].as_or_throw<PrimExpr>())) {
	if (cond.value()) {
	return this->VisitExpr(op->args[1]);
	} else {
	return this->VisitExpr(op->args[2]);
	}
	}
	}
	return Parent::VisitExpr_(op);
	}

	Expr VisitExpr_(const BufferLoadNode* op) override { return Parent::VisitExpr_(op); }

	// eliminate useless stores
	Stmt VisitStmt_(const BufferStoreNode* op) override {
	BufferStore store = Parent::VisitStmt_(op).as_or_throw<BufferStore>();
	if (const BufferLoadNode* load = store->value.as<BufferLoadNode>()) {
	if (load->buffer.same_as(store->buffer) && ArrayDeepEqual(load->indices, store->indices) &&
	tirx::ExprDeepEqual()(load->buffer->elem_offset, store->buffer->elem_offset) &&
	ArrayDeepEqual(load->buffer->shape, store->buffer->shape) &&
	ArrayDeepEqual(load->buffer->strides, store->buffer->strides)) {
	return Evaluate(0);
	}
	}
	return store;
	}

	private:
	bool ArrayDeepEqual(const ffi::Array<PrimExpr>& lhs, const ffi::Array<PrimExpr>& rhs) {
	if (lhs.size() != rhs.size()) {
	return false;
	}
	for (size_t i = 0; i < lhs.size(); i++) {
	if (!tirx::ExprDeepEqual()(lhs[i], rhs[i])) {
	return false;
	}
	}
	return true;
	}

	/* \brief Internal utility for checking conditionals
	*
	* Substitutes any known Bind values and then simplifies with the analyzer.
	*/
	ffi::Optional<bool> ProveCondition(PrimExpr condition) const {
	condition = Substitute(condition, non_inlined_bindings_);
	condition = analyzer_->Simplify(condition);
	if (const int64_t* as_int = as_const_int(condition)) {
	return *as_int != 0;
	} else {
	return std::nullopt;
	}
	}

	StmtSimplifyConfig config_;

	// Pure Bind values kept for substitution into assert conditions.
	// Grows monotonically under SSA — no scope-based cleanup required.
	ffi::Map<Var, PrimExpr> non_inlined_bindings_;
	};

	} // namespace arith

	namespace tirx {

	PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer) {
	return arith::StmtSimplifier::Apply(std::move(func), analyzer);
	}

	namespace transform {

	Pass StmtSimplify() {
	auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
	arith::Analyzer analyzer;
	auto cfg = ctx->GetConfig<arith::StmtSimplifyConfig>("tirx.StmtSimplify");

	return arith::StmtSimplifier::Apply(f, analyzer, cfg);
	};
	return CreatePrimFuncPass(pass_func, 0, "tirx.StmtSimplify", {});
	}

	TVM_FFI_STATIC_INIT_BLOCK() {
	namespace refl = tvm::ffi::reflection;
	refl::GlobalDef().def("tirx.transform.StmtSimplify", StmtSimplify);
	}

	} // namespace transform
	} // namespace tirx
	} // namespace tvm