src/relax/transform/convert_layout.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 src/relax/transform/convert_layout.cc
  * \brief Automatic layout conversion pass, especially for axis swapping.
  */

 #include <tvm/ffi/cast.h>
 #include <tvm/ffi/extra/serialization.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/relax/expr_functor.h>
 #include <tvm/relax/nested_msg.h>
 #include <tvm/relax/op_attr_types.h>
 #include <tvm/relax/transform.h>
 #include <tvm/tirx/index_map.h>

 #include "../op/tensor/manipulate.h"
 #include "infer_layout_utils.h"
 #include "utils.h"

 namespace tvm {
 namespace relax {

 using tirx::IndexMap;
 using tirx::SLayout;
 using LayoutCb = tvm::relax::transform::LayoutCb;

 /*!
  * \brief Main logic to convert the layout of conv2d. Other ops
  * can adapt to such layout conversion following conv2d accordingly.
  *
  * Structurally speaking, a Relax function is composed of a series of VarBinding and
  * MatchCast. And a specific class of VarBindings is the basic unit we want to rewrite.
  * Formally, they are of the form:
  *
  * var = Call(Op, [args], attrs)
  *
  * where Op is a specific op we want to rewrite, and attrs is the attributes of the op.
  * var and args are all exprs with type Tensor or Tuple of Tensors. They might
  * be vars, constants, or Tuple of vars and constants.
  *
  * We register the layout inference function for each op (FRelaxInferLayout), which accepts the
  * current call, the desired layout of conv2d ops, and the layout map of previous vars. The result
  * of the layout inference function is contained in an InferLayoutOutput object, which contains 3
  * fields: input_layouts, output_layouts, and attr, which represents the expected input layout,
  * output_layout and converted attrs of the new op call.
  *
  * The rewrite pass does the rewriting in a single forward pass, where for each Call(Op),
  * we collect the current SLayout of each input var, and let the InferLayout function to infer the
  * desired layout of the output. The rewriter will use these info to convert
  * the layout of inputs and attrs of the op call, and note down the new layout of the output.
  *
  * The desired layout of conv2d ops is a map from the name of the op to the desired layout of the
  * desired feature map, weight and output. For example, if we want to convert the layout of conv2d
  * from NCHW to NHWC, we can set the desired layout of conv2d to be {"conv2d": ["NHWC", "OHWI"]}.
  *
  * The way we represent the layout of a var is a NLayout object, which is a nested tuple of SLayout.
  * The incoming layout of the module will be set as the default layout (We use ABCD... as the
  * default) Note that for operators like conv, pool, people typically use NHWC to refer to the axes.
  * But to be generic and support more operators, we use ABCD... to refer to the axes.
  *
  * Note that currently the layout conversion of conv2d only support axis swapping, such as NCHW to
  * NWHC. Packed layout such as NCHW to NCHW4c is not supported now.
  */
 class LayoutConvertMutator : public ExprMutator {
  public:
   explicit LayoutConvertMutator(
       const ffi::Map<ffi::String, ffi::Array<ffi::String>>& desired_layouts, LayoutCb layout_cb)
       : desired_layouts_(desired_layouts), layout_cb_(layout_cb) {}

  private:
   ffi::Array<int64_t> LayoutToIntegers(const SLayout& layout) {
     ffi::Array<int64_t> ret;
     LayoutDecision src = InitialLayoutDecision(layout.ndim());
     for (size_t i = 0; i < layout.ndim(); ++i) {
       ret.push_back(static_cast<int64_t>(src->layout.IndexOf(layout[i])));
     }
     return ret;
   }

   IndexMap LayoutIndexMap(int ndim, const SLayout& src_layout, const SLayout& desired_layout) {
     tirx::SBijectiveLayout todesired(src_layout, desired_layout);
     ffi::Optional<IndexMap> inverse_index_map;

     ffi::Array<tvm::tirx::Var> initial_indices;
     ffi::Array<PrimExpr> initial_indices_expr;
     initial_indices.reserve(ndim);
     for (int i = 0; i < ndim; ++i) {
       auto var = tvm::tirx::Var("i" + std::to_string(i), PrimType::Int(32));
       initial_indices.push_back(var);
       initial_indices_expr.push_back(var.as_or_throw<PrimExpr>());
     }
     ffi::Array<PrimExpr> desired_shape = todesired.ForwardIndex(initial_indices_expr);
     return IndexMap(initial_indices.Map(
                         [](tvm::tirx::Var var) { return var.as_or_throw<tvm::tirx::PrimVar>(); }),
                     desired_shape, std::move(inverse_index_map));
   }

   Expr RewriteExpr(const Expr& expr, const NLayout& to) {
     auto fvisitleaf = [&](const Expr& expr, std::array<NLayout, 2> layouts) -> Expr {
       NLayout from = layouts[0], to = layouts[1];
       if (NLayoutEqual()(from, to) || layouts[0].LeafValue()->layout.name() == "") return expr;
       // If not both from and to are unknown, then none of them can be unknown.
       TVM_FFI_ICHECK(!NLayoutEqual()(from, LayoutDecision::InitUnknownDim()) &&
                      !NLayoutEqual()(to, LayoutDecision::InitUnknownDim()))
           << "Cannot convert when exactly one of the layouts is unknown";
       const auto* tensor = GetTypeAs<TensorTypeNode>(expr);
       TVM_FFI_ICHECK(tensor != nullptr) << "Expect a tensor, but got: " << expr;

       if (from.LeafValue()->layout.ndim() == to.LeafValue()->layout.ndim()) {
         SLayout axes = TransposeLike(InitialLayoutDecision(tensor->ndim)->layout,
                                      from.LeafValue()->layout, to.LeafValue()->layout);
         return permute_dims(expr, LayoutToIntegers(axes));
       } else {
         auto index_map = LayoutIndexMap(from.LeafValue()->layout.ndim(), from.LeafValue()->layout,
                                         to.LeafValue()->layout);
         ffi::ObjectPtr<LayoutTransformAttrs> attrs = ffi::make_object<LayoutTransformAttrs>();
         attrs->index_map = ffi::FromJSONGraph(ffi::ToJSONGraph(index_map)).as_or_throw<IndexMap>();
         const Op& layout_transform_op_ = Op::Get("relax.layout_transform");
         auto ret_expr =
             Call(Type::Missing(), layout_transform_op_, {expr}, Attrs{std::move(attrs)}, {});
         return ret_expr;
       }
     };
     return TransformTupleLeaf<LayoutDecision>(
         VarReplacer::Replace(expr, var_remap_),
         std::array<NLayout, 2>({GetNLayout(var_layout_map_, expr), to}), fvisitleaf);
   }

   ffi::Array<Expr> RewriteArgs(const ffi::Array<Expr>& args, const ffi::Array<NLayout>& to) {
     // The `ffi::Array<Expr> args` array contains both tensor and
     // non-tensor arguments, where the `ffi::Array<NLayout> to` array only
     // contains tensor arguments.  The number of tensor arguments in
     // `args` should match the full extent of `to`.

     TVM_FFI_ICHECK_LE(to.size(), args.size());

     std::vector<Expr> new_args;
     for (size_t i = 0; i < args.size(); ++i) {
       Expr arg = args[i];
       if (i < to.size()) {
         arg = RewriteExpr(arg, to[i]);
       }
       new_args.push_back(arg);
     }

     return new_args;
   }

   void VisitBinding(const Binding& binding) final {
     // Emit the binding
     ExprMutator::VisitBinding(binding);
     // The layout is default to be initial if not rewritten.
     if (var_layout_map_.find(binding->var) == var_layout_map_.end()) {
       var_layout_map_[binding->var] = InitialNLayout(binding->var);
     }
   }

   Expr VisitVars_(const Var& var) {
     // We encounter a var use outside of inferrable regions, we rewrite it to initial layout.
     return RewriteExpr(var, InitialNLayout(var));
   }

   Expr VisitExpr_(const VarNode* op) final { return VisitVars_(ffi::GetRef<Var>(op)); }

   bool HasUnknownDimTensor(const NLayout& nlayout) {
     bool find = false;
     auto fvisit = [&](const LayoutDecision& layout) {
       find = find | (NLayoutEqual()(layout, LayoutDecision::InitUnknownDim()));
     };
     ForEachLeaf<LayoutDecision>(nlayout, fvisit);
     return find;
   }

   bool HasUnknownDimTensor(const ffi::Array<Expr>& args) {
     for (const auto& arg : args) {
       if (IsNestedTensor(arg)) {
         if (HasUnknownDimTensor(GetNLayout(var_layout_map_, arg))) {
           return true;
         }
       }
     }
     return false;
   }

   ffi::Optional<InferLayoutOutput> GetInferLayoutInfo(
       const CallNode* call_node,
       const ffi::Map<ffi::String, ffi::Array<ffi::String>>& desired_layouts,
       const LayoutCb& layout_cb, const VarLayoutMap& var_layout_map) {
     const OpNode* op_node = call_node->op.as<OpNode>();
     if (op_node == nullptr) return std::nullopt;
     Op op = ffi::GetRef<Op>(op_node).as_or_throw<Op>();
     const auto attr_map = Op::GetAttrMap<FRelaxInferLayout>("FRelaxInferLayout");
     if (attr_map.count(op) && !HasUnknownDimTensor(call_node->args)) {
       // If the op has FRelaxInferLayout, and all the input tensors have known ndim
       FRelaxInferLayout f = attr_map[op];
       auto call = ffi::GetRef<Call>(call_node);
       if (layout_cb != nullptr) {
         auto custom_layouts = layout_cb(call);
         return f(call, custom_layouts, var_layout_map);
       } else {
         return f(call, desired_layouts, var_layout_map);
       }
     } else {
       // Otherwise, we use the default policy.
       return std::nullopt;
     }
   }

   void VisitBinding_(const VarBindingNode* binding, const CallNode* call_node) final {
     ffi::Optional<InferLayoutOutput> res =
         GetInferLayoutInfo(call_node, desired_layouts_, layout_cb_, var_layout_map_);
     ffi::ObjectPtr<CallNode> new_call = ffi::make_object<CallNode>(*call_node);
     new_call->ty = Type::Missing();
     if (!res.has_value() ||
         (!IsNestedTensor(binding->var) && !binding->var->IsInstance<DataflowVarNode>())) {
       // Default policy: use the initial layout.
       // When we don't have the infer layout info, or it's a non-tensor global var binding.
       std::vector<NLayout> input_layout;
       for (const auto& arg : call_node->args) {
         input_layout.push_back(InitialNLayout(arg));
       }
       ffi::Array<Expr> new_args = RewriteArgs(call_node->args, std::move(input_layout));
       new_call->args = std::move(new_args);
       ReEmitBinding(binding, builder_->Normalize(Call(new_call)));
       // update the layout map
       var_layout_map_[binding->var] = InitialNLayout(binding->var);
     } else {
       // Convert the layout according to the inferred layout output.
       ffi::Array<Expr> new_args = RewriteArgs(call_node->args, res.value()->input_layouts);
       for (const auto& [i, arg] : res.value()->new_args) {
         new_args.Set(i->value, arg);
       }
       new_call->args = std::move(new_args);

       new_call->attrs = std::move(res.value()->new_attrs);
       Expr cur_call = builder_->Normalize(Call(new_call));
       if (binding->var->IsInstance<DataflowVarNode>()) {
         // Dataflow var, we emit the rewritten call.
         ReEmitBinding(binding, cur_call);
         // update the layout map
         var_layout_map_[binding->var] = res.value()->output_layouts[0];
       } else {
         // Global var (tensor), we rewrite it to initial layout
         TVM_FFI_ICHECK(IsNestedTensor(binding->var));
         if (!NLayoutEqual()(res.value()->output_layouts[0], InitialNLayout(binding->var))) {
           Var new_var = builder_->Emit(cur_call);
           var_layout_map_[new_var] = res.value()->output_layouts[0];
           cur_call = builder_->Normalize(RewriteExpr(new_var, InitialNLayout(binding->var)));
         }
         ReEmitBinding(binding, cur_call);
         // update the layout map
         var_layout_map_[binding->var] = InitialNLayout(binding->var);
       }
     }
   }

   void VisitBinding_(const VarBindingNode* binding, const TupleNode* val) final {
     std::vector<NLayout> input_layout;
     for (const auto& field : val->fields) {
       if (binding->var->IsInstance<DataflowVarNode>()) {
         // Df var: Use the current realized layout to group the tuple;
         input_layout.push_back(GetNLayout(var_layout_map_, field));
       } else {
         // Global var: Use the initial layout to group the tuple;
         input_layout.push_back(InitialNLayout(field));
       }
     }
     ffi::Array<Expr> new_fields = RewriteArgs(val->fields, std::move(input_layout));
     if (IsNestedTensor(binding->var)) {
       ReEmitBinding(binding, builder_->Normalize(Tuple(new_fields)));
       var_layout_map_[binding->var] = input_layout;
     }
   }

   void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val) final {
     NLayout input_layout = binding->var->IsInstance<DataflowVarNode>()
                                ? GetNLayout(var_layout_map_, val->tuple)
                                : InitialNLayout(val->tuple);
     ReEmitBinding(binding, builder_->Normalize(
                                TupleGetItem(RewriteExpr(val->tuple, input_layout), val->index)));
     // update the layout map
     var_layout_map_[binding->var] = input_layout.NestedArray()[val->index];
   }

   void VisitBinding_(const MatchCastNode* binding) final {
     if (!binding->var->IsInstance<DataflowVarNode>()) {
       ExprMutator::VisitBinding_(binding);
       return;
     }
     NLayout from_layout = InitialNLayout(binding->value);
     NLayout input_layout = GetNLayout(var_layout_map_, binding->value);
     auto fvisitleaf = [&](const Type& ty, std::array<NLayout, 2> layouts) -> Type {
       NLayout from = layouts[0], to = layouts[1];
       if (NLayoutEqual()(from, to)) return ty;
       // If not both from and to are unknown, then none of them can be unknown.
       TVM_FFI_ICHECK(!NLayoutEqual()(from, LayoutDecision::InitUnknownDim()) &&
                      !NLayoutEqual()(to, LayoutDecision::InitUnknownDim()))
           << "Cannot convert when exactly one of the layouts is unknown";
       const TensorTypeNode* tensor_ty = ty.as<TensorTypeNode>();
       TVM_FFI_ICHECK(tensor_ty != nullptr) << "We can not set layout for non-tensor struct";
       if (!tensor_ty->shape.has_value()) return ty;
       const ShapeExprNode* shape = tensor_ty->shape.value().as<ShapeExprNode>();
       if (shape == nullptr) return ty;
       TVM_FFI_ICHECK_EQ(shape->values.size(), to.LeafValue()->layout.ndim());
       std::vector<PrimExpr> new_shape;
       for (size_t i = 0; i < shape->values.size(); ++i) {
         new_shape.push_back(
             shape->values[from.LeafValue()->layout.IndexOf(to.LeafValue()->layout[i])]);
       }
       VDevice vdev = tensor_ty->vdevice.value_or(VDevice());
       return TensorType(ShapeExpr(new_shape), tensor_ty->dtype, vdev, tensor_ty->span);
     };
     Type new_ty = TransformTupleLeaf<LayoutDecision>(
         binding->ty, std::array<NLayout, 2>({from_layout, input_layout}), fvisitleaf);
     // re-emit old binding if nothing changes
     if (new_ty.same_as(binding->ty)) {
       builder_->EmitNormalized(ffi::GetRef<MatchCast>(binding));
     } else {
       Var new_var = builder_->EmitMatchCast(RewriteExpr(binding->value, input_layout), new_ty);
       var_layout_map_[binding->var] = input_layout;
       this->var_remap_[binding->var] = new_var;
     }
   }

   std::unordered_map<Var, NLayout> var_layout_map_;
   ffi::Map<ffi::String, ffi::Array<ffi::String>> desired_layouts_;
   LayoutCb layout_cb_;
 };  // namespace relax

 DataflowBlock ConvertLayoutPass(const DataflowBlock& df_block,
                                 ffi::Map<ffi::String, ffi::Array<ffi::String>> desired_layouts,
                                 LayoutCb layout_cb) {
   LayoutConvertMutator mutator(desired_layouts, layout_cb);
   return mutator.VisitBindingBlock(df_block).as_or_throw<DataflowBlock>();
 }

 namespace transform {

 Pass ConvertLayout(ffi::Map<ffi::String, ffi::Array<ffi::String>> desired_layouts,
                    LayoutCb layout_cb) {
   ffi::TypedFunction<DataflowBlock(DataflowBlock, IRModule, PassContext)> pass_func =
       [=](DataflowBlock df_block, IRModule m, PassContext pc) {
         return ConvertLayoutPass(df_block, desired_layouts, layout_cb);
       };
   return CreateDataflowBlockPass(pass_func, 0, "ConvertLayout", {});
 }

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

 }  // namespace transform
 }  // namespace relax
 }  // 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 src/relax/transform/convert_layout.cc
	* \brief Automatic layout conversion pass, especially for axis swapping.
	*/

	#include <tvm/ffi/cast.h>
	#include <tvm/ffi/extra/serialization.h>
	#include <tvm/ffi/reflection/registry.h>
	#include <tvm/relax/expr_functor.h>
	#include <tvm/relax/nested_msg.h>
	#include <tvm/relax/op_attr_types.h>
	#include <tvm/relax/transform.h>
	#include <tvm/tirx/index_map.h>

	#include "../op/tensor/manipulate.h"
	#include "infer_layout_utils.h"
	#include "utils.h"

	namespace tvm {
	namespace relax {

	using tirx::IndexMap;
	using tirx::SLayout;
	using LayoutCb = tvm::relax::transform::LayoutCb;

	/*!
	* \brief Main logic to convert the layout of conv2d. Other ops
	* can adapt to such layout conversion following conv2d accordingly.
	*
	* Structurally speaking, a Relax function is composed of a series of VarBinding and
	* MatchCast. And a specific class of VarBindings is the basic unit we want to rewrite.
	* Formally, they are of the form:
	*
	* var = Call(Op, [args], attrs)
	*
	* where Op is a specific op we want to rewrite, and attrs is the attributes of the op.
	* var and args are all exprs with type Tensor or Tuple of Tensors. They might
	* be vars, constants, or Tuple of vars and constants.
	*
	* We register the layout inference function for each op (FRelaxInferLayout), which accepts the
	* current call, the desired layout of conv2d ops, and the layout map of previous vars. The result
	* of the layout inference function is contained in an InferLayoutOutput object, which contains 3
	* fields: input_layouts, output_layouts, and attr, which represents the expected input layout,
	* output_layout and converted attrs of the new op call.
	*
	* The rewrite pass does the rewriting in a single forward pass, where for each Call(Op),
	* we collect the current SLayout of each input var, and let the InferLayout function to infer the
	* desired layout of the output. The rewriter will use these info to convert
	* the layout of inputs and attrs of the op call, and note down the new layout of the output.
	*
	* The desired layout of conv2d ops is a map from the name of the op to the desired layout of the
	* desired feature map, weight and output. For example, if we want to convert the layout of conv2d
	* from NCHW to NHWC, we can set the desired layout of conv2d to be {"conv2d": ["NHWC", "OHWI"]}.
	*
	* The way we represent the layout of a var is a NLayout object, which is a nested tuple of SLayout.
	* The incoming layout of the module will be set as the default layout (We use ABCD... as the
	* default) Note that for operators like conv, pool, people typically use NHWC to refer to the axes.
	* But to be generic and support more operators, we use ABCD... to refer to the axes.
	*
	* Note that currently the layout conversion of conv2d only support axis swapping, such as NCHW to
	* NWHC. Packed layout such as NCHW to NCHW4c is not supported now.
	*/
	class LayoutConvertMutator : public ExprMutator {
	public:
	explicit LayoutConvertMutator(
	const ffi::Map<ffi::String, ffi::Array<ffi::String>>& desired_layouts, LayoutCb layout_cb)
	: desired_layouts_(desired_layouts), layout_cb_(layout_cb) {}

	private:
	ffi::Array<int64_t> LayoutToIntegers(const SLayout& layout) {
	ffi::Array<int64_t> ret;
	LayoutDecision src = InitialLayoutDecision(layout.ndim());
	for (size_t i = 0; i < layout.ndim(); ++i) {
	ret.push_back(static_cast<int64_t>(src->layout.IndexOf(layout[i])));
	}
	return ret;
	}

	IndexMap LayoutIndexMap(int ndim, const SLayout& src_layout, const SLayout& desired_layout) {
	tirx::SBijectiveLayout todesired(src_layout, desired_layout);
	ffi::Optional<IndexMap> inverse_index_map;

	ffi::Array<tvm::tirx::Var> initial_indices;
	ffi::Array<PrimExpr> initial_indices_expr;
	initial_indices.reserve(ndim);
	for (int i = 0; i < ndim; ++i) {
	auto var = tvm::tirx::Var("i" + std::to_string(i), PrimType::Int(32));
	initial_indices.push_back(var);
	initial_indices_expr.push_back(var.as_or_throw<PrimExpr>());
	}
	ffi::Array<PrimExpr> desired_shape = todesired.ForwardIndex(initial_indices_expr);
	return IndexMap(initial_indices.Map(
	[](tvm::tirx::Var var) { return var.as_or_throw<tvm::tirx::PrimVar>(); }),
	desired_shape, std::move(inverse_index_map));
	}

	Expr RewriteExpr(const Expr& expr, const NLayout& to) {
	auto fvisitleaf = [&](const Expr& expr, std::array<NLayout, 2> layouts) -> Expr {
	NLayout from = layouts[0], to = layouts[1];
	if (NLayoutEqual()(from, to) \|\| layouts[0].LeafValue()->layout.name() == "") return expr;
	// If not both from and to are unknown, then none of them can be unknown.
	TVM_FFI_ICHECK(!NLayoutEqual()(from, LayoutDecision::InitUnknownDim()) &&
	!NLayoutEqual()(to, LayoutDecision::InitUnknownDim()))
	<< "Cannot convert when exactly one of the layouts is unknown";
	const auto* tensor = GetTypeAs<TensorTypeNode>(expr);
	TVM_FFI_ICHECK(tensor != nullptr) << "Expect a tensor, but got: " << expr;

	if (from.LeafValue()->layout.ndim() == to.LeafValue()->layout.ndim()) {
	SLayout axes = TransposeLike(InitialLayoutDecision(tensor->ndim)->layout,
	from.LeafValue()->layout, to.LeafValue()->layout);
	return permute_dims(expr, LayoutToIntegers(axes));
	} else {
	auto index_map = LayoutIndexMap(from.LeafValue()->layout.ndim(), from.LeafValue()->layout,
	to.LeafValue()->layout);
	ffi::ObjectPtr<LayoutTransformAttrs> attrs = ffi::make_object<LayoutTransformAttrs>();
	attrs->index_map = ffi::FromJSONGraph(ffi::ToJSONGraph(index_map)).as_or_throw<IndexMap>();
	const Op& layout_transform_op_ = Op::Get("relax.layout_transform");
	auto ret_expr =
	Call(Type::Missing(), layout_transform_op_, {expr}, Attrs{std::move(attrs)}, {});
	return ret_expr;
	}
	};
	return TransformTupleLeaf<LayoutDecision>(
	VarReplacer::Replace(expr, var_remap_),
	std::array<NLayout, 2>({GetNLayout(var_layout_map_, expr), to}), fvisitleaf);
	}

	ffi::Array<Expr> RewriteArgs(const ffi::Array<Expr>& args, const ffi::Array<NLayout>& to) {
	// The `ffi::Array<Expr> args` array contains both tensor and
	// non-tensor arguments, where the `ffi::Array<NLayout> to` array only
	// contains tensor arguments. The number of tensor arguments in
	// `args` should match the full extent of `to`.

	TVM_FFI_ICHECK_LE(to.size(), args.size());

	std::vector<Expr> new_args;
	for (size_t i = 0; i < args.size(); ++i) {
	Expr arg = args[i];
	if (i < to.size()) {
	arg = RewriteExpr(arg, to[i]);
	}
	new_args.push_back(arg);
	}

	return new_args;
	}

	void VisitBinding(const Binding& binding) final {
	// Emit the binding
	ExprMutator::VisitBinding(binding);
	// The layout is default to be initial if not rewritten.
	if (var_layout_map_.find(binding->var) == var_layout_map_.end()) {
	var_layout_map_[binding->var] = InitialNLayout(binding->var);
	}
	}

	Expr VisitVars_(const Var& var) {
	// We encounter a var use outside of inferrable regions, we rewrite it to initial layout.
	return RewriteExpr(var, InitialNLayout(var));
	}

	Expr VisitExpr_(const VarNode* op) final { return VisitVars_(ffi::GetRef<Var>(op)); }

	bool HasUnknownDimTensor(const NLayout& nlayout) {
	bool find = false;
	auto fvisit = [&](const LayoutDecision& layout) {
	find = find \| (NLayoutEqual()(layout, LayoutDecision::InitUnknownDim()));
	};
	ForEachLeaf<LayoutDecision>(nlayout, fvisit);
	return find;
	}

	bool HasUnknownDimTensor(const ffi::Array<Expr>& args) {
	for (const auto& arg : args) {
	if (IsNestedTensor(arg)) {
	if (HasUnknownDimTensor(GetNLayout(var_layout_map_, arg))) {
	return true;
	}
	}
	}
	return false;
	}

	ffi::Optional<InferLayoutOutput> GetInferLayoutInfo(
	const CallNode* call_node,
	const ffi::Map<ffi::String, ffi::Array<ffi::String>>& desired_layouts,
	const LayoutCb& layout_cb, const VarLayoutMap& var_layout_map) {
	const OpNode* op_node = call_node->op.as<OpNode>();
	if (op_node == nullptr) return std::nullopt;
	Op op = ffi::GetRef<Op>(op_node).as_or_throw<Op>();
	const auto attr_map = Op::GetAttrMap<FRelaxInferLayout>("FRelaxInferLayout");
	if (attr_map.count(op) && !HasUnknownDimTensor(call_node->args)) {
	// If the op has FRelaxInferLayout, and all the input tensors have known ndim
	FRelaxInferLayout f = attr_map[op];
	auto call = ffi::GetRef<Call>(call_node);
	if (layout_cb != nullptr) {
	auto custom_layouts = layout_cb(call);
	return f(call, custom_layouts, var_layout_map);
	} else {
	return f(call, desired_layouts, var_layout_map);
	}
	} else {
	// Otherwise, we use the default policy.
	return std::nullopt;
	}
	}

	void VisitBinding_(const VarBindingNode* binding, const CallNode* call_node) final {
	ffi::Optional<InferLayoutOutput> res =
	GetInferLayoutInfo(call_node, desired_layouts_, layout_cb_, var_layout_map_);
	ffi::ObjectPtr<CallNode> new_call = ffi::make_object<CallNode>(*call_node);
	new_call->ty = Type::Missing();
	if (!res.has_value() \|\|
	(!IsNestedTensor(binding->var) && !binding->var->IsInstance<DataflowVarNode>())) {
	// Default policy: use the initial layout.
	// When we don't have the infer layout info, or it's a non-tensor global var binding.
	std::vector<NLayout> input_layout;
	for (const auto& arg : call_node->args) {
	input_layout.push_back(InitialNLayout(arg));
	}
	ffi::Array<Expr> new_args = RewriteArgs(call_node->args, std::move(input_layout));
	new_call->args = std::move(new_args);
	ReEmitBinding(binding, builder_->Normalize(Call(new_call)));
	// update the layout map
	var_layout_map_[binding->var] = InitialNLayout(binding->var);
	} else {
	// Convert the layout according to the inferred layout output.
	ffi::Array<Expr> new_args = RewriteArgs(call_node->args, res.value()->input_layouts);
	for (const auto& [i, arg] : res.value()->new_args) {
	new_args.Set(i->value, arg);
	}
	new_call->args = std::move(new_args);

	new_call->attrs = std::move(res.value()->new_attrs);
	Expr cur_call = builder_->Normalize(Call(new_call));
	if (binding->var->IsInstance<DataflowVarNode>()) {
	// Dataflow var, we emit the rewritten call.
	ReEmitBinding(binding, cur_call);
	// update the layout map
	var_layout_map_[binding->var] = res.value()->output_layouts[0];
	} else {
	// Global var (tensor), we rewrite it to initial layout
	TVM_FFI_ICHECK(IsNestedTensor(binding->var));
	if (!NLayoutEqual()(res.value()->output_layouts[0], InitialNLayout(binding->var))) {
	Var new_var = builder_->Emit(cur_call);
	var_layout_map_[new_var] = res.value()->output_layouts[0];
	cur_call = builder_->Normalize(RewriteExpr(new_var, InitialNLayout(binding->var)));
	}
	ReEmitBinding(binding, cur_call);
	// update the layout map
	var_layout_map_[binding->var] = InitialNLayout(binding->var);
	}
	}
	}

	void VisitBinding_(const VarBindingNode* binding, const TupleNode* val) final {
	std::vector<NLayout> input_layout;
	for (const auto& field : val->fields) {
	if (binding->var->IsInstance<DataflowVarNode>()) {
	// Df var: Use the current realized layout to group the tuple;
	input_layout.push_back(GetNLayout(var_layout_map_, field));
	} else {
	// Global var: Use the initial layout to group the tuple;
	input_layout.push_back(InitialNLayout(field));
	}
	}
	ffi::Array<Expr> new_fields = RewriteArgs(val->fields, std::move(input_layout));
	if (IsNestedTensor(binding->var)) {
	ReEmitBinding(binding, builder_->Normalize(Tuple(new_fields)));
	var_layout_map_[binding->var] = input_layout;
	}
	}

	void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val) final {
	NLayout input_layout = binding->var->IsInstance<DataflowVarNode>()
	? GetNLayout(var_layout_map_, val->tuple)
	: InitialNLayout(val->tuple);
	ReEmitBinding(binding, builder_->Normalize(
	TupleGetItem(RewriteExpr(val->tuple, input_layout), val->index)));
	// update the layout map
	var_layout_map_[binding->var] = input_layout.NestedArray()[val->index];
	}

	void VisitBinding_(const MatchCastNode* binding) final {
	if (!binding->var->IsInstance<DataflowVarNode>()) {
	ExprMutator::VisitBinding_(binding);
	return;
	}
	NLayout from_layout = InitialNLayout(binding->value);
	NLayout input_layout = GetNLayout(var_layout_map_, binding->value);
	auto fvisitleaf = [&](const Type& ty, std::array<NLayout, 2> layouts) -> Type {
	NLayout from = layouts[0], to = layouts[1];
	if (NLayoutEqual()(from, to)) return ty;
	// If not both from and to are unknown, then none of them can be unknown.
	TVM_FFI_ICHECK(!NLayoutEqual()(from, LayoutDecision::InitUnknownDim()) &&
	!NLayoutEqual()(to, LayoutDecision::InitUnknownDim()))
	<< "Cannot convert when exactly one of the layouts is unknown";
	const TensorTypeNode* tensor_ty = ty.as<TensorTypeNode>();
	TVM_FFI_ICHECK(tensor_ty != nullptr) << "We can not set layout for non-tensor struct";
	if (!tensor_ty->shape.has_value()) return ty;
	const ShapeExprNode* shape = tensor_ty->shape.value().as<ShapeExprNode>();
	if (shape == nullptr) return ty;
	TVM_FFI_ICHECK_EQ(shape->values.size(), to.LeafValue()->layout.ndim());
	std::vector<PrimExpr> new_shape;
	for (size_t i = 0; i < shape->values.size(); ++i) {
	new_shape.push_back(
	shape->values[from.LeafValue()->layout.IndexOf(to.LeafValue()->layout[i])]);
	}
	VDevice vdev = tensor_ty->vdevice.value_or(VDevice());
	return TensorType(ShapeExpr(new_shape), tensor_ty->dtype, vdev, tensor_ty->span);
	};
	Type new_ty = TransformTupleLeaf<LayoutDecision>(
	binding->ty, std::array<NLayout, 2>({from_layout, input_layout}), fvisitleaf);
	// re-emit old binding if nothing changes
	if (new_ty.same_as(binding->ty)) {
	builder_->EmitNormalized(ffi::GetRef<MatchCast>(binding));
	} else {
	Var new_var = builder_->EmitMatchCast(RewriteExpr(binding->value, input_layout), new_ty);
	var_layout_map_[binding->var] = input_layout;
	this->var_remap_[binding->var] = new_var;
	}
	}

	std::unordered_map<Var, NLayout> var_layout_map_;
	ffi::Map<ffi::String, ffi::Array<ffi::String>> desired_layouts_;
	LayoutCb layout_cb_;
	}; // namespace relax

	DataflowBlock ConvertLayoutPass(const DataflowBlock& df_block,
	ffi::Map<ffi::String, ffi::Array<ffi::String>> desired_layouts,
	LayoutCb layout_cb) {
	LayoutConvertMutator mutator(desired_layouts, layout_cb);
	return mutator.VisitBindingBlock(df_block).as_or_throw<DataflowBlock>();
	}

	namespace transform {

	Pass ConvertLayout(ffi::Map<ffi::String, ffi::Array<ffi::String>> desired_layouts,
	LayoutCb layout_cb) {
	ffi::TypedFunction<DataflowBlock(DataflowBlock, IRModule, PassContext)> pass_func =
	[=](DataflowBlock df_block, IRModule m, PassContext pc) {
	return ConvertLayoutPass(df_block, desired_layouts, layout_cb);
	};
	return CreateDataflowBlockPass(pass_func, 0, "ConvertLayout", {});
	}

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

	} // namespace transform
	} // namespace relax
	} // namespace tvm