tqchen commented on a change in pull request #6190: URL: https://github.com/apache/incubator-tvm/pull/6190#discussion_r465446085
########## File path: src/auto_scheduler/feature.cc ########## @@ -0,0 +1,1568 @@ +/* + * 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 auto_scheduler/feature.cc + * \brief Feature extraction for the cost model + */ + +#include <tvm/arith/analyzer.h> +#include <tvm/auto_scheduler/feature.h> +#include <tvm/auto_scheduler/measure.h> +#include <tvm/auto_scheduler/measure_record.h> +#include <tvm/runtime/registry.h> +#include <tvm/te/operation.h> +#include <tvm/te/schedule_pass.h> +#include <tvm/tir/analysis.h> +#include <tvm/tir/op_attr_types.h> +#include <tvm/tir/stmt_functor.h> +#include <tvm/tir/transform.h> + +#include <algorithm> +#include <cmath> +#include <numeric> +#include <unordered_map> +#include <vector> + +#include "utils.h" + +namespace tvm { +// import the function from driver_api.cc +extern void GetBinds(const Array<te::Tensor>& args, bool compact, + const std::unordered_map<te::Tensor, tir::Buffer>& binds, + Map<te::Tensor, tir::Buffer>* out_binds, Array<ObjectRef>* out_arg_list); +} // namespace tvm + +namespace tvm { +namespace auto_scheduler { + +using namespace tvm::tir; +using arith::Analyzer; +using arith::ConstIntBound; + +template <class T> +using BufferMap = std::unordered_map<Buffer, T, ObjectHash, ObjectEqual>; + +// The number of samples to extract for arithmetic intensity curves +static const int ARITH_INTENSITY_CURVE_SAMPLE_N = 10; + +// Annotation position encoding +enum class AnnotationPosType : int { + kPosNone = 0, + kPosInnerSpatial = 1, + kPosMiddleSpatial = 2, + kPosOuterSpatial = 3, + kPosInnerReduce = 4, + kPosMiddleReduce = 5, + kPosOuterReduce = 6, + kPosMixed = 7 +}; + +// Buffer access type +enum class BufferAccessType : int { kRead = 0, kWrite = 1, kReadWrite = 2, kUnknownRW = 3 }; + +// Accesses to a buffer +struct BufferAccess { + BufferAccessType acc_type{BufferAccessType::kUnknownRW}; + std::vector<std::vector<PrimExpr>> indices; +}; + +// Data reuse type +enum class ReuseType : int { kLoopMultipleRead = 0, kSerialMultipleReadWrite = 1, kNoReuse = 2 }; + +// Feature for an access of a buffer +struct BufferAccessFeature { + std::string buffer_name; // The name of the buffer + BufferAccessType acc_type; // The type of the access + float bytes; // touched memory in bytes + float unique_bytes; // touched unique memory in bytes + float lines; // touched cache lines + float unique_lines; // touched unique cache lines + ReuseType reuse_type; // type of data reuse + float reuse_dis_iter; // reuse distance in iterator number + float reuse_dis_bytes; // reuse distance in total touched bytes + float reuse_ct; // reuse times + float bytes_d_reuse_ct; // bytes / reuse_ct + float unique_bytes_d_reuse_ct; // unique_bytes / reuse_ct + float lines_d_reuse_ct; // lines / reuse_ct + float unique_lines_d_reuse_ct; // unique_lines / reuse_ct + float stride; // The stride in access +}; + +// Feature set of a BufferStore statement +struct FeatureSet { + // compute feature + float float_mad; // The number of float MAD (Multiply–add) ops + float float_addsub; // The number of float add and sub ops + float float_mul; // The number of float multiply ops + float float_divmod; // The number of float div and mod ops + float float_cmp; // The number of float comparison ops + float float_math_func; // The number of float math func calls + float float_other_func; // The number of other float func calls + float int_mad; // The number of integer MAD (Multiply–add) ops + float int_addsub; // The number of integer add and sub ops + float int_mul; // The number of float multiply ops + float int_divmod; // The number of float div and mod ops + float int_cmp; // The number of float comparison ops + float int_math_func; // The number of float math func calls + float int_other_func; // The number of other float func calls + float bool_op; // The number of bool ops + float select_op; // The number of select ops + float vec_num; // The number of vectorized iterators + float vec_prod; // The product of the lengths of vectorized iterators + float vec_len; // The length of the innermost vectorized iterator + AnnotationPosType vec_type; // The type of vectorizatoin position + float unroll_num; // The number of unrolled iterators + float unroll_prod; // The product of the lengths of vectorized iterators + float unroll_len; // The length of the innermost unrolled iterator + AnnotationPosType unroll_type; // The type of unroll position + float parallel_num; // The number of paralleled iterators + float parallel_prod; // The product of the lengths of paralleled iterators + float parallel_len; // The length of the innermost paralleled iterators + AnnotationPosType parallel_type; // The type of parallel position + float is_gpu; // Whether it is a GPU task + float blockIdx_x_len; // The length of blockIdx.x + float blockIdx_y_len; // The length of blockIdx.y + float blockIdx_z_len; // The length of blockIdx.z + float threadIdx_x_len; // The length of threadIdx.x + float threadIdx_y_len; // The length of threadIdx.y + float threadIdx_z_len; // The length of threadIdx.z + float vthread_len; // The length of virtual thread + + // Points sampled from the arithmetic intensity curve. + float arith_intensity_curve[ARITH_INTENSITY_CURVE_SAMPLE_N]; + + // Buffer access feature (per buffer) + std::vector<BufferAccessFeature> access_feas; + + // Allocation feature + float alloc_size; // The size of allocated buffer in bytes + float alloc_outer_prod; // The product of lenghts of loops outside the scope of the allocation + float alloc_inner_prod; // The product of lenghts of loops inside the score of the allocation + float alloc_prod; // alloc_outer_prod * alloc_inner_prod + + // Overall feature + float outer_prod; // The product of lenghts of outer loops + float num_loops; // The number of outer loops + float auto_unroll_max_step; // The value of pragma "auto_unroll_max_step" +}; + +// Return whether a var is in an expr +bool VarInExpr(const Var& var, const PrimExpr& expr) { + bool find = false; + + PostOrderVisit(expr, [&find, &var](const ObjectRef& node) { + if (find) { + return; + } + + if (const VarNode* op = node.as<VarNode>()) { + if (op == var.get()) { + find = true; + } + } + }); + + return find; +} + +// Get position encoding for annotation +AnnotationPosType GetAnnotationPosEncoding(const Var& var, const Array<PrimExpr>& spatial_args, + const Array<IterVar>& axis, + const Array<IterVar>& reduce_axis) { + // Try to match spatial args first + size_t find_i = 0; + size_t find_ct = 0; + for (size_t i = 0; i < spatial_args.size(); ++i) { + if (VarInExpr(var, spatial_args[i])) { + find_i = i; + find_ct += 1; + } + } + + if (find_ct == 0) { + // If it is not found in spacial args, then it is a reduce iterator. + // Use name to match + const std::string& var_name = var->name_hint; + for (size_t i = 0; i < reduce_axis.size(); ++i) { + if (var_name.find(reduce_axis[i]->var->name_hint) != std::string::npos) { + find_i = i; + find_ct++; + } + } + if (find_ct >= 1) { + if (find_i == 0) { + return AnnotationPosType::kPosInnerReduce; + } else if (find_i == reduce_axis.size() - 1) { + return AnnotationPosType::kPosOuterReduce; + } else { + return AnnotationPosType::kPosMiddleReduce; + } + } else { + // If the axis is not found in both spatial args and reduce axis, + // then this stage must compute_at somewhere under this aixs and this axis is simplified out + // We assume it is an outer spatial + return AnnotationPosType::kPosOuterSpatial; + } + } else if (find_ct == 1) { + if (find_i == spatial_args.size() - 1) { + return AnnotationPosType::kPosInnerSpatial; + } else if (find_i == 0) { + return AnnotationPosType::kPosOuterSpatial; + } else { + return AnnotationPosType::kPosMiddleSpatial; + } + } else { + return AnnotationPosType::kPosMixed; + } +} + +// Return the extent of a for loop +int64_t GetLoopExtent(const ForNode* node) { + auto pint = node->extent.as<IntImmNode>(); + if (pint != nullptr) { + return pint->value; + } else { + return 1; + } +} + +// Count math ops in an expr +class MathOpCounter : public StmtExprVisitor { + public: +#define VisitBinary(Type, float_ct, int_ct) \ + void VisitExpr_(const Type* op) final { \ + if (op->a.dtype().is_float()) { \ + float_ct++; \ + } else { \ + int_ct++; \ + } \ + StmtExprVisitor::VisitExpr_(op); \ + } + + VisitBinary(AddNode, float_addsub, int_addsub); + VisitBinary(SubNode, float_addsub, int_addsub); + VisitBinary(MulNode, float_mul, int_mul); + VisitBinary(DivNode, float_divmod, int_divmod); + VisitBinary(ModNode, float_divmod, int_divmod); + VisitBinary(FloorDivNode, float_divmod, int_divmod); + VisitBinary(FloorModNode, float_divmod, int_divmod); + VisitBinary(MaxNode, float_cmp, int_cmp); + VisitBinary(MinNode, float_cmp, int_cmp); + VisitBinary(EQNode, float_cmp, int_cmp); + VisitBinary(NENode, float_cmp, int_cmp); + VisitBinary(LTNode, float_cmp, int_cmp); + VisitBinary(LENode, float_cmp, int_cmp); + VisitBinary(GTNode, float_cmp, int_cmp); + VisitBinary(GENode, float_cmp, int_cmp); + + void VisitExpr_(const AndNode* op) final { + bool_op++; + StmtExprVisitor::VisitExpr_(op); + } + void VisitExpr_(const OrNode* op) final { + bool_op++; + StmtExprVisitor::VisitExpr_(op); + } + void VisitExpr_(const NotNode* op) final { + bool_op++; + StmtExprVisitor::VisitExpr_(op); + } + void VisitExpr_(const SelectNode* op) final { + select_op++; + StmtExprVisitor::VisitExpr_(op); + } + + void VisitExpr_(const CallNode* op) final { + auto* pop = op->op.as<OpNode>(); + CHECK(pop != nullptr); + auto effect_kind = op_call_effect_[GetRef<Op>(pop)]; + bool is_pure = + effect_kind == CallEffectKind::kPure || effect_kind == CallEffectKind::kExprAnnotation; + + if (is_pure) { + if (op->dtype.is_float()) { + float_math_func++; + } else { + int_math_func++; + } + } else { + if (op->dtype.is_float()) { + float_other_func++; + } else { + int_other_func++; + } + } + StmtExprVisitor::VisitExpr_(op); + } + + // todo(lmzheng): detect mad + size_t float_mad{0}, float_addsub{0}, float_mul{0}, float_divmod{0}, float_cmp{0}, Review comment: for each field, document by comment, use one field per line ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
