alter_op_layout.cc 13.4 KB
Newer Older
1 2 3 4 5 6 7 8
/*
 * 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
9
 *
10
 *   http://www.apache.org/licenses/LICENSE-2.0
11
 *
12 13 14 15 16 17 18 19
 * 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.
 */

20
/*!
21
 * Copyright (c) 2019 by Contributors
22 23 24 25 26
 * \file alter_op_layout.cc
 * \brief Alternate the layouts of operators or replace primitive operators with
          other expressions. This pass can be used for computing convolution in
          custom layouts or other general weight pre-transformation.
 */
Zhi committed
27 28
#include <tvm/relay/analysis.h>
#include <tvm/relay/transform.h>
29 30
#include <tvm/relay/op_attr_types.h>
#include <tvm/relay/attrs/transform.h>
31
#include <tvm/relay/transform.h>
32
#include <tvm/operation.h>
33 34 35 36
#include <tuple>
#include <vector>
#include <functional>
#include <string>
37 38
#include <utility>
#include <unordered_map>
39 40

#include "alter_op_layout.h"
41
#include "pattern_util.h"
42 43 44 45 46 47 48

namespace tvm {
namespace relay {

namespace alter_op_layout {

// Make a transform CallNode
49 50 51 52 53
/* Performs 2 operations
 * 1) If src_layout ndim is smaller then dst_layout, expand_dim is inserted to match the dim size.
 *    For example, src_layout = C, dst_layout = NCHW16c. The src is expanded to NHWC.
 * 2) Call layout transform with new src layout.
 */
54
Expr TransformLayout(Expr raw, Layout src_layout, Layout dst_layout) {
55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77
  if (src_layout.Equals(dst_layout)) {
    return raw;
  }

  // 1) Check if the shape lengths are different. If yes, expand dims.
  Expr input_expr = raw;
  Layout new_src_layout = src_layout;
  if (src_layout.ndim_primal() < dst_layout.ndim_primal()) {
    int num_new_axis = dst_layout.ndim_primal() - src_layout.ndim_primal();
    new_src_layout = src_layout.ExpandPrimal(dst_layout);
    input_expr = MakeExpandDims(input_expr, 0, num_new_axis);
    if (new_src_layout.Equals(dst_layout)) {
      return input_expr;
    }
  }

  // 2) Insert layout transform on the transformed src.
  CHECK(new_src_layout.defined() && dst_layout.defined())
      << "Cannot insert layout transform because there are undefined layouts";
  CHECK(BijectiveLayoutNode::make(new_src_layout, dst_layout).defined())
      << "Cannot insert layout transform because there are inconvertible layouts: "
      << new_src_layout << " v.s. " << dst_layout;
  return MakeLayoutTransform(input_expr, new_src_layout.name(), dst_layout.name());
78 79 80 81 82 83 84
}

// Memorize layout transform so we can reuse internal transformed nodes
class TransformMemorizerNode : public Node {
 public:
  // map from (Expr, src_layout, dst_layout) to transformed Expr
  using TransformKey = std::tuple<const Node*, std::string, std::string>;
85
struct key_hash : public std::function<std::size_t(TransformKey)> {
86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208
    std::size_t operator()(const TransformKey& k) const {
      return dmlc::HashCombine<std::string>(dmlc::HashCombine<std::string>(
              std::hash<const Node*>()(std::get<0>(k)), std::get<1>(k)), (std::get<2>(k)));
    }
  };

  std::unordered_map<TransformKey, Expr, key_hash> memo;
  static constexpr const char *_type_key = "relay.alter_op_layout.TransformMemorizerNode";
  TVM_DECLARE_NODE_TYPE_INFO(TransformMemorizerNode, Node);
};

class TransformMemorizer : public NodeRef {
 public:
  TransformMemorizer() {}
  explicit TransformMemorizer(NodePtr<Node> n) : NodeRef(n) {}

  TransformMemorizerNode* operator->() {
    return static_cast<TransformMemorizerNode*>(node_.get());
  }

  // Transform layout with memorizer
  Expr Transform(Expr raw, const Layout& src_layout, const Layout& dst_layout) {
    if (src_layout.Equals(dst_layout)) { return raw; }

    std::tuple<const Node*, std::string, std::string> key =
        std::make_tuple<>(raw.get(), src_layout.name(), dst_layout.name());
    auto& memo = operator->()->memo;

    auto iter = memo.find(key);
    if (iter != memo.end()) {
      return iter->second;
    } else {
      Expr transform = TransformLayout(raw, src_layout, dst_layout);
      memo[key] = transform;
      return transform;
    }
  }

  using ContainerType = TransformMemorizerNode;
};


// TempExprNode during layout transform
// Instance of this expr will be Realized to normal expr ultimately
class LayoutAlternatedExprNode : public TempExprNode {
 public:
  Expr value;
  Layout old_layout;
  Layout new_layout;
  TransformMemorizer memorizer;

  Expr Realize() const final {
    // NOTE: use a copy to discard the "const" qualifier
    TransformMemorizer tmp_memorizer = memorizer;
    // fallback to old layout
    return tmp_memorizer.Transform(value, new_layout, old_layout);
  }

  void VisitAttrs(AttrVisitor *v) final {
    v->Visit("value", &value);
    v->Visit("old_layout", &old_layout);
    v->Visit("new_layout", &new_layout);
  }

  static constexpr const char *_type_key = "relay.alter_op_layout.LayoutAlternatedExprNode";
  TVM_DECLARE_NODE_TYPE_INFO(LayoutAlternatedExprNode, TempExprNode);
};

RELAY_DEFINE_NODE_REF(LayoutAlternatedExpr, LayoutAlternatedExprNode, TempExpr);

// Call registered FInferCorrectLayout of an op.
// Parameters are the same as the parameters for FInferCorrectLayout
// Returns inferred_input_layout, inferred_output_layout, success
std::tuple<Array<Layout>, Array<Layout>, bool> CallInfer(
    const Call& call,
    const Array<Layout>& new_in_layouts,
    const Array<Layout>& old_in_layouts,
    const Array<Array<IndexExpr> > &old_in_shapes) {
  static auto finfer_layout = Op::GetAttr<FInferCorrectLayout>("FInferCorrectLayout");

  Op op = Downcast<Op>(call->op);
  if (finfer_layout.count(op)) {
    Array<Array<Layout> > inferred_layouts;
    inferred_layouts = finfer_layout[op](call->attrs, new_in_layouts,
                                         old_in_layouts, old_in_shapes);
    CHECK_EQ(inferred_layouts.size(), 2)
      << "FInferCorrectLayout should return an array with size of 2";
    for (auto x : inferred_layouts) {
      for (auto y : x) {
        if (!y.defined()) {  // inference fails
          return std::make_tuple<>(Array<Layout>(nullptr), Array<Layout>(nullptr), false);
        }
      }
    }
    return std::make_tuple<>(inferred_layouts[0], inferred_layouts[1], true);
  } else {
    return std::make_tuple<>(Array<Layout>(nullptr), Array<Layout>(nullptr), false);
  }
}

// Call registered FTVMAlterOpLayout of an op
// Returns the altered expression
Call CallAlter(const Call& ref_call,
               const std::vector<Expr>& new_args) {
  static auto falter_layout = Op::GetAttr<FTVMAlterOpLayout>("FTVMAlterOpLayout");
  Op op = Downcast<Op>(ref_call->op);

  Expr new_e;
  bool modified = false;
  if (falter_layout.count(op)) {
    tvm::Array<tvm::Tensor> tinfos;
    for (auto expr : ref_call->args) {
      auto ttype = expr->type_as<TensorTypeNode>();
      tinfos.push_back(tvm::placeholder(ttype->shape, ttype->dtype));
    }
    Expr altered_value = falter_layout[op](ref_call->attrs, new_args, tinfos);
    if (altered_value.defined()) {
      new_e = altered_value;
      modified = true;
    }
  }
  if (!modified) {
    new_e = CallNode::make(ref_call->op, new_args,
209
                           ref_call->attrs);
210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226
  }

  const CallNode *new_call = new_e.as<CallNode>();
  CHECK(new_call) << "Can only replace the original operator with another call node";
  return GetRef<Call>(new_call);
}

Expr AlterOpLayoutRewrite(const Call &ref_call,
                          const Array<Expr> &new_args,
                          const NodeRef& ctx) {
  std::vector<LayoutAlternatedExpr> inputs;
  std::vector<Expr> normal_new_args;
  Array<Array<IndexExpr> > input_shapes;

  // NOTE: discard the "const" qualifier
  TransformMemorizer memorizer = Downcast<TransformMemorizer>(ctx);

227 228 229 230 231 232 233 234 235 236 237 238 239 240 241
  // fill incomplete state and flatten tuple
  auto push_back_one_arg = [&inputs, memorizer](Expr arg) {
    // We always expect LayoutAlternatedExpr.
    // This is used to convert the normal Expr to LayoutAlternatedExpr.
    if (const LayoutAlternatedExprNode *inp = arg.as<LayoutAlternatedExprNode>()) {
      inputs.push_back(GetRef<LayoutAlternatedExpr>(inp));
      return inp->value;
    } else {
      auto inode = make_node<LayoutAlternatedExprNode>();
      inode->value = arg;
      inode->memorizer = memorizer;
      inputs.push_back(LayoutAlternatedExpr(inode));
      return arg;
    }
  };
242

243 244
  for (auto new_arg : new_args) {
    // NOTE: do not support nested tuple
245 246
    if (new_arg->is_type<TupleNode>()) {
      Tuple tuple_new_arg = Downcast<Tuple>(new_arg);
247
      std::vector<Expr> fields;
248
      for (auto x : tuple_new_arg->fields) {
249 250
        Expr tmp = push_back_one_arg(x);
        fields.push_back(tmp);
251
      }
252
      normal_new_args.push_back(TupleNode::make(fields));
253
    } else {
254 255
      Expr tmp = push_back_one_arg(new_arg);
      normal_new_args.push_back(tmp);
256 257 258 259 260 261 262 263 264 265 266
    }
  }

  // old_in, new_in = state[inputs]
  Array<Layout> old_in, old_out, new_in, new_out, new_in2;
  for (auto inp : inputs) {
    old_in.push_back(inp->old_layout);
    new_in.push_back(inp->new_layout);
  }

  for (auto arg : ref_call->args) {
267
    if (arg->is_type<TupleNode>()) {  // flatten tuple
268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310
      Tuple tuple_arg = Downcast<Tuple>(arg);
      for (auto x : tuple_arg->fields) {
        input_shapes.push_back(x->type_as<TensorTypeNode>()->shape);
      }
    } else {
      input_shapes.push_back(arg->type_as<TensorTypeNode>()->shape);
    }
  }

  // old_in, old_out = op.infer(old_in)
  bool success = false;
  std::tie(old_in, old_out, success) = CallInfer(ref_call,
                                                 Array<Layout>(nullptr),
                                                 old_in, input_shapes);
  if (!success) { return Expr(nullptr); }
  CHECK_EQ(old_in.size(), new_in.size());

  // if new_in == 'undef':  new_in = old_in
  for (size_t i = 0; i < new_in.size(); ++i) {
    if (!new_in[i].defined()) {
      new_in.Set(i, old_in[i]);
    }
  }

  // new_op = alter(op)
  Call new_call = CallAlter(ref_call, normal_new_args);

  // new_in2, new_out = op.infer(new_in)
  if (new_call->op->is_type<OpNode>()) {
    success = false;
    std::tie(new_in2, new_out, success) = CallInfer(new_call, new_in, old_in, input_shapes);
    if (!success) { return Expr(nullptr); }
  } else {
    return Expr(nullptr);
  }

  CHECK_EQ(new_out.size(), old_out.size())
    << "The number of output nodes should keep the same during alter_op_layout";
  CHECK_EQ(new_in.size(), new_in2.size())
    << "The number of input nodes should keep the same during alter_op_layout";

  // if (new_in != new_in2): insert transform (new_in -> new_in2)
  Array<Expr> transformed_args;
311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326
  size_t pt = 0;
  for (auto arg : new_call->args) {
    if (arg->is_type<TupleNode>()) {  // unflatten tuple
      Tuple tuple_arg = Downcast<Tuple>(arg);
      std::vector<Expr> transformed_tuple_arg;
      for (auto arg_item : tuple_arg->fields) {
          transformed_tuple_arg.push_back(
                  memorizer.Transform(arg_item, new_in[pt], new_in2[pt]));
          pt++;
      }
      transformed_args.push_back(TupleNode::make(transformed_tuple_arg));
    } else {
      transformed_args.push_back(
              memorizer.Transform(arg, new_in[pt], new_in2[pt]));
      pt++;
    }
327
  }
328
  CHECK_EQ(pt, inputs.size());
329 330

  // state[node] = (old_out, new_out)
331
  // (handle tuple output)
332 333
  if (ref_call->checked_type()->is_type<TupleTypeNode>()) {
    Expr tuple_output = CallNode::make(new_call->op, transformed_args,
334
                                       new_call->attrs);
335 336 337 338 339 340 341 342 343 344 345 346 347 348
    Array<Expr> fields;
    for (size_t i = 0; i < new_out.size(); ++i) {
      auto rnode = make_node<LayoutAlternatedExprNode>();
      rnode->value = TupleGetItemNode::make(tuple_output, i);
      rnode->old_layout = old_out[i];
      rnode->new_layout = new_out[i];
      rnode->memorizer = memorizer;
      fields.push_back(Expr(rnode));
    }
    return TupleNode::make(fields);
  } else {
    auto rnode = make_node<LayoutAlternatedExprNode>();
    CHECK_EQ(new_out.size(), 1);
    rnode->value = CallNode::make(new_call->op, transformed_args,
349
                                  new_call->attrs);
350 351 352 353 354 355 356
    rnode->old_layout = old_out[0];
    rnode->new_layout = new_out[0];
    rnode->memorizer = memorizer;
    return Expr(rnode);
  }
}

357 358 359
// Limiations:
// 1. the altered op should have the same number of arguments as the previous one
// 2. do not support nested tuple arguments
360
Expr AlterOpLayout(const Expr& expr) {
361 362 363 364 365
  TransformMemorizer transformMemorizer(make_node<TransformMemorizerNode>());
  auto fcontext = [&](const Call& call) -> NodeRef{
    return transformMemorizer;
  };

366 367 368
  return ForwardRewrite(expr, AlterOpLayoutRewrite, fcontext);
}

369 370
}  // namespace alter_op_layout

371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386
namespace transform {

Pass AlterOpLayout() {
  runtime::TypedPackedFunc<Function(Function, Module, PassContext)> pass_func =
    [=](Function f, Module m, PassContext pc) {
      return Downcast<Function>(relay::alter_op_layout::AlterOpLayout(f));
  };
  return CreateFunctionPass(pass_func, 3, "AlterOpLayout",
                            {ir::StringImm::make("InferType")});
}

TVM_REGISTER_API("relay._transform.AlterOpLayout")
.set_body_typed(AlterOpLayout);

}  // namespace transform

387 388
}  // namespace relay
}  // namespace tvm