Skip to content

Commit

Permalink
Recover dropped commits
Browse files Browse the repository at this point in the history
[TVMScript] B4: If branch support (tlc-pack#263)
B8: Local Function Support  (tlc-pack#258)
[TVMScript] B3: Type annotation checks (tlc-pack#256)
[TVMScript][Parser] B1: Dataflow block (tlc-pack#252)
[TVMScript] B2: match shape support (tlc-pack#251)
[TVMScript] B6/B7: Symbolic shape and var shadowing  (tlc-pack#245)
[TVMScript] B5: Support relax op (tlc-pack#244)
[TVMScript] B0: Call_tir support (tlc-pack#243)
enhance parser error reporting (tlc-pack#242)
[TVMScript] A1: Relax Parser infra (tlc-pack#240)
update ci image versions. (tlc-pack#241)
[TVMScript] B2-4: TIR IRBuilder (tlc-pack#239)
[TVMScript] A0: Relax IRBuilder infra (tlc-pack#235)
[TVMScript] B5-6: TIR IRBuilder (tlc-pack#231)
[TVMScript] B1: IRBuilder (tlc-pack#228)
[TVMScript] New Parser: Part C (tlc-pack#218)
[TVMScript] New Parser: Part A (tlc-pack#221)
[TVMScript] New Parser: Part B (tlc-pack#217)

Not recovered:
[Pass] Separate ApplyHistoryBest from tuning passes (tlc-pack#226)
[Bugfix] Couple of bug fixes to run TVM-gen code together with BYOC (tlc-pack#249)

co-authored-by: Yuchen Jin <yuchenj@cs.washington.edu>
co-authored-by: Siyuan Feng <Hzfengsy@sjtu.edu.cn>
co-authored-by: Ruihang Lai <ruihangl@cs.cmu.edu>
  • Loading branch information
4 people committed Jan 29, 2023
1 parent 451abda commit 72d6f4f
Show file tree
Hide file tree
Showing 62 changed files with 4,635 additions and 353 deletions.
10 changes: 0 additions & 10 deletions include/tvm/relax/transform.h
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,6 @@
#define TVM_RELAX_TRANSFORM_H_

#include <tvm/ir/transform.h>
#include <tvm/meta_schedule/apply_history_best.h>
#include <tvm/meta_schedule/database.h>
#include <tvm/relax/expr.h>

namespace tvm {
Expand Down Expand Up @@ -119,14 +117,6 @@ TVM_DLL Pass CanonicalizeBindings();
*/
TVM_DLL Pass Normalize();

/*!
* \brief Apply the best schedule from tuning database.
*
* \return The Pass.
*/
TVM_DLL Pass MetaScheduleApplyHistoryBest(const tvm::meta_schedule::Database& database,
Target target);

/*!
* \brief Bind params of function of the module to constant tensors.
*
Expand Down
11 changes: 8 additions & 3 deletions include/tvm/script/ir_builder/ir/frame.h
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,17 @@ namespace ir {
*/
class IRModuleFrameNode : public IRBuilderFrameNode {
public:
Array<GlobalVar> global_vars;
Array<BaseFunc> functions;
/*! \brief A map from string names to global variables that ensures global uniqueness. */
Map<String, GlobalVar> global_var_map;
/*!
* \brief A map from GlobalVar to all global functions.
* \note Only defined functions are in the map, while declared functions are not included.
*/
Map<GlobalVar, BaseFunc> functions;

void VisitAttrs(tvm::AttrVisitor* v) {
IRBuilderFrameNode::VisitAttrs(v);
v->Visit("global_vars", &global_vars);
v->Visit("global_vars", &global_var_map);
v->Visit("functions", &functions);
}

Expand Down
15 changes: 15 additions & 0 deletions include/tvm/script/ir_builder/ir/ir.h
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,21 @@ namespace ir {
*/
TVM_DLL IRModuleFrame IRModule();

/*!
* \brief Declare a Function without given the specific function implementation.
* \note It is usually used in cross-function call. And we can specify the function by `DefFunction`
* \param func_name The function unique name.
* \return The corresponding GlobalVar.
*/
TVM_DLL GlobalVar DeclFunction(const String& func_name);

/*!
* \brief Define the function which is declared before.
* \param func_name The function unique name.
* \param func The given function implementation
*/
TVM_DLL void DefFunction(const String& func_name, const BaseFunc& func);

} // namespace ir
} // namespace ir_builder
} // namespace script
Expand Down
283 changes: 283 additions & 0 deletions include/tvm/script/ir_builder/relax/frame.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,283 @@
/*
* 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.
*/
#ifndef TVM_SCRIPT_IR_BUILDER_RELAX_FRAME_H_
#define TVM_SCRIPT_IR_BUILDER_RELAX_FRAME_H_

#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
#include <tvm/script/ir_builder/base.h>
#include <tvm/script/ir_builder/ir/frame.h>
#include <tvm/script/ir_builder/ir/ir.h>

namespace tvm {
namespace script {
namespace ir_builder {
namespace relax {

/*! \brief The base ir_builder frame for the relax dialect. */
class RelaxFrameNode : public IRBuilderFrameNode {
public:
void VisitAttrs(tvm::AttrVisitor* v) { IRBuilderFrameNode::VisitAttrs(v); }

static constexpr const char* _type_key = "script.ir_builder.relax.RelaxFrame";
TVM_DECLARE_BASE_OBJECT_INFO(RelaxFrameNode, IRBuilderFrameNode);
};

class RelaxFrame : public IRBuilderFrame {
public:
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(RelaxFrame, IRBuilderFrame, RelaxFrameNode);

protected:
RelaxFrame() = default;
};

/*! \brief The base ir_builder frame for frames with SeqExpr
i.e. Functions, If branches
*/
class SeqExprFrameNode : public RelaxFrameNode {
public:
/*! \brief The binding blocks inside the frame. */
Array<tvm::relax::BindingBlock> binding_blocks;
/*! \brief The frame output expr. `NullOpt` when undefined. */
Optional<tvm::relax::Expr> output;

void VisitAttrs(tvm::AttrVisitor* v) {
RelaxFrameNode::VisitAttrs(v);
v->Visit("binding_blocks", &binding_blocks);
v->Visit("output", &output);
}

static constexpr const char* _type_key = "script.ir_builder.relax.SeqExprFrame";
TVM_DECLARE_BASE_OBJECT_INFO(SeqExprFrameNode, RelaxFrameNode);

public:
void ExitWithScope() override;
};

class SeqExprFrame : public RelaxFrame {
public:
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(SeqExprFrame, RelaxFrame, SeqExprFrameNode);
};

/*! \brief The ir_builder frame for the relax function. */
class FunctionFrameNode : public SeqExprFrameNode {
public:
/*!
* \brief The function name.
* \note The name will not be specified in constructor, so it is "Optional",
* However, we must specify the name by `R.func_name` before exit this frame.
*/
Optional<String> name;
/*! \brief The function params. */
Array<tvm::relax::Var> params;
/*!
* \brief The function return type.
* \note Usually the function return type can be deduced by the function body.
* But we can use this field to specify a more "accurate" return type.
* i.e. If the `ret_type` is None, try to use the deduced type from body
* If the `ret_type` is not None, check the deduced type is a base type of the given one.
*/
Optional<Type> ret_type;
/*! \brief The function attributes. */
Map<String, ObjectRef> attrs;
/*! \brief The block builder to create Relax function. */
tvm::relax::BlockBuilder block_builder;

void VisitAttrs(tvm::AttrVisitor* v) {
SeqExprFrameNode::VisitAttrs(v);
v->Visit("name", &name);
v->Visit("params", &params);
v->Visit("ret_type", &ret_type);
v->Visit("attrs", &attrs);
v->Visit("binding_blocks", &binding_blocks);
v->Visit("output", &output);
// `block_builder` is not visited.
}

static constexpr const char* _type_key = "script.ir_builder.relax.FunctionFrame";
TVM_DECLARE_FINAL_OBJECT_INFO(FunctionFrameNode, SeqExprFrameNode);

public:
void ExitWithScope() final;
};

class FunctionFrame : public SeqExprFrame {
public:
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(FunctionFrame, SeqExprFrame, FunctionFrameNode);
};

/*! \brief The ir_builder frame for relax binding blocks. */
class BlockFrameNode : public RelaxFrameNode {
public:
/*! \brief The flag that indicates whether the block is a dataflow block. */
bool is_dataflow;
/*! \brief The variables emitted in this block. */
Array<tvm::relax::Var> emitted_vars;
/*!
* \brief (Only used for a dataflow block.) A boolean indicating if the dataflow block is ended of
* construction. If it is true, any new binding trying to be emitted into this block will cause an
* error.
*/
bool block_ended;

void VisitAttrs(tvm::AttrVisitor* v) {
RelaxFrameNode::VisitAttrs(v);
v->Visit("is_dataflow", &is_dataflow);
v->Visit("emitted_vars", &emitted_vars);
v->Visit("block_ended", &block_ended);
}

static constexpr const char* _type_key = "script.ir_builder.relax.BlockFrame";
TVM_DECLARE_FINAL_OBJECT_INFO(BlockFrameNode, RelaxFrameNode);

public:
void EnterWithScope() final;
void ExitWithScope() final;
};

class BlockFrame : public RelaxFrame {
public:
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(BlockFrame, RelaxFrame, BlockFrameNode);
};

/*!
* \brief A frame that represents if statement.
*
* \sa IfFrame
*/
class IfFrameNode : public RelaxFrameNode {
public:
/*! \brief The condition of the if statement. */
tvm::relax::Expr condition;
/*! \brief The Bindings in the true branch. */
Optional<tvm::relax::Expr> then_expr;
/*! \brief The Bindings in the false branch. */
Optional<tvm::relax::Expr> else_expr;
/*! \brief The Binding var. */
tvm::relax::Var var;
/*! \brief The binding var name. */
String var_name;

void VisitAttrs(tvm::AttrVisitor* v) {
RelaxFrameNode::VisitAttrs(v);
v->Visit("condition", &condition);
v->Visit("then_expr", &then_expr);
v->Visit("else_expr", &else_expr);
v->Visit("var", &var);
v->Visit("var_name", &var_name);
}

static constexpr const char* _type_key = "script.ir_builder.relax.IfFrame";
TVM_DECLARE_FINAL_OBJECT_INFO(IfFrameNode, RelaxFrameNode);

public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};

/*!
* \brief Managed reference to IfFrameNode.
*
* \sa IfFrameNode
*/
class IfFrame : public RelaxFrame {
public:
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IfFrame, RelaxFrame, IfFrameNode);
};

/*!
* \brief A frame that represents then.
*
* \sa ThenFrame
*/
class ThenFrameNode : public SeqExprFrameNode {
public:
static constexpr const char* _type_key = "script.ir_builder.relax.ThenFrame";
TVM_DECLARE_FINAL_OBJECT_INFO(ThenFrameNode, SeqExprFrameNode);

public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};

/*!
* \brief Managed reference to ThenFrameNode.
*
* \sa ThenFrameNode
*/
class ThenFrame : public SeqExprFrame {
public:
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(ThenFrame, SeqExprFrame, ThenFrameNode);
};

/*!
* \brief A frame that represents else.
*
* \sa ElseFrame
*/
class ElseFrameNode : public SeqExprFrameNode {
public:
static constexpr const char* _type_key = "script.ir_builder.relax.ElseFrame";
TVM_DECLARE_FINAL_OBJECT_INFO(ElseFrameNode, SeqExprFrameNode);

public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};

/*!
* \brief Managed reference to ElseFrameNode.
*
* \sa ElseFrameNode
*/
class ElseFrame : public SeqExprFrame {
public:
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(ElseFrame, SeqExprFrame, ElseFrameNode);
};

} // namespace relax
} // namespace ir_builder
} // namespace script
} // namespace tvm

#endif // TVM_SCRIPT_IR_BUILDER_RELAX_FRAME_H_
Loading

0 comments on commit 72d6f4f

Please sign in to comment.