You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/09/26 05:49:28 UTC
[GitHub] [tvm] wrongtest opened a new pull request #9121: [TIR] add loop partition hint pragma
wrongtest opened a new pull request #9121:
URL: https://github.com/apache/tvm/pull/9121
Currently `LoopPartition` pass will try to partition loops assiociated with condition in likely tag, it would be great if developers can take control of which loop to partition, no-matter whether the condition to eliminate is "likely" tagged.
The PR add a pragma attr key `loop_partition_hint`, which can be tagged explicitly in schedule phace. The loop partition pass will consider all arith comparation conditions for hinted loop var.
Below are two examples of how explicit controlled loop partition benefits, the target is on Ubuntu20.08 i7-7700, with llvm version 11.0 :
- For max pooling with padding inlined, which create conditional buffer accesses
```python
data = te.placeholder([1, 128, 56, 56], name="x")
out = topi.nn.pool2d(data, kernel=[5, 5], stride=[1, 1], padding=[2, 2, 2, 2], pool_type="max", dilation=[1, 1], layout="NCHW")
pad = out.op.input_tensors[0]
x = tvm.nd.array(np.random.randint(0, 64, [1, 128, 56, 56]).astype("float32"))
def test(do_partition):
s = te.create_schedule([out.op])
s[pad].compute_inline()
n, c, h, w = s[out].op.axis
if do_partition:
s[out].pragma(h, "loop_partition_hint")
s[out].pragma(w, "loop_partition_hint")
with tvm.ir.transform.PassContext(config={"tir.LoopPartition": {"partition_const_loop": True}}):
f = tvm.build(s, [data, out], "llvm")
y = tvm.nd.array(np.zeros([1, 128, 56, 56]).astype("float32"))
f(x, y)
result = y.asnumpy()
print(f.get_source("asm"))
evaluator = f.time_evaluator(f.entry_name, tvm.cpu(), number=1000)
print("partition=%s: %.3f millisecs" % (do_partition, evaluator(x, y).mean * 1000))
return result
r1 = test(do_partition=False)
r2 = test(do_partition=True)
testing.assert_allclose(r1, r2, rtol=1e-5)
```
The performance I get:
- no loop partition: 3.708 millisecs
- with loop partition: 0.975 millisecs
- For tiled matmul following TVM tensor expression tutorial, but with shape not divided by tiling factor. The tir split do not create a likely condition for it now.
```python
M, N, K = 1025, 1025, 1025
dtype = "float32"
dev = tvm.cpu()
a = tvm.nd.array(np.random.rand(M, K).astype(dtype), dev)
b = tvm.nd.array(np.random.rand(K, N).astype(dtype), dev)
k = te.reduce_axis((0, K), "k")
A = te.placeholder((M, K), name="A")
B = te.placeholder((K, N), name="B")
C = te.compute((M, N), lambda x, y: te.sum(A[x, k] * B[k, y], axis=k), name="C")
f = te.create_prim_func([A, B, C])
s = tvm.tir.Schedule(f)
def evaluate_operation(s, target, optimization):
with tvm.ir.transform.PassContext(config={"tir.LoopPartition": {"partition_const_loop": True}}):
print(tvm.lower(s.mod["main"], [], simple_mode=True))
func = tvm.build(s.mod["main"], [], target=target, name="mmult")
assert func
c = tvm.nd.array(np.zeros((M, N), dtype=dtype), dev)
func(a, b, c)
evaluator = func.time_evaluator(func.entry_name, dev, number=10)
mean_time = evaluator(a, b, c).mean
print("%s: %f" % (optimization, mean_time))
# no opt
evaluate_operation(s, target="llvm", optimization="none")
# tiling and vectorize
x, y, k = s.get_loops(s.get_block("C"))
xo, xi = s.split(x, factors=[None, 32])
yo, yi = s.split(y, factors=[None, 32])
ko, ki = s.split(k, factors=[None, 4])
s.reorder(xo, yo, ko, ki, xi, yi)
s.vectorize(yi)
evaluate_operation(s, target="llvm", optimization="blocking")
# loop partition
def pragma(s, rv, key):
sref = s.get_sref(rv)
loop = sref.stmt
new_loop = tvm.tir.For(loop.loop_var, loop.min, loop.extent, loop.kind, loop.body, annotations={key: 1})
s.state.replace(sref, new_loop)
pragma(s, xo, "pragma_loop_partition_hint")
pragma(s, yo, "pragma_loop_partition_hint")
evaluate_operation(s, target="llvm", optimization="loop_partition")
```
The performance I get:
- no opt: 3.708 millisecs: 1.374402
- with tiling + vectorize: 0.843930
- with tiling + vectorize + loop partition: 0.272183
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #9121:
URL: https://github.com/apache/tvm/pull/9121#issuecomment-928244294
CC @areusch @Hzfengsy @vinx13 @zxybazh @ZihengJiang would you guys review this PR? looks like it's relevant to some of our previous discussion
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Hzfengsy commented on a change in pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on a change in pull request #9121:
URL: https://github.com/apache/tvm/pull/9121#discussion_r717139926
##########
File path: tests/python/unittest/test_tir_transform_loop_partition.py
##########
@@ -538,6 +538,27 @@ def test_simple_rfactor():
assert not tvm.ir.structural_equal(stmt1.body, stmt2.body)
+def test_explicit_partition_hint():
+ A = te.placeholder((16,), name="A")
+ B = te.placeholder((16,), name="B")
+ C = te.compute((32,), lambda i: te.if_then_else(i < 16, A[i], B[i]), name="C")
+ s = te.create_schedule(C.op)
+ s.normalize()
+ s[C].pragma(s[C].op.axis[0], "loop_partition_hint")
+ bounds = tvm.te.schedule.InferBound(s)
+ stmt1 = tvm.te.schedule.ScheduleOps(s, bounds)
+
+ mod1 = tvm.IRModule.from_expr(tvm.tir.PrimFunc([], stmt1))
+ stmt1 = tvm.tir.transform.Simplify()(mod1)["main"].body
+
+ with tvm.transform.PassContext(config={"tir.LoopPartition": {"partition_const_loop": True}}):
+ mod2 = tvm.tir.transform.LoopPartition()(mod1)
+ stmt2 = tvm.tir.transform.Simplify()(mod2)["main"].body
+
+ # make sure loop partition actually did something
+ assert not tvm.ir.structural_equal(stmt1.body, stmt2.body)
Review comment:
It's better to assert `structural_equal` to the expected result.
##########
File path: include/tvm/tir/stmt.h
##########
@@ -1339,6 +1339,12 @@ constexpr const char* hand_threaded = "hand_threaded";
* if (mask & 2) the write region should be detected.
*/
constexpr const char* script_parsing_detect_access = "tir.script_parsing_detect_access";
+
+/*!
+ * \brief Mark that the loop should be partitioned.
+ */
+constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint";
Review comment:
```suggestion
constexpr const char* pragma_loop_partition_hint = "tir.pragma_loop_partition_hint";
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Hzfengsy merged pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
Hzfengsy merged pull request #9121:
URL: https://github.com/apache/tvm/pull/9121
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #9121:
URL: https://github.com/apache/tvm/pull/9121#issuecomment-928244294
CC @areusch @Hzfengsy @vinx13 @zxybazh @ZihengJiang would you guys review this PR? looks like it's relevant to some of our previous discussion
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Hzfengsy commented on a change in pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on a change in pull request #9121:
URL: https://github.com/apache/tvm/pull/9121#discussion_r717139926
##########
File path: tests/python/unittest/test_tir_transform_loop_partition.py
##########
@@ -538,6 +538,27 @@ def test_simple_rfactor():
assert not tvm.ir.structural_equal(stmt1.body, stmt2.body)
+def test_explicit_partition_hint():
+ A = te.placeholder((16,), name="A")
+ B = te.placeholder((16,), name="B")
+ C = te.compute((32,), lambda i: te.if_then_else(i < 16, A[i], B[i]), name="C")
+ s = te.create_schedule(C.op)
+ s.normalize()
+ s[C].pragma(s[C].op.axis[0], "loop_partition_hint")
+ bounds = tvm.te.schedule.InferBound(s)
+ stmt1 = tvm.te.schedule.ScheduleOps(s, bounds)
+
+ mod1 = tvm.IRModule.from_expr(tvm.tir.PrimFunc([], stmt1))
+ stmt1 = tvm.tir.transform.Simplify()(mod1)["main"].body
+
+ with tvm.transform.PassContext(config={"tir.LoopPartition": {"partition_const_loop": True}}):
+ mod2 = tvm.tir.transform.LoopPartition()(mod1)
+ stmt2 = tvm.tir.transform.Simplify()(mod2)["main"].body
+
+ # make sure loop partition actually did something
+ assert not tvm.ir.structural_equal(stmt1.body, stmt2.body)
Review comment:
It's better to assert `structural_equal` to the expected result.
##########
File path: include/tvm/tir/stmt.h
##########
@@ -1339,6 +1339,12 @@ constexpr const char* hand_threaded = "hand_threaded";
* if (mask & 2) the write region should be detected.
*/
constexpr const char* script_parsing_detect_access = "tir.script_parsing_detect_access";
+
+/*!
+ * \brief Mark that the loop should be partitioned.
+ */
+constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint";
Review comment:
```suggestion
constexpr const char* pragma_loop_partition_hint = "tir.pragma_loop_partition_hint";
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] wrongtest commented on pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
wrongtest commented on pull request #9121:
URL: https://github.com/apache/tvm/pull/9121#issuecomment-929903485
@Hzfengsy hi~ I resolve the comment issues，can you kindly take another round of review：）？
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Hzfengsy commented on pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on pull request #9121:
URL: https://github.com/apache/tvm/pull/9121#issuecomment-929971822
Thanks, @wrongtest
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Hzfengsy merged pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
Hzfengsy merged pull request #9121:
URL: https://github.com/apache/tvm/pull/9121
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] wrongtest commented on a change in pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
wrongtest commented on a change in pull request #9121:
URL: https://github.com/apache/tvm/pull/9121#discussion_r717195678
##########
File path: include/tvm/tir/stmt.h
##########
@@ -1339,6 +1339,12 @@ constexpr const char* hand_threaded = "hand_threaded";
* if (mask & 2) the write region should be detected.
*/
constexpr const char* script_parsing_detect_access = "tir.script_parsing_detect_access";
+
+/*!
+ * \brief Mark that the loop should be partitioned.
+ */
+constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint";
Review comment:
I encounter the problem that the `tir.` prefix seems not compatible with `s[X].pragma(axis, key)`, which add an "pragma_" prefix to key string.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] wrongtest commented on pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
wrongtest commented on pull request #9121:
URL: https://github.com/apache/tvm/pull/9121#issuecomment-929903485
@Hzfengsy hi~ I resolve the comment issues，can you kindly take another round of review：）？
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Hzfengsy commented on pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on pull request #9121:
URL: https://github.com/apache/tvm/pull/9121#issuecomment-929971822
Thanks, @wrongtest
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] wrongtest commented on a change in pull request #9121: [TIR] add loop partition hint pragma
Posted by GitBox <gi...@apache.org>.
wrongtest commented on a change in pull request #9121:
URL: https://github.com/apache/tvm/pull/9121#discussion_r717195678
##########
File path: include/tvm/tir/stmt.h
##########
@@ -1339,6 +1339,12 @@ constexpr const char* hand_threaded = "hand_threaded";
* if (mask & 2) the write region should be detected.
*/
constexpr const char* script_parsing_detect_access = "tir.script_parsing_detect_access";
+
+/*!
+ * \brief Mark that the loop should be partitioned.
+ */
+constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint";
Review comment:
I encounter the problem that the `tir.` prefix seems not compatible with `s[X].pragma(axis, key)`, which add an "pragma_" prefix to key string.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org