elemwise.cc 32.2 KB
Newer Older
1 2 3 4 5 6 7 8
/*!
 *  Copyright (c) 2017 by Contributors
 * \file elemwise.cc
 * \brief Elemenwise operators
 */
#include <nnvm/op.h>
#include <nnvm/node.h>
#include <nnvm/op_attr_types.h>
9
#include <nnvm/compiler/op_attr_types.h>
10
#include <nnvm/compiler/util.h>
11
#include <nnvm/top/tensor.h>
12
#include <cmath>
13 14
#include "../op_common.h"
#include "../elemwise_op_common.h"
15 16 17
#include "topi/broadcast.h"
#include "topi/elemwise.h"
#include "topi/tags.h"
18
#include "../../compiler/compile_engine.h"
19 20 21

namespace nnvm {
namespace top {
22

23 24
using namespace tvm;
using namespace nnvm::compiler;
25

26 27 28 29 30 31 32 33 34 35
// undefined op
NNVM_REGISTER_ELEMWISE_UNARY_OP(__undef__)
.describe(R"code(undefined op.

Used to produce invalide node during optimization.

)code" NNVM_ADD_FILELINE)
.set_num_outputs(1)
.set_num_inputs(0);

36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83
// floor
NNVM_REGISTER_ELEMWISE_UNARY_OP(floor)
.describe(R"code(Take floor input array, computed element-wise.
)code" NNVM_ADD_FILELINE)
.set_support_level(3)
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::floor(inputs[0]) };
});

// ceil
NNVM_REGISTER_ELEMWISE_UNARY_OP(ceil)
.describe(R"code(Take ceil input array, computed element-wise.
)code" NNVM_ADD_FILELINE)
.set_support_level(3)
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::ceil(inputs[0]) };
});

// trunc
NNVM_REGISTER_ELEMWISE_UNARY_OP(trunc)
.describe(R"code(Take truncated value of the input, element-wise.
)code" NNVM_ADD_FILELINE)
.set_support_level(3)
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::trunc(inputs[0]) };
});

// round
NNVM_REGISTER_ELEMWISE_UNARY_OP(round)
.describe(R"code(Round elements of the input to nearest integer.
)code" NNVM_ADD_FILELINE)
.set_support_level(3)
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::round(inputs[0]) };
});

84 85 86 87 88 89 90 91 92 93 94 95
// abs
NNVM_REGISTER_ELEMWISE_UNARY_OP(abs)
.describe(R"code(Take absolute value of elements of the input.
)code" NNVM_ADD_FILELINE)
.set_support_level(3)
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::abs(inputs[0]) };
});

96 97 98 99 100
// sigmoid
NNVM_REGISTER_ELEMWISE_UNARY_OP(sigmoid)
.describe(R"code(Computes sigmoid.

.. math::
101
  Y = 1 / (1 + exp(-X))
102 103

)code" NNVM_ADD_FILELINE)
104
.set_support_level(1)
105 106 107 108 109 110
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::sigmoid(inputs[0]) };
})
111 112 113 114 115 116 117 118 119 120 121 122 123 124
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds) {
    // y = 1 / (1 + exp(-n0))
    // grad_0 = grad_y * y * (1 - y)
    NodeEntry sub0 = MakeNode("elemwise_mul", n->attrs.name + "_grad_sub_0",
                              {ograds[0], NodeEntry{n, 0, 0}});
    NodeEntry sub1 = MakeNode("__rsub_scalar__", n->attrs.name + "_grad_sub_1",
                              {NodeEntry{n, 0, 0}}, {{"scalar", "1"}});
    return std::vector<NodeEntry>{
      MakeNode("elemwise_mul", n->attrs.name + "_grad_0",
               {sub0, sub1})
    };
});
125 126 127

// tanh
NNVM_REGISTER_ELEMWISE_UNARY_OP(tanh)
128
.describe(R"code(Computes hyperbolic tangent.
129 130

.. math::
131
   Y = sinh(X) / cosh(X)
132 133

)code" NNVM_ADD_FILELINE)
134
.set_support_level(1)
135 136 137 138 139 140
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::tanh(inputs[0]) };
})
141 142 143 144 145 146 147 148 149 150 151 152 153 154
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds) {
    // y = sinh(n0) / cosh(n0)
    // grad_0 = grad_y * (1 - y^2)
    NodeEntry sub0 = MakeNode("elemwise_mul", n->attrs.name + "_grad_sub_0",
                              {NodeEntry{n, 0, 0}, NodeEntry{n, 0, 0}});
    NodeEntry sub1 = MakeNode("__rsub_scalar__", n->attrs.name + "_grad_sub_1",
                              {sub0}, {{"scalar", "1"}});
    return std::vector<NodeEntry>{
      MakeNode("elemwise_mul", n->attrs.name + "_grad_0",
               {ograds[0], sub1})
    };
});
155 156 157 158 159 160 161 162 163

// exp
NNVM_REGISTER_ELEMWISE_UNARY_OP(exp)
.describe(R"code(Returns the exp input array, computed element-wise.

.. math::
   exp(x)

)code" NNVM_ADD_FILELINE)
164
.set_support_level(1)
165 166 167 168 169 170
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::exp(inputs[0]) };
})
171 172 173 174 175 176 177 178 179 180
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds) {
    // y = exp(n0)
    // grad_0 = grad_y * y
    return std::vector<NodeEntry>{
      MakeNode("elemwise_mul", n->attrs.name + "_grad_0",
               {ograds[0], NodeEntry{n, 0, 0}})
    };
});
181 182 183 184 185 186 187 188 189

// log
NNVM_REGISTER_ELEMWISE_UNARY_OP(log)
.describe(R"code(Returns the log input array, computed element-wise.

.. math::
   log(x)

)code" NNVM_ADD_FILELINE)
190
.set_support_level(1)
191 192 193 194 195 196
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::log(inputs[0]) };
})
197 198 199 200 201 202 203 204 205 206
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds) {
    // y = log(n0)
    // grad_0 = grad_y / n0
    return std::vector<NodeEntry>{
      MakeNode("elemwise_div", n->attrs.name + "_grad_0",
               {ograds[0], n->inputs[0]})
    };
});
207

208 209 210 211 212 213 214 215
// sqrt
NNVM_REGISTER_ELEMWISE_UNARY_OP(sqrt)
.describe(R"code(Returns the sqrt input array, computed element-wise.

.. math::
   \sqrt(x)

)code" NNVM_ADD_FILELINE)
216
.set_support_level(1)
217 218 219 220 221 222
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::sqrt(inputs[0]) };
})
223 224 225 226 227 228 229 230 231 232 233 234
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds) {
    // y = sqrt(n0)
    // grad_0 = grad_y / (2 * y)
    NodeEntry sub0 = MakeNode("__mul_scalar__", n->attrs.name + "_grad_sub_0",
                              {NodeEntry{n, 0, 0}}, {{"scalar", "2"}});
    return std::vector<NodeEntry>{
      MakeNode("elemwise_div", n->attrs.name + "_grad_0",
             {ograds[0], sub0})
    };
});
235

236 237 238 239 240 241
// binary ops

NNVM_REGISTER_ELEMWISE_BINARY_OP(elemwise_add)
.describe(R"code(Element-wise add

)code")
242
.set_support_level(1)
243 244 245 246
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
247
      return Array<Tensor>{ topi::add(inputs[0], inputs[1]) };
248
  })
249 250 251 252 253 254
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = n0 + n1
    // grad_0 = grad_y
    // grad_1 = grad_y
Yao Wang committed
255 256 257 258
    return std::vector<NodeEntry>{ MakeNode("copy", n->attrs.name + "_grad_0",
                                            {ograds[0]}),
                                   MakeNode("copy", n->attrs.name + "_grad_0",
                                            {ograds[0]}) };
259
});
260 261 262 263 264

NNVM_REGISTER_ELEMWISE_BINARY_OP(elemwise_sub)
.describe(R"code(Element-wise substraction

)code"  NNVM_ADD_FILELINE)
265
.set_support_level(1)
266 267 268 269
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
270
    return Array<Tensor>{ topi::subtract(inputs[0], inputs[1]) };
271
})
272 273 274 275 276 277 278 279 280 281 282
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = n0 - n1
    // grad_0 = grad_y
    // grad_1 = - grad_y
    return std::vector<NodeEntry>{
      ograds[0],
      MakeNode("negative", n->attrs.name + "_grad_1", {ograds[0]}),
    };
});
283 284 285 286 287

NNVM_REGISTER_ELEMWISE_BINARY_OP(elemwise_mul)
.describe(R"code(Element-wise multiplication

)code"  NNVM_ADD_FILELINE)
288
.set_support_level(1)
289 290 291 292
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
293
      return Array<Tensor>{ topi::multiply(inputs[0], inputs[1]) };
294
})
295 296 297 298 299 300 301 302 303 304 305 306 307
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = n0 * n1
    // grad_0 = grad_y * n1
    // grad_1 = grad_y * n0
    return std::vector<NodeEntry>{
      MakeNode("elemwise_mul", n->attrs.name + "_grad_0",
               {ograds[0], n->inputs[1]}),
      MakeNode("elemwise_mul", n->attrs.name + "_grad_1",
               {ograds[0], n->inputs[0]})
    };
});
308 309 310 311 312

NNVM_REGISTER_ELEMWISE_BINARY_OP(elemwise_div)
.describe(R"code(Element-wise multiplication

)code"  NNVM_ADD_FILELINE)
313
.set_support_level(1)
314 315 316 317
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
318
      return Array<Tensor>{ topi::divide(inputs[0], inputs[1]) };
319
})
320 321 322 323 324 325 326 327
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = n0 / n1
    // grad_0 = grad_y / n1
    // grad_1 = - grad_y * n0 / n1^2
    NodeEntry sub0 = MakeNode("elemwise_mul", n->attrs.name + "_grad_sub_0",
                              {ograds[0], n->inputs[0]});
Yao Wang committed
328 329
    NodeEntry sub1 = MakeNode("negative", n->attrs.name + "_grad_sub_1",
                              {sub0});
330 331 332 333 334 335 336 337 338
    NodeEntry sub2 = MakeNode("elemwise_mul", n->attrs.name + "_grad_sub_2",
                              {n->inputs[1], n->inputs[1]});
    return std::vector<NodeEntry>{
      MakeNode("elemwise_div", n->attrs.name + "_grad_0",
               {ograds[0], n->inputs[1]}),
      MakeNode("elemwise_div", n->attrs.name + "_grad_1",
               {sub1, sub2})
    };
});
339

340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363
NNVM_REGISTER_ELEMWISE_BINARY_OP(elemwise_mod)
  .describe(R"code(Element-wise modulo

)code" NNVM_ADD_FILELINE)
.set_support_level(1)
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::mod(inputs[0], inputs[1]) };
});

NNVM_REGISTER_ELEMWISE_BINARY_OP(elemwise_pow)
  .describe(R"code(Element-wise power

)code" NNVM_ADD_FILELINE)
.set_support_level(1)
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::power(inputs[0], inputs[1]) };
});

364 365 366 367 368
// negative
NNVM_REGISTER_ELEMWISE_UNARY_OP(negative)
.describe(R"code(Elemenwise numeric negative

)code"  NNVM_ADD_FILELINE)
369
.set_support_level(3)
370 371 372 373 374 375
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::negative(inputs[0]) };
})
376 377 378 379 380 381 382 383 384
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = - n0
    // grad_0 = - grad_y
    return std::vector<NodeEntry>{
      MakeNode("negative", n->attrs.name + "_grad_0", {ograds[0]}),
    };
});
385 386 387 388 389 390

// copy
NNVM_REGISTER_ELEMWISE_UNARY_OP(copy)
.describe(R"code(Copy tensor to another one.

)code"  NNVM_ADD_FILELINE)
391
.set_support_level(3)
392 393 394 395 396 397
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
      return Array<Tensor>{ topi::identity(inputs[0]) };
})
398 399 400 401 402
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = copy(n0)
    // grad_0 = grad_y
Yao Wang committed
403 404
    return std::vector<NodeEntry>{ MakeNode("copy", n->attrs.name + "_grad_0",
                                            {ograds[0]}) };
405 406 407
});

DMLC_REGISTER_PARAMETER(InitOpParam);
408 409
DMLC_REGISTER_PARAMETER(InitOpWithScalarParam);
DMLC_REGISTER_PARAMETER(FillValueParam);
410

Yao Wang committed
411 412
// full
NNVM_REGISTER_INIT_OP(full)
413 414 415
.describe(R"code(Fill array with scalar value

)code"  NNVM_ADD_FILELINE)
416 417 418 419 420 421
.set_attr_parser(ParamParser<InitOpWithScalarParam>)
.set_attr<FGetAttrDict>(
  "FGetAttrDict", ParamGetAttrDict<InitOpWithScalarParam>)
.add_arguments(InitOpWithScalarParam::__FIELDS__())
.set_attr<FInferShape>("FInferShape", ZeroShape<InitOpWithScalarParam>)
.set_attr<FInferType>("FInferType", ZeroType<InitOpWithScalarParam>)
422
.set_attr<FCorrectLayout>("FCorrectLayout", ZeroLayout)
423 424 425 426 427 428 429 430 431 432
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    const InitOpWithScalarParam& param = nnvm::get<InitOpWithScalarParam>(attrs.parsed);
    Array<Expr> shape = ShapeToArray(param.shape);
    Type dtype = GetTVMType(param.dtype);
    Expr fill_value = tvm::make_const(dtype, param.fill_value);
    return Array<Tensor>{ topi::full(shape, dtype, fill_value) };
})
Yao Wang committed
433
.set_support_level(4);
434

Yao Wang committed
435 436 437 438
NNVM_REGISTER_INIT_OP(zeros)
.describe(R"code(Fill target with zeros

)code"  NNVM_ADD_FILELINE)
439 440 441 442 443 444
.set_attr_parser(ParamParser<InitOpParam>)
.set_attr<FGetAttrDict>(
  "FGetAttrDict", ParamGetAttrDict<InitOpParam>)
.add_arguments(InitOpParam::__FIELDS__())
.set_attr<FInferShape>("FInferShape", ZeroShape<InitOpParam>)
.set_attr<FInferType>("FInferType", ZeroType<InitOpParam>)
445
.set_attr<FCorrectLayout>("FCorrectLayout", ZeroLayout)
446 447 448 449 450 451 452 453 454 455
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    const InitOpParam& param = nnvm::get<InitOpParam>(attrs.parsed);
    Array<Expr> shape = ShapeToArray(param.shape);
    Type dtype = GetTVMType(param.dtype);
    Expr fill_value = tvm::make_const(dtype, 0);
    return Array<Tensor>{ topi::full(shape, dtype, fill_value) };
})
Yao Wang committed
456
.set_support_level(4);
Yao Wang committed
457 458 459 460 461

NNVM_REGISTER_INIT_OP(ones)
.describe(R"code(Fill target with ones

)code"  NNVM_ADD_FILELINE)
462 463 464 465 466 467
.set_attr_parser(ParamParser<InitOpParam>)
.set_attr<FGetAttrDict>(
  "FGetAttrDict", ParamGetAttrDict<InitOpParam>)
.add_arguments(InitOpParam::__FIELDS__())
.set_attr<FInferShape>("FInferShape", ZeroShape<InitOpParam>)
.set_attr<FInferType>("FInferType", ZeroType<InitOpParam>)
468
.set_attr<FCorrectLayout>("FCorrectLayout", ZeroLayout)
469 470 471 472 473 474 475 476 477 478
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    const InitOpParam& param = nnvm::get<InitOpParam>(attrs.parsed);
    Array<Expr> shape = ShapeToArray(param.shape);
    Type dtype = GetTVMType(param.dtype);
    Expr fill_value = tvm::make_const(dtype, 1);
    return Array<Tensor>{ topi::full(shape, dtype, fill_value) };
})
Yao Wang committed
479
.set_support_level(4);
Yao Wang committed
480 481

// full_like
482 483
NNVM_REGISTER_INIT_LIKE_OP(full_like)
.describe(R"code(Return an scalar value array with the same shape and type
484 485 486
as the input array

)code"  NNVM_ADD_FILELINE)
487 488 489
.add_arguments(FillValueParam::__FIELDS__())
.set_attr_parser(ParamParser<FillValueParam>)
.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<FillValueParam>)
490 491 492 493 494 495 496 497
.set_attr<FTVMCompute>(
    "FTVMCompute", [](const NodeAttrs& attrs,
                      const Array<Tensor>& inputs,
                      const Array<Tensor>& out_info) {
      const FillValueParam& param = nnvm::get<FillValueParam>(attrs.parsed);
      const Expr fill_value = tvm::make_const(out_info[0]->dtype, param.fill_value);
      return Array<Tensor> { topi::full_like(inputs[0], fill_value) };
})
Yao Wang committed
498
.set_support_level(4);
Yao Wang committed
499

500
NNVM_REGISTER_INIT_LIKE_OP(zeros_like)
Yao Wang committed
501 502 503 504
.describe(R"code(Return an array of zeros with the same shape and type
as the input array.

)code")
505 506 507 508 509 510 511
.set_attr<FTVMCompute>(
    "FTVMCompute", [](const NodeAttrs& attrs,
                      const Array<Tensor>& inputs,
                      const Array<Tensor>& out_info) {
      return Array<Tensor> { topi::full_like(inputs[0],
                                             tvm::make_const(out_info[0]->dtype, 0)) };
})
Yao Wang committed
512
.set_support_level(4);
Yao Wang committed
513

514
NNVM_REGISTER_INIT_LIKE_OP(ones_like)
Yao Wang committed
515 516 517 518
.describe(R"code(Return an array of ones with the same shape and type
as the input array.

)code")
519 520 521 522 523 524 525
.set_attr<FTVMCompute>(
    "FTVMCompute", [](const NodeAttrs& attrs,
                      const Array<Tensor>& inputs,
                      const Array<Tensor>& out_info) {
      return Array<Tensor> { topi::full_like(inputs[0],
                                             tvm::make_const(out_info[0]->dtype, 1)) };
})
Yao Wang committed
526
.set_support_level(4);
527 528 529 530

// unary scalar op
DMLC_REGISTER_PARAMETER(ScalarParam);

531 532 533 534 535 536
#define NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(op)                        \
  NNVM_REGISTER_ELEMWISE_UNARY_OP(op)                                   \
  .add_arguments(ScalarParam::__FIELDS__())                             \
  .set_attr_parser(ParamParser<ScalarParam>)                            \
  .set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<ScalarParam>)

537 538 539 540 541 542 543 544 545 546
inline Tensor binary_scalar_op(const NodeAttrs& attrs,
                               const Tensor& x,
                               std::function<Expr(Expr, Expr)> f) {
  const ScalarParam& param = nnvm::get<ScalarParam>(attrs.parsed);
  auto scalar_val = static_cast<float>(param.scalar);
  return compute(x->shape, [&](const Array<Var>& i) {
    auto scalar_const = make_const(x->dtype, scalar_val);
    return f(x(i), scalar_const);
    }, "tensor", topi::kElementWise);
}
547 548

NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(__add_scalar__)
549 550 551
.describe(R"code(Tensor add scalar

)code"  NNVM_ADD_FILELINE)
552
.set_support_level(3)
553 554 555 556 557 558 559
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    return Array<Tensor>{ binary_scalar_op(attrs, inputs[0],
      [](Expr x, Expr y) { return x + y; }) };
})
560 561 562
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
Yao Wang committed
563 564
    return std::vector<NodeEntry>{ MakeNode("copy", n->attrs.name + "_grad_0",
                                            {ograds[0]}) };
565
});
566

567
NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(__sub_scalar__)
568 569 570
.describe(R"code(Tensor substract scalar

)code"  NNVM_ADD_FILELINE)
571
.set_support_level(3)
572 573 574 575 576 577 578
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    return Array<Tensor>{ binary_scalar_op(attrs, inputs[0],
      [](Expr x, Expr y) { return x - y; }) };
})
579 580 581 582 583
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    return std::vector<NodeEntry>{ograds[0]};
});
584

585
NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(__rsub_scalar__)
586 587 588
.describe(R"code(scalar substract Tensor

)code"  NNVM_ADD_FILELINE)
589
.set_support_level(3)
590 591 592 593 594 595 596
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    return Array<Tensor>{ binary_scalar_op(attrs, inputs[0],
      [](Expr x, Expr y) { return y - x; }) };
})
597 598 599 600 601 602 603
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    return std::vector<NodeEntry>{
      MakeNode("negative", n->attrs.name + "_grad_0", {ograds[0]})
    };
});
604

605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637

NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(__lshift_scalar__)
.describe(R"code(Tensor left shift by scalar

)code"  NNVM_ADD_FILELINE)
.set_support_level(3)
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    const ScalarParam& param = nnvm::get<ScalarParam>(attrs.parsed);
    int scalar_val = static_cast<int>(param.scalar);
    return Array<Tensor>{
      topi::left_shift(inputs[0],
                       make_const(inputs[0]->dtype, scalar_val))};
    });

NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(__rshift_scalar__)
.describe(R"code(Tensor right shift by scalar

)code"  NNVM_ADD_FILELINE)
.set_support_level(3)
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    const ScalarParam& param = nnvm::get<ScalarParam>(attrs.parsed);
    int scalar_val = static_cast<int>(param.scalar);
    return Array<Tensor>{
      topi::right_shift(inputs[0],
                        make_const(inputs[0]->dtype, scalar_val))};
  });

638
NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(__mul_scalar__)
639 640 641
.describe(R"code(Tensor multiplies scalar

)code"  NNVM_ADD_FILELINE)
642
.set_support_level(3)
643 644 645 646 647 648 649
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    return Array<Tensor>{ binary_scalar_op(attrs, inputs[0],
      [](Expr x, Expr y) { return x * y; }) };
})
650 651 652 653 654 655 656 657 658 659
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = n0 * scalar
    // grad_0 = grad_y * scalar
    return std::vector<NodeEntry>{
      MakeNode("__mul_scalar__", n->attrs.name + "_grad_0",
               {ograds[0]}, {{"scalar", n->attrs.dict["scalar"]}})
    };
});
660

661
NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(__div_scalar__)
662 663 664
.describe(R"code(Tensor divides scalar

)code"  NNVM_ADD_FILELINE)
665
.set_support_level(3)
666 667 668 669 670 671 672
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    return Array<Tensor>{ binary_scalar_op(attrs, inputs[0],
      [](Expr x, Expr y) { return x / y; }) };
})
673 674 675 676 677 678 679 680 681 682
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = n0 / scalar
    // grad_0 = grad_y / scalar
    return std::vector<NodeEntry>{
      MakeNode("__div_scalar__", n->attrs.name + "_grad_0",
               {ograds[0]}, {{"scalar", n->attrs.dict["scalar"]}})
    };
});
683

684
NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(__rdiv_scalar__)
685 686 687
.describe(R"code(scalar divides Tensor

)code"  NNVM_ADD_FILELINE)
688
.set_support_level(3)
689 690 691 692 693 694 695
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    return Array<Tensor>{ binary_scalar_op(attrs, inputs[0],
      [](Expr x, Expr y) { return y / x; }) };
})
696 697 698 699 700 701
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = scalar / n0
    // grad_0 = - grad_y * scalar / n0^2
    NodeEntry sub0 = MakeNode("__mul_scalar__", n->attrs.name + "_grad_sub_0",
Yao Wang committed
702 703 704 705
                              {ograds[0]},
                              {{"scalar", n->attrs.dict["scalar"]}});
    NodeEntry sub1 = MakeNode("negative", n->attrs.name + "_grad_sub_1",
                              {sub0});
706 707 708 709 710 711 712
    NodeEntry sub2 = MakeNode("elemwise_mul", n->attrs.name + "_grad_sub_2",
                              {n->inputs[0], n->inputs[0]});
    return std::vector<NodeEntry>{
      MakeNode("elemwise_div", n->attrs.name + "_grad_0",
               {sub1, sub2})
    };
});
713

714
NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(__pow_scalar__)
715 716 717
.describe(R"code(Tensor power scalar

)code"  NNVM_ADD_FILELINE)
718
.set_support_level(3)
719 720 721 722 723 724 725
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    return Array<Tensor>{ binary_scalar_op(attrs, inputs[0],
      [](Expr x, Expr y) { return tvm::pow(x, y); }) };
})
726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = n0^scalar
    // grad_0 = grad_y * scalar * n0^(scalar - 1)
    double scalar = std::stod(n->attrs.dict["scalar"]);
    NodeEntry sub0 = MakeNode("__pow_scalar__", n->attrs.name + "_grad_sub_0",
                              {n->inputs[0]},
                              {{"scalar", std::to_string(scalar - 1)}});
    NodeEntry sub1 = MakeNode("__mul_scalar__", n->attrs.name + "_grad_sub_1",
                              {ograds[0]},
                              {{"scalar", std::to_string(scalar)}});
    return std::vector<NodeEntry>{
      MakeNode("elemwise_mul", n->attrs.name + "_grad_0",
               {sub0, sub1})
    };
});
743

744
NNVM_REGISTER_ELEMWISE_BINARY_SCALAR(__rpow_scalar__)
745 746 747
.describe(R"code(scalar power Tensor

)code"  NNVM_ADD_FILELINE)
748
.set_support_level(3)
749 750 751 752 753 754 755
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    return Array<Tensor>{ binary_scalar_op(attrs, inputs[0],
      [](Expr x, Expr y) { return tvm::pow(y, x); }) };
})
756 757 758 759 760 761 762 763 764 765 766 767 768 769 770
.set_attr<FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = scalar^n0
    // grad_0 = grad_y * scalar^n0 * log(scalar)
    double num = std::stod(n->attrs.dict["scalar"]);
    NodeEntry sub0 = MakeNode("__mul_scalar__", n->attrs.name + "_grad_sub_0",
                              {NodeEntry{n, 0, 0}},
                              {{"scalar", std::to_string(std::log(num))}});
    return std::vector<NodeEntry>{
      MakeNode("__mul_symbol__", n->attrs.name + "_grad_0",
               {ograds[0], sub0})
    };
});

771
DMLC_REGISTER_PARAMETER(ElementWiseReduceParam);
772

773
NNVM_REGISTER_ELEMWISE_REDUCE_OP(elemwise_sum)
Yao Wang committed
774 775 776
.describe(R"code(Adds all input arguments element-wise.

)code"  NNVM_ADD_FILELINE)
777 778 779 780 781 782 783 784
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    const ElementWiseReduceParam& param = nnvm::get<ElementWiseReduceParam>(attrs.parsed);
    CHECK_EQ(param.num_args, inputs.size()) << """Compute definition of elemwise sum""";
    return Array<Tensor>{ topi::elemwise_sum(inputs) };
})
785 786 787 788 789 790
.set_attr<nnvm::FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    CHECK_EQ(ograds.size(), 1);
    std::vector<NodeEntry> ret;
    for (size_t i = 0; i < n->inputs.size(); i++) {
Yao Wang committed
791
      ret.push_back(MakeNode("copy", n->attrs.name + "_grad_0", {ograds[0]}));
792 793
    }
    return ret;
Yao Wang committed
794 795
  })
.set_support_level(4);
796 797 798 799 800 801 802 803

NNVM_REGISTER_ELEMWISE_UNARY_OP(block_grad)
.describe(R"code(Blocks gradient computation for input.

)code" NNVM_ADD_FILELINE)
.set_attr<nnvm::FInplaceIdentity>(
  "FInplaceIdentity", [](const NodeAttrs& attrs){
    return std::vector<bool>{true};
Yao Wang committed
804
})
Yao Wang committed
805 806
.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
.set_support_level(4);
807 808 809 810 811 812 813 814 815 816 817 818 819

DMLC_REGISTER_PARAMETER(IndicatorParam);

// indicator function
NNVM_REGISTER_INDICATOR_OP(greater)
.describe(R"code(Greater function that returns a mask tensor
with 1.0 if (left > right), otherwise 0.0 element-wise.

)code" NNVM_ADD_FILELINE)
.add_argument("lhs", "Tensor", "First input")
.add_argument("rhs", "Tensor", "Second input")
.set_num_inputs(2)
.set_attr<nnvm::FInferShape>("FInferShape", ElemwiseShape<2, 1>)
820 821 822 823 824 825
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    return Array<Tensor>{ topi::cast(topi::greater(inputs[0], inputs[1]), out_info[0]->dtype) };
})
Yao Wang committed
826
.set_support_level(4);
827 828 829 830 831 832 833 834 835 836 837


NNVM_REGISTER_INDICATOR_OP(less)
  .describe(R"code(Less function that returns a mask tensor
with 1.0 if (left < right), otherwise 0.0 element-wise.

)code" NNVM_ADD_FILELINE)
.add_argument("lhs", "Tensor", "First input")
.add_argument("rhs", "Tensor", "Second input")
.set_num_inputs(2)
.set_attr<nnvm::FInferShape>("FInferShape", ElemwiseShape<2, 1>)
838 839 840 841 842 843
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    return Array<Tensor>{ topi::cast(topi::less(inputs[0], inputs[1]), out_info[0]->dtype) };
})
Yao Wang committed
844
.set_support_level(4);
845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870

NNVM_REGISTER_INDICATOR_OP(_max_mask)
  .describe(R"code(Function that returns a mask tensor
with 1.0 if the value is maximum over given axes, otherwise 0.0 element-wise.

)code" NNVM_ADD_FILELINE)
.add_argument("data", "Tensor", "Input")
.set_num_inputs(1)
.add_arguments(IndicatorParam::__FIELDS__())
.set_attr_parser(ParamParser<IndicatorParam>)
.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<IndicatorParam>)
.set_attr<nnvm::FInferShape>("FInferShape", ElemwiseShape<1, 1>)
.set_support_level(1);

NNVM_REGISTER_INDICATOR_OP(_min_mask)
  .describe(R"code(Function that returns a mask tensor
with 1.0 if the value is minimum over given axes, otherwise 0.0 element-wise.

)code" NNVM_ADD_FILELINE)
.add_argument("data", "Tensor", "Input")
.set_num_inputs(1)
.add_arguments(IndicatorParam::__FIELDS__())
.set_attr_parser(ParamParser<IndicatorParam>)
.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<IndicatorParam>)
.set_attr<nnvm::FInferShape>("FInferShape", ElemwiseShape<1, 1>)
.set_support_level(1);
871

Yao Wang committed
872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889

DMLC_REGISTER_PARAMETER(ClipParam);

NNVM_REGISTER_OP(clip)
.describe(R"doc(Clips (limits) the values in an array.
Given an interval, values outside the interval are clipped to the interval edges.
Clipping ``x`` between `a_min` and `a_x` would be::
   clip(x, a_min, a_max) = max(min(x, a_max), a_min))
Example::
    x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    clip(x,1,8) = [ 1.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  8.]
)doc" NNVM_ADD_FILELINE)
.set_num_inputs(1)
.set_num_outputs(1)
.set_attr_parser(ParamParser<ClipParam>)
.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<ClipParam>)
.set_attr<nnvm::FInferShape>("FInferShape", ElemwiseShape<1, 1>)
.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)
890
.set_attr<nnvm::FCorrectLayout>("FCorrectLayout", ElemwiseFixedLayoutUnknownOut<1, 1>)
Yao Wang committed
891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940
.set_attr<FTVMCompute>(
  "FTVMCompute", [](const NodeAttrs& attrs,
                    const Array<Tensor>& inputs,
                    const Array<Tensor>& out_info) {
    const ClipParam params = get<ClipParam>(attrs.parsed);
    return Array<Tensor>{
      topi::clip(inputs[0], tvm::make_const(tvm::Float(32), params.a_min),
                 tvm::make_const(tvm::Float(32), params.a_max)) };
  })
.add_argument("data", "NDArray-or-Symbol", "Input array.")
.add_arguments(ClipParam::__FIELDS__())
.set_attr<nnvm::FGradient>(
  "FGradient", [](const NodePtr& n,
                  const std::vector<NodeEntry>& ograds){
    // y = clip(x, a_min, a_max)
    // min_mask = greater_equal(x, a_min*ones_like(x))
    //          => ones_like(x) - less(x, a_min)
    // max_mask = less_equal(x, a_max*ones_like(x))
    //          => ones_like(x) - greater(x, a_max)
    // grad_x = min_mask * max_mask * grad_y
    CHECK_EQ(ograds.size(), 1);

    NodeEntry sub0 = MakeNode("ones_like", n->attrs.name + "_grad_sub_0",
                              {n->inputs[0]});
    // min_mask
    NodeEntry sub1 = MakeNode("__mul_scalar__", n->attrs.name + "_grad_sub_1",
                              {sub0}, {{"scalar", n->attrs.dict["a_min"]}});
    NodeEntry sub2 = MakeNode("less", n->attrs.name + "_grad_sub_2",
                              {n->inputs[0], sub1});
    NodeEntry sub3 = MakeNode("elemwise_sub", n->attrs.name + "_grad_sub_3",
                              {sub0, sub2});

    // max_mask
    NodeEntry sub4 = MakeNode("__mul_scalar__", n->attrs.name + "_grad_sub_4",
                              {sub0}, {{"scalar", n->attrs.dict["a_max"]}});
    NodeEntry sub5 = MakeNode("greater", n->attrs.name + "_grad_sub_5",
                              {n->inputs[0], sub4});
    NodeEntry sub6 = MakeNode("elemwise_sub", n->attrs.name + "_grad_sub_6",
                              {sub0, sub5});

    // min_mask * max_mask
    NodeEntry sub7 = MakeNode("elemwise_mul", n->attrs.name + "_grad_sub_7",
                              {sub3, sub6});
    return std::vector<NodeEntry>{
      MakeNode("elemwise_mul", n->attrs.name + "_grad",
               {sub7, ograds[0]})
    };
  })
.set_support_level(4);

941 942
}  // namespace top
}  // namespace nnvm