Annotate kernel params #289
Conversation
chelini
force-pushed
the
lchelini/annotate
branch
2 times, most recently
from
1057012 to
20604de
Compare
src/enzyme_ad/jax/Passes/Passes.td
Outdated
| 2. Block dimension | ||
| 3. Block index. | ||
| Additionally, set all the following attributes for all the kernel pointers: | ||
| 1. align 16 |
There was a problem hiding this comment.
technically I looked it up and we can assume 128 byte aligned: https://github.com/openxla/xla/blob/109a8ff382f2d174ee13504fb4f140a745cb25a2/xla/service/gpu/gpu_constants.h#L41
chelini
force-pushed
the
lchelini/annotate
branch
2 times, most recently
from
3caf315 to
0dba906
Compare
| auto callee = symTable.lookup<FunctionOpInterface>(symbolName); | ||
| if (!callee) | ||
| return WalkResult::advance(); | ||
| MLIRContext *ctx = callee->getContext(); |
There was a problem hiding this comment.
Something I realize here. this (and the above) is wrong. A callee could have multiple callers. You need to conservatively min the dereferencable size over all callers (and same for the range above, though that's a max)
There was a problem hiding this comment.
Mind taking another look? thanks!
chelini
force-pushed
the
lchelini/annotate
branch
2 times, most recently
from
e0733ea to
fb7f715
Compare
| if (matchPattern(maybeCst, m_ConstantInt(&intValue))) | ||
| target->setAttr("range", LLVM::ConstantRangeAttr::get( | ||
| ctx, 32, 0, intValue.getSExtValue())); | ||
| if (matchPattern(maybeCst, m_ConstantInt(&intValue))) { |
There was a problem hiding this comment.
I don't think this is the best way to do things.
If, for example, there is already range metadata (that is weaker), we won't do the update.
I think the thing to do here is to
- get all kernelcall ops
- then get all called functions
- for each called function
a) loop over all callers, getting the max over all ranges/etc
b) set the derived values/aliasing/etc
There was a problem hiding this comment.
I did not understand your previous comment then. What do you mean by "weaker range"? Can you provide an example?
There was a problem hiding this comment.
Yeah supposed that the existing kernel code before this pass has a range of 1-100000, but all the callers have a range 1-100. We would be taking the max over the previous setting and all the actual callers
No description provided.