TVM 原始碼閱讀PASS — VectorizeLoop-有解無憂

本文地址：https://www.cnblogs.com/wanger-sjtu/p/17501119.html

VectorizeLoop這個PASS就是對標記為ForKind::kVectorized的For回圈做向量化處理，并對For回圈中的陳述句涉及到的變數，替換為Ramp，以便于在Codegen的程序中生成相關的向量化運算的指令，

VectorizeLoop這個PASS的入口函式如下，只有在打開enable_vectorize=true的情況下載才會被啟用，否則VectorizeSkipper會把ForKind::kVectorized的For回圈替換為普通回圈，

Pass VectorizeLoop(bool enable_vectorize) {
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
    auto* n = f.CopyOnWrite();
    if (enable_vectorize) {
      n->body = LoopVectorizer()(std::move(n->body));
    } else {
      n->body = VectorizeSkipper()(std::move(n->body));
    }
    return f;
  };
  return CreatePrimFuncPass(pass_func, 0, "tir.VectorizeLoop", {});
}

下面就以UT中的幾個例子，介紹一下原始碼實作，

vectorize_loop

dtype = "int64"
n = te.var("n")
ib = tvm.tir.ir_builder.create()
A = ib.pointer("float32", name="A")

with ib.for_range(0, n) as i:
 with ib.for_range(0, 4, kind="vectorize") as j:
     A[i*4+j] += tvm.tir.const(1, A.dtype)
stmt = ib.get()
assert isinstance(stmt.body, tvm.tir.For)
mod = tvm.IRModule.from_expr(tvm.tir.PrimFunc([A, n], stmt))
stmt = tvm.tir.transform.VectorizeLoop()(mod)["main"].body

上面的這個代碼完成的是，向量加法，長度為4n的向量A，對每個元素+1，

# before
for (i, 0, n) {
  vectorized (j, 0, 4) {
    A[((i*4) + j)] = (A[((i*4) + j)] + 1f)
  }
}
# after
for (i, 0, n) {
  A[ramp((i*4), 1, 4)] = (A[ramp((i*4), 1, 4)] + x4(1f))
}

可以看到在經過VectorizeLoop的PASS以后，內層的回圈消掉了，替換成為了一個Ramp的向量指令，這個在CPU中會被替換為SIMD指令（neon，AVX等）

PASS流程

在向量化的處理的PASS中是在LoopVectorizer中處理的，處理For回圈部分，

class LoopVectorizer : public StmtMutator {
 public:
  Stmt VisitStmt_(const ForNode* op) final {
    if (op->kind == ForKind::kVectorized) {
      ICHECK(is_zero(op->min));
      auto* extent_as_int = op->extent.as<IntImmNode>();
      if (!extent_as_int || extent_as_int->value < 1) {
        LOG(FATAL) << "Failed to vectorize loop with extent " << op->extent;
      }
      return Vectorizer(op->loop_var, static_cast<int>(extent_as_int->value))(op->body);
    } else {
      return StmtMutator::VisitStmt_(op);
    }
  }
};

當遇到需要向量化的節點時，首先記錄回圈變數和范圍，這個在后續替換相應的Load和Store操作為Ramp時用到，然后就到了Vectorizer部分，遍歷For回圈體，修改相應的stmt，

Vectorizer(Var var, int var_lanes) : var_(var), var_lanes_(var_lanes) {
    ramp_ = Ramp(0, 1, var_lanes);
}

在Vectorizer中對不同的PrimExpr、Stmt做了多載，這里不逐一介紹，就以上面的向量加計算，介紹一下用到的函式以及流程，

首先看一下這里的上面sch的For的回圈內的計算邏輯：

 A[((i*4) + j)] = (A[((i*4) + j)] + 1f)

因為TVM中，Stmt的表達可以視為一個DSL的語言，訪問的時候也是按照深度優先的策略遍歷的AST，這里把上面的計算程序簡單表示為一個AST的語法樹，然后再分析一下流程中呼叫的各個函式是如何處理的，

從上面的AST的示意圖可以看出來，對于上面的sch，依次訪問了BufferStoreNode、Add Mul、BufferLoadNode 等，這里就以這幾個Node的處理介紹一下向量化的程序，

所謂向量化的程序就是把這個標記為kVectorized的標量回圈操作映射到向量化的操作，對于上面的例子來說就是把所有關于j的訪問映射為RampNode，以便于后續處理可以正確生成相應的指令，

BufferStoreNode

BufferStoreNode中有三部分：

buffer——寫入的buffer
value——待寫入的值或者運算式
indices——寫入buffer的坐標
這里的目的就是修改value和indices中的內容，
對于indices，是在這里完成的，最終通過MapHelper依次訪問了indices的運算式，

auto fmutate = [this](const PrimExpr& index) { return this->VisitExpr(index); };
Array<PrimExpr> indices = op->indices.Map(fmutate);

對于value 則是直接遍歷，

PrimExpr value = https://www.cnblogs.com/wanger-sjtu/archive/2023/06/24/this->VisitExpr(op->value);

AddNode

對于AddNode和SubNode 都會走到AddSubVec這個模板函式，
這個函式里面首先會遍歷左右運算式，

PrimExpr a = this->VisitExpr(op->a);
PrimExpr b = this->VisitExpr(op->b);
if (a.same_as(op->a) && b.same_as(op->b)) {
 return GetRef<PrimExpr>(op);
} else {
int lanes = std::max(a.dtype().lanes(), b.dtype().lanes());
if (lanes != 1) {
 const RampNode* b_ramp = b.as<RampNode>();
 const RampNode* a_ramp = a.as<RampNode>();
 if (a.dtype().lanes() == 1 && b_ramp) {
   return Ramp(fcompute(a, b_ramp->base),
		 fcompute(make_zero(b_ramp->stride.dtype()), b_ramp->stride), b_ramp->lanes);
 }
 if (b.dtype().lanes() == 1 && a_ramp) {
   return Ramp(fcompute(a_ramp->base, b), a_ramp->stride, a_ramp->lanes);
 }
}
return fcompute(BroadcastTo(a, lanes), BroadcastTo(b, lanes));

如果遍歷之后沒有變化，就直接回傳了，而對于這里的我們需要計算的是

((i*4) + j)

j 是需要向量化的坐標，i*4 是沒有變化的，遍歷以后a沒變化，b變成了T.Ramp(0, 1, 4) 這時候lanes=4，會走到第一個if分支，回傳的是新構造的RampNode

 T.Ramp(i * 4, 1, 4)

其他的分支也類似，比如：

A[i * 4 + j] + T.float32(1)
// --- after ---
A[i * 4:i * 4 + 4]   T.float32(1)

這里會把a、b broadcast為一個向量再做計算，

VarNode

對于這里的VarNode判斷就比較簡單了，如果匹配到的是需要向量化的變數，就回傳建構式中構造的RampNode，否則就回傳，其他的操作，暫時略過，

Var var = GetRef<Var>(op);
if (var.same_as(var_)) {
 return ramp_;
}
// ...
else {
 return std::move(var);
}

MulNode

PrimExpr a = this->VisitExpr(op->a);
PrimExpr b = this->VisitExpr(op->b);
if (a.same_as(op->a) && b.same_as(op->b)) {
return GetRef<PrimExpr>(op);
} else {
int lanes = std::max(a.dtype().lanes(), b.dtype().lanes());
if (lanes != 1) {
 const RampNode* b_ramp = b.as<RampNode>();
 const RampNode* a_ramp = a.as<RampNode>();
 if (a_ramp && b.dtype().lanes() == 1 && analyzer_.CanProve(b > 0)) {
   return Ramp(a_ramp->base * b, a_ramp->stride * b, a_ramp->lanes);
 }
 if (b_ramp && a.dtype().lanes() == 1 && analyzer_.CanProve(a > 0)) {
   return Ramp(b_ramp->base * a, b_ramp->stride * a, b_ramp->lanes);
 }
}
return Mul(BroadcastTo(a, lanes), BroadcastTo(b, lanes));
}
return BinaryVec<Mul>(op);

這里的處理邏輯與Add基本一致，只是在計算RampNode的時候有點區別，

轉載請註明出處，本文鏈接：https://www.uj5u.com/qita/555872.html

標籤：其他

上一篇：怎么讓英文大預言模型支持中文？（一）構建自己的tokenization

下一篇：返回列表