tensor.h 30.8 KB
Newer Older
1
// Copyright 2019-2020 Yan Yan
tusimple's avatar
tusimple committed
2
3
4
5
6
7
8
9
10
11
12
13
14
//
// Licensed 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.

15
16
17
18
19
20
21
22
23
/*
tv::Tensor is a lightweight header-only tensor container
without template and annoying dependencies. no algorithm is implemented.
it should only be used when you want a no-template simple container but
dont want to link with libtorch.

If you can use libtorch, dont use tv::Tensor.
*/

tusimple's avatar
tusimple committed
24
#pragma once
25
#include "mp_helper.h"
tusimple's avatar
tusimple committed
26
#include "tensorview.h"
27
28
#include <cstring>
#include <iomanip>
tusimple's avatar
tusimple committed
29
#include <memory>
30
31
32
#include <type_traits>
#ifdef TV_CUDA
#include <cuda_fp16.h>
tusimple's avatar
tusimple committed
33
34
35
36
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
#endif

37
38
namespace tv {
enum DType {
tusimple's avatar
tusimple committed
39
40
41
42
43
44
45
46
47
48
49
50
51
52
  float32,
  int32,
  int16,
  int8,
  float64,
  bool_,
  uint8,
  float16,
  int64,
  uint16,
  uint32,
  uint64
};

53
namespace detail {
tusimple's avatar
tusimple committed
54

Yan Yan's avatar
Yan Yan committed
55
56
57
58
using dtype_collection_t =
    tv::mp_list_c<int, float32, int32, int16, int8, float64, bool_, uint8,
                  float16, int64, uint16, uint32, uint64>;

yanyan's avatar
yanyan committed
59
#ifdef TV_CUDA
60
61
62
using all_tensor_types_t =
    std::tuple<float, double, int8_t, int16_t, int32_t, int64_t, uint8_t,
               uint16_t, uint32_t, uint64_t, bool>;
yanyan's avatar
yanyan committed
63
64
65
66
67
#else
using all_tensor_types_t =
    std::tuple<float, double, int8_t, int16_t, int32_t, int64_t, uint8_t,
               uint16_t, uint32_t, uint64_t, bool>;
#endif
68
69

template <typename T> class TensorStorage {
tusimple's avatar
tusimple committed
70
public:
71
72
73
74
  TensorStorage(size_t size, int device = -1, bool managed = false,
                bool pinned = false)
      : mSize(size), device_(device), managed_(managed), pinned_(pinned) {
    if (size == 0) {
tusimple's avatar
tusimple committed
75
      mPtr = nullptr;
76
77
78
79
80
    } else {
      if (device == -1) {
        if (pinned_) {
#ifdef TV_CUDA
          checkCudaErrors(cudaMallocHost(&mPtr, size * sizeof(T)));
tusimple's avatar
tusimple committed
81
#else
82
          TV_THROW_INVALID_ARG("you need to define TV_CUDA to use pinned");
tusimple's avatar
tusimple committed
83
#endif
84
85
86
87
88
        } else {
          mPtr = new T[size];
        }
      } else {
#ifdef TV_CUDA
yanyan's avatar
yanyan committed
89
90
        // we should select device in external
        /*
tusimple's avatar
tusimple committed
91
92
        int deviceCount;
        cudaGetDeviceCount(&deviceCount);
93
94
95
        if (device >= deviceCount) {
          TV_THROW_INVALID_ARG("you provide device ", device,
                               " but you only have ", deviceCount, " device.");
tusimple's avatar
tusimple committed
96
97
        }
        cudaSetDevice(device);
yanyan's avatar
yanyan committed
98
        */
99
        if (managed) {
tusimple's avatar
tusimple committed
100
          checkCudaErrors(cudaMallocManaged(&this->mPtr, size * sizeof(T)));
101
        } else {
tusimple's avatar
tusimple committed
102
103
104
          checkCudaErrors(cudaMalloc(&mPtr, size * sizeof(T)));
        }
#else
105
        TV_THROW_INVALID_ARG("don't compiled with cuda");
tusimple's avatar
tusimple committed
106
107
108
109
110
111
112
#endif
      }
    }
  }
  TensorStorage(T *ptr, size_t size, int device)
      : mSize(size), mPtr(ptr), from_blob_(true), device_(device) {}

113
114
  virtual ~TensorStorage() {
    if (empty()) {
tusimple's avatar
tusimple committed
115
116
      return;
    }
117
    if (from_blob_) {
tusimple's avatar
tusimple committed
118
119
      return;
    }
120
121
122
123
    if (device_ == -1) {
      if (pinned_) {
#ifdef TV_CUDA
        cudaFreeHost(mPtr);
tusimple's avatar
tusimple committed
124
#endif
125
126
127
128
129
      } else {
        delete[] mPtr;
      }
    } else {
#ifdef TV_CUDA
tusimple's avatar
tusimple committed
130
131
132
133
134
135
136
137
138
139
140
141
      cudaFree(mPtr);
#endif
    }
  };

  inline size_t size() const { return mSize; }

  T *data() { return mPtr; }
  const T *data() const { return mPtr; }

  bool empty() const { return mPtr == nullptr || mSize == 0; }
  bool managed() const { return managed_; }
142
143
  bool pinned() const { return pinned_; }

tusimple's avatar
tusimple committed
144
  int device() const { return device_; }
145
146
  void zero_() {
    if (device_ == -1) {
tusimple's avatar
tusimple committed
147
148
      std::memset(data(), 0, mSize);
      // std::fill(data(), data() + mSize, 0);
149
150
    } else {
#ifdef TV_CUDA
tusimple's avatar
tusimple committed
151
152
      checkCudaErrors(cudaMemset(data(), 0, mSize / sizeof(T)));
#else
153
      TV_THROW_INVALID_ARG("don't compiled with cuda");
tusimple's avatar
tusimple committed
154
155
156
157
158
159
#endif
    }
  }

private:
  size_t mSize = 0;
160
  T *mPtr = nullptr;
tusimple's avatar
tusimple committed
161
  bool from_blob_ = false;
162
  int device_ = -1;
tusimple's avatar
tusimple committed
163
  bool managed_ = false;
164
  bool pinned_ = false;
tusimple's avatar
tusimple committed
165
166
};

167
168
template <typename T> size_t sizeof_dtype(T dtype) {
  switch (dtype) {
tusimple's avatar
tusimple committed
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
  case float32:
    return sizeof(float);
  case int8:
    return sizeof(int8_t);
  case int16:
    return sizeof(int16_t);
  case int32:
    return sizeof(int32_t);
  case float64:
    return sizeof(double);
  case int64:
    return sizeof(int64_t);
  case bool_:
    return sizeof(bool);
  case uint8:
    return sizeof(uint8_t);
  case uint16:
    return sizeof(uint16_t);
  case uint32:
    return sizeof(uint32_t);
  case uint64:
    return sizeof(uint64_t);
  case float16:
192
    return 2;
tusimple's avatar
tusimple committed
193
194
195
196
197
198
  default:
    TV_THROW_RT_ERR("unsupported dtype");
  }
  return 0;
}

199
200
template <typename T> std::string typeString(T t) {
  switch (t) {
tusimple's avatar
tusimple committed
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
  case DType::bool_:
    return "bool";
  case DType::float32:
    return "float32";
  case DType::int8:
    return "int8";
  case DType::int16:
    return "int16";
  case DType::int32:
    return "int32";
  case DType::float64:
    return "float64";
  case DType::int64:
    return "int64";
  case DType::uint8:
    return "uint8";
  case DType::uint16:
    return "uint16";
  case DType::uint32:
    return "uint32";
  case DType::uint64:
    return "uint64";
  case DType::float16:
    return "half";
  default:
    return "";
  }
}

230
231
232
233
234
235
236
237
238
239
240
template <typename T> struct TypeToDtypeTraits;

template <> struct TypeToDtypeTraits<int32_t> {
  static constexpr DType dtype = int32;
};

#ifdef TV_CUDA
template <> struct TypeToDtypeTraits<__half> {
  static constexpr DType dtype = float16;
};
#endif
tusimple's avatar
tusimple committed
241

242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
template <> struct TypeToDtypeTraits<float> {
  static constexpr DType dtype = float32;
};
template <> struct TypeToDtypeTraits<double> {
  static constexpr DType dtype = float64;
};
template <> struct TypeToDtypeTraits<int16_t> {
  static constexpr DType dtype = int16;
};
template <> struct TypeToDtypeTraits<int8_t> {
  static constexpr DType dtype = int8;
};
template <> struct TypeToDtypeTraits<int64_t> {
  static constexpr DType dtype = int64;
};
template <> struct TypeToDtypeTraits<uint8_t> {
  static constexpr DType dtype = uint8;
};
template <> struct TypeToDtypeTraits<uint16_t> {
  static constexpr DType dtype = uint16;
};
template <> struct TypeToDtypeTraits<uint32_t> {
  static constexpr DType dtype = uint32;
};
template <> struct TypeToDtypeTraits<uint64_t> {
  static constexpr DType dtype = uint64;
};
template <> struct TypeToDtypeTraits<bool> {
  static constexpr DType dtype = bool_;
};
template <> struct TypeToDtypeTraits<const int32_t> {
tusimple's avatar
tusimple committed
273
274
275
  static constexpr DType dtype = int32;
};

276
277
#ifdef TV_CUDA
template <> struct TypeToDtypeTraits<const __half> {
tusimple's avatar
tusimple committed
278
279
280
281
  static constexpr DType dtype = float16;
};
#endif

282
template <> struct TypeToDtypeTraits<const float> {
tusimple's avatar
tusimple committed
283
284
  static constexpr DType dtype = float32;
};
285
template <> struct TypeToDtypeTraits<const double> {
tusimple's avatar
tusimple committed
286
287
  static constexpr DType dtype = float64;
};
288
template <> struct TypeToDtypeTraits<const int16_t> {
tusimple's avatar
tusimple committed
289
290
  static constexpr DType dtype = int16;
};
291
template <> struct TypeToDtypeTraits<const int8_t> {
tusimple's avatar
tusimple committed
292
293
  static constexpr DType dtype = int8;
};
294
template <> struct TypeToDtypeTraits<const int64_t> {
tusimple's avatar
tusimple committed
295
296
  static constexpr DType dtype = int64;
};
297
template <> struct TypeToDtypeTraits<const uint8_t> {
tusimple's avatar
tusimple committed
298
299
  static constexpr DType dtype = uint8;
};
300
template <> struct TypeToDtypeTraits<const uint16_t> {
tusimple's avatar
tusimple committed
301
302
  static constexpr DType dtype = uint16;
};
303
template <> struct TypeToDtypeTraits<const uint32_t> {
tusimple's avatar
tusimple committed
304
305
  static constexpr DType dtype = uint32;
};
306
template <> struct TypeToDtypeTraits<const uint64_t> {
tusimple's avatar
tusimple committed
307
308
  static constexpr DType dtype = uint64;
};
309
310
311
template <> struct TypeToDtypeTraits<const bool> {
  static constexpr DType dtype = bool_;
};
tusimple's avatar
tusimple committed
312
313
314

} // namespace detail

315
template <class T> constexpr DType type_v = detail::TypeToDtypeTraits<T>::dtype;
tusimple's avatar
tusimple committed
316

317
template <class... Ts, typename F> bool dispatch_noexcept(DType t, F &&f) {
318
319
320
  static_assert(sizeof...(Ts) > 0, "you need to provide at least one type");
  bool notFound = true;
  mp_for_each<mp_list<Ts...>>([=, &notFound, &f](auto I) {
Yan Yan's avatar
Yan Yan committed
321
    if (type_v<decltype(I)> == t && notFound) {
322
323
324
325
      std::forward<F>(f)(decltype(I)());
      notFound = false;
    }
  });
326
327
328
329
330
  return !notFound;
}

template <class... Ts, typename F> void dispatch(DType t, F &&f) {
  if (!dispatch_noexcept<Ts...>(t, std::forward<F>(f))) {
331
332
333
334
335
336
337
338
339
340
341
342
343
344
    std::stringstream ss;
    mp_for_each<mp_list<Ts...>>([=, &ss](auto I) {
      ss << detail::TypeToString<decltype(I)>::value << " ";
    });
    TV_THROW_RT_ERR("unknown type", detail::typeString(t),
                    ", available:", ss.str());
  }
}

template <typename T, T... Is, typename F> void dispatch_scalar(T idx, F &&f) {
  static_assert(sizeof...(Is) > 0,
                "you need to provide at least one candidate");
  bool notFound = true;
  mp_for_each<mp_list_c<T, Is...>>([=, &notFound, &f](auto I) {
Yan Yan's avatar
Yan Yan committed
345
    if (T(I) == idx && notFound) {
346
347
348
349
350
351
352
353
354
355
356
      std::forward<F>(f)(I);
      notFound = false;
    }
  });
  if (notFound) {
    std::stringstream ss;
    mp_for_each<mp_list_c<T, Is...>>([=, &ss](auto I) { ss << T(I) << " "; });
    TV_THROW_RT_ERR("unknown value", idx, ", available:", ss.str());
  }
}

357
template <int... Is, typename F> bool dispatch_int_noexcept(int idx, F &&f) {
358
359
360
361
  static_assert(sizeof...(Is) > 0,
                "you need to provide at least one candidate");
  bool notFound = true;
  mp_for_each<mp_list_c<int, Is...>>([=, &notFound, &f](auto I) {
Yan Yan's avatar
Yan Yan committed
362
363
364
365
366
    if (decltype(I)::value == idx && notFound) {
      std::forward<F>(f)(I);
      notFound = false;
    }
  });
367
  return !notFound;
Yan Yan's avatar
Yan Yan committed
368
369
370
}

template <int... Is, typename F, class BinaryPredicate>
371
bool dispatch_int_noexcept(int idx, BinaryPredicate p, F &&f) {
Yan Yan's avatar
Yan Yan committed
372
373
374
375
376
  static_assert(sizeof...(Is) > 0,
                "you need to provide at least one candidate");
  bool notFound = true;
  mp_for_each<mp_list_c<int, Is...>>([=, &notFound, &f](auto I) {
    if (p(idx, decltype(I)::value) && notFound) {
377
378
379
380
      std::forward<F>(f)(I);
      notFound = false;
    }
  });
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
  return !notFound;
}

template <int... Is, typename F> void dispatch_int(int idx, F &&f) {
  if (!dispatch_int_noexcept<Is...>(idx, std::forward<F>(f))) {
    std::stringstream ss;
    mp_for_each<mp_list_c<int, Is...>>(
        [=, &ss](auto I) { ss << decltype(I)::value << " "; });
    TV_THROW_RT_ERR("unknown value", idx, ", available:", ss.str());
  }
}

template <int... Is, typename F, class BinaryPredicate>
void dispatch_int(int idx, BinaryPredicate p, F &&f) {
  // BinaryPredicate: BinaryPredicate(idx, candidate)
  if (!dispatch_int_noexcept<Is...>(idx, p, std::forward<F>(f))) {
397
398
    std::stringstream ss;
    mp_for_each<mp_list_c<int, Is...>>(
Yan Yan's avatar
Yan Yan committed
399
        [=, &ss](auto I) { ss << decltype(I)::value << " "; });
400
401
402
403
    TV_THROW_RT_ERR("unknown value", idx, ", available:", ss.str());
  }
}

yanyan's avatar
yanyan committed
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
// Ts is pack of mp_list_c
template <class... Ts, typename Iterator, typename F>
bool dispatch_container_noexcept(Iterator begin, Iterator end, F &&f) {
  static_assert(sizeof...(Ts) > 0,
                "you need to provide at least one candidate");
  bool notFound = true;
  mp_for_each<mp_list<Ts...>>([=, &notFound, &f](auto I) {
    using val_lst_t = decltype(I);
    auto val_lst_size = mp_size<val_lst_t>::value;
    bool equal = true;
    std::size_t count = 0;
    auto iter = begin;
    mp_for_each<val_lst_t>([&](auto E) {
      if (iter == end || !equal) {
        return;
      }
      if (count >= val_lst_size) {
        TV_THROW_INVALID_ARG("iterator length invalid:", val_lst_size);
      }
      constexpr auto c = decltype(E)::value;
      if (c != *iter) {
        equal = false;
      }
      ++count;
      std::advance(iter, 1);
    });
    if (count != val_lst_size || iter != end) {
      equal = false;
    }
    if (equal && notFound) {
      std::forward<F>(f)(I);
      notFound = false;
    }
  });

  return !notFound;
}

template <class... Ts, typename Iterator, typename F>
void dispatch_container(Iterator begin, Iterator end, F &&f) {
  if (!dispatch_container_noexcept<Ts...>(begin, end, std::forward<F>(f))) {
    std::stringstream ss;
    ss << "unknown value [";
    for (auto iter = begin; iter != end; std::advance(iter, 1)) {
      ss << *iter << ",";
    }
    ss << "], available: ";
    mp_for_each<mp_list<Ts...>>([=, &ss](auto I) {
      ss << "[";
      mp_for_each<decltype(I)>(
          [=, &ss](auto E) { ss << decltype(E)::value << ","; });
      ss << "]";
    });
    TV_THROW_RT_ERR(ss.str());
  }
}

461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
/*
template <int... Is, typename F> void dispatch_int(int idx, F &&f) {
  return dispatch_scalar<int, Is...>(idx, f);
}
*/

template <class T> struct Dispatch;

template <template <class...> class T, class... Args>
struct Dispatch<T<Args...>> {
  template <typename F> inline void operator()(DType t, F &&f) {
    return dispatch<Args...>(t, std::forward<F>(f));
  }
};

yanyan's avatar
yanyan committed
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
template <class T> struct DispatchContainer;

template <template <class...> class T, class... Args>
struct DispatchContainer<T<Args...>> {
  template <typename Iterator, typename F>
  inline void operator()(Iterator begin, Iterator end, F &&f) {
    return dispatch_container<Args...>(begin, end, std::forward<F>(f));
  }
};

template <class T> struct DispatchContainerNoexcept;

template <template <class...> class T, class... Args>
struct DispatchContainerNoexcept<T<Args...>> {
  template <typename Iterator, typename F>
  inline bool operator()(Iterator begin, Iterator end, F &&f) {
    return dispatch_container_noexcept<Args...>(begin, end, std::forward<F>(f));
  }
};

496
497
template <class T> struct DispatchInt;

498
499
500
501
502
503
// Args should be std::integral_constant<int, value>
// you need to use type_container<std::integral_constant<int, value>...>
// as template parameter of DispatchInt.
// tv::mp_list_c is ok.
template <template <class...> class T, class... Args>
struct DispatchInt<T<Args...>> {
504
  template <typename F> inline void operator()(int t, F &&f) {
505
    return dispatch_int<Args::value...>(t, std::forward<F>(f));
506
  }
507
508
509
  template <typename F, typename BinaryPredicate>
  inline void operator()(int t, BinaryPredicate p, F &&f) {
    return dispatch_int<Args::value...>(t, p, std::forward<F>(f));
Yan Yan's avatar
Yan Yan committed
510
  }
511
};
Yan Yan's avatar
Yan Yan committed
512

513
514
515
516
constexpr size_t kTensorMaxDim = 10;
using TensorShape = ShapeBase<kTensorMaxDim, int64_t>;

struct Tensor {
tusimple's avatar
tusimple committed
517
  Tensor() {}
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
  Tensor(TensorShape shape, TensorShape stride, DType dtype, int device = -1,
         bool pinned = false, bool managed = false)
      : dtype_(dtype) {
    TV_ASSERT_INVALID_ARG(!shape.empty(), "dont support empty shape");
    storage_ = std::make_shared<detail::TensorStorage<uint8_t>>(
        shape.size() * detail::sizeof_dtype(dtype), device, managed, pinned);
    shape_ = shape;
    stride_ = stride;
  }

  Tensor(TensorShape shape, DType dtype, int device = -1, bool pinned = false,
         bool managed = false)
      : dtype_(dtype) {
    TV_ASSERT_INVALID_ARG(!shape.empty(), "dont support empty shape");
    storage_ = std::make_shared<detail::TensorStorage<uint8_t>>(
        shape.size() * detail::sizeof_dtype(dtype), device, managed, pinned);
    shape_ = shape;
    stride_ = shape.stride_rowmajor();
  }
  Tensor(void *ptr, TensorShape shape, TensorShape stride, DType dtype,
         int device = -1)
      : dtype_(dtype) {
    TV_ASSERT_INVALID_ARG(!shape.empty(), "dont support empty shape");
tusimple's avatar
tusimple committed
541
    storage_ = std::make_shared<detail::TensorStorage<uint8_t>>(
542
543
        reinterpret_cast<uint8_t *>(ptr),
        shape.size() * detail::sizeof_dtype(dtype), device);
tusimple's avatar
tusimple committed
544
    shape_ = shape;
545
    stride_ = stride;
tusimple's avatar
tusimple committed
546
  }
547
548
549
  Tensor(void *ptr, TensorShape shape, DType dtype, int device = -1)
      : dtype_(dtype) {
    TV_ASSERT_INVALID_ARG(!shape.empty(), "dont support empty shape");
tusimple's avatar
tusimple committed
550
551
552
553
    storage_ = std::make_shared<detail::TensorStorage<uint8_t>>(
        reinterpret_cast<uint8_t *>(ptr),
        shape.size() * detail::sizeof_dtype(dtype), device);
    shape_ = shape;
554
    stride_ = shape.stride_rowmajor();
tusimple's avatar
tusimple committed
555
556
  }

557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
  Tensor(const void *ptr, TensorShape shape, TensorShape stride, DType dtype,
         int device = -1)
      : dtype_(dtype), writeable_(false) {
    TV_ASSERT_INVALID_ARG(!shape.empty(), "dont support empty shape");
    storage_ = std::make_shared<detail::TensorStorage<uint8_t>>(
        reinterpret_cast<uint8_t *>(const_cast<void *>(ptr)),
        shape.size() * detail::sizeof_dtype(dtype), device);
    shape_ = shape;
    stride_ = stride;
  }
  Tensor(const void *ptr, TensorShape shape, DType dtype, int device = -1)
      : dtype_(dtype), writeable_(false) {
    TV_ASSERT_INVALID_ARG(!shape.empty(), "dont support empty shape");
    storage_ = std::make_shared<detail::TensorStorage<uint8_t>>(
        reinterpret_cast<uint8_t *>(const_cast<void *>(ptr)),
        shape.size() * detail::sizeof_dtype(dtype), device);
    shape_ = shape;
    stride_ = shape.stride_rowmajor();
  }

  Tensor(std::initializer_list<int32_t> init)
      : Tensor({int(init.size())}, tv::int32) {
    std::copy(init.begin(), init.end(), data<int32_t>());
  }
  Tensor(std::initializer_list<int64_t> init)
      : Tensor({int(init.size())}, tv::int64) {
    std::copy(init.begin(), init.end(), data<int64_t>());
  }
  Tensor(std::initializer_list<float> init)
      : Tensor({int(init.size())}, tv::float32) {
    std::copy(init.begin(), init.end(), data<float>());
  }
  Tensor(std::initializer_list<double> init)
      : Tensor({int(init.size())}, tv::float64) {
    std::copy(init.begin(), init.end(), data<double>());
  }

  template <typename T, int Rank = -1,
            template <class> class PtrTraits = DefaultPtrTraits,
            typename Tindex = int,
            typename std::enable_if<(Rank > 0), int>::type = 0>
  TensorView<T, Rank, PtrTraits, Tindex> tview() {
    using tv_shape_t =
        typename TensorView<T, Rank, PtrTraits, Tindex>::tv_shape_t;
    writable_check();
    static_assert(Rank == -1 || Rank > 0, "error");
tusimple's avatar
tusimple committed
603
    TV_ASSERT_RT_ERR(dtype_ == type_v<T>, "error");
604
605
606
607
608
609
610
    tv_shape_t shape(Rank), stride(Rank);
    for (int i = 0; i < Rank; ++i) {
      shape[i] = shape_[i];
      stride[i] = stride_[i];
    }
    return TensorView<T, Rank, PtrTraits, Tindex>(
        reinterpret_cast<T *>(data<T>()), shape, stride);
tusimple's avatar
tusimple committed
611
  }
612
613
614
615
616
617
618
  template <typename T, int Rank = -1,
            template <class> class PtrTraits = DefaultPtrTraits,
            typename Tindex = int,
            typename std::enable_if<Rank == -1, int>::type = 0>
  TensorView<T, Rank, PtrTraits, Tindex> tview() {
    writable_check();
    static_assert(Rank == -1 || Rank > 0, "error");
tusimple's avatar
tusimple committed
619
    TV_ASSERT_RT_ERR(dtype_ == type_v<T>, "error");
620
    ShapeBase<TV_MAX_DIM, Tindex> shape(ndim()), stride(ndim());
yanyan's avatar
yanyan committed
621
    for (size_t i = 0; i < ndim(); ++i) {
622
623
624
625
626
      shape[i] = shape_[i];
      stride[i] = stride_[i];
    }
    return TensorView<T, Rank, PtrTraits, Tindex>(
        reinterpret_cast<T *>(data<T>()), shape, stride);
tusimple's avatar
tusimple committed
627
  }
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662

  template <typename T, int Rank = -1,
            template <class> class PtrTraits = DefaultPtrTraits,
            typename Tindex = int,
            typename std::enable_if<(Rank > 0), int>::type = 0>
  TensorView<const std::remove_const_t<T>, Rank, PtrTraits, Tindex>
  tview() const {
    static_assert(Rank == -1 || Rank > 0, "error");
    if (Rank > 0) {
      TV_ASSERT_RT_ERR(Rank == ndim(), "error");
    }
    TV_ASSERT_RT_ERR(dtype_ == type_v<T>, "error");

    ShapeBase<Rank == -1 ? TV_MAX_DIM : Rank, Tindex> shape(Rank), stride(Rank);
    for (int i = 0; i < Rank; ++i) {
      shape[i] = shape_[i];
      stride[i] = stride_[i];
    }
    return TensorView<const std::remove_const_t<T>, Rank, PtrTraits, Tindex>(
        reinterpret_cast<const std::remove_const_t<T> *>(data<T>()), shape,
        stride);
  }
  template <typename T, int Rank = -1,
            template <class> class PtrTraits = DefaultPtrTraits,
            typename Tindex = int,
            typename std::enable_if<Rank == -1, int>::type = 0>
  TensorView<const std::remove_const_t<T>, Rank, PtrTraits, Tindex>
  tview() const {
    static_assert(Rank == -1 || Rank > 0, "error");
    if (Rank > 0) {
      TV_ASSERT_RT_ERR(Rank == ndim(), "error");
    }
    TV_ASSERT_RT_ERR(dtype_ == type_v<T>, "error");

    ShapeBase<TV_MAX_DIM, Tindex> shape(ndim()), stride(ndim());
yanyan's avatar
yanyan committed
663
    for (int i = 0; i < int(ndim()); ++i) {
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
      shape[i] = shape_[i];
      stride[i] = stride_[i];
    }
    return TensorView<const std::remove_const_t<T>, Rank, PtrTraits, Tindex>(
        reinterpret_cast<const std::remove_const_t<T> *>(data<T>()), shape,
        stride);
  }

  template <class... Inds> Tensor view(Inds... newShapes) const {
    static_assert(sizeof...(newShapes) > 0, "dont support empty for now");
    TensorShape shape{int(newShapes)...};
    bool found_minus_1 = false;
    for (size_t i = 0; i < shape.ndim(); ++i) {
      if (!found_minus_1) {
        if (shape[i] == -1) {
          shape[i] = 1;
          shape[i] = size() / shape.size();
          found_minus_1 = true;
        } else {
          TV_ASSERT_INVALID_ARG(shape[i] > 0,
                                "shape except -1 must larger than 0");
        }
      } else {
        TV_ASSERT_INVALID_ARG(shape[i] > 0, "multiple -1 in your argument.");
      }
    }
    TV_ASSERT_RT_ERR(shape.size() == size(), "error");
    Tensor res(*this);
    res.shape_ = shape;
    res.stride_ = shape.stride_rowmajor();
    return res;
  }

  Tensor view(TensorShape shape) const {
    TV_ASSERT_RT_ERR(shape.size() == size(), "error");
    Tensor res(*this);
    res.shape_ = shape;
    res.stride_ = shape.stride_rowmajor();
    return res;
  }

yanyan's avatar
yanyan committed
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
  Tensor operator[](int64_t index) {
    TV_ASSERT_INVALID_ARG(ndim() > 1, "error");
    if (index < 0) {
      index += dim(0);
    }
    TV_ASSERT_INVALID_ARG(index < dim(0), "error");
    Tensor res = Tensor();
    res.storage_ = storage_;
    res.shape_ = shape_.subshape(1);
    res.offset_ = offset_ + index * stride_[0];
    res.stride_ = stride_.subshape(1);
    res.writeable_ = writeable_;
    return res;
  }

720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
  Tensor squeeze() const { return view(shape_.squeeze()); }

  Tensor squeeze(int axis) const {
    if (axis < 0) {
      axis = ndim() + axis;
    }
    return view(shape_.squeeze(axis));
  }

  Tensor unsqueeze(int axis) const {
    if (axis < 0) {
      axis = ndim() + axis;
    }
    return view(shape_.unsqueeze(axis));
  }

  bool pinned() const { return storage_->pinned(); }

  Tensor slice_first_axis(int start, int end) const {
    TV_ASSERT_INVALID_ARG(contiguous_, "only support contiguous for now");
    if (start < 0) {
      start = shape_[0] + start;
    }
    if (end < 0) {
      end = shape_[0] + end;
    }
    TV_ASSERT_INVALID_ARG(start < shape_[0], "start must small than dim 0");
    TV_ASSERT_INVALID_ARG(start < end, "start must small than end");
    size_t new_offset = start * shape_.prod(1) * itemsize();
    Tensor res(*this);
    TensorShape newshape(shape_);
    newshape[0] = end - start;
    res.shape_ = newshape;
    res.stride_ = stride_;
    res.offset_ = new_offset;
    return res;
  }

tusimple's avatar
tusimple committed
758
759
760
  bool empty() const { return storage_->empty(); }
  DType dtype() const { return dtype_; }
  int device() const { return storage_->device(); }
761
762
763
  size_t ndim() const { return shape_.ndim(); }

  const TensorShape &shape() const { return shape_; }
yanyan's avatar
yanyan committed
764
  const TensorShape &sizes() const { return shape_; }
765
766
767
768
769
770
771
772
773
774
  const TensorShape &stride() const { return stride_; }

  int dim(int idx) const {
    if (idx < 0) {
      TV_ASSERT_RT_ERR(shape_.size() + idx < shape_.size(), idx, shape_);
      return shape_[shape_.size() + idx];
    } else {
      TV_ASSERT_RT_ERR(idx < int(shape_.size()), idx, shape_);
      return shape_[idx];
    }
tusimple's avatar
tusimple committed
775
  }
776
777
  const uint8_t *raw_data() const { return storage_->data() + offset_; }
  size_t raw_size() const { return size() * itemsize(); }
tusimple's avatar
tusimple committed
778
  size_t size() const { return shape_.size(); }
yanyan's avatar
yanyan committed
779
  size_t size(int64_t idx) const { return dim(idx); }
780
781
782
  size_t itemsize() const { return detail::sizeof_dtype(dtype_); }
  Tensor &zero_() {
    writable_check();
tusimple's avatar
tusimple committed
783
784
785
    storage_->zero_();
    return *this;
  }
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
  uint8_t *raw_data() {
    writable_check();
    return storage_->data() + offset_;
  }
  template <typename T> Tensor &fill_(T value) {
    writable_check();
    TV_ASSERT_RT_ERR(device() == -1, "error");
    Dispatch<detail::all_tensor_types_t>()(dtype_, [&](auto I) {
      using Treal = decltype(I);
      if (std::is_convertible<T, Treal>::value) {
        auto ptr = reinterpret_cast<Treal *>(raw_data());
        std::fill(ptr, ptr + size(), Treal(value));
      } else {
        TV_THROW_INVALID_ARG("not convertable from", type_s<T>, "to",
                             type_s<Treal>);
      }
    });
tusimple's avatar
tusimple committed
803
804
805
    return *this;
  }

806
  template <typename T> T *data() {
tusimple's avatar
tusimple committed
807
    TV_ASSERT_RT_ERR(dtype_ == type_v<T>, "error");
808
    writable_check();
tusimple's avatar
tusimple committed
809
810
811
    return reinterpret_cast<T *>(raw_data());
  }

812
  template <typename T> const T *data() const {
tusimple's avatar
tusimple committed
813
814
815
816
    TV_ASSERT_RT_ERR(dtype_ == type_v<T>, "error");
    return reinterpret_cast<const T *>(raw_data());
  }

yanyan's avatar
yanyan committed
817
818
819
820
821
822
823
824
825
826
  template <typename T> T *data_ptr() { return data<T>(); }

  template <typename T> const T *data_ptr() const { return data<T>(); }

  void *data_ptr() { return reinterpret_cast<void *>(raw_data()); }

  const void *data_ptr() const {
    return reinterpret_cast<const void *>(raw_data());
  }

827
828
829
  void copy_(const Tensor &tensor) {
    writable_check();
    TV_ASSERT_INVALID_ARG(contiguous_, "only support contiguous for now");
tusimple's avatar
tusimple committed
830
831
    TV_ASSERT_RT_ERR(!empty() && !tensor.empty(), "must not empty");
    TV_ASSERT_RT_ERR(size() == tensor.size(), "must have same size");
832
833
834
835
836
    TV_ASSERT_RT_ERR(dtype() == tensor.dtype(), "must have same dtype",
                     detail::typeString(dtype()),
                     detail::typeString(tensor.dtype()));
    if (device() == -1 && tensor.device() == -1) {
#ifdef TV_CUDA
tusimple's avatar
tusimple committed
837
838
839
840
841
842
843
844
      host2host(storage_->data(), tensor.raw_data(),
                size() * detail::sizeof_dtype(dtype_));
#else
      std::copy(tensor.raw_data(),
                tensor.raw_data() + size() * detail::sizeof_dtype(dtype_),
                storage_->data());
#endif
    }
845
846
#ifdef TV_CUDA
    else if (device() >= 0 && tensor.device() == -1) {
tusimple's avatar
tusimple committed
847
848
      host2dev(storage_->data(), tensor.raw_data(),
               size() * detail::sizeof_dtype(dtype_));
849
    } else if (device() == -1 && tensor.device() >= 0) {
tusimple's avatar
tusimple committed
850
851
      dev2host(storage_->data(), tensor.raw_data(),
               size() * detail::sizeof_dtype(dtype_));
852
    } else if (device() >= 0 && tensor.device() >= 0) {
tusimple's avatar
tusimple committed
853
854
855
856
      dev2dev(storage_->data(), tensor.raw_data(),
              size() * detail::sizeof_dtype(dtype_));
    }
#endif
857
858
    else {
      TV_THROW_RT_ERR("only support cpu tensor");
tusimple's avatar
tusimple committed
859
860
861
    }
  }

862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
#ifdef TV_CUDA
  void copy_(const Tensor &tensor, cudaStream_t stream) {
    writable_check();
    TV_ASSERT_INVALID_ARG(contiguous_, "only support contiguous for now");
    TV_ASSERT_RT_ERR(!empty() && !tensor.empty(), "must not empty");
    TV_ASSERT_RT_ERR(size() == tensor.size(), "must have same size");
    TV_ASSERT_RT_ERR(dtype() == tensor.dtype(), "must have same dtype",
                     detail::typeString(dtype()),
                     detail::typeString(tensor.dtype()));
    if (device() == -1 && tensor.device() == -1) {
      host2host(storage_->data(), tensor.raw_data(),
                size() * detail::sizeof_dtype(dtype_), stream);
    } else if (device() >= 0 && tensor.device() == -1) {
      host2dev(storage_->data(), tensor.raw_data(),
               size() * detail::sizeof_dtype(dtype_), stream);
    } else if (device() == -1 && tensor.device() >= 0) {
      dev2host(storage_->data(), tensor.raw_data(),
               size() * detail::sizeof_dtype(dtype_), stream);
    } else if (device() >= 0 && tensor.device() >= 0) {
      dev2dev(storage_->data(), tensor.raw_data(),
              size() * detail::sizeof_dtype(dtype_), stream);
    } else {
      TV_THROW_RT_ERR("only support cpu tensor");
tusimple's avatar
tusimple committed
885
    }
886
887
888
889
890
891
892
893
894
  }
#endif

  Tensor cpu() const {
    if (storage_->device() == -1) {
      // cpu() should always copy tensor.
      return clone();
    }
    Tensor res(shape_, stride_, dtype_, -1, storage_->managed());
tusimple's avatar
tusimple committed
895
896
897
898
    res.copy_(*this);
    return res;
  }

899
900
901
  template <typename T> void copy_(const TensorView<T> &tensor, int device) {
    writable_check();
    TV_ASSERT_INVALID_ARG(contiguous_, "only support contiguous for now");
tusimple's avatar
tusimple committed
902
903
904
905
    Tensor src = from_blob(tensor, device);
    return copy_(src);
  }

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
941
942
943
  Tensor &operator=(const Tensor &tensor) {
    dtype_ = tensor.dtype_;
    storage_ = tensor.storage_;
    shape_ = tensor.shape_;
    writeable_ = tensor.writeable_;
    offset_ = tensor.offset_;
    stride_ = tensor.stride_;
    return *this;
  }

  Tensor(const Tensor &tensor) {
    dtype_ = tensor.dtype_;
    storage_ = tensor.storage_;
    shape_ = tensor.shape_;
    writeable_ = tensor.writeable_;
    offset_ = tensor.offset_;
    stride_ = tensor.stride_;
  }

  Tensor clone(bool pinned = false) const {
    TV_ASSERT_RT_ERR(!empty(), "clone a empty tensor");
    TV_ASSERT_INVALID_ARG(contiguous_, "only support contiguous for now");
    Tensor newtensor(shape_, stride_, dtype_, device(), pinned,
                     storage_->managed());
    newtensor.copy_(*this);
    return newtensor;
  }

  Tensor astype(DType dtype) {
    if (dtype == dtype_) {
      return clone();
    }
    TV_ASSERT_INVALID_ARG(device() == -1, "only support cpu tensor");
    TV_ASSERT_INVALID_ARG(!empty(), "can't be used in empty tensor");
    TV_ASSERT_INVALID_ARG(contiguous_, "only support contiguous for now");
    auto tensor = Tensor();
    Dispatch<detail::all_tensor_types_t>()(dtype, [&](auto Idst) {
      using Tdst = decltype(Idst);
yanyan's avatar
yanyan committed
944
      Dispatch<detail::all_tensor_types_t>()(this->dtype_, [&](auto Icur) {
945
946
        using Tcur = decltype(Icur);
        if (std::is_convertible<Tcur, Tdst>::value) {
yanyan's avatar
yanyan committed
947
          auto ptr = this->data<Tcur>();
yanyan's avatar
yanyan committed
948
949
          tensor = Tensor(this->shape_, this->stride_, dtype, this->device(),
                          this->pinned(), this->storage_->managed());
yanyan's avatar
yanyan committed
950
          std::copy(ptr, ptr + this->size(), tensor.data<Tdst>());
951
952
953
954
955
956
957
958
959
960
961
962
963
        } else {
          TV_THROW_INVALID_ARG("not convertable from", type_s<Tcur>, "to",
                               type_s<Tdst>);
        }
      });
    });
    return tensor;
  }

  template <class... Ts, typename F> inline void dispatch(F &&f) {
    return tv::dispatch<Ts...>(dtype_, std::forward<F>(f));
  }

tusimple's avatar
tusimple committed
964
protected:
965
966
967
968
969
  inline void writable_check() {
    TV_ASSERT_RT_ERR(writeable_,
                     "you cant do non-const operation when not writable");
  }

tusimple's avatar
tusimple committed
970
971
  DType dtype_;
  std::shared_ptr<detail::TensorStorage<uint8_t>> storage_;
972
973
974
975
976
977
978
  TensorShape shape_;
  size_t offset_ = 0;
  TensorShape stride_;

private:
  bool writeable_ = true;
  bool contiguous_ = true;
tusimple's avatar
tusimple committed
979
980
};

981
982
983
984
985
986
987
988
989
990
991
template <typename Os> Os &operator<<(Os &os, const Tensor &tensor) {
  TV_ASSERT_INVALID_ARG(tensor.device() == -1, "must be cpu tensor");
  Dispatch<detail::all_tensor_types_t>()(tensor.dtype(), [&](auto I) {
    using T = decltype(I);
    std::stringstream ss;
    if (std::is_same<T, float>::value || std::is_same<T, double>::value) {
      ss << std::setprecision(4);
    }
    os << tensor.tview<T, -1, DefaultPtrTraits, int64_t>().repr(ss);
  });
  return os;
tusimple's avatar
tusimple committed
992
993
}

994
995
inline Tensor from_blob(void *ptr, TensorShape shape, DType dtype, int device) {
  return Tensor(ptr, shape, dtype, device);
tusimple's avatar
tusimple committed
996
997
}

998
999
1000
inline Tensor from_blob(const void *ptr, TensorShape shape, DType dtype,
                        int device) {
  return Tensor(ptr, shape, dtype, device);
tusimple's avatar
tusimple committed
1001
1002
1003
}

} // namespace tv