elem.h 1.53 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
/*!
 * \file tl/op/elem.h
 * \brief Define elment-wise operators.
 *
 */

#ifndef TVM_TL_OP_ELEM_H_
#define TVM_TL_OP_ELEM_H_

#include "op.h"
#include "parallel.h"

namespace tvm {
namespace tl {

using namespace tir;

class Copy : public Operator {
19
public:
20
  Copy(Array<PrimExpr> args, BufferMap vmap);
21
22
  Stmt Lower(const LowerArgs &T, arith::Analyzer *analyzer) const final;
  LayoutMap InferLayout(const LayoutInferArgs &T, InferLevel level) final;
23

24
  static const Op &Get();
25

26
27
28
protected:
  Stmt LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const;
  Stmt LowerLDSMCopy(const LowerArgs &T, arith::Analyzer *analyzer) const;
29

30
  For MakeSIMTLoop(arith::Analyzer *analyzer) const;
31
32
33
34
  Array<IterVar> MakeIterVars() const;

  // ivs: itervars returned by MakeIterVars()
  // src_dst: 0 for src_indices, 1 for dst_indices
35
  Array<PrimExpr> MakeIndices(const Array<IterVar> &ivs, int src_dst) const;
36

37
  PrimExpr MakePredicate(arith::Analyzer *analyzer, const Array<IterVar> &ivs,
38
39
40
41
42
43
44
45
46
47
48
49
                         Array<PrimExpr> extents, int src_dst) const;

  Array<PrimExpr> args_;

  Buffer src, dst;
  Array<Range> src_range, dst_range;
  IntImm coalesced_width;

  std::unique_ptr<ParallelOp> par_op_;
};

class Fill : public Operator {
50
public:
51
  Fill(Array<PrimExpr> args, BufferMap vmap);
52
53
  Stmt Lower(const LowerArgs &T, arith::Analyzer *analyzer) const final;
  static const Op &Get();
54

55
56
private:
  For MakeSIMTLoop(arith::Analyzer *analyzer) const;
57
58
59
60
  tir::Buffer dst;
  PrimExpr value;
};

61
62
} // namespace tl
} // namespace tvm
63

64
#endif //  TVM_TL_OP_ELEM_H_