16aed94ff5
This code in this branch assumes the `src` operand in `insert_slice_async` always aliases the result, which shouldn't hold for generally cases but is just a workaround to make the pipeline pass work. I'm also working on the complete analysis in another [branch](https://github.com/openai/triton-mlir/tree/keren/analyze-slice).
66 lines
1.9 KiB
C++
66 lines
1.9 KiB
C++
#include "triton/Analysis/Alias.h"
|
|
#include "triton/Analysis/Utility.h"
|
|
#include "triton/Dialect/TritonGPU/IR/Dialect.h"
|
|
|
|
namespace mlir {
|
|
|
|
AliasInfo AliasInfo::join(const AliasInfo &lhs, const AliasInfo &rhs) {
|
|
if (lhs == rhs)
|
|
return lhs;
|
|
AliasInfo ret;
|
|
for (auto value : lhs.allocs) {
|
|
ret.insert(value);
|
|
}
|
|
for (auto value : rhs.allocs) {
|
|
ret.insert(value);
|
|
}
|
|
return ret;
|
|
}
|
|
|
|
ChangeResult SharedMemoryAliasAnalysis::visitOperation(
|
|
Operation *op, ArrayRef<LatticeElement<AliasInfo> *> operands) {
|
|
AliasInfo aliasInfo;
|
|
bool pessimistic = true;
|
|
if (maybeSharedAllocationOp(op)) {
|
|
// These ops may allocate a new shared memory buffer.
|
|
auto result = op->getResult(0);
|
|
if (isSharedEncoding(result)) {
|
|
// FIXME(Keren): extract and insert are always alias for now
|
|
if (auto extractSliceOp = dyn_cast<triton::gpu::ExtractSliceOp>(op)) {
|
|
// extract_slice %src, %index
|
|
aliasInfo = AliasInfo(operands[0]->getValue());
|
|
} else if (auto insertSliceOp =
|
|
dyn_cast<triton::gpu::InsertSliceAsyncOp>(op)) {
|
|
// insert_slice_async %src, %dst, %index
|
|
aliasInfo = AliasInfo(operands[1]->getValue());
|
|
} else {
|
|
aliasInfo.insert(result);
|
|
}
|
|
pessimistic = false;
|
|
}
|
|
}
|
|
|
|
if (pessimistic) {
|
|
return markAllPessimisticFixpoint(op->getResults());
|
|
}
|
|
// Join all latice elements
|
|
ChangeResult result = ChangeResult::NoChange;
|
|
for (Value value : op->getResults()) {
|
|
result |= getLatticeElement(value).join(aliasInfo);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
AliasResult SharedMemoryAliasAnalysis::alias(Value lhs, Value rhs) {
|
|
// TODO: implement
|
|
return AliasResult::MayAlias;
|
|
}
|
|
|
|
ModRefResult SharedMemoryAliasAnalysis::getModRef(Operation *op,
|
|
Value location) {
|
|
// TODO: implement
|
|
return ModRefResult::getModAndRef();
|
|
}
|
|
|
|
} // namespace mlir
|