1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 * \file elemwise.h
22 * \brief Elementwise op constructions
23 */
24#ifndef TVM_TOPI_ELEMWISE_H_
25#define TVM_TOPI_ELEMWISE_H_
26
27#include <tvm/tir/builtin.h>
28#include <tvm/tir/expr.h>
29#include <tvm/topi/tags.h>
30
31#include <algorithm>
32#include <string>
33
34#include "broadcast.h"
35
36namespace tvm {
37namespace topi {
38
39using namespace tvm::te;
40
41// Unary intrinsic operators
42#define TOPI_DECLARE_UNARY_OP(OpName) \
43 inline Tensor OpName(const Tensor& x, std::string name = "T_" #OpName, \
44 std::string tag = kElementWise) { \
45 return compute( \
46 x->shape, [&](const Array<Var>& i) { return ::tvm::OpName(x(i)); }, name, tag); \
47 }
48
49TOPI_DECLARE_UNARY_OP(exp);
50TOPI_DECLARE_UNARY_OP(erf);
51TOPI_DECLARE_UNARY_OP(sigmoid);
52TOPI_DECLARE_UNARY_OP(sqrt);
53TOPI_DECLARE_UNARY_OP(log);
54TOPI_DECLARE_UNARY_OP(log2);
55TOPI_DECLARE_UNARY_OP(log10);
56TOPI_DECLARE_UNARY_OP(floor);
57TOPI_DECLARE_UNARY_OP(ceil);
58TOPI_DECLARE_UNARY_OP(round);
59TOPI_DECLARE_UNARY_OP(trunc);
60TOPI_DECLARE_UNARY_OP(abs);
61TOPI_DECLARE_UNARY_OP(cos);
62TOPI_DECLARE_UNARY_OP(cosh);
63TOPI_DECLARE_UNARY_OP(tan);
64TOPI_DECLARE_UNARY_OP(sin);
65TOPI_DECLARE_UNARY_OP(sinh);
66TOPI_DECLARE_UNARY_OP(acos);
67TOPI_DECLARE_UNARY_OP(acosh);
68TOPI_DECLARE_UNARY_OP(asin);
69TOPI_DECLARE_UNARY_OP(asinh);
70TOPI_DECLARE_UNARY_OP(atan);
71TOPI_DECLARE_UNARY_OP(atanh);
72TOPI_DECLARE_UNARY_OP(isnan);
73TOPI_DECLARE_UNARY_OP(tanh);
74TOPI_DECLARE_UNARY_OP(isfinite);
75TOPI_DECLARE_UNARY_OP(isinf);
76
77/*!
78 * \brief Fast_tanh_float implementation from Eigen
79 * https://github.com/eigenteam/eigen-git-mirror/blob/master/Eigen/src/Core/MathFunctionsImpl.h#L26
80 */
81inline Tensor fast_tanh_float(const Tensor& in, std::string name, std::string tag) {
82 // Clamp the inputs to the range [-9, 9] since anything outside
83 // this range is +/-1.0f in single-precision.
84 auto x = maximum(make_const(in->dtype, -9.0), minimum(make_const(in->dtype, 9.0), in));
85
86 // The monomial coefficients of the numerator polynomial (odd).
87 auto alpha_1 = make_const(in->dtype, 4.89352455891786e-03);
88 auto alpha_3 = make_const(in->dtype, 6.37261928875436e-04);
89 auto alpha_5 = make_const(in->dtype, 1.48572235717979e-05);
90 auto alpha_7 = make_const(in->dtype, 5.12229709037114e-08);
91 auto alpha_9 = make_const(in->dtype, -8.60467152213735e-11);
92 auto alpha_11 = make_const(in->dtype, 2.00018790482477e-13);
93 auto alpha_13 = make_const(in->dtype, -2.76076847742355e-16);
94
95 // The monomial coefficients of the denominator polynomial (even).
96 auto beta_0 = make_const(in->dtype, 4.89352518554385e-03);
97 auto beta_2 = make_const(in->dtype, 2.26843463243900e-03);
98 auto beta_4 = make_const(in->dtype, 1.18534705686654e-04);
99 auto beta_6 = make_const(in->dtype, 1.19825839466702e-06);
100
101 return compute(
102 x->shape,
103 [&](const Array<Var>& i) {
104 auto x2 = x(i) * x(i);
105 auto p = x2 * alpha_13 + alpha_11;
106 p = x2 * p + alpha_9;
107 p = x2 * p + alpha_7;
108 p = x2 * p + alpha_5;
109 p = x2 * p + alpha_3;
110 p = x2 * p + alpha_1;
111 p = x(i) * p;
112
113 auto q = x2 * beta_6 + beta_4;
114 q = x2 * q + beta_2;
115 q = x2 * q + beta_0;
116 return p / q;
117 },
118 name, tag);
119}
120
121/*!
122 * \brief Creates an operation that returns hyperbolic tanh of a given tensor
123 *
124 * \param x The input tensor
125 * \param name The name of the operation
126 * \param tag The tag to mark the operation
127 *
128 * \return A Tensor whose op member is tanh
129 */
130inline Tensor fast_tanh(const Tensor& x, std::string name = "T_fast_tanh",
131 std::string tag = kElementWise) {
132 if (x->dtype == DataType::Float(32)) {
133 // invoke fast_tanh_float implementation
134 return fast_tanh_float(x, name, tag);
135 } else {
136 // fallback to default implementation
137 return compute(
138 x->shape, [&](const Array<Var>& i) { return ::tvm::tanh(x(i)); }, name, tag);
139 }
140}
141
142/*!
143 * \brief Creates an operation that returns identity of a given tensor
144 *
145 * \param x The input tensor
146 * \param name The name of the operation
147 * \param tag The tag to mark the operation
148 *
149 * \return A Tensor whose op member is the identity operation
150 */
151inline Tensor identity(const Tensor& x, std::string name = "T_identity",
152 std::string tag = kElementWise) {
153 return compute(
154 x->shape, [&](const Array<Var>& i) { return x(i); }, name, tag);
155}
156
157/*!
158 * \brief Creates an operation that returns the negation of a given tensor
159 *
160 * \param x The input tensor
161 * \param name The name of the operation
162 * \param tag The tag to mark the operation
163 *
164 * \return A Tensor whose op member is the negation operation
165 */
166inline Tensor negative(const Tensor& x, std::string name = "T_negative",
167 std::string tag = kElementWise) {
168 return compute(
169 x->shape, [&](const Array<Var>& i) { return -x(i); }, name, tag);
170}
171
172/*!
173 * \brief Creates an operation that returns the logical NOT of a given tensor
174 *
175 * \param x The input tensor
176 * \param name The name of the operation
177 * \param tag The tag to mark the operation
178 *
179 * \return A Tensor whose op member is the logical NOT operation
180 */
181inline Tensor logical_not(const Tensor& x, std::string name = "T_logical_not",
182 std::string tag = kElementWise) {
183 return compute(
184 x->shape, [&](const Array<Var>& i) { return !x(i); }, name, tag);
185}
186
187/*!
188 * \brief Creates an operation that returns the bitwise NOT of a given tensor
189 *
190 * \param x The input tensor
191 * \param name The name of the operation
192 * \param tag The tag to mark the operation
193 *
194 * \return A Tensor whose op member is the bitwise NOT operation
195 */
196inline Tensor bitwise_not(const Tensor& x, std::string name = "T_bitwise_not",
197 std::string tag = kElementWise) {
198 return compute(
199 x->shape, [&](const Array<Var>& i) { return ~x(i); }, name, tag);
200}
201
202/*!
203 * \brief Returns the sign of the tensor
204 *
205 * \param x The input tensor
206 * \param name The name of the operation
207 * \param tag The tag to mark the operation
208 *
209 * \return A Tensor whose op member is the sign
210 */
211inline Tensor sign(const Tensor& x, std::string name = "T_sign", std::string tag = kElementWise) {
212 return compute(
213 x->shape,
214 [&](const Array<Var>& i) {
215 PrimExpr zero = make_zero(x->dtype);
216 PrimExpr one = make_const(x->dtype, 1);
217 PrimExpr minus_one = make_const(x->dtype, -1);
218 auto s1 = tvm::tir::Select((x(i) < zero), minus_one, zero);
219 auto s2 = tvm::tir::Select((x(i) > zero), one, s1);
220 return s2;
221 },
222 name, tag);
223}
224
225/*!
226 * \brief Creates an operation that returns rsqrt of a given tensor
227 *
228 * \param x The input tensor
229 * \param name The name of the operation
230 * \param tag The tag to mark the operation
231 *
232 * \return A Tensor whose op member is the rsqrt operation
233 */
234inline Tensor rsqrt(const Tensor& x, std::string name = "tensor", std::string tag = kElementWise) {
235 return compute(
236 x->shape,
237 [&](const Array<Var>& i) {
238 PrimExpr one = make_const(x->dtype, 1);
239 return one / tvm::sqrt(x(i));
240 },
241 name, tag);
242}
243
244/*!
245 * \brief Creates an operation that clips each element of a tensor to
246 * the interval [a_min, a_max]
247 *
248 * \param x The input tensor
249 * \param a_min The inclusive lower bound of the interval
250 * \param a_max The inclusive upper bound of the interval
251 * \param name The name of the operation
252 * \param tag The tag to mark the operation
253 *
254 * \return A Tensor whose op member is the clip operation
255 */
256inline Tensor clip(const Tensor& x, const PrimExpr& a_min, const PrimExpr& a_max,
257 std::string name = "T_clip", std::string tag = kElementWise) {
258 return compute(
259 x->shape,
260 [&](const Array<Var>& i) {
261 auto min_val = tvm::cast(x->dtype, a_min);
262 auto max_val = tvm::cast(x->dtype, a_max);
263 return tvm::max(tvm::min(x(i), max_val), min_val); // NOLINT(*)
264 },
265 name, tag);
266}
267
268/*!
269 * \brief Cast each element of x to the given type. If expr is
270 * scalar and type is a corresponding vector type, a
271 * Broadcast is generated, otherwise a Cast is generated.
272 *
273 * \param x The input tensor
274 * \param type The type to cast to
275 * \param name The name of the operation
276 * \param tag The tag to mark the operation
277 *
278 * \return A Tensor whose op member is the cast operation
279 */
280inline Tensor cast(const Tensor& x, DataType type, std::string name = "T_cast",
281 std::string tag = kElementWise) {
282 return compute(
283 x->shape,
284 [&](const Array<Var>& i) -> PrimExpr {
285 auto expr = x(i);
286 if (expr.dtype().code() == type.code() && expr.dtype().bits() == type.bits()) {
287 if (expr.dtype().lanes() == type.lanes()) {
288 return expr;
289 } else if (expr.dtype().lanes() == 1 && type.lanes() > 1) {
290 return tvm::tir::Broadcast(expr, type.lanes());
291 }
292 }
293
294 return tvm::cast(type, x(i));
295 },
296 name, tag);
297}
298
299/*!
300 * \brief Reinterpret each element of x to the given type.
301
302 * \param x The input tensor
303 * \param type The type to cast to
304 * \param name The name of the operation
305 * \param tag The tag to mark the operation
306 *
307 * \return A Tensor whose op member is the reinterpret operation
308 */
309inline Tensor reinterpret(const Tensor& x, DataType type, std::string name = "tensor",
310 std::string tag = kElementWise) {
311 return compute(
312 x->shape,
313 [&](const Array<Var>& i) {
314 return tvm::tir::Call(type, tvm::tir::builtin::reinterpret(), {x(i)});
315 },
316 name, tag);
317}
318
319/*!
320 * \brief Creates an operation that sum each element of a tensor
321 *
322 * \param xs The input tensor array
323 * \param name The name of the operation
324 * \param tag The tag to mark the operation
325 *
326 * \return A Tensor whose op member is the sum operation
327 */
328inline Tensor elemwise_sum(const Array<Tensor>& xs, std::string name = "T_elemwise_sum",
329 std::string tag = kElementWise) {
330 ICHECK_GT(xs.size(), 0) << "elemwise sum must have at least one input tensor.";
331 return compute(
332 xs[0]->shape,
333 [&](const Array<Var>& i) {
334 auto sum_expr = xs[0](i);
335 for (size_t j = 1; j < xs.size(); j++) {
336 sum_expr = sum_expr + xs[j](i);
337 }
338 return sum_expr;
339 },
340 name, tag);
341}
342
343/*!
344 * \brief Creates an operation that fill a tensor with fill_value
345 *
346 * \param shape The shape of a tensor
347 * \param dtype The Type of fill_value
348 * \param fill_value The value to be filled
349 * \param name The name of the operation
350 * \param tag The tag to mark the operation
351 *
352 * \return A Tensor whose op member is the full operation
353 */
354inline Tensor full(const Array<PrimExpr>& shape, DataType dtype, const PrimExpr fill_value,
355 std::string name = "T_full", std::string tag = kElementWise) {
356 PrimExpr ev = cast(dtype, fill_value);
357 if (!ev.defined()) {
358 LOG(ERROR) << "Can't cast fill_value to " << dtype;
359 }
360 return compute(
361 shape, [&](const Array<Var>& i) { return ev; }, name, tag);
362}
363
364/*!
365 * \brief Creates an operation that construct a tensor with same shape as input tensor,
366 * then fill a tensor with fill_value
367 *
368 * \param x The input tensor
369 * \param fill_value The value to be filled
370 * \param name The name of the operation
371 * \param tag The tag to mark the operation
372 *
373 * \return A Tensor whose op memeber is the full_like operation
374 */
375inline Tensor full_like(const Tensor& x, const PrimExpr fill_value,
376 std::string name = "T_full_like", std::string tag = kElementWise) {
377 PrimExpr ev = cast(x->dtype, fill_value);
378 return compute(
379 x->shape, [&](const Array<Var>& i) { return ev; }, name, tag);
380}
381
382/*!
383 * \brief Fast exponential function implementation
384 *
385 * \param _x The input tensor
386 * \param name The name of the operation
387 * \param tag The tag to mark the operation
388 *
389 * \return A Tensor whose op member is exponent operation
390 *
391 * \note Function computes:
392 * log2(e^x) = x * log2(e) * log2(2) =>
393 * log2(e^x) = log2(2^(x*log2(e))) =>
394 * e^x = 2^(x*log2(e))
395 * Splitting power x*log2(e) into integer and fractional parts:
396 * e^(n+f) = e^n * e^f
397 * n = floor(x*log2(e) + 1/2)
398 * f = x - n * ln(2)
399 * exp(x) = 2^n * exp(y)
400 * Approximation for fractional part:
401 * y = exp(f) = 1 + 2 * P(x**2)/(Q(x**2) - P(x**2))
402 */
403inline Tensor fast_exp_float32(const Tensor& _x, std::string name, std::string tag) {
404 auto x_hi = make_const(DataType::Float(32), 88.3762626647950f);
405 auto x_lo = make_const(DataType::Float(32), -88.3762626647949f);
406 auto log2e = make_const(DataType::Float(32), 1.44269504088896341f);
407 auto ln2 = make_const(DataType::Float(32), 0.6931471805599453f);
408 PrimExpr p[6] = {make_const(DataType::Float(32), 1.9875691500E-4f),
409 make_const(DataType::Float(32), 1.3981999507E-3f),
410 make_const(DataType::Float(32), 8.3334519073E-3f),
411 make_const(DataType::Float(32), 4.1665795894E-2f),
412 make_const(DataType::Float(32), 1.6666665459E-1f),
413 make_const(DataType::Float(32), 5.0000001201E-1f)};
414 auto one = make_const(DataType::Float(32), 1.0f);
415 auto one_half = make_const(DataType::Float(32), 0.5f);
416 auto b = make_const(DataType::Float(32), 127.0f);
417
418 return compute(
419 _x->shape,
420 [&](const Array<Var>& i) {
421 // clamp x
422 auto x = ::tvm::max(::tvm::min(_x(i), x_hi), x_lo);
423 // integer part
424 auto n = ::tvm::floor(x * log2e + one_half);
425 // fractional part
426 auto f = x - n * ln2;
427 auto y =
428 (((((p[0] * f + p[1]) * f + p[2]) * f + p[3]) * f + p[4]) * f + p[5]) * f * f + f + one;
429 // Return 2^m * exp(r).
430 auto ef =
431 tvm::reinterpret(DataType::Float(32), ::tvm::cast(DataType::Int(32), n + b) << 23);
432 return ::tvm::max(ef * y, _x(i)); // NOLINT(*)
433 },
434 name, tag);
435}
436
437/*!
438 * \brief Fast exponential function implementation
439 *
440 * \param x The input tensor
441 * \param name The name of the operation
442 * \param tag The tag to mark the operation
443 *
444 * \return A Tensor whose op member is exponent operation
445 *
446 */
447inline Tensor fast_exp(const Tensor& x, std::string name = "T_fast_exp",
448 std::string tag = kElementWise) {
449 if (x->dtype == DataType::Float(32)) {
450 auto ret = fast_exp_float32(x, name, tag);
451 return ret;
452 } else {
453 return compute(
454 x->shape, [&](const Array<Var>& i) { return ::tvm::exp(x(i)); }, name, tag);
455 }
456}
457
458/*!
459 * \brief Fast_erf_float expression from Eigen
460 * https://github.com/eigenteam/eigen-git-mirror/blob/master/unsupported/Eigen/src/SpecialFunctions/SpecialFunctionsImpl.h#L290
461 * \param arg The input expression.
462 * \param bits The number of bits in the type.
463 */
464inline PrimExpr fast_erf_float_expr(PrimExpr arg, int bits) {
465 auto plus_4 = make_const(DataType::Float(bits), 4.f);
466 auto minus_4 = make_const(DataType::Float(bits), -4.f);
467
468 // The monomial coefficients of the numerator polynomial (odd).
469 auto alpha_1 = make_const(DataType::Float(bits), -1.60960333262415e-02f);
470 auto alpha_3 = make_const(DataType::Float(bits), -2.95459980854025e-03f);
471 auto alpha_5 = make_const(DataType::Float(bits), -7.34990630326855e-04f);
472 auto alpha_7 = make_const(DataType::Float(bits), -5.69250639462346e-05f);
473 auto alpha_9 = make_const(DataType::Float(bits), -2.10102402082508e-06f);
474 auto alpha_11 = make_const(DataType::Float(bits), 2.77068142495902e-08f);
475 auto alpha_13 = make_const(DataType::Float(bits), -2.72614225801306e-10f);
476
477 // The monomial coefficients of the denominator polynomial (even).
478 auto beta_0 = make_const(DataType::Float(bits), -1.42647390514189e-02f);
479 auto beta_2 = make_const(DataType::Float(bits), -7.37332916720468e-03f);
480 auto beta_4 = make_const(DataType::Float(bits), -1.68282697438203e-03f);
481 auto beta_6 = make_const(DataType::Float(bits), -2.13374055278905e-04f);
482 auto beta_8 = make_const(DataType::Float(bits), -1.45660718464996e-05f);
483
484 // clamp x
485 auto x = tvm::max(tvm::min(arg, plus_4), minus_4);
486 auto x2 = x * x;
487
488 // Evaluate the numerator polynomial p.
489 auto p = x2 * alpha_13 + alpha_11;
490 p = x2 * p + alpha_9;
491 p = x2 * p + alpha_7;
492 p = x2 * p + alpha_5;
493 p = x2 * p + alpha_3;
494 p = x2 * p + alpha_1;
495 p = x * p;
496
497 // Evaluate the denominator polynomial p.
498 auto q = x2 * beta_8 + beta_6;
499 q = x2 * q + beta_4;
500 q = x2 * q + beta_2;
501 q = x2 * q + beta_0;
502
503 return p / q;
504}
505
506/*!
507 * \brief Fast_erf_float expression from Eigen
508 */
509inline Tensor fast_erf_float32(const Tensor& data, std::string name, std::string tag) {
510 return compute(
511 data->shape, [&](const Array<Var>& i) { return fast_erf_float_expr(data(i), 32); }, name,
512 tag);
513}
514
515/*!
516 * \brief Fast_erf_float expression from Eigen for float16.
517 */
518inline Tensor fast_erf_float16(const Tensor& data, std::string name, std::string tag) {
519 return compute(
520 data->shape, [&](const Array<Var>& i) { return fast_erf_float_expr(data(i), 16); }, name,
521 tag);
522}
523
524/*!
525 * \brief Fast erf implementation
526 *
527 * \param x The input tensor
528 * \param name The name of the operation
529 * \param tag The tag to mark the operation
530 *
531 * \return A Tensor whose op member is erf operation
532 */
533inline Tensor fast_erf(const Tensor& x, std::string name = "T_fast_erf",
534 std::string tag = kElementWise) {
535 if (x->dtype == DataType::Float(32)) {
536 auto ret = fast_erf_float32(x, name, tag);
537 return ret;
538 } else if (x->dtype == DataType::Float(16)) {
539 auto ret = fast_erf_float16(x, name, tag);
540 return ret;
541 } else {
542 return topi::erf(x);
543 }
544}
545
546} // namespace topi
547} // namespace tvm
548#endif // TVM_TOPI_ELEMWISE_H_
549