taco/src/lower/mode_format_dense.cpp at 1fae902e77413e8ee41682ae067f9195e60b514c · tensor-compiler/taco · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
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
84
85
86
87
88
89
90
#include "taco/lower/mode_format_dense.h"

using namespace std;
using namespace taco::ir;

namespace taco {

DenseModeFormat::DenseModeFormat() : DenseModeFormat(true, true, false) {
}

DenseModeFormat::DenseModeFormat(const bool isOrdered, const bool isUnique, const bool isZeroless) :
    ModeFormatImpl("dense", true, isOrdered, isUnique, false, true, isZeroless, false,
                   false, true, true, false) {
}

ModeFormat DenseModeFormat::copy(
    std::vector<ModeFormat::Property> properties) const {
  bool isOrdered = this->isOrdered;
  bool isUnique = this->isUnique;
  bool isZeroless = this->isZeroless;
  for (const auto property : properties) {
    switch (property) {
      case ModeFormat::ORDERED:
        isOrdered = true;
        break;
      case ModeFormat::NOT_ORDERED:
        isOrdered = false;
        break;
      case ModeFormat::UNIQUE:
        isUnique = true;
        break;
      case ModeFormat::NOT_UNIQUE:
        isUnique = false;
        break;
      case ModeFormat::ZEROLESS:
        isZeroless = true;
        break;
      case ModeFormat::NOT_ZEROLESS:
        isZeroless = false;
        break;
      default:
        break;
    }
  }
  return ModeFormat(std::make_shared<DenseModeFormat>(isOrdered, isUnique, isZeroless));
}

ModeFunction DenseModeFormat::locate(ir::Expr parentPos,
                                   std::vector<ir::Expr> coords,
                                   Mode mode) const {
  Expr pos = Add::make(Mul::make(parentPos, getWidth(mode)), coords.back());
  return ModeFunction(Stmt(), {pos, true});
}

Stmt DenseModeFormat::getInsertCoord(Expr p,
    const std::vector<Expr>& i, Mode mode) const {
  return Stmt();
}

Expr DenseModeFormat::getWidth(Mode mode) const {
  return (mode.getSize().isFixed() && mode.getSize().getSize() < 16) ?
         (int)mode.getSize().getSize() :
         getSizeArray(mode.getModePack());
}

Stmt DenseModeFormat::getInsertInitCoords(Expr pBegin,
    Expr pEnd, Mode mode) const {
  return Stmt();
}

Stmt DenseModeFormat::getInsertInitLevel(Expr szPrev, Expr sz,
    Mode mode) const {
  return Stmt();
}

Stmt DenseModeFormat::getInsertFinalizeLevel(Expr szPrev,
    Expr sz, Mode mode) const {
  return Stmt();
}

vector<Expr> DenseModeFormat::getArrays(Expr tensor, int mode,
                                        int level) const {
  return {GetProperty::make(tensor, TensorProperty::Dimension, mode)};
}

Expr DenseModeFormat::getSizeArray(ModePack pack) const {
  return pack.getArray(0);
}

}