You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2023/02/22 12:56:51 UTC
[tvm] branch unity updated: [Unity][Analysis] TIR pattern kind analysis for multi-buffer write block (#14075)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 7e792d29af [Unity][Analysis] TIR pattern kind analysis for multi-buffer write block (#14075)
7e792d29af is described below
commit 7e792d29af7c924177b66b6527c338cc058656d9
Author: Ruihang Lai <ru...@cs.cmu.edu>
AuthorDate: Wed Feb 22 07:56:45 2023 -0500
[Unity][Analysis] TIR pattern kind analysis for multi-buffer write block (#14075)
This PR supports TIR pattern kind analysis for TIR blocks which write
to multiple buffers, which is helpful for normalization operators like
layernorm, groupnorm, etc.
Prior to this PR, the analyzer does not support a blocks which write to
multiple buffers. On seeing such a block, the analyzer simply sets the
analysis result to "opaque". With this PR, on seeing a block which
writes multiple buffers, the analyzer will check if all the BufferStores
have the same indices. And it will only set the result to "opaque" when
the BufferStores have different indices.
By doing this, the analysis works for common cases where a block may
write to multiple buffers, like layernorm or groupnorm.
Besides the unit test for the analysis itself, this PR also adds a unit
test for FuseOps pass, make sure that a "layernorm + relu" pattern can
be fused together.
---
src/relax/analysis/tir_op_pattern_kind.cc | 14 +--
.../test_transform_annotate_tir_op_pattern.py | 27 ++++++
tests/python/relax/test_transform_fuse_ops.py | 102 +++++++++++++++++++++
3 files changed, 136 insertions(+), 7 deletions(-)
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc
index b7ac8faddd..dfa073fd9c 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -50,9 +50,9 @@ class PatternKindAnalyzer : public StmtExprVisitor {
}
void VisitStmt_(const BufferStoreNode* op) final {
- // We only support one buffer store in a block (ususally generated by TE compute)
+ // We only support one buffer store in a block (usually generated by TE compute)
// If we have already seen buffer store in the current block, classify as Opaque.
- if (store_.defined()) {
+ if (store_.defined() && !IsSameArray(op->indices, store_.value()->indices)) {
kind_ = relay::kOpaque;
return;
}
@@ -85,7 +85,7 @@ class PatternKindAnalyzer : public StmtExprVisitor {
for (const BufferLoad& load : loads_) {
// Since elemwise is stricter than broadcast and broadcast is stricter than injective,
// while the order amount enums: kElemWise < kBroadcast < kInjective.
- // We can simpily use `std::max` to detect these three patterns.
+ // We can simply use `std::max` to detect these three patterns.
// E.g Here is only one store node but two load nodes, like C[i, j] = A[i, j] + B[i]
// Buffer C and A are elemwise but C and B are broadcast. So the whole block follows
// broadcast pattern.
@@ -190,7 +190,7 @@ class PatternKindAnalyzer : public StmtExprVisitor {
continue;
}
- // Try to find the i-th load indice in the store indices.
+ // Try to find the i-th load index in the store indices.
while (j < ndim_store_buf && !store->indices[j].same_as(load->indices[i])) {
++j;
}
@@ -205,10 +205,10 @@ class PatternKindAnalyzer : public StmtExprVisitor {
/*!
* \brief Checking the load indices and store indices follows injective pattern.
- * It's injective pattern iff all load indice vars are in the store indices, no matter orders.
+ * It's injective pattern iff all load index vars are in the store indices, no matter orders.
* Note that we only support store indices are direct vars so far, which can be enhance later.
* E.g. A[i, j] = B[j, i] is injective.
- * A[i, j] = B[i - j] is injective since the load indice vars are only i, j
+ * A[i, j] = B[i - j] is injective since the load index vars are only i, j
*/
static bool IsInjectivePattern(const BufferStore& store, const BufferLoad& load) {
std::unordered_set<const tir::VarNode*> vars;
@@ -307,7 +307,7 @@ class PatternKindAnalyzer : public StmtExprVisitor {
private:
/*!
* \brief The BufferStore node in the current block.
- * \note We only support one BufferStore node in a block (ususally generated by TE compute)
+ * \note We only support one BufferStore node in a block (usually generated by TE compute)
*/
Optional<BufferStore> store_;
/*! \brief The BufferLoad nodes in the current block. */
diff --git a/tests/python/relax/test_transform_annotate_tir_op_pattern.py b/tests/python/relax/test_transform_annotate_tir_op_pattern.py
index 23ce49a7c2..5fc8c99367 100644
--- a/tests/python/relax/test_transform_annotate_tir_op_pattern.py
+++ b/tests/python/relax/test_transform_annotate_tir_op_pattern.py
@@ -356,5 +356,32 @@ def test_multiple_bufer_stores_fallback():
assert new_mod["cumsum"].attrs["op_pattern"] == OpPatternKind.kOpaque
+def test_sum_sqsum():
+ @tvm.script.ir_module
+ class Module:
+ @T.prim_func
+ def sum_sqsum(
+ A: T.Buffer((32, 64), "float32"),
+ vsum: T.Buffer((32,), "float32"),
+ sqsum: T.Buffer((32,), "float32"),
+ ):
+ for ax0, k0 in T.grid(32, 64):
+ with T.block("block"):
+ v_ax0, v_k0 = T.axis.remap("SR", [ax0, k0])
+ T.reads(A[v_ax0, v_k0])
+ T.writes(vsum[v_ax0], sqsum[v_ax0])
+ with T.init():
+ vsum[v_ax0] = T.float32(0)
+ sqsum[v_ax0] = T.float32(0)
+ v_vsum: T.float32 = vsum[v_ax0] + A[v_ax0, v_k0]
+ v_sqsum: T.float32 = sqsum[v_ax0] + A[v_ax0, v_k0] * A[v_ax0, v_k0]
+ vsum[v_ax0] = v_vsum
+ sqsum[v_ax0] = v_sqsum
+
+ mod = Module
+ new_mod = relax.transform.AnnotateTIROpPattern()(mod)
+ assert new_mod["sum_sqsum"].attrs["op_pattern"] == OpPatternKind.kCommReduce
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_transform_fuse_ops.py b/tests/python/relax/test_transform_fuse_ops.py
index d38e582981..a7a6066c4b 100644
--- a/tests/python/relax/test_transform_fuse_ops.py
+++ b/tests/python/relax/test_transform_fuse_ops.py
@@ -854,5 +854,107 @@ def test_edge_with_call_dps_packed():
_check(Module, Module)
+def test_layer_norm_silu():
+ # fmt: off
+ @I.ir_module
+ class Module:
+ @R.function
+ def main(x: R.Tensor((1, 512, 64, 64), "float32"), mean: R.Tensor((64, 64), "float32"), var: R.Tensor((64, 64), "float32")):
+ with R.dataflow():
+ gv0 = R.call_tir(layer_norm, (x, mean, var), out_sinfo=R.Tensor((1, 512, 64, 64)))
+ gv1 = R.call_tir(relu, gv0, out_sinfo=R.Tensor((1, 512, 64, 64), "float32"))
+ R.output(gv1)
+ return gv1
+
+ @T.prim_func
+ def layer_norm(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), gamma: T.Buffer((T.int64(64), T.int64(64)), "float32"), beta: T.Buffer((T.int64(64), T.int64(64)), "float32"), T_layer_norm: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")):
+ rxplaceholder_red_temp_v0 = T.alloc_buffer([T.int64(64), T.int64(64)], dtype="float32")
+ rxplaceholder_red_temp_v1 = T.alloc_buffer([T.int64(64), T.int64(64)], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
+ with T.block("rxplaceholder_red_temp"):
+ ax0, ax1, k2, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3])
+ T.reads(A[ax0, ax1, k2, k3])
+ T.writes(rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1])
+ with T.init():
+ rxplaceholder_red_temp_v0[ax0, ax1] = T.float32(0)
+ rxplaceholder_red_temp_v1[ax0, ax1] = T.float32(0)
+ v_rxplaceholder_red_temp_v0: T.float32 = rxplaceholder_red_temp_v0[ax0, ax1] + A[ax0, ax1, k2, k3]
+ v_rxplaceholder_red_temp_v1: T.float32 = rxplaceholder_red_temp_v1[ax0, ax1] + A[ax0, ax1, k2, k3] * A[ax0, ax1, k2, k3]
+ rxplaceholder_red_temp_v0[ax0, ax1] = v_rxplaceholder_red_temp_v0
+ rxplaceholder_red_temp_v1[ax0, ax1] = v_rxplaceholder_red_temp_v1
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
+ with T.block("T_layer_norm"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(A[ax0, ax1, ax2, ax3], rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1], gamma[ax2, ax3], beta[ax2, ax3])
+ T.writes(T_layer_norm[ax0, ax1, ax2, ax3])
+ T_layer_norm[ax0, ax1, ax2, ax3] = (A[ax0, ax1, ax2, ax3] - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.05)) * T.rsqrt(rxplaceholder_red_temp_v1[ax0, ax1] * T.float32(0.05) - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.05) * (rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.05)) + T.float32(1e-05), dtype="float32") * gamma[ax2, ax3] + beta[ax2, ax3]
+
+ @T.prim_func
+ def relu(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), B: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")):
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
+ with T.block("relu"):
+ v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(A[v_i0, v_i1, v_i2, v_i3])
+ T.writes(B[v_i0, v_i1, v_i2, v_i3])
+ B[v_i0, v_i1, v_i2, v_i3] = T.max(A[v_i0, v_i1, v_i2, v_i3], T.float32(0))
+
+ @I.ir_module
+ class Expected:
+ @T.prim_func
+ def layer_norm(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), gamma: T.Buffer((T.int64(64), T.int64(64)), "float32"), beta: T.Buffer((T.int64(64), T.int64(64)), "float32"), T_layer_norm: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")):
+ T.func_attr({"op_pattern": 4})
+ # with T.block("root"):
+ rxplaceholder_red_temp_v0 = T.alloc_buffer((T.int64(64), T.int64(64)))
+ rxplaceholder_red_temp_v1 = T.alloc_buffer((T.int64(64), T.int64(64)))
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
+ with T.block("rxplaceholder_red_temp"):
+ ax0, ax1, k2, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3])
+ T.reads(A[ax0, ax1, k2, k3])
+ T.writes(rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1])
+ with T.init():
+ rxplaceholder_red_temp_v0[ax0, ax1] = T.float32(0)
+ rxplaceholder_red_temp_v1[ax0, ax1] = T.float32(0)
+ v_rxplaceholder_red_temp_v0: T.float32 = rxplaceholder_red_temp_v0[ax0, ax1] + A[ax0, ax1, k2, k3]
+ v_rxplaceholder_red_temp_v1: T.float32 = rxplaceholder_red_temp_v1[ax0, ax1] + A[ax0, ax1, k2, k3] * A[ax0, ax1, k2, k3]
+ rxplaceholder_red_temp_v0[ax0, ax1] = v_rxplaceholder_red_temp_v0
+ rxplaceholder_red_temp_v1[ax0, ax1] = v_rxplaceholder_red_temp_v1
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
+ with T.block("T_layer_norm"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(A[ax0, ax1, ax2, ax3], rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1], gamma[ax2, ax3], beta[ax2, ax3])
+ T.writes(T_layer_norm[ax0, ax1, ax2, ax3])
+ T_layer_norm[ax0, ax1, ax2, ax3] = (A[ax0, ax1, ax2, ax3] - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.050000000000000003)) * T.rsqrt(rxplaceholder_red_temp_v1[ax0, ax1] * T.float32(0.050000000000000003) - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.050000000000000003) * (rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.050000000000000003)) + T.float32(1.0000000000000001e-05)) * gamma[ax2, ax3] + beta[ax2, ax3]
+
+ @T.prim_func
+ def relu(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), B: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")):
+ T.func_attr({"op_pattern": 0})
+ # with T.block("root"):
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)):
+ with T.block("relu"):
+ v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(A[v_i0, v_i1, v_i2, v_i3])
+ T.writes(B[v_i0, v_i1, v_i2, v_i3])
+ B[v_i0, v_i1, v_i2, v_i3] = T.max(A[v_i0, v_i1, v_i2, v_i3], T.float32(0))
+
+ @R.function
+ def fused_layer_norm_relu(x: R.Tensor((1, 512, 64, 64), dtype="float32"), mean: R.Tensor((64, 64), dtype="float32"), var: R.Tensor((64, 64), dtype="float32")) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
+ R.func_attr({"Primitive": 1})
+ with R.dataflow():
+ gv0 = R.call_tir(layer_norm, (x, mean, var), out_sinfo=R.Tensor((1, 512, 64, 64)))
+ gv = R.call_tir(relu, (gv0,), out_sinfo=R.Tensor((1, 512, 64, 64), dtype="float32"))
+ R.output(gv)
+ return gv
+
+ @R.function
+ def main(x: R.Tensor((1, 512, 64, 64), dtype="float32"), mean: R.Tensor((64, 64), dtype="float32"), var: R.Tensor((64, 64), dtype="float32")) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
+ with R.dataflow():
+ gv: R.Tensor((1, 512, 64, 64), dtype="float32") = fused_layer_norm_relu(x, mean, var)
+ R.output(gv)
+ return gv
+ # fmt: on
+
+ _check(Module, Expected)
+
+
if __name__ == "__main__":
tvm.testing.main()