-
Notifications
You must be signed in to change notification settings - Fork 3.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[PASS] InstrumentBoundCheckers pass #2079
Conversation
0240240
to
ff7d970
Compare
Thanks for the contribution, please request reviews |
Hi @tqchen can you please review the patch. |
I would love to and will do it when I have time:) but because I cannot personally review every PRs, please also request reviews from other reviewers :) |
@tqchen got it :) |
One thing I'm curious about - have you considered instead (or in addition) adding an option for generating code for LLVM backends with address sanitizer? It seems pretty simple to add (https://github.com/halide/Halide/blob/1d471591077939fda9d9cbb77f92c6322cf28206/src/CodeGen_LLVM.cpp#L1094-L1104), and would catch these OOB accesses, along with a lot more kinds of bugs (https://clang.llvm.org/docs/AddressSanitizer.html). I noticed you've contributed to LLVM's sanitizers in the past, so I wonder if you agree this is useful here? |
@ajtulloch thanks for point this out.
So, IMHO that every architecture which TVM supports and every new architecture would require to custom ASan's runtime, but otherwise if we need just x86/x64 it should be ok. |
Overall this is very good work @denis0x0D! Does it support bound checking on vectorized loops? |
ff7d970
to
649d4aa
Compare
@tqchen @tmoreau89 thanks for review, I've updated patch regarding to your suggestions. |
great thank you for addressing my comments @denis0x0D! |
src/pass/bound_checker.h
Outdated
namespace tvm { | ||
namespace ir { | ||
struct BoundCheckerManager { | ||
std::unordered_map<const Variable *, Array<Expr>> mem_to_buffer; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The use of global singleton to pass information between passes seems to be dangerous, as:
- The const Variable* can get de-allocated and re-allocated to a different variable
- It makes the pass less self-contained, need to depend on context information.
I would try two possible ways to alleviate this:
- Do the bound checking during storage flatten
- Embed the mem_to_buffer information inside the IR itself, using AttrStmt, like https://github.com/dmlc/tvm/blob/master/include/tvm/ir.h#L208
The use of global singleton to pass information between passes seems to be dangerous, because:
I would recommend two possible ways to alleviate this:
|
@tqchen @tmoreau89 thanks for review. In other case, IMHO, it seems like a normal practice to pass some context depended info between passes in other compilers, for example GCC UBSAN check for object size, at first walk on control flow graph for every basic block and insert internal call, if needed, with info like ptr to OBJECT, builtin __builtin_object_size which helps to get actual size - https://github.com/gcc-mirror/gcc/blob/master/gcc/ubsan.c#L2213 |
649d4aa
to
6f8192e
Compare
@tqchen @tmoreau89 @yzhliu I've updated patch. Now it based on AttrStmt, could you please review it. |
src/pass/storage_flatten.cc
Outdated
@@ -96,11 +96,15 @@ class StorageFlattener : public IRMutator { | |||
return this->Mutate(op->body); | |||
} else if (op->attr_key == attr::opengl_stage_scope) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can we pass create_bound_attributes as argument of storage flatten?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ok, fixed, but this also required to change public api for Storage Flatten pass and some tests.
src/pass/storage_flatten.cc
Outdated
// The size of cacheline | ||
int cache_line_size_; | ||
// The current stage is an OpenGL shader. | ||
bool is_opengl_{false}; | ||
// Whether to mark load/store with theirs bounds. | ||
bool create_bound_attributes{false}; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Google C style, private member names end with underscore _
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks, fixed.
@xqdan @sgrechanik-h @ZihengJiang can you also help review this PR? |
To follow up the discussion @denis0x0D
The global context object probably will cause a hard time for the programmer to reason about things because it depends on the exact order of the pass. Imagine if we do another pass in parallel and both of them write to the same global singleton context, then there would be some problem. When possible, use an explicit context info object or embed the information into the IR will make the reasoning easier, since the pass writer only has to consider how to implement the exact signature of the function. |
6f8192e
to
57116f6
Compare
thanks! for inject copy intrin pass, the reason of failure is the if stmt inserted, we need to support assert if in this pass. we can fix this issue in another PR. |
Tensorize is in scheduleOps pass, which is before storage flatten. inject copy is another codegen pass which is after flatten, looks not hard to work with this pr. but for tensorize, we need more discussion. |
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + 1) | ||
|
||
def test_in_bounds_loop_partition_basic_llvm(): | ||
n = tvm.var('n') |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can you add another case for const loop? also enable cfg.partition_const_loop, so we can verify this PR for tail part. we may need to check the IR generated of tail part after InstrumentBoundCheckers pass.
https://github.com/dmlc/tvm/blob/master/python/tvm/build_module.py#L353
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@xqdan Added new test cases on const loop with check for TVM ir.
b4f049e
to
1f5cba3
Compare
I've added more tests cases. The condition for the branch should always be correct, because it relies on index from load/store. So the algorithm is straightforward -
Example of the condition for spliting const loop :
@xqdan thanks for example
@tmoreau89 @xqdan It looks like the support for VTA backed is more complicated then I was thinking.
which inserts before VTALoadBuffer*, and checks the out of the bound, but I'm not sure about this.
@tqchen @tmoreau89 @ZihengJiang @yzhliu @xqdan @sgrechanik-h Anyway should we create another PR to support Bound Checkers for VTA ? |
LGTM. thanks for your work! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Some quick comments, it is possible for us to have zero-dimensional tensors, we do need to confirm if we handle this case correctly
include/tvm/ir_pass.h
Outdated
Stmt StorageFlatten(Stmt stmt, | ||
Map<Tensor, Buffer> extern_buffer, | ||
int cache_line_size); | ||
Stmt StorageFlatten(Stmt stmt, Map<Tensor, Buffer> extern_buffer, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit, one line per argument
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
thanks, fixed. BTW should TVM has its own .clang-format file ?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There is a proposal in #1732 , we generally use Google C style, but haven't yet have our own .clang-format
@@ -429,6 +452,30 @@ class StorageFlattener : public IRMutator { | |||
} | |||
} | |||
}; | |||
|
|||
bool ShapeIsValid(const Array<Expr> &shape) { | |||
if (!shape.size()) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
it is possible for us to have zero-dimensional tensors(scalar), in which case shape.size() == 0
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, I believe that it's right way to skip this type of tensors and do not instrument, please fix me if I'm wrong.
I've added some tests cases. The compute op could be more complicated, at first I was thinking if the store accesses buffer with shape == 0, we should skip annotations and instrumentations, but there are could be example where store accesses zero shape buffer, but load accesses not zero shape.
I've added test cases for this situation.
c6083d6#diff-fb278cb819f24c9c0369504de1bc2e01R351
c6083d6#diff-f904c24c2b3bba4f50a565961755b153R434
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I see, can you add a comment that zero-dimensional tensor does not need boundary check, so it can give context future readers.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Of course, added.
1f5cba3
to
aa0745c
Compare
@xqdan thanks for review.
@tqchen thanks for review.
I've added test cases for this situation. |
c6083d6
to
683c6aa
Compare
src/api/api_pass.cc
Outdated
TVM_REGISTER_API("ir_pass.StorageFlatten") | ||
.set_body([](TVMArgs args, TVMRetValue *ret) { | ||
*ret = StorageFlatten(args[0], args[1], args[2], | ||
args.size() == 3 ? false : args[3]); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The following would be better (to avoid duplicating the default value in two places)
if (args.size() <= 3) {
*ret = StorageFlatten(args[0], args[1], args[2]);
} else {
*ret = StorageFlatten(args[0], args[1], args[2], args[3]);
}
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for review, updated.
683c6aa
to
8e75e11
Compare
src/pass/bound_checker.cc
Outdated
} | ||
|
||
bool IndexIsValid(const Expr &index) const { | ||
if (!index.defined()) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
for if spans multiple lines, enclose with {}, or put return in the same line
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
thanks for review, updated.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @denis0x0D , I think it is good modulo a nit comment
cc @tmoreau89 |
The pass which instruments checkers before memory accesses (load/store). This allows to handle invalid memory accesses. The patch is related to issue: https://discuss.tvm.ai/t/array-bounds-checking/944
8e75e11
to
2317282
Compare
The pass which instruments checkers before memory accesses (load/store). This allows to handle invalid memory accesses. The patch is related to issue: https://discuss.tvm.ai/t/array-bounds-checking/944
The pass which instruments checkers before memory accesses (load/store). This allows to handle invalid memory accesses. The patch is related to issue: https://discuss.tvm.ai/t/array-bounds-checking/944
The pass which instruments checkers before memory accesses (load/store). This allows to handle invalid memory accesses. The patch is related to issue: https://discuss.tvm.ai/t/array-bounds-checking/944
The patch is related to issue:
https://discuss.tvm.ai/t/array-bounds-checking/944
Lets imagine simple example with out of bounds access:
We will get IR like this:
This is actually out of the bounds access, this code could cause a segfault or access to "junk" memory since TMV has native runtime.
Could we set for example, in debug mode, an instrumentation before every memory access with simple check and assert. So at the runtime we can handle it, and make sure our model does not access any invalid memory.
So, I wrote the pass, which instruments bound check before potentially invalid memory access.
The pass could be enabled by option instrument_bound_checkers=True.
The idea is simple,
at first we associate every buffer_var with actual Buffer shape, as far as I understood, in term of TVM buffer_var is like a pointer to actual Buffer, and then in special pass insert check before memory access.
Also new buffer could be created - we check it.
Some buffers could increase size by allocate - we check it.
Could have situation when optimization pass insert an intrinsic tvm.if_then_else - we check it.
I think it could solve the problem mentioned in the issue.
Could someone please review and give me any feedback.
Thanks.