Better implementation of MergeComputeNodesPass
Validate Operations / validate-operations (push) Waiting to run
Validate Operations / validate-operations (push) Waiting to run
This commit is contained in:
NiccoloN
@@ -9,6 +9,7 @@ add_pim_library(OMPimCommon
|
||||
IR/LoopUtils.cpp
|
||||
IR/ShapeUtils.cpp
|
||||
IR/ShapingUtils.cpp
|
||||
IR/StaticIntSequence.cpp
|
||||
IR/SubviewUtils.cpp
|
||||
IR/TensorSliceUtils.cpp
|
||||
IR/WeightUtils.cpp
|
||||
|
||||
@@ -10,6 +10,32 @@ using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
ConstantPool::ConstantPool(Operation *constantAnchor, OpBuilder &builder)
|
||||
: anchor(constantAnchor), block(getConstantInsertionBlock(constantAnchor)),
|
||||
builder(builder) {
|
||||
for (Operation &op : *block)
|
||||
if (auto constant = dyn_cast<arith::ConstantOp>(&op))
|
||||
cache.try_emplace(
|
||||
std::make_pair(constant.getType(), constant.getValue()),
|
||||
constant.getResult());
|
||||
}
|
||||
|
||||
Value ConstantPool::getIndex(int64_t value) {
|
||||
return get(builder.getIndexType(), builder.getIndexAttr(value));
|
||||
}
|
||||
|
||||
Value ConstantPool::get(Type type, Attribute value) {
|
||||
auto key = std::make_pair(type, value);
|
||||
if (Value existing = cache.lookup(key))
|
||||
return existing;
|
||||
OpBuilder::InsertionGuard guard(builder);
|
||||
builder.setInsertionPointToStart(block);
|
||||
Value constant = arith::ConstantOp::create(
|
||||
builder, anchor->getLoc(), type, cast<TypedAttr>(value)).getResult();
|
||||
cache.try_emplace(key, constant);
|
||||
return constant;
|
||||
}
|
||||
|
||||
static std::optional<int64_t> getIndexConstantValue(arith::ConstantOp constantOp) {
|
||||
if (!constantOp.getType().isIndex())
|
||||
return std::nullopt;
|
||||
|
||||
@@ -6,10 +6,26 @@
|
||||
#include "mlir/IR/Value.h"
|
||||
#include "mlir/Transforms/FoldUtils.h"
|
||||
|
||||
#include "llvm/ADT/DenseMap.h"
|
||||
|
||||
#include <optional>
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
class ConstantPool {
|
||||
public:
|
||||
ConstantPool(mlir::Operation *constantAnchor, mlir::OpBuilder &builder);
|
||||
|
||||
mlir::Value getIndex(int64_t value);
|
||||
mlir::Value get(mlir::Type type, mlir::Attribute value);
|
||||
|
||||
private:
|
||||
mlir::Operation *anchor;
|
||||
mlir::Block *block;
|
||||
mlir::OpBuilder &builder;
|
||||
llvm::DenseMap<std::pair<mlir::Type, mlir::Attribute>, mlir::Value> cache;
|
||||
};
|
||||
|
||||
mlir::Block* getConstantInsertionBlock(mlir::Operation* anchorOp);
|
||||
|
||||
mlir::Value
|
||||
|
||||
@@ -0,0 +1,539 @@
|
||||
#include "mlir/Dialect/Arith/IR/Arith.h"
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
|
||||
#include "llvm/ADT/STLExtras.h"
|
||||
#include "llvm/Support/ErrorHandling.h"
|
||||
#include "llvm/Support/MathExtras.h"
|
||||
|
||||
#include <limits>
|
||||
|
||||
#include "AffineUtils.hpp"
|
||||
#include "ConstantUtils.hpp"
|
||||
#include "StaticIntSequence.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace {
|
||||
|
||||
static bool getAffineValue(int64_t base, int64_t step, size_t index,
|
||||
int64_t &value) {
|
||||
if (index > static_cast<size_t>(std::numeric_limits<int64_t>::max()))
|
||||
return false;
|
||||
int64_t scaled;
|
||||
return !llvm::MulOverflow(static_cast<int64_t>(index), step, scaled)
|
||||
&& !llvm::AddOverflow(base, scaled, value);
|
||||
}
|
||||
|
||||
static FailureOr<SmallVector<int64_t>> getI64Values(Operation *op,
|
||||
StringRef name) {
|
||||
Attribute attr = op->getAttr(name);
|
||||
if (!attr)
|
||||
return op->emitOpError() << "is missing " << name << " metadata",
|
||||
failure();
|
||||
if (auto scalar = dyn_cast<IntegerAttr>(attr))
|
||||
return SmallVector<int64_t> {scalar.getInt()};
|
||||
if (auto array = dyn_cast<DenseI64ArrayAttr>(attr))
|
||||
return SmallVector<int64_t>(array.asArrayRef());
|
||||
auto elements = dyn_cast<DenseIntElementsAttr>(attr);
|
||||
auto type = elements ? dyn_cast<RankedTensorType>(elements.getType())
|
||||
: RankedTensorType();
|
||||
if (!elements || !type || type.getRank() != 1
|
||||
|| !type.getElementType().isInteger(64))
|
||||
return op->emitOpError() << "has invalid " << name << " metadata",
|
||||
failure();
|
||||
SmallVector<int64_t> values;
|
||||
values.reserve(elements.getNumElements());
|
||||
for (APInt value : elements.getValues<APInt>())
|
||||
values.push_back(value.getSExtValue());
|
||||
return values;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
StaticIntSequence StaticIntSequence::uniform(int64_t value, size_t count) {
|
||||
assert(count != 0 && "empty static integer sequence");
|
||||
StaticIntSequence result;
|
||||
result.kind = StaticIntSequenceKind::Uniform;
|
||||
result.count = count;
|
||||
result.base = value;
|
||||
return result;
|
||||
}
|
||||
|
||||
StaticIntSequence StaticIntSequence::affine(int64_t base, int64_t step,
|
||||
size_t count) {
|
||||
assert(count != 0 && "empty static integer sequence");
|
||||
int64_t last;
|
||||
assert(getAffineValue(base, step, count - 1, last)
|
||||
&& "overflowing static affine sequence");
|
||||
if (count == 1 || step == 0)
|
||||
return uniform(base, count);
|
||||
StaticIntSequence result;
|
||||
result.kind = StaticIntSequenceKind::Affine;
|
||||
result.count = count;
|
||||
result.base = base;
|
||||
result.step = step;
|
||||
return result;
|
||||
}
|
||||
|
||||
StaticIntSequence StaticIntSequence::runLengthEncoded(
|
||||
ArrayRef<int64_t> runs, size_t count) {
|
||||
assert(count != 0 && runs.size() % 2 == 0
|
||||
&& "invalid run-length encoded sequence");
|
||||
StaticIntSequence result;
|
||||
result.kind = StaticIntSequenceKind::RunLengthEncoded;
|
||||
result.count = count;
|
||||
result.data.assign(runs);
|
||||
return result;
|
||||
}
|
||||
|
||||
StaticIntSequence StaticIntSequence::fromValues(ArrayRef<int64_t> values) {
|
||||
assert(!values.empty() && "empty static integer sequence");
|
||||
if (llvm::all_equal(values))
|
||||
return uniform(values.front(), values.size());
|
||||
int64_t step;
|
||||
bool isAffine = !llvm::SubOverflow(values[1], values[0], step);
|
||||
for (size_t index = 1; isAffine && index < values.size(); ++index) {
|
||||
int64_t difference;
|
||||
isAffine = !llvm::SubOverflow(values[index], values[index - 1],
|
||||
difference)
|
||||
&& difference == step;
|
||||
}
|
||||
if (isAffine)
|
||||
return affine(values.front(), step, values.size());
|
||||
|
||||
SmallVector<int64_t> runs;
|
||||
for (int64_t value : values) {
|
||||
if (!runs.empty() && runs[runs.size() - 2] == value) {
|
||||
++runs.back();
|
||||
continue;
|
||||
}
|
||||
runs.push_back(value);
|
||||
runs.push_back(1);
|
||||
}
|
||||
if (runs.size() < values.size())
|
||||
return runLengthEncoded(runs, values.size());
|
||||
StaticIntSequence result;
|
||||
result.kind = StaticIntSequenceKind::Dense;
|
||||
result.count = values.size();
|
||||
result.data.assign(values);
|
||||
return result;
|
||||
}
|
||||
|
||||
int64_t StaticIntSequence::valueAt(size_t index) const {
|
||||
assert(index < count && "static integer sequence index out of range");
|
||||
if (kind == StaticIntSequenceKind::Uniform)
|
||||
return base;
|
||||
if (kind == StaticIntSequenceKind::Affine) {
|
||||
int64_t value;
|
||||
bool valid = getAffineValue(base, step, index, value);
|
||||
assert(valid && "overflowing static affine sequence");
|
||||
return value;
|
||||
}
|
||||
if (kind == StaticIntSequenceKind::Dense)
|
||||
return data[index];
|
||||
for (size_t run = 0; run < data.size(); run += 2) {
|
||||
size_t length = static_cast<size_t>(data[run + 1]);
|
||||
if (index < length)
|
||||
return data[run];
|
||||
index -= length;
|
||||
}
|
||||
llvm_unreachable("malformed run-length encoded sequence");
|
||||
}
|
||||
|
||||
std::optional<size_t> StaticIntSequence::find(int64_t value, size_t begin,
|
||||
size_t length) const {
|
||||
assert(begin <= count && length <= count - begin
|
||||
&& "invalid static integer sequence search");
|
||||
if (length == 0)
|
||||
return std::nullopt;
|
||||
size_t end = begin + length;
|
||||
if (kind == StaticIntSequenceKind::Uniform)
|
||||
return value == base ? std::optional<size_t>(begin) : std::nullopt;
|
||||
if (kind == StaticIntSequenceKind::Affine) {
|
||||
int64_t delta;
|
||||
if (llvm::SubOverflow(value, base, delta) || delta % step != 0)
|
||||
return std::nullopt;
|
||||
int64_t index = delta / step;
|
||||
return index >= static_cast<int64_t>(begin)
|
||||
&& index < static_cast<int64_t>(end)
|
||||
? std::optional<size_t>(index)
|
||||
: std::nullopt;
|
||||
}
|
||||
if (kind == StaticIntSequenceKind::Dense) {
|
||||
ArrayRef<int64_t> selected = ArrayRef(data).slice(begin, length);
|
||||
auto found = llvm::find(selected, value);
|
||||
return found == selected.end()
|
||||
? std::nullopt
|
||||
: std::optional<size_t>(begin + (found - selected.begin()));
|
||||
}
|
||||
size_t runBegin = 0;
|
||||
for (size_t run = 0; run < data.size(); run += 2) {
|
||||
size_t runEnd = runBegin + static_cast<size_t>(data[run + 1]);
|
||||
if (runEnd > begin && runBegin < end && data[run] == value)
|
||||
return std::max(begin, runBegin);
|
||||
if (runBegin >= end)
|
||||
break;
|
||||
runBegin = runEnd;
|
||||
}
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
StaticIntSequence StaticIntSequence::slice(size_t begin, size_t length) const {
|
||||
assert(length != 0 && begin <= count - length && "invalid sequence slice");
|
||||
if (kind == StaticIntSequenceKind::Uniform)
|
||||
return uniform(base, length);
|
||||
if (kind == StaticIntSequenceKind::Affine)
|
||||
return affine(valueAt(begin), step, length);
|
||||
if (kind == StaticIntSequenceKind::Dense)
|
||||
return fromValues(ArrayRef(data).slice(begin, length));
|
||||
SmallVector<int64_t> runs;
|
||||
size_t end = begin + length;
|
||||
forEachEqualRun([&](int64_t value, size_t runBegin, size_t runCount) {
|
||||
size_t selectedBegin = std::max(begin, runBegin);
|
||||
size_t selectedEnd = std::min(end, runBegin + runCount);
|
||||
if (selectedBegin >= selectedEnd)
|
||||
return;
|
||||
if (!runs.empty() && runs[runs.size() - 2] == value)
|
||||
runs.back() += selectedEnd - selectedBegin;
|
||||
else {
|
||||
runs.push_back(value);
|
||||
runs.push_back(selectedEnd - selectedBegin);
|
||||
}
|
||||
});
|
||||
if (runs.size() == 2)
|
||||
return uniform(runs.front(), length);
|
||||
if (runs.size() < length)
|
||||
return runLengthEncoded(runs, length);
|
||||
SmallVector<int64_t> values;
|
||||
for (size_t run = 0; run < runs.size(); run += 2)
|
||||
values.append(runs[run + 1], runs[run]);
|
||||
return fromValues(values);
|
||||
}
|
||||
|
||||
StaticIntSequence StaticIntSequence::remap(ArrayRef<unsigned> indices) const {
|
||||
assert(!indices.empty() && "empty static integer sequence remap");
|
||||
SmallVector<int64_t> values;
|
||||
values.reserve(indices.size());
|
||||
for (unsigned index : indices)
|
||||
values.push_back(valueAt(index));
|
||||
return fromValues(values);
|
||||
}
|
||||
|
||||
bool StaticIntSequence::operator==(const StaticIntSequence& other) const {
|
||||
return kind == other.kind && count == other.count && base == other.base
|
||||
&& step == other.step && data == other.data;
|
||||
}
|
||||
|
||||
llvm::hash_code StaticIntSequence::hash() const {
|
||||
return llvm::hash_combine(kind, count, base, step,
|
||||
llvm::hash_combine_range(data.begin(), data.end()));
|
||||
}
|
||||
|
||||
void StaticIntSequence::forEachEqualRun(
|
||||
llvm::function_ref<void(int64_t, size_t, size_t)> callback) const {
|
||||
if (kind == StaticIntSequenceKind::Uniform) {
|
||||
callback(base, 0, count);
|
||||
return;
|
||||
}
|
||||
if (kind == StaticIntSequenceKind::RunLengthEncoded) {
|
||||
size_t begin = 0;
|
||||
for (size_t run = 0; run < data.size(); run += 2) {
|
||||
size_t runCount = static_cast<size_t>(data[run + 1]);
|
||||
callback(data[run], begin, runCount);
|
||||
begin += runCount;
|
||||
}
|
||||
return;
|
||||
}
|
||||
size_t begin = 0;
|
||||
while (begin < count) {
|
||||
int64_t value = valueAt(begin);
|
||||
size_t end = begin + 1;
|
||||
while (end < count && valueAt(end) == value)
|
||||
++end;
|
||||
callback(value, begin, end - begin);
|
||||
begin = end;
|
||||
}
|
||||
}
|
||||
|
||||
void StaticIntSequenceChain::append(const StaticIntSequence &sequence,
|
||||
size_t begin, size_t length) {
|
||||
assert(length != 0 && begin <= sequence.size() - length
|
||||
&& "invalid static integer sequence chain slice");
|
||||
if (!slices.empty()) {
|
||||
StaticIntSequenceSlice &last = slices.back();
|
||||
if (last.sequence == &sequence && last.begin + last.count == begin) {
|
||||
last.count += length;
|
||||
count += length;
|
||||
return;
|
||||
}
|
||||
auto affinePart = [](const StaticIntSequenceSlice &slice,
|
||||
int64_t &base, int64_t &step) {
|
||||
base = slice.sequence->valueAt(slice.begin);
|
||||
if (slice.count == 1) {
|
||||
step = 0;
|
||||
return true;
|
||||
}
|
||||
return !llvm::SubOverflow(slice.sequence->valueAt(slice.begin + 1),
|
||||
base, step)
|
||||
&& (slice.sequence->kind == StaticIntSequenceKind::Uniform
|
||||
|| slice.sequence->kind == StaticIntSequenceKind::Affine);
|
||||
};
|
||||
StaticIntSequenceSlice next {&sequence, begin, length};
|
||||
int64_t leftBase, leftStep, rightBase, rightStep, expected;
|
||||
if (affinePart(last, leftBase, leftStep)
|
||||
&& affinePart(next, rightBase, rightStep)
|
||||
&& (last.count == 1 || length == 1 || leftStep == rightStep)) {
|
||||
int64_t step = last.count == 1 ? rightStep : leftStep;
|
||||
if (getAffineValue(leftBase, step, last.count, expected)
|
||||
&& expected == rightBase) {
|
||||
owned.push_back(std::make_unique<StaticIntSequence>(
|
||||
StaticIntSequence::affine(leftBase, step, last.count + length)));
|
||||
last = {owned.back().get(), 0, last.count + length};
|
||||
count += length;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
slices.push_back({&sequence, begin, length});
|
||||
count += length;
|
||||
}
|
||||
|
||||
void StaticIntSequenceChain::append(StaticIntSequence sequence) {
|
||||
size_t length = sequence.size();
|
||||
owned.push_back(std::make_unique<StaticIntSequence>(std::move(sequence)));
|
||||
append(*owned.back(), 0, length);
|
||||
}
|
||||
|
||||
int64_t StaticIntSequenceChain::valueAt(size_t index) const {
|
||||
assert(index < count && "static integer sequence chain index out of range");
|
||||
for (const StaticIntSequenceSlice &slice : slices) {
|
||||
if (index < slice.count)
|
||||
return slice.sequence->valueAt(slice.begin + index);
|
||||
index -= slice.count;
|
||||
}
|
||||
llvm_unreachable("malformed static integer sequence chain");
|
||||
}
|
||||
|
||||
void StaticIntSequenceChain::forEachSegment(llvm::function_ref<void(
|
||||
const StaticIntSequence &, size_t, size_t)> callback) const {
|
||||
for (const StaticIntSequenceSlice &slice : slices)
|
||||
callback(*slice.sequence, slice.begin, slice.count);
|
||||
}
|
||||
|
||||
void StaticIntSequenceChain::forEachEqualRun(
|
||||
llvm::function_ref<void(int64_t, size_t, size_t)> callback) const {
|
||||
std::optional<int64_t> pendingValue;
|
||||
size_t pendingBegin = 0, pendingCount = 0, chainBegin = 0;
|
||||
auto flush = [&] {
|
||||
if (pendingValue)
|
||||
callback(*pendingValue, pendingBegin, pendingCount);
|
||||
};
|
||||
for (const StaticIntSequenceSlice &slice : slices) {
|
||||
size_t sliceEnd = slice.begin + slice.count;
|
||||
slice.sequence->forEachEqualRun(
|
||||
[&](int64_t value, size_t runBegin, size_t runCount) {
|
||||
size_t begin = std::max(slice.begin, runBegin);
|
||||
size_t end = std::min(sliceEnd, runBegin + runCount);
|
||||
if (begin >= end)
|
||||
return;
|
||||
size_t selectedCount = end - begin;
|
||||
size_t globalBegin = chainBegin + begin - slice.begin;
|
||||
if (pendingValue && *pendingValue == value
|
||||
&& pendingBegin + pendingCount == globalBegin) {
|
||||
pendingCount += selectedCount;
|
||||
return;
|
||||
}
|
||||
flush();
|
||||
pendingValue = value;
|
||||
pendingBegin = globalBegin;
|
||||
pendingCount = selectedCount;
|
||||
});
|
||||
chainBegin += slice.count;
|
||||
}
|
||||
flush();
|
||||
}
|
||||
|
||||
StaticIntSequence StaticIntSequenceChain::canonicalize() const {
|
||||
assert(count != 0 && "empty static integer sequence chain");
|
||||
int64_t first = valueAt(0);
|
||||
bool uniform = true;
|
||||
forEachEqualRun([&](int64_t value, size_t, size_t) {
|
||||
uniform &= value == first;
|
||||
});
|
||||
if (uniform)
|
||||
return StaticIntSequence::uniform(first, count);
|
||||
|
||||
int64_t step = 0, previous = first;
|
||||
bool affine = true, haveStep = false;
|
||||
size_t position = 0;
|
||||
forEachSegment([&](const StaticIntSequence &sequence, size_t begin,
|
||||
size_t length) {
|
||||
if (!affine)
|
||||
return;
|
||||
for (size_t index = 0; index < length; ++index) {
|
||||
int64_t value = sequence.valueAt(begin + index);
|
||||
if (position++ == 0) {
|
||||
previous = value;
|
||||
continue;
|
||||
}
|
||||
if (!haveStep) {
|
||||
affine = !llvm::SubOverflow(value, previous, step);
|
||||
haveStep = true;
|
||||
} else if (haveStep) {
|
||||
int64_t difference;
|
||||
affine = !llvm::SubOverflow(value, previous, difference)
|
||||
&& difference == step;
|
||||
}
|
||||
previous = value;
|
||||
if (!affine)
|
||||
break;
|
||||
}
|
||||
});
|
||||
if (affine && haveStep)
|
||||
return StaticIntSequence::affine(first, step, count);
|
||||
|
||||
SmallVector<int64_t> runs;
|
||||
forEachEqualRun([&](int64_t value, size_t, size_t runCount) {
|
||||
runs.push_back(value);
|
||||
runs.push_back(runCount);
|
||||
});
|
||||
if (runs.size() < count)
|
||||
return StaticIntSequence::runLengthEncoded(runs, count);
|
||||
SmallVector<int64_t> values;
|
||||
values.reserve(count);
|
||||
for (size_t run = 0; run < runs.size(); run += 2)
|
||||
values.append(runs[run + 1], runs[run]);
|
||||
return StaticIntSequence::fromValues(values);
|
||||
}
|
||||
|
||||
int64_t StaticIntSequenceChainCursor::value() const {
|
||||
assert(!done() && "static integer sequence chain cursor is done");
|
||||
const StaticIntSequenceSlice ¤t = chain.slices[slice];
|
||||
return current.sequence->valueAt(current.begin + offset);
|
||||
}
|
||||
|
||||
void StaticIntSequenceChainCursor::advance() {
|
||||
assert(!done() && "static integer sequence chain cursor is done");
|
||||
if (++offset != chain.slices[slice].count)
|
||||
return;
|
||||
offset = 0;
|
||||
++slice;
|
||||
}
|
||||
|
||||
void setStaticIntSequenceAttr(Operation *op, StringRef name,
|
||||
const StaticIntSequence &sequence,
|
||||
size_t logicalCount) {
|
||||
assert(sequence.size() == logicalCount && logicalCount != 0
|
||||
&& "invalid static integer metadata count");
|
||||
SmallVector<int64_t> values;
|
||||
StringRef encoding;
|
||||
switch (sequence.kind) {
|
||||
case StaticIntSequenceKind::Uniform:
|
||||
encoding = "uniform";
|
||||
values.push_back(sequence.base);
|
||||
break;
|
||||
case StaticIntSequenceKind::Affine:
|
||||
encoding = "affine";
|
||||
values = {sequence.base, sequence.step};
|
||||
break;
|
||||
case StaticIntSequenceKind::RunLengthEncoded:
|
||||
encoding = "rle";
|
||||
values = sequence.data;
|
||||
break;
|
||||
case StaticIntSequenceKind::Dense:
|
||||
encoding = "dense";
|
||||
values = sequence.data;
|
||||
break;
|
||||
}
|
||||
OpBuilder builder(op);
|
||||
auto type = RankedTensorType::get(
|
||||
{static_cast<int64_t>(values.size())}, builder.getI64Type());
|
||||
op->setAttr(name, DenseIntElementsAttr::get(type, values));
|
||||
if (sequence.kind != StaticIntSequenceKind::Dense)
|
||||
op->setAttr((name + "_encoding").str(), builder.getStringAttr(encoding));
|
||||
}
|
||||
|
||||
FailureOr<StaticIntSequence> getStaticIntSequenceAttr(
|
||||
Operation *op, StringRef name, size_t logicalCount) {
|
||||
if (logicalCount == 0)
|
||||
return op->emitOpError() << "has zero logical count for " << name,
|
||||
failure();
|
||||
auto values = getI64Values(op, name);
|
||||
if (failed(values))
|
||||
return failure();
|
||||
auto encoding = op->getAttrOfType<StringAttr>((name + "_encoding").str());
|
||||
if (!encoding) {
|
||||
if (values->size() != logicalCount)
|
||||
return op->emitOpError() << "has invalid dense " << name << " count",
|
||||
failure();
|
||||
return StaticIntSequence::fromValues(*values);
|
||||
}
|
||||
if (encoding.getValue() == "uniform") {
|
||||
if (values->size() != 1)
|
||||
return op->emitOpError() << "has invalid uniform " << name,
|
||||
failure();
|
||||
return StaticIntSequence::uniform(values->front(), logicalCount);
|
||||
}
|
||||
if (encoding.getValue() == "affine") {
|
||||
int64_t last;
|
||||
if (values->size() != 2
|
||||
|| !getAffineValue((*values)[0], (*values)[1], logicalCount - 1, last))
|
||||
return op->emitOpError() << "has invalid affine " << name,
|
||||
failure();
|
||||
return StaticIntSequence::affine((*values)[0], (*values)[1], logicalCount);
|
||||
}
|
||||
if (encoding.getValue() == "rle") {
|
||||
size_t count = 0;
|
||||
if (values->empty() || values->size() % 2 != 0)
|
||||
return op->emitOpError() << "has invalid RLE " << name, failure();
|
||||
for (size_t index = 1; index < values->size(); index += 2) {
|
||||
if ((*values)[index] <= 0
|
||||
|| static_cast<uint64_t>((*values)[index]) > logicalCount - count)
|
||||
return op->emitOpError() << "has invalid RLE " << name, failure();
|
||||
count += (*values)[index];
|
||||
}
|
||||
if (count != logicalCount)
|
||||
return op->emitOpError() << "has mismatched RLE " << name << " count",
|
||||
failure();
|
||||
return StaticIntSequence::runLengthEncoded(*values, count);
|
||||
}
|
||||
if (encoding.getValue() == "dense") {
|
||||
if (values->size() != logicalCount)
|
||||
return op->emitOpError() << "has invalid dense " << name << " count",
|
||||
failure();
|
||||
return StaticIntSequence::fromValues(*values);
|
||||
}
|
||||
return op->emitOpError() << "has unknown " << name << " encoding",
|
||||
failure();
|
||||
}
|
||||
|
||||
Value emitStaticIntLookup(const StaticIntSequence& sequence, Value position,
|
||||
Operation* constantAnchor,
|
||||
ConstantPool& constants, OpBuilder& builder,
|
||||
Location loc) {
|
||||
if (sequence.getKind() == StaticIntSequenceKind::Uniform)
|
||||
return constants.getIndex(sequence.valueAt(0));
|
||||
if (sequence.getKind() == StaticIntSequenceKind::Affine) {
|
||||
Value scaled = affineMulConst(builder, loc, position,
|
||||
sequence.valueAt(1) - sequence.valueAt(0),
|
||||
constantAnchor);
|
||||
return affineAddConst(builder, loc, scaled, sequence.valueAt(0),
|
||||
constantAnchor);
|
||||
}
|
||||
SmallVector<int64_t> values;
|
||||
values.reserve(sequence.size());
|
||||
sequence.forEachEqualRun([&](int64_t value, size_t, size_t count) {
|
||||
values.append(count, value);
|
||||
});
|
||||
auto type = RankedTensorType::get(
|
||||
{static_cast<int64_t>(values.size())}, builder.getI64Type());
|
||||
Value table = constants.get(type,
|
||||
DenseElementsAttr::get(type, ArrayRef<int64_t>(values)));
|
||||
Value selected = tensor::ExtractOp::create(
|
||||
builder, loc, table, ValueRange {position});
|
||||
return arith::IndexCastOp::create(
|
||||
builder, loc, builder.getIndexType(), selected);
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir
|
||||
@@ -0,0 +1,117 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/IR/Builders.h"
|
||||
|
||||
#include "llvm/ADT/ArrayRef.h"
|
||||
#include "llvm/ADT/FunctionExtras.h"
|
||||
#include "llvm/ADT/Hashing.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
#include "llvm/ADT/StringRef.h"
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
class ConstantPool;
|
||||
|
||||
enum class StaticIntSequenceKind {
|
||||
Uniform,
|
||||
Affine,
|
||||
RunLengthEncoded,
|
||||
Dense
|
||||
};
|
||||
|
||||
class StaticIntSequence {
|
||||
public:
|
||||
static StaticIntSequence fromValues(llvm::ArrayRef<int64_t> values);
|
||||
static StaticIntSequence uniform(int64_t value, size_t count);
|
||||
static StaticIntSequence affine(int64_t base, int64_t step, size_t count);
|
||||
|
||||
size_t size() const { return count; }
|
||||
int64_t valueAt(size_t index) const;
|
||||
std::optional<size_t> find(int64_t value, size_t begin, size_t length) const;
|
||||
StaticIntSequence slice(size_t begin, size_t count) const;
|
||||
StaticIntSequence remap(llvm::ArrayRef<unsigned> indices) const;
|
||||
StaticIntSequenceKind getKind() const { return kind; }
|
||||
|
||||
bool operator==(const StaticIntSequence& other) const;
|
||||
llvm::hash_code hash() const;
|
||||
|
||||
void forEachEqualRun(
|
||||
llvm::function_ref<void(int64_t, size_t, size_t)> callback) const;
|
||||
|
||||
private:
|
||||
friend class StaticIntSequenceChain;
|
||||
friend void setStaticIntSequenceAttr(mlir::Operation *, llvm::StringRef,
|
||||
const StaticIntSequence &, size_t);
|
||||
friend mlir::FailureOr<StaticIntSequence>
|
||||
getStaticIntSequenceAttr(mlir::Operation *, llvm::StringRef, size_t);
|
||||
|
||||
static StaticIntSequence runLengthEncoded(
|
||||
llvm::ArrayRef<int64_t> runs, size_t count);
|
||||
StaticIntSequenceKind kind = StaticIntSequenceKind::Dense;
|
||||
size_t count = 0;
|
||||
int64_t base = 0;
|
||||
int64_t step = 0;
|
||||
llvm::SmallVector<int64_t> data;
|
||||
};
|
||||
|
||||
struct StaticIntSequenceSlice {
|
||||
const StaticIntSequence *sequence = nullptr;
|
||||
size_t begin = 0;
|
||||
size_t count = 0;
|
||||
};
|
||||
|
||||
class StaticIntSequenceChain {
|
||||
public:
|
||||
void append(const StaticIntSequence &sequence, size_t begin, size_t count);
|
||||
void append(StaticIntSequence sequence);
|
||||
size_t size() const { return count; }
|
||||
int64_t valueAt(size_t index) const;
|
||||
void forEachSegment(llvm::function_ref<void(
|
||||
const StaticIntSequence &, size_t, size_t)> callback) const;
|
||||
void forEachEqualRun(
|
||||
llvm::function_ref<void(int64_t, size_t, size_t)> callback) const;
|
||||
StaticIntSequence canonicalize() const;
|
||||
|
||||
private:
|
||||
friend class StaticIntSequenceChainCursor;
|
||||
llvm::SmallVector<StaticIntSequenceSlice> slices;
|
||||
llvm::SmallVector<std::unique_ptr<StaticIntSequence>> owned;
|
||||
size_t count = 0;
|
||||
};
|
||||
|
||||
class StaticIntSequenceChainCursor {
|
||||
public:
|
||||
explicit StaticIntSequenceChainCursor(const StaticIntSequenceChain &chain)
|
||||
: chain(chain) {}
|
||||
|
||||
bool done() const { return slice == chain.slices.size(); }
|
||||
int64_t value() const;
|
||||
void advance();
|
||||
|
||||
private:
|
||||
const StaticIntSequenceChain &chain;
|
||||
size_t slice = 0;
|
||||
size_t offset = 0;
|
||||
};
|
||||
|
||||
void setStaticIntSequenceAttr(mlir::Operation *op, llvm::StringRef name,
|
||||
const StaticIntSequence &sequence,
|
||||
size_t logicalCount);
|
||||
|
||||
mlir::FailureOr<StaticIntSequence>
|
||||
getStaticIntSequenceAttr(mlir::Operation *op, llvm::StringRef name,
|
||||
size_t logicalCount);
|
||||
|
||||
mlir::Value emitStaticIntLookup(const StaticIntSequence& sequence,
|
||||
mlir::Value position,
|
||||
mlir::Operation* constantAnchor,
|
||||
ConstantPool& constants,
|
||||
mlir::OpBuilder& builder,
|
||||
mlir::Location loc);
|
||||
|
||||
} // namespace onnx_mlir
|
||||
@@ -68,6 +68,23 @@ Value insertStaticSlice(
|
||||
.getResult();
|
||||
}
|
||||
|
||||
Value extractMixedSliceOrIdentity(RewriterBase &rewriter,
|
||||
Location loc,
|
||||
Value source,
|
||||
RankedTensorType resultType,
|
||||
const MixedSliceGeometry &geometry) {
|
||||
return extractStaticSliceOrIdentity(rewriter, loc, source, resultType,
|
||||
geometry.offsets, geometry.sizes,
|
||||
geometry.strides);
|
||||
}
|
||||
|
||||
Value insertMixedSlice(OpBuilder &builder, Location loc, Value source,
|
||||
Value dest, const MixedSliceGeometry &geometry) {
|
||||
return tensor::InsertSliceOp::create(builder, loc, source, dest,
|
||||
geometry.offsets, geometry.sizes,
|
||||
geometry.strides);
|
||||
}
|
||||
|
||||
FailureOr<Value> addLeadingUnitTensorDimension(OpBuilder& builder, Location loc, Value value) {
|
||||
auto type = dyn_cast<RankedTensorType>(value.getType());
|
||||
if (!type || !type.hasStaticShape())
|
||||
|
||||
@@ -8,6 +8,12 @@
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
struct MixedSliceGeometry {
|
||||
llvm::SmallVector<mlir::OpFoldResult> offsets;
|
||||
llvm::SmallVector<mlir::OpFoldResult> sizes;
|
||||
llvm::SmallVector<mlir::OpFoldResult> strides;
|
||||
};
|
||||
|
||||
mlir::Value extractAxisSlice(
|
||||
mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value source, int64_t axis, int64_t offset, int64_t size);
|
||||
|
||||
@@ -25,6 +31,18 @@ mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
|
||||
mlir::Value dest,
|
||||
llvm::ArrayRef<mlir::OpFoldResult> offsets);
|
||||
|
||||
mlir::Value extractMixedSliceOrIdentity(mlir::RewriterBase &rewriter,
|
||||
mlir::Location loc,
|
||||
mlir::Value source,
|
||||
mlir::RankedTensorType resultType,
|
||||
const MixedSliceGeometry &geometry);
|
||||
|
||||
mlir::Value insertMixedSlice(mlir::OpBuilder &builder,
|
||||
mlir::Location loc,
|
||||
mlir::Value source,
|
||||
mlir::Value dest,
|
||||
const MixedSliceGeometry &geometry);
|
||||
|
||||
mlir::FailureOr<mlir::Value>
|
||||
addLeadingUnitTensorDimension(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value);
|
||||
|
||||
|
||||
@@ -11,7 +11,12 @@ add_pim_library(SpatialOps
|
||||
Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
|
||||
Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
|
||||
Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp
|
||||
Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp
|
||||
Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp
|
||||
Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp
|
||||
Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp
|
||||
Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp
|
||||
Transforms/MergeComputeNodes/DeferredResultRealization.cpp
|
||||
Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
|
||||
Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
|
||||
Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
|
||||
|
||||
@@ -0,0 +1,658 @@
|
||||
#include "llvm/ADT/MapVector.h"
|
||||
|
||||
#include "DeferredBoundaryPlanning.hpp"
|
||||
#include "DeferredCommunicationScheduling.hpp"
|
||||
#include "DeferredProjectionAnalysis.hpp"
|
||||
#include "DeferredTransferPlanning.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
using namespace mlir;
|
||||
namespace {
|
||||
|
||||
enum class BoundaryEventKind { Send, Receive };
|
||||
|
||||
struct BoundaryEvent {
|
||||
BoundaryEventKind kind = BoundaryEventKind::Send;
|
||||
ScheduledTransferSlice slice;
|
||||
LaneSet activeLanes;
|
||||
TransferEmissionSignature emission;
|
||||
};
|
||||
|
||||
struct SequenceNode {
|
||||
unsigned previous = 0;
|
||||
unsigned instruction = 0;
|
||||
};
|
||||
|
||||
struct IntervalClass {
|
||||
LaneSet lanes;
|
||||
unsigned sequence = 0;
|
||||
};
|
||||
|
||||
enum class SemanticKind {
|
||||
Send,
|
||||
Availability,
|
||||
Result
|
||||
};
|
||||
|
||||
struct SemanticKey {
|
||||
SemanticKind kind = SemanticKind::Send;
|
||||
TransferEmissionSignature emission;
|
||||
DeferredExchangePlan* exchange = nullptr;
|
||||
RequirementCoordinate coordinate;
|
||||
Type fragmentType;
|
||||
|
||||
bool operator==(const SemanticKey& other) const {
|
||||
if (kind != other.kind)
|
||||
return false;
|
||||
if (kind == SemanticKind::Send)
|
||||
return emission == other.emission;
|
||||
if (kind == SemanticKind::Result)
|
||||
return exchange == other.exchange;
|
||||
return exchange == other.exchange && coordinate == other.coordinate
|
||||
&& fragmentType == other.fragmentType;
|
||||
}
|
||||
};
|
||||
|
||||
struct PendingToken {
|
||||
LaneSet lanes;
|
||||
unsigned semantic = 0;
|
||||
std::variant<unsigned, LocalAvailabilityFamily*, DeferredExchangePlan*> value;
|
||||
};
|
||||
|
||||
struct SequenceCursor {
|
||||
LaneSet lanes;
|
||||
unsigned position = 0;
|
||||
};
|
||||
|
||||
struct CanonicalAction {
|
||||
SemanticKey key;
|
||||
SmallVector<ScheduledTransferSlice> slices;
|
||||
SmallVector<LocalAvailabilityFamily*> locals;
|
||||
LaneSet instructionLanes;
|
||||
LaneSet receiveLanes;
|
||||
LaneSet localLanes;
|
||||
DeferredExchangePlan* result = nullptr;
|
||||
};
|
||||
|
||||
struct BoundaryWork {
|
||||
BoundaryProgram program;
|
||||
SmallVector<BoundaryEvent> events;
|
||||
};
|
||||
|
||||
static unsigned getBoundaryIndex(SmallVectorImpl<BoundaryWork>& boundaries,
|
||||
DenseMap<BoundaryKey, unsigned>& indices,
|
||||
BoundaryKey key) {
|
||||
if (auto it = indices.find(key); it != indices.end())
|
||||
return it->second;
|
||||
unsigned index = boundaries.size();
|
||||
indices[key] = index;
|
||||
BoundaryWork work;
|
||||
work.program.key = key;
|
||||
boundaries.push_back(std::move(work));
|
||||
return index;
|
||||
}
|
||||
|
||||
static unsigned getResultBoundary(DeferredExchangePlan& exchange,
|
||||
const ScheduledCommunicationPlan& schedule) {
|
||||
std::optional<unsigned> step;
|
||||
for (const ScheduledTransferSlice& slice : schedule.slices)
|
||||
if (slice.family->requirement->exchange == &exchange)
|
||||
step = std::max(step.value_or(0), slice.targetInsertionStep);
|
||||
return step.value_or(exchange.consumerStep);
|
||||
}
|
||||
|
||||
static size_t hashSemanticKey(const SemanticKey& key) {
|
||||
if (key.kind == SemanticKind::Send)
|
||||
return llvm::hash_combine(key.kind,
|
||||
key.emission.scheduled,
|
||||
key.emission.payload.getAsOpaquePointer(),
|
||||
key.emission.fragmentType.getAsOpaquePointer(),
|
||||
key.emission.hasGraphLane,
|
||||
key.emission.sourceIsBatch);
|
||||
if (key.kind == SemanticKind::Result)
|
||||
return llvm::hash_combine(key.kind, key.exchange);
|
||||
return llvm::hash_combine(key.kind, key.exchange,
|
||||
key.coordinate.leafIndex,
|
||||
key.coordinate.selectedPosition,
|
||||
key.fragmentType.getAsOpaquePointer());
|
||||
}
|
||||
|
||||
static unsigned internSemantic(const SemanticKey& key,
|
||||
SmallVectorImpl<SemanticKey>& keys,
|
||||
DenseMap<size_t, SmallVector<unsigned>>& byHash) {
|
||||
size_t hash = hashSemanticKey(key);
|
||||
for (unsigned candidate : byHash.lookup(hash))
|
||||
if (keys[candidate] == key)
|
||||
return candidate;
|
||||
unsigned id = keys.size();
|
||||
keys.push_back(key);
|
||||
byHash[hash].push_back(id);
|
||||
return id;
|
||||
}
|
||||
|
||||
static SemanticKey getEventKey(const BoundaryEvent& event) {
|
||||
if (event.kind == BoundaryEventKind::Send) {
|
||||
SemanticKey key;
|
||||
key.kind = SemanticKind::Send;
|
||||
key.emission = event.emission;
|
||||
return key;
|
||||
}
|
||||
RequirementFamily& requirement = *event.slice.family->requirement;
|
||||
SemanticKey key;
|
||||
key.kind = SemanticKind::Availability;
|
||||
key.exchange = requirement.exchange;
|
||||
key.coordinate = requirement.coordinate;
|
||||
key.fragmentType = requirement.publicationFragmentType;
|
||||
return key;
|
||||
}
|
||||
|
||||
static SemanticKey getLocalKey(LocalAvailabilityFamily& local) {
|
||||
RequirementFamily& requirement = *local.requirement;
|
||||
SemanticKey key;
|
||||
key.kind = SemanticKind::Availability;
|
||||
key.exchange = requirement.exchange;
|
||||
key.coordinate = requirement.coordinate;
|
||||
key.fragmentType = requirement.publicationFragmentType;
|
||||
return key;
|
||||
}
|
||||
|
||||
static SmallVector<ScheduledTransferSlice> intersectReceiveSlice(const ScheduledTransferSlice& slice,
|
||||
const LaneSet& lanes) {
|
||||
LaneInterval family = slice.family->targetLanes.intervals().front();
|
||||
unsigned sliceBegin = family.begin + slice.familyOffset;
|
||||
LaneSet active = LaneSet::range(sliceBegin, sliceBegin + slice.transferCount).intersect(lanes);
|
||||
SmallVector<ScheduledTransferSlice> result;
|
||||
for (LaneInterval selected : active.intervals()) {
|
||||
ScheduledTransferSlice part = slice;
|
||||
part.familyOffset += selected.begin - sliceBegin;
|
||||
part.transferCount = selected.end - selected.begin;
|
||||
result.push_back(part);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
static void mergeSequenceCursors(SmallVectorImpl<SequenceCursor>& states) {
|
||||
DenseMap<unsigned, unsigned> byPosition;
|
||||
SmallVector<SequenceCursor> merged;
|
||||
for (SequenceCursor& state : states) {
|
||||
auto [it, inserted] = byPosition.try_emplace(state.position, merged.size());
|
||||
if (inserted)
|
||||
merged.push_back(std::move(state));
|
||||
else
|
||||
merged[it->second].lanes = merged[it->second].lanes.unite(state.lanes);
|
||||
}
|
||||
states = std::move(merged);
|
||||
}
|
||||
|
||||
static LogicalResult appendReplayToken(const PendingToken& token,
|
||||
const LaneSet& lanes,
|
||||
CanonicalAction& action,
|
||||
ArrayRef<BoundaryEvent> events) {
|
||||
if (auto eventId = std::get_if<unsigned>(&token.value)) {
|
||||
const BoundaryEvent& event = events[*eventId];
|
||||
if (event.kind == BoundaryEventKind::Send) {
|
||||
action.slices.push_back(event.slice);
|
||||
action.instructionLanes =
|
||||
action.instructionLanes.unite(event.activeLanes);
|
||||
}
|
||||
else {
|
||||
llvm::append_range(action.slices, intersectReceiveSlice(event.slice, lanes));
|
||||
action.receiveLanes = action.receiveLanes.unite(lanes);
|
||||
}
|
||||
return success();
|
||||
}
|
||||
if (auto local = std::get_if<LocalAvailabilityFamily*>(&token.value)) {
|
||||
if (!llvm::is_contained(action.locals, *local))
|
||||
action.locals.push_back(*local);
|
||||
action.localLanes = action.localLanes.unite(lanes);
|
||||
return success();
|
||||
}
|
||||
action.result = *std::get_if<DeferredExchangePlan*>(&token.value);
|
||||
return success();
|
||||
}
|
||||
|
||||
static FailureOr<SmallVector<CanonicalAction, 0>> collectCanonicalActions(
|
||||
ArrayRef<unsigned> sequence, ArrayRef<PendingToken> tokens,
|
||||
ArrayRef<SemanticKey> semantics, ArrayRef<BoundaryEvent> events,
|
||||
const LaneSet& lanes) {
|
||||
SmallVector<CanonicalAction, 0> actions;
|
||||
for (unsigned semantic : sequence) {
|
||||
CanonicalAction action;
|
||||
action.key = semantics[semantic];
|
||||
actions.push_back(std::move(action));
|
||||
}
|
||||
SmallVector<SequenceCursor> states {
|
||||
{lanes, 0}
|
||||
};
|
||||
for (const PendingToken& token : tokens) {
|
||||
SmallVector<SequenceCursor> next;
|
||||
for (const SequenceCursor& state : states) {
|
||||
LaneSet intersection = state.lanes.intersect(token.lanes);
|
||||
LaneSet difference = state.lanes.subtract(token.lanes);
|
||||
if (!difference.empty())
|
||||
next.push_back({difference, state.position});
|
||||
if (intersection.empty())
|
||||
continue;
|
||||
if (state.position == sequence.size()) {
|
||||
next.push_back({intersection, state.position});
|
||||
continue;
|
||||
}
|
||||
if (state.position >= sequence.size() || sequence[state.position] != token.semantic
|
||||
|| failed(appendReplayToken(token, intersection, actions[state.position], events)))
|
||||
return failure();
|
||||
next.push_back({intersection, state.position + 1});
|
||||
}
|
||||
mergeSequenceCursors(next);
|
||||
states = std::move(next);
|
||||
}
|
||||
if (llvm::any_of(states, [&](const SequenceCursor& state) { return state.position != sequence.size(); }))
|
||||
return failure();
|
||||
return actions;
|
||||
}
|
||||
|
||||
static bool matchesAssembly(ArrayRef<CanonicalAction> actions, size_t begin, SmallVectorImpl<unsigned>& entryOrder) {
|
||||
if (begin >= actions.size())
|
||||
return false;
|
||||
const CanonicalAction& first = actions[begin];
|
||||
DeferredExchangePlan* exchange = first.key.exchange;
|
||||
auto& assembly = exchange->program.insertAssembly;
|
||||
if (first.key.kind != SemanticKind::Availability || !first.locals.empty()
|
||||
|| !assembly || assembly->entries.empty()
|
||||
|| begin + assembly->entries.size() > actions.size())
|
||||
return false;
|
||||
SmallVector<bool> matched(assembly->entries.size());
|
||||
for (size_t offset = 0; offset < assembly->entries.size(); ++offset) {
|
||||
const CanonicalAction& action = actions[begin + offset];
|
||||
if (action.key.kind != SemanticKind::Availability || !action.locals.empty()
|
||||
|| action.key.exchange != exchange || action.slices.empty())
|
||||
return false;
|
||||
std::optional<unsigned> entry;
|
||||
for (auto [entryIndex, candidate] : llvm::enumerate(assembly->entries))
|
||||
if (!matched[entryIndex] && action.key.coordinate == candidate.coordinate) {
|
||||
entry = entryIndex;
|
||||
break;
|
||||
}
|
||||
if (!entry)
|
||||
return false;
|
||||
matched[*entry] = true;
|
||||
entryOrder.push_back(*entry);
|
||||
}
|
||||
if (assembly->entries.size() > 1
|
||||
&& llvm::any_of(ArrayRef(assembly->entries).drop_front(), [&](const DeferredInsertAssemblyEntryTemplate& entry) {
|
||||
return entry.sourceTransform != assembly->entries.front().sourceTransform
|
||||
|| entry.sourceType != assembly->entries.front().sourceType;
|
||||
}))
|
||||
return false;
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool matchesProjectionAssembly(ArrayRef<CanonicalAction> actions,
|
||||
size_t begin,
|
||||
const LaneSet& lanes,
|
||||
unsigned& leafIndex,
|
||||
SmallVectorImpl<unsigned>& positions) {
|
||||
if (begin >= actions.size() || lanes.empty())
|
||||
return false;
|
||||
const CanonicalAction& first = actions[begin];
|
||||
DeferredExchangePlan* exchange = first.key.exchange;
|
||||
if (first.key.kind != SemanticKind::Availability || !first.locals.empty()
|
||||
|| !exchange || exchange->program.insertAssembly
|
||||
|| first.key.coordinate.leafIndex >= exchange->program.leaves.size())
|
||||
return false;
|
||||
leafIndex = first.key.coordinate.leafIndex;
|
||||
const DeferredProjectionLeafTemplate& leaf = exchange->program.leaves[leafIndex];
|
||||
if (leaf.form != DeferredLeafForm::DirectSource)
|
||||
return false;
|
||||
unsigned representative = lanes.intervals().front().begin;
|
||||
unsigned positionCount = 0;
|
||||
unsigned requirementCount = 0;
|
||||
Type fragmentType;
|
||||
for (RequirementFamily& requirement : exchange->requirements) {
|
||||
if (requirement.coordinate.leafIndex != leafIndex || !requirement.targetLanes.contains(representative))
|
||||
continue;
|
||||
++requirementCount;
|
||||
positionCount = std::max(positionCount, requirement.coordinate.selectedPosition + 1);
|
||||
if (fragmentType && fragmentType != requirement.publicationFragmentType)
|
||||
return false;
|
||||
fragmentType = requirement.publicationFragmentType;
|
||||
}
|
||||
auto fragment = dyn_cast<RankedTensorType>(fragmentType);
|
||||
if (positionCount < 2 || requirementCount != positionCount || !fragment
|
||||
|| leaf.reconstructedType.getRank() != fragment.getRank() + 1
|
||||
|| leaf.reconstructedType.getDimSize(0) != positionCount
|
||||
|| leaf.reconstructedType.getShape().drop_front() != fragment.getShape())
|
||||
return false;
|
||||
SmallVector<bool> seen(positionCount);
|
||||
for (size_t offset = 0; begin + offset < actions.size(); ++offset) {
|
||||
const CanonicalAction& action = actions[begin + offset];
|
||||
if (action.key.kind != SemanticKind::Availability || !action.locals.empty()
|
||||
|| action.key.exchange != exchange
|
||||
|| action.key.coordinate.leafIndex != leafIndex || action.key.fragmentType != fragmentType)
|
||||
break;
|
||||
unsigned position = action.key.coordinate.selectedPosition;
|
||||
if (action.slices.empty() || position >= positionCount || seen[position])
|
||||
return false;
|
||||
seen[position] = true;
|
||||
positions.push_back(position);
|
||||
}
|
||||
if (positions.size() < 2)
|
||||
return false;
|
||||
for (RequirementFamily& requirement : exchange->requirements) {
|
||||
if (requirement.coordinate.leafIndex != leafIndex || !requirement.targetLanes.contains(representative)
|
||||
|| seen[requirement.coordinate.selectedPosition])
|
||||
continue;
|
||||
bool local = llvm::any_of(exchange->local, [&](const LocalAvailabilityFamily& availability) {
|
||||
return availability.requirement == &requirement && availability.targetLanes.contains(representative);
|
||||
});
|
||||
if (!local)
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static BoundaryInstructionList materializeInstructions(ArrayRef<CanonicalAction> actions,
|
||||
const LaneSet& lanes) {
|
||||
BoundaryInstructionList result;
|
||||
auto& instructions = result.instructions;
|
||||
for (size_t index = 0; index < actions.size();) {
|
||||
const CanonicalAction& action = actions[index];
|
||||
if (action.key.kind == SemanticKind::Send) {
|
||||
EmitSendRun run;
|
||||
run.lanes = action.instructionLanes;
|
||||
do {
|
||||
llvm::append_range(run.slices, actions[index].slices);
|
||||
run.lanes = run.lanes.unite(actions[index].instructionLanes);
|
||||
++index;
|
||||
}
|
||||
while (index < actions.size() && actions[index].key.kind == SemanticKind::Send
|
||||
&& actions[index].key.emission == action.key.emission);
|
||||
instructions.push_back(std::move(run));
|
||||
continue;
|
||||
}
|
||||
SmallVector<unsigned> assemblyEntries;
|
||||
if (matchesAssembly(actions, index, assemblyEntries)) {
|
||||
EmitReceiveAssemblyRun run;
|
||||
run.lanes = lanes;
|
||||
run.assemblyEntries = assemblyEntries;
|
||||
run.entryOffsets.push_back(0);
|
||||
for (size_t offset = 0; offset < assemblyEntries.size(); ++offset) {
|
||||
llvm::append_range(run.slices, actions[index + offset].slices);
|
||||
run.entryOffsets.push_back(run.slices.size());
|
||||
}
|
||||
instructions.push_back(std::move(run));
|
||||
index += assemblyEntries.size();
|
||||
continue;
|
||||
}
|
||||
unsigned projectionLeaf = 0;
|
||||
SmallVector<unsigned> projectionPositions;
|
||||
if (matchesProjectionAssembly(actions, index, lanes, projectionLeaf, projectionPositions)) {
|
||||
EmitReceiveAssemblyRun run;
|
||||
run.lanes = lanes;
|
||||
run.projectionLeaf = projectionLeaf;
|
||||
run.assemblyEntries = projectionPositions;
|
||||
run.entryOffsets.push_back(0);
|
||||
for (size_t offset = 0; offset < projectionPositions.size(); ++offset) {
|
||||
llvm::append_range(run.slices, actions[index + offset].slices);
|
||||
run.entryOffsets.push_back(run.slices.size());
|
||||
}
|
||||
instructions.push_back(std::move(run));
|
||||
index += projectionPositions.size();
|
||||
continue;
|
||||
}
|
||||
if (action.key.kind == SemanticKind::Availability) {
|
||||
ResolveAvailability availability;
|
||||
availability.exchange = action.key.exchange;
|
||||
availability.coordinate = action.key.coordinate;
|
||||
availability.fragmentType = action.key.fragmentType;
|
||||
if (!action.slices.empty()) {
|
||||
EmitReceiveRun receive;
|
||||
receive.slices = action.slices;
|
||||
receive.entryOffsets = {0, receive.slices.size()};
|
||||
receive.lanes = action.receiveLanes;
|
||||
availability.alternatives.push_back(
|
||||
{receive.lanes, AvailabilitySource(std::move(receive))});
|
||||
}
|
||||
if (!action.locals.empty()) {
|
||||
MaterializeLocalFamily local {action.locals, action.localLanes};
|
||||
availability.alternatives.push_back(
|
||||
{local.lanes, AvailabilitySource(std::move(local))});
|
||||
}
|
||||
instructions.push_back(std::move(availability));
|
||||
++index;
|
||||
continue;
|
||||
}
|
||||
instructions.push_back(ProduceDeferredResult {action.result, lanes});
|
||||
++index;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
static void addTokenToClasses(const LaneSet& active,
|
||||
unsigned semantic,
|
||||
SmallVectorImpl<IntervalClass>& classes,
|
||||
SmallVectorImpl<SequenceNode>& nodes,
|
||||
DenseMap<std::pair<unsigned, unsigned>, unsigned>& interned) {
|
||||
SmallVector<IntervalClass> next;
|
||||
for (const IntervalClass& current : classes) {
|
||||
LaneSet intersection = current.lanes.intersect(active);
|
||||
LaneSet difference = current.lanes.subtract(active);
|
||||
if (!difference.empty())
|
||||
next.push_back({difference, current.sequence});
|
||||
if (intersection.empty())
|
||||
continue;
|
||||
auto key = std::make_pair(current.sequence, semantic);
|
||||
auto [it, inserted] = interned.try_emplace(key, nodes.size());
|
||||
if (inserted)
|
||||
nodes.push_back({current.sequence, semantic});
|
||||
next.push_back({intersection, it->second});
|
||||
}
|
||||
classes = std::move(next);
|
||||
}
|
||||
|
||||
struct SequenceClass {
|
||||
LaneSet lanes;
|
||||
SmallVector<unsigned> sequence;
|
||||
};
|
||||
|
||||
static SmallVector<SequenceClass>
|
||||
buildSequenceClasses(unsigned laneCount, ArrayRef<PendingToken> tokens) {
|
||||
SmallVector<IntervalClass> classes {
|
||||
{LaneSet::all(laneCount), 0}
|
||||
};
|
||||
SmallVector<SequenceNode> nodes {
|
||||
{0, 0}
|
||||
};
|
||||
DenseMap<std::pair<unsigned, unsigned>, unsigned> interned;
|
||||
for (const PendingToken& token : tokens)
|
||||
addTokenToClasses(token.lanes, token.semantic, classes, nodes, interned);
|
||||
|
||||
SmallVector<SequenceClass> result;
|
||||
DenseMap<unsigned, unsigned> classBySequence;
|
||||
SmallVector<unsigned> classSequences;
|
||||
for (const IntervalClass& interval : classes) {
|
||||
auto [it, inserted] = classBySequence.try_emplace(
|
||||
interval.sequence, result.size());
|
||||
if (inserted) {
|
||||
result.push_back({interval.lanes, {}});
|
||||
classSequences.push_back(interval.sequence);
|
||||
}
|
||||
else
|
||||
result[it->second].lanes = result[it->second].lanes.unite(interval.lanes);
|
||||
}
|
||||
for (auto [behavior, sequence] : llvm::zip_equal(result, classSequences)) {
|
||||
for (unsigned node = sequence; node != 0; node = nodes[node].previous)
|
||||
behavior.sequence.push_back(nodes[node].instruction);
|
||||
std::reverse(behavior.sequence.begin(), behavior.sequence.end());
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
static SmallVector<DeferredExchangePlan*>
|
||||
getProducedExchanges(const BoundaryInstructionList& list) {
|
||||
SmallVector<DeferredExchangePlan*> result;
|
||||
for (const BoundaryInstruction& instruction : list.instructions)
|
||||
if (auto produced = std::get_if<ProduceDeferredResult>(&instruction))
|
||||
result.push_back(produced->exchange);
|
||||
return result;
|
||||
}
|
||||
|
||||
static LogicalResult buildCanonicalBoundary(
|
||||
BoundaryProgram& boundary, ArrayRef<BoundaryEvent> events,
|
||||
ArrayRef<PendingToken> tokens, ArrayRef<SemanticKey> semantics) {
|
||||
unsigned laneCount = boundary.key.scheduled->cores.size();
|
||||
SmallVector<SequenceClass> classes =
|
||||
buildSequenceClasses(laneCount, tokens);
|
||||
if (classes.empty())
|
||||
return failure();
|
||||
size_t prefix = classes.front().sequence.size();
|
||||
for (const SequenceClass& behavior : ArrayRef(classes).drop_front()) {
|
||||
prefix = std::min(prefix, behavior.sequence.size());
|
||||
size_t index = 0;
|
||||
while (index < prefix
|
||||
&& behavior.sequence[index] == classes.front().sequence[index])
|
||||
++index;
|
||||
prefix = index;
|
||||
}
|
||||
auto prefixActions = collectCanonicalActions(
|
||||
ArrayRef(classes.front().sequence).take_front(prefix), tokens, semantics,
|
||||
events, LaneSet::all(laneCount));
|
||||
if (failed(prefixActions))
|
||||
return failure();
|
||||
boundary.root = materializeInstructions(*prefixActions, LaneSet::all(laneCount));
|
||||
if (classes.size() == 1)
|
||||
return success();
|
||||
|
||||
auto dispatch = std::make_unique<LaneDispatch>();
|
||||
SmallVector<int64_t> classIds(laneCount);
|
||||
for (auto [classId, behavior] : llvm::enumerate(classes)) {
|
||||
for (LaneInterval interval : behavior.lanes.intervals())
|
||||
for (unsigned lane = interval.begin; lane < interval.end; ++lane)
|
||||
classIds[lane] = classId;
|
||||
auto actions = collectCanonicalActions(
|
||||
behavior.sequence, tokens, semantics, events, behavior.lanes);
|
||||
if (failed(actions))
|
||||
return failure();
|
||||
dispatch->branches.push_back(materializeInstructions(
|
||||
ArrayRef(*actions).drop_front(prefix), behavior.lanes));
|
||||
auto produced = getProducedExchanges(dispatch->branches.back());
|
||||
if (classId == 0)
|
||||
dispatch->producedExchanges = std::move(produced);
|
||||
else if (produced != dispatch->producedExchanges)
|
||||
return failure();
|
||||
}
|
||||
dispatch->branchByLane = StaticIntSequence::fromValues(classIds);
|
||||
boundary.root.instructions.push_back(std::move(dispatch));
|
||||
return success();
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransferPlan& transfers,
|
||||
const ScheduledCommunicationPlan& schedule) {
|
||||
DeferredBoundaryPlan result;
|
||||
SmallVector<BoundaryWork> boundaries;
|
||||
DenseMap<BoundaryKey, unsigned> boundaryIndices;
|
||||
for (const ScheduledTransferSlice& slice : schedule.slices) {
|
||||
ExternalTransferFamily& family = *slice.family;
|
||||
unsigned sourceIndex =
|
||||
getBoundaryIndex(boundaries, boundaryIndices,
|
||||
{family.sourceScheduled, slice.sourceInsertionStep});
|
||||
unsigned targetIndex =
|
||||
getBoundaryIndex(boundaries, boundaryIndices,
|
||||
{family.targetScheduled, slice.targetInsertionStep});
|
||||
unsigned sourceLane = family.requirement->producer->scheduledLane;
|
||||
boundaries[sourceIndex].events.push_back({BoundaryEventKind::Send,
|
||||
slice,
|
||||
LaneSet::range(sourceLane, sourceLane + 1),
|
||||
getTransferEmissionSignature(family)});
|
||||
LaneInterval familyLanes = family.targetLanes.intervals().front();
|
||||
unsigned targetBegin = familyLanes.begin + slice.familyOffset;
|
||||
boundaries[targetIndex].events.push_back({BoundaryEventKind::Receive,
|
||||
slice,
|
||||
LaneSet::range(targetBegin, targetBegin + slice.transferCount),
|
||||
getTransferEmissionSignature(family)});
|
||||
}
|
||||
|
||||
DenseMap<BoundaryKey, SmallVector<LocalAvailabilityFamily*>> locals;
|
||||
DenseMap<BoundaryKey, SmallVector<DeferredExchangePlan*>> exchanges;
|
||||
for (const std::unique_ptr<DeferredExchangePlan>& exchange : transfers.exchanges) {
|
||||
unsigned step = getResultBoundary(*exchange, schedule);
|
||||
BoundaryKey key {exchange->target, step};
|
||||
getBoundaryIndex(boundaries, boundaryIndices, key);
|
||||
for (LocalAvailabilityFamily& local : exchange->local)
|
||||
locals[key].push_back(&local);
|
||||
exchanges[key].push_back(exchange.get());
|
||||
auto resultPlan = buildDeferredResultPlan(*exchange);
|
||||
if (failed(resultPlan))
|
||||
return exchange->deferred.emitOpError("cannot evaluate deferred result lane functions"), failure();
|
||||
result.results.push_back(std::move(*resultPlan));
|
||||
}
|
||||
|
||||
DenseMap<ScheduledInfo*, unsigned> scheduledOrder;
|
||||
for (auto [index, scheduled] : llvm::enumerate(transfers.scheduled))
|
||||
scheduledOrder[&scheduled] = index;
|
||||
llvm::stable_sort(boundaries, [&](const BoundaryWork& lhs, const BoundaryWork& rhs) {
|
||||
return std::tie(scheduledOrder[lhs.program.key.scheduled], lhs.program.key.insertionStep)
|
||||
< std::tie(scheduledOrder[rhs.program.key.scheduled], rhs.program.key.insertionStep);
|
||||
});
|
||||
for (BoundaryWork& work : boundaries) {
|
||||
BoundaryProgram& boundary = work.program;
|
||||
SmallVector<PendingToken> tokens;
|
||||
SmallVector<SemanticKey> semantics;
|
||||
DenseMap<size_t, SmallVector<unsigned>> semanticsByHash;
|
||||
DenseMap<unsigned, SmallVector<LocalAvailabilityFamily*>> localsBySemantic;
|
||||
SmallPtrSet<LocalAvailabilityFamily*, 8> emittedLocals;
|
||||
for (LocalAvailabilityFamily* local : locals[boundary.key]) {
|
||||
unsigned semantic =
|
||||
internSemantic(getLocalKey(*local), semantics, semanticsByHash);
|
||||
localsBySemantic[semantic].push_back(local);
|
||||
}
|
||||
SmallVector<unsigned> eventSemantics;
|
||||
DenseMap<unsigned, SmallVector<unsigned>> receivesBySemantic;
|
||||
for (auto [eventId, event] : llvm::enumerate(work.events)) {
|
||||
unsigned semantic =
|
||||
internSemantic(getEventKey(event), semantics, semanticsByHash);
|
||||
eventSemantics.push_back(semantic);
|
||||
if (event.kind == BoundaryEventKind::Receive)
|
||||
receivesBySemantic[semantic].push_back(eventId);
|
||||
}
|
||||
llvm::SmallDenseSet<unsigned, 8> emittedReceives;
|
||||
for (auto [eventId, event] : llvm::enumerate(work.events)) {
|
||||
unsigned semantic = eventSemantics[eventId];
|
||||
if (event.kind == BoundaryEventKind::Send) {
|
||||
tokens.push_back({LaneSet::all(
|
||||
boundary.key.scheduled->cores.size()),
|
||||
semantic, static_cast<unsigned>(eventId)});
|
||||
continue;
|
||||
}
|
||||
if (!emittedReceives.insert(semantic).second)
|
||||
continue;
|
||||
for (unsigned receive : receivesBySemantic[semantic])
|
||||
tokens.push_back({work.events[receive].activeLanes, semantic, receive});
|
||||
for (LocalAvailabilityFamily* local : localsBySemantic[semantic]) {
|
||||
tokens.push_back({local->targetLanes, semantic, local});
|
||||
emittedLocals.insert(local);
|
||||
}
|
||||
localsBySemantic.erase(semantic);
|
||||
}
|
||||
for (LocalAvailabilityFamily* local : locals[boundary.key])
|
||||
if (!emittedLocals.contains(local)) {
|
||||
unsigned semantic =
|
||||
internSemantic(getLocalKey(*local), semantics, semanticsByHash);
|
||||
tokens.push_back({local->targetLanes, semantic, local});
|
||||
}
|
||||
for (DeferredExchangePlan* exchange : exchanges[boundary.key]) {
|
||||
SemanticKey key;
|
||||
key.kind = SemanticKind::Result;
|
||||
key.exchange = exchange;
|
||||
unsigned semantic = internSemantic(key, semantics, semanticsByHash);
|
||||
tokens.push_back({LaneSet::all(exchange->targetLaneCount), semantic, exchange});
|
||||
}
|
||||
if (failed(buildCanonicalBoundary(boundary, work.events, tokens, semantics)))
|
||||
return boundary.key.scheduled->op->emitOpError("cannot construct canonical boundary program"), failure();
|
||||
result.boundaries.push_back(std::move(boundary));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
@@ -0,0 +1,102 @@
|
||||
#pragma once
|
||||
|
||||
#include "DeferredCommunicationScheduling.hpp"
|
||||
#include "DeferredResultRealization.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
struct DeferredTransferPlan;
|
||||
|
||||
struct BoundaryKey {
|
||||
ScheduledInfo* scheduled = nullptr;
|
||||
unsigned insertionStep = 0;
|
||||
|
||||
bool operator==(const BoundaryKey& other) const {
|
||||
return scheduled == other.scheduled
|
||||
&& insertionStep == other.insertionStep;
|
||||
}
|
||||
};
|
||||
|
||||
struct EmitSendRun {
|
||||
llvm::SmallVector<ScheduledTransferSlice> slices;
|
||||
LaneSet lanes;
|
||||
};
|
||||
struct EmitReceiveRun {
|
||||
llvm::SmallVector<ScheduledTransferSlice> slices;
|
||||
llvm::SmallVector<size_t> entryOffsets;
|
||||
LaneSet lanes;
|
||||
};
|
||||
struct EmitReceiveAssemblyRun {
|
||||
llvm::SmallVector<ScheduledTransferSlice> slices;
|
||||
llvm::SmallVector<size_t> entryOffsets;
|
||||
llvm::SmallVector<unsigned> assemblyEntries;
|
||||
std::optional<unsigned> projectionLeaf;
|
||||
LaneSet lanes;
|
||||
};
|
||||
struct MaterializeLocalFamily {
|
||||
llvm::SmallVector<LocalAvailabilityFamily*> families;
|
||||
LaneSet lanes;
|
||||
};
|
||||
using AvailabilitySource =
|
||||
std::variant<EmitReceiveRun, MaterializeLocalFamily>;
|
||||
struct AvailabilityAlternative {
|
||||
LaneSet lanes;
|
||||
AvailabilitySource source;
|
||||
};
|
||||
struct ResolveAvailability {
|
||||
DeferredExchangePlan* exchange = nullptr;
|
||||
RequirementCoordinate coordinate;
|
||||
mlir::Type fragmentType;
|
||||
llvm::SmallVector<AvailabilityAlternative, 2> alternatives;
|
||||
};
|
||||
struct ProduceDeferredResult {
|
||||
DeferredExchangePlan* exchange = nullptr;
|
||||
LaneSet lanes;
|
||||
};
|
||||
|
||||
struct BoundaryInstructionList;
|
||||
struct LaneDispatch;
|
||||
using BoundaryInstruction = std::variant<EmitSendRun, ResolveAvailability,
|
||||
EmitReceiveAssemblyRun, ProduceDeferredResult,
|
||||
std::unique_ptr<LaneDispatch>>;
|
||||
struct BoundaryInstructionList {
|
||||
llvm::SmallVector<BoundaryInstruction, 0> instructions;
|
||||
};
|
||||
struct LaneDispatch {
|
||||
StaticIntSequence branchByLane = StaticIntSequence::uniform(0, 1);
|
||||
llvm::SmallVector<BoundaryInstructionList> branches;
|
||||
llvm::SmallVector<DeferredExchangePlan*> producedExchanges;
|
||||
};
|
||||
struct BoundaryProgram {
|
||||
BoundaryKey key;
|
||||
BoundaryInstructionList root;
|
||||
};
|
||||
|
||||
struct DeferredBoundaryPlan {
|
||||
llvm::SmallVector<BoundaryProgram, 0> boundaries;
|
||||
llvm::SmallVector<DeferredResultPlan, 0> results;
|
||||
};
|
||||
|
||||
mlir::FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransferPlan& transfers,
|
||||
const ScheduledCommunicationPlan& schedule);
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
|
||||
namespace llvm {
|
||||
template <>
|
||||
struct DenseMapInfo<onnx_mlir::spatial::BoundaryKey> {
|
||||
static onnx_mlir::spatial::BoundaryKey getEmptyKey() {
|
||||
return {DenseMapInfo<onnx_mlir::spatial::ScheduledInfo*>::getEmptyKey(), 0};
|
||||
}
|
||||
static onnx_mlir::spatial::BoundaryKey getTombstoneKey() {
|
||||
return {DenseMapInfo<onnx_mlir::spatial::ScheduledInfo*>::getTombstoneKey(), 0};
|
||||
}
|
||||
static unsigned getHashValue(const onnx_mlir::spatial::BoundaryKey& key) {
|
||||
return hash_combine(key.scheduled, key.insertionStep);
|
||||
}
|
||||
static bool isEqual(const onnx_mlir::spatial::BoundaryKey& lhs,
|
||||
const onnx_mlir::spatial::BoundaryKey& rhs) {
|
||||
return lhs == rhs;
|
||||
}
|
||||
};
|
||||
} // namespace llvm
|
||||
@@ -0,0 +1,824 @@
|
||||
#include "mlir/Dialect/Arith/IR/Arith.h"
|
||||
#include "mlir/Dialect/SCF/IR/SCF.h"
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
|
||||
#include "DeferredBoundaryRealization.hpp"
|
||||
#include "DeferredResultRealization.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
using namespace mlir;
|
||||
namespace {
|
||||
|
||||
struct LogicalTransferMetadataView {
|
||||
StaticIntSequenceChain channels;
|
||||
StaticIntSequenceChain parents;
|
||||
StaticIntSequenceChain parentCounts;
|
||||
StaticIntSequenceChain sourceCores;
|
||||
StaticIntSequenceChain targetCores;
|
||||
StaticIntSequenceChain targetLanes;
|
||||
StaticIntSequenceChain localOffsets;
|
||||
|
||||
size_t size() const { return channels.size(); }
|
||||
};
|
||||
|
||||
static void appendMetadata(const ScheduledTransferSlice& slice,
|
||||
LogicalTransferMetadataView& metadata) {
|
||||
ExternalTransferFamily& family = *slice.family;
|
||||
LaneInterval familyLanes = family.targetLanes.intervals().front();
|
||||
LaneInterval requirementLanes = family.requirement->targetLanes.intervals().front();
|
||||
size_t count = slice.transferCount;
|
||||
size_t familyIndex = slice.familyOffset;
|
||||
size_t targetLane = familyLanes.begin + familyIndex;
|
||||
metadata.channels.append(family.channelIds, familyIndex, count);
|
||||
metadata.parents.append(StaticIntSequence::uniform(
|
||||
family.requirement->exchange->exchangeId, count));
|
||||
metadata.parentCounts.append(StaticIntSequence::uniform(
|
||||
family.requirement->exchange->externalTransferCount, count));
|
||||
metadata.sourceCores.append(family.sourceCores, familyIndex, count);
|
||||
metadata.targetCores.append(family.targetCores, familyIndex, count);
|
||||
metadata.targetLanes.append(
|
||||
StaticIntSequence::affine(targetLane, 1, count));
|
||||
if (family.requirement->producerLocalOffsets)
|
||||
metadata.localOffsets.append(
|
||||
*family.requirement->producerLocalOffsets,
|
||||
targetLane - requirementLanes.begin, count);
|
||||
else
|
||||
metadata.localOffsets.append(StaticIntSequence::uniform(0, count));
|
||||
}
|
||||
|
||||
static LogicalTransferMetadataView
|
||||
buildMetadataView(ArrayRef<ScheduledTransferSlice> slices) {
|
||||
LogicalTransferMetadataView metadata;
|
||||
for (const ScheduledTransferSlice& slice : slices)
|
||||
appendMetadata(slice, metadata);
|
||||
return metadata;
|
||||
}
|
||||
|
||||
static void setLogicalTransferMetadata(
|
||||
Operation* op, const LogicalTransferMetadataView& metadata) {
|
||||
size_t logicalCount = metadata.size();
|
||||
OpBuilder builder(op);
|
||||
if (logicalCount == 1) {
|
||||
op->setAttr("raptor.exchange_id",
|
||||
builder.getI64IntegerAttr(metadata.channels.valueAt(0)));
|
||||
op->setAttr("raptor.channel_id",
|
||||
builder.getI64IntegerAttr(metadata.channels.valueAt(0)));
|
||||
op->setAttr("raptor.parent_exchange_id",
|
||||
builder.getI64IntegerAttr(metadata.parents.valueAt(0)));
|
||||
op->setAttr("raptor.parent_transfer_count",
|
||||
builder.getI64IntegerAttr(metadata.parentCounts.valueAt(0)));
|
||||
op->setAttr("raptor.source_core",
|
||||
builder.getI64IntegerAttr(metadata.sourceCores.valueAt(0)));
|
||||
op->setAttr("raptor.target_core",
|
||||
builder.getI64IntegerAttr(metadata.targetCores.valueAt(0)));
|
||||
return;
|
||||
}
|
||||
op->setAttr("raptor.batch_transfer_count", builder.getI64IntegerAttr(logicalCount));
|
||||
setStaticIntSequenceAttr(op, "raptor.batch_channel_ids",
|
||||
metadata.channels.canonicalize(), logicalCount);
|
||||
setStaticIntSequenceAttr(op, "raptor.batch_source_cores",
|
||||
metadata.sourceCores.canonicalize(), logicalCount);
|
||||
setStaticIntSequenceAttr(op, "raptor.batch_target_cores",
|
||||
metadata.targetCores.canonicalize(), logicalCount);
|
||||
setStaticIntSequenceAttr(op, "raptor.batch_parent_exchange_ids",
|
||||
metadata.parents.canonicalize(), logicalCount);
|
||||
setStaticIntSequenceAttr(op, "raptor.batch_parent_transfer_counts",
|
||||
metadata.parentCounts.canonicalize(), logicalCount);
|
||||
}
|
||||
|
||||
static Value
|
||||
lookup(ArrayRef<int64_t> table, Value position, Operation* anchor, DeferredEmissionContext& context, Location loc) {
|
||||
return emitStaticIntLookup(
|
||||
StaticIntSequence::fromValues(table), position, anchor, context.constants, context.rewriter, loc);
|
||||
}
|
||||
|
||||
static OpFoldResult lookupGeometry(
|
||||
ArrayRef<int64_t> table, Value position, Operation* anchor, DeferredEmissionContext& context, Location loc) {
|
||||
StaticIntSequence sequence = StaticIntSequence::fromValues(table);
|
||||
if (sequence.getKind() == StaticIntSequenceKind::Uniform)
|
||||
return context.rewriter.getIndexAttr(sequence.valueAt(0));
|
||||
return emitStaticIntLookup(sequence, position, anchor, context.constants, context.rewriter, loc);
|
||||
}
|
||||
|
||||
static FailureOr<Value> materializeSendPayload(const RequirementFamily& requirement,
|
||||
Value localOffset,
|
||||
DeferredEmissionContext& context,
|
||||
Location loc) {
|
||||
Value payload = requirement.producer->payload;
|
||||
if (payload.getType() == requirement.publicationFragmentType)
|
||||
return payload;
|
||||
auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
|
||||
auto fragmentType = dyn_cast<RankedTensorType>(requirement.publicationFragmentType);
|
||||
if (!payloadType || !fragmentType || !requirement.graphLanes || payloadType.getRank() != fragmentType.getRank() + 1)
|
||||
return failure();
|
||||
MixedSliceGeometry geometry;
|
||||
geometry.offsets.assign(payloadType.getRank(), context.rewriter.getIndexAttr(0));
|
||||
geometry.sizes.push_back(context.rewriter.getIndexAttr(1));
|
||||
geometry.strides.assign(payloadType.getRank(), context.rewriter.getIndexAttr(1));
|
||||
geometry.offsets.front() = localOffset;
|
||||
for (int64_t dimension : fragmentType.getShape())
|
||||
geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
|
||||
SmallVector<int64_t> unitShape {1};
|
||||
llvm::append_range(unitShape, fragmentType.getShape());
|
||||
auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType());
|
||||
Value unit = extractMixedSliceOrIdentity(context.rewriter, loc, payload, unitType, geometry);
|
||||
return removeLeadingUnitTensorDimension(context.rewriter, loc, unit, fragmentType);
|
||||
}
|
||||
|
||||
static LogicalResult
|
||||
emitSendRun(const EmitSendRun& run, Value lane, unsigned laneCount, DeferredEmissionContext& context) {
|
||||
SmallVector<LogicalTransferMetadataView, 0> metadataByLane(laneCount);
|
||||
for (const ScheduledTransferSlice& slice : run.slices) {
|
||||
unsigned sourceLane = slice.family->requirement->producer->scheduledLane;
|
||||
appendMetadata(slice, metadataByLane[sourceLane]);
|
||||
}
|
||||
LogicalTransferMetadataView logical = buildMetadataView(run.slices);
|
||||
size_t actionCount = 0;
|
||||
for (const LogicalTransferMetadataView& laneMetadata : metadataByLane)
|
||||
actionCount = std::max(actionCount, laneMetadata.size());
|
||||
SmallVector<int64_t> channelTable(
|
||||
actionCount * laneCount, logical.channels.valueAt(0));
|
||||
SmallVector<int64_t> sourceTable(
|
||||
actionCount * laneCount, logical.sourceCores.valueAt(0));
|
||||
SmallVector<int64_t> targetTable(
|
||||
actionCount * laneCount, logical.targetCores.valueAt(0));
|
||||
SmallVector<int64_t> offsetTable(
|
||||
actionCount * laneCount, logical.localOffsets.valueAt(0));
|
||||
SmallVector<int64_t> counts(laneCount);
|
||||
for (unsigned sourceLane = 0; sourceLane < laneCount; ++sourceLane) {
|
||||
const LogicalTransferMetadataView& source = metadataByLane[sourceLane];
|
||||
counts[sourceLane] = source.size();
|
||||
for (size_t action = 0; action < source.size(); ++action) {
|
||||
size_t index = action * laneCount + sourceLane;
|
||||
channelTable[index] = source.channels.valueAt(action);
|
||||
sourceTable[index] = source.sourceCores.valueAt(action);
|
||||
targetTable[index] = source.targetCores.valueAt(action);
|
||||
offsetTable[index] = source.localOffsets.valueAt(action);
|
||||
}
|
||||
}
|
||||
ExternalTransferFamily& firstFamily = *run.slices.front().family;
|
||||
RequirementFamily& requirement = *firstFamily.requirement;
|
||||
Operation* anchor = requirement.exchange->deferred;
|
||||
Location loc = requirement.exchange->deferred.getLoc();
|
||||
auto emitOne = [&](Value position) -> LogicalResult {
|
||||
Value localOffset = lookup(offsetTable, position, anchor, context, loc);
|
||||
auto payload = materializeSendPayload(requirement, localOffset, context, loc);
|
||||
if (failed(payload))
|
||||
return failure();
|
||||
auto send = SpatChannelSendOp::create(context.rewriter,
|
||||
loc,
|
||||
lookup(channelTable, position, anchor, context, loc),
|
||||
lookup(sourceTable, position, anchor, context, loc),
|
||||
lookup(targetTable, position, anchor, context, loc),
|
||||
*payload);
|
||||
setLogicalTransferMetadata(send, logical);
|
||||
return success();
|
||||
};
|
||||
Value runtimeLane = lane ? lane : context.constants.getIndex(0);
|
||||
if (actionCount == 1)
|
||||
return emitOne(runtimeLane);
|
||||
bool uniformCount = true;
|
||||
std::optional<int64_t> firstCount;
|
||||
for (LaneInterval interval : run.lanes.intervals())
|
||||
for (unsigned activeLane = interval.begin; activeLane < interval.end; ++activeLane) {
|
||||
if (!firstCount)
|
||||
firstCount = counts[activeLane];
|
||||
uniformCount &= counts[activeLane] == *firstCount;
|
||||
}
|
||||
Value count =
|
||||
uniformCount ? context.constants.getIndex(*firstCount) : lookup(counts, runtimeLane, anchor, context, loc);
|
||||
auto loop = buildNormalizedScfFor(context.rewriter,
|
||||
loc,
|
||||
context.constants.getIndex(0),
|
||||
count,
|
||||
context.constants.getIndex(1),
|
||||
ValueRange {},
|
||||
[&](OpBuilder&, Location, Value index, ValueRange, SmallVectorImpl<Value>&) {
|
||||
Value base = affineMulConst(context.rewriter, loc, index, laneCount, anchor);
|
||||
Value position = arith::AddIOp::create(context.rewriter, loc, base, runtimeLane);
|
||||
return emitOne(position);
|
||||
});
|
||||
return success(succeeded(loop));
|
||||
}
|
||||
|
||||
static FailureOr<Value>
|
||||
emitReceiveValue(const EmitReceiveRun& run, Value lane, unsigned laneCount, DeferredEmissionContext& context) {
|
||||
LogicalTransferMetadataView metadata = buildMetadataView(run.slices);
|
||||
RequirementFamily& requirement = *run.slices.front().family->requirement;
|
||||
Operation* anchor = requirement.exchange->deferred;
|
||||
Location loc = requirement.exchange->deferred.getLoc();
|
||||
SmallVector<int64_t> channelTable(laneCount, metadata.channels.valueAt(0));
|
||||
SmallVector<int64_t> sourceTable(laneCount, metadata.sourceCores.valueAt(0));
|
||||
SmallVector<int64_t> targetTable(laneCount, metadata.targetCores.valueAt(0));
|
||||
for (size_t index = 0; index < metadata.size(); ++index) {
|
||||
unsigned targetLane = metadata.targetLanes.valueAt(index);
|
||||
channelTable[targetLane] = metadata.channels.valueAt(index);
|
||||
sourceTable[targetLane] = metadata.sourceCores.valueAt(index);
|
||||
targetTable[targetLane] = metadata.targetCores.valueAt(index);
|
||||
}
|
||||
Value position = lane ? lane : context.constants.getIndex(0);
|
||||
auto receive = SpatChannelReceiveOp::create(context.rewriter,
|
||||
loc,
|
||||
requirement.publicationFragmentType,
|
||||
lookup(channelTable, position, anchor, context, loc),
|
||||
lookup(sourceTable, position, anchor, context, loc),
|
||||
lookup(targetTable, position, anchor, context, loc));
|
||||
setLogicalTransferMetadata(receive, metadata);
|
||||
return receive.getOutput();
|
||||
}
|
||||
|
||||
static LogicalResult
|
||||
emitConditionalSendRun(const EmitSendRun& run, Value lane, unsigned laneCount, DeferredEmissionContext& context) {
|
||||
if (run.lanes.size() == laneCount)
|
||||
return emitSendRun(run, lane, laneCount, context);
|
||||
if (!lane)
|
||||
return failure();
|
||||
SmallVector<int64_t> active(laneCount);
|
||||
for (LaneInterval interval : run.lanes.intervals())
|
||||
for (unsigned index = interval.begin; index < interval.end; ++index)
|
||||
active[index] = 1;
|
||||
Location loc = run.slices.front().family->requirement->exchange->deferred.getLoc();
|
||||
Value selected = lookup(active, lane, run.slices.front().family->sourceScheduled->op, context, loc);
|
||||
Value condition =
|
||||
arith::CmpIOp::create(context.rewriter, loc, arith::CmpIPredicate::ne, selected, context.constants.getIndex(0));
|
||||
auto conditional = scf::IfOp::create(context.rewriter, loc, TypeRange {}, condition, false);
|
||||
Block& block = conditional.getThenRegion().front();
|
||||
auto yield = cast<scf::YieldOp>(block.getTerminator());
|
||||
OpBuilder::InsertionGuard guard(context.rewriter);
|
||||
context.rewriter.setInsertionPoint(yield);
|
||||
return emitSendRun(run, lane, laneCount, context);
|
||||
}
|
||||
|
||||
static FailureOr<Value> materializeLocalValue(const MaterializeLocalFamily& local,
|
||||
Value lane,
|
||||
unsigned laneCount,
|
||||
DeferredEmissionContext& context) {
|
||||
RequirementFamily& reference = *local.families.front()->requirement;
|
||||
Value fragment = reference.producer->payload;
|
||||
if (fragment.getType() != reference.publicationFragmentType) {
|
||||
SmallVector<int64_t> offsets(laneCount);
|
||||
for (LocalAvailabilityFamily* family : local.families) {
|
||||
RequirementFamily& requirement = *family->requirement;
|
||||
LaneInterval requirementLanes = requirement.targetLanes.intervals().front();
|
||||
for (LaneInterval interval : family->targetLanes.intervals())
|
||||
for (unsigned targetLane = interval.begin; targetLane < interval.end; ++targetLane)
|
||||
offsets[targetLane] = requirement.producerLocalOffsets->valueAt(targetLane - requirementLanes.begin);
|
||||
}
|
||||
Location loc = reference.exchange->deferred.getLoc();
|
||||
Value position = lane ? lane : context.constants.getIndex(0);
|
||||
auto materialized = materializeSendPayload(
|
||||
reference, lookup(offsets, position, reference.exchange->deferred, context, loc), context, loc);
|
||||
if (failed(materialized))
|
||||
return failure();
|
||||
fragment = *materialized;
|
||||
}
|
||||
return fragment;
|
||||
}
|
||||
|
||||
static const DeferredResultPlan* findResultPlan(ArrayRef<DeferredResultPlan> results, DeferredExchangePlan* exchange) {
|
||||
auto it = llvm::find_if(results, [&](const DeferredResultPlan& result) { return result.exchange == exchange; });
|
||||
return it == results.end() ? nullptr : &*it;
|
||||
}
|
||||
|
||||
static FailureOr<Value> applyAssemblyTransform(Value source,
|
||||
const DeferredInsertAssemblyEntryTemplate& entry,
|
||||
DeferredEmissionContext& context,
|
||||
Location loc) {
|
||||
auto sourceType = dyn_cast<RankedTensorType>(source.getType());
|
||||
auto targetType = entry.sourceType;
|
||||
if (!sourceType)
|
||||
return failure();
|
||||
if (sourceType == targetType)
|
||||
return source;
|
||||
if (targetType.getRank() == sourceType.getRank() + 1 && targetType.getDimSize(0) == 1
|
||||
&& targetType.getShape().drop_front() == sourceType.getShape())
|
||||
return addLeadingUnitTensorDimension(context.rewriter, loc, source);
|
||||
if (sourceType.getRank() != targetType.getRank() + 1 || sourceType.getDimSize(0) != 1
|
||||
|| sourceType.getShape().drop_front() != targetType.getShape())
|
||||
return failure();
|
||||
return removeLeadingUnitTensorDimension(context.rewriter, loc, source, targetType);
|
||||
}
|
||||
|
||||
static FailureOr<Value> emitProjectionAssemblyRun(const EmitReceiveAssemblyRun& run,
|
||||
Value lane,
|
||||
unsigned laneCount,
|
||||
DeferredEmissionContext& context) {
|
||||
if (!run.projectionLeaf || run.slices.empty() || run.entryOffsets.size() != run.assemblyEntries.size() + 1)
|
||||
return failure();
|
||||
DeferredExchangePlan* exchange = run.slices.front().family->requirement->exchange;
|
||||
unsigned leafIndex = *run.projectionLeaf;
|
||||
if (leafIndex >= exchange->program.leaves.size())
|
||||
return failure();
|
||||
const DeferredProjectionLeafTemplate& leaf = exchange->program.leaves[leafIndex];
|
||||
size_t entryCount = run.assemblyEntries.size();
|
||||
LogicalTransferMetadataView metadata = buildMetadataView(run.slices);
|
||||
SmallVector<int64_t> channelTable(
|
||||
entryCount * laneCount, metadata.channels.valueAt(0));
|
||||
SmallVector<int64_t> sourceTable(
|
||||
entryCount * laneCount, metadata.sourceCores.valueAt(0));
|
||||
SmallVector<int64_t> targetTable(
|
||||
entryCount * laneCount, metadata.targetCores.valueAt(0));
|
||||
SmallVector<int64_t> positions(entryCount * laneCount);
|
||||
for (size_t entry = 0; entry < entryCount; ++entry) {
|
||||
ArrayRef<ScheduledTransferSlice> slices =
|
||||
ArrayRef(run.slices).slice(run.entryOffsets[entry], run.entryOffsets[entry + 1] - run.entryOffsets[entry]);
|
||||
LogicalTransferMetadataView item = buildMetadataView(slices);
|
||||
for (size_t index = 0; index < item.size(); ++index) {
|
||||
size_t tableIndex = entry * laneCount + item.targetLanes.valueAt(index);
|
||||
channelTable[tableIndex] = item.channels.valueAt(index);
|
||||
sourceTable[tableIndex] = item.sourceCores.valueAt(index);
|
||||
targetTable[tableIndex] = item.targetCores.valueAt(index);
|
||||
positions[tableIndex] = run.assemblyEntries[entry];
|
||||
}
|
||||
}
|
||||
Location loc = exchange->deferred.getLoc();
|
||||
Value initial = tensor::EmptyOp::create(
|
||||
context.rewriter, loc, leaf.reconstructedType.getShape(), leaf.reconstructedType.getElementType());
|
||||
auto emitEntry = [&](Value entry, Value current) -> FailureOr<Value> {
|
||||
Value tableIndex = entry;
|
||||
if (lane) {
|
||||
Value base = affineMulConst(context.rewriter, loc, entry, laneCount, exchange->deferred);
|
||||
tableIndex = arith::AddIOp::create(context.rewriter, loc, base, lane);
|
||||
}
|
||||
Type fragmentType = run.slices.front().family->requirement->publicationFragmentType;
|
||||
auto receive = SpatChannelReceiveOp::create(context.rewriter,
|
||||
loc,
|
||||
fragmentType,
|
||||
lookup(channelTable, tableIndex, exchange->deferred, context, loc),
|
||||
lookup(sourceTable, tableIndex, exchange->deferred, context, loc),
|
||||
lookup(targetTable, tableIndex, exchange->deferred, context, loc));
|
||||
setLogicalTransferMetadata(receive, metadata);
|
||||
auto source = addLeadingUnitTensorDimension(context.rewriter, loc, receive.getOutput());
|
||||
if (failed(source))
|
||||
return failure();
|
||||
auto sourceType = cast<RankedTensorType>(source->getType());
|
||||
MixedSliceGeometry geometry;
|
||||
geometry.offsets.assign(sourceType.getRank(), context.rewriter.getIndexAttr(0));
|
||||
geometry.offsets.front() = lookupGeometry(positions, tableIndex, exchange->deferred, context, loc);
|
||||
for (int64_t dimension : sourceType.getShape())
|
||||
geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
|
||||
geometry.strides.assign(sourceType.getRank(), context.rewriter.getIndexAttr(1));
|
||||
return insertMixedSlice(context.rewriter, loc, *source, current, geometry);
|
||||
};
|
||||
auto loop = buildNormalizedScfFor(
|
||||
context.rewriter,
|
||||
loc,
|
||||
context.constants.getIndex(0),
|
||||
context.constants.getIndex(entryCount),
|
||||
context.constants.getIndex(1),
|
||||
ValueRange {initial},
|
||||
[&](OpBuilder&, Location, Value entry, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) -> LogicalResult {
|
||||
auto value = emitEntry(entry, iterArgs.front());
|
||||
if (failed(value))
|
||||
return failure();
|
||||
yielded.push_back(*value);
|
||||
return success();
|
||||
});
|
||||
if (failed(loop))
|
||||
return failure();
|
||||
Value result = loop->results.front();
|
||||
context.projectionAssemblies[{exchange, leafIndex}] = result;
|
||||
return result;
|
||||
}
|
||||
|
||||
static FailureOr<Value> emitAssemblyRun(const EmitReceiveAssemblyRun& run,
|
||||
Value lane,
|
||||
unsigned laneCount,
|
||||
ArrayRef<DeferredResultPlan> results,
|
||||
DeferredEmissionContext& context) {
|
||||
DeferredExchangePlan* exchange = run.slices.front().family->requirement->exchange;
|
||||
const DeferredResultPlan* resultPlan = findResultPlan(results, exchange);
|
||||
if (!resultPlan)
|
||||
return failure();
|
||||
auto& assembly = exchange->program.insertAssembly;
|
||||
size_t entryCount = run.entryOffsets.empty() ? 0 : run.entryOffsets.size() - 1;
|
||||
if (!assembly || assembly->entries.size() != entryCount)
|
||||
return failure();
|
||||
if (run.assemblyEntries.size() != entryCount)
|
||||
return failure();
|
||||
LogicalTransferMetadataView metadata = buildMetadataView(run.slices);
|
||||
SmallVector<int64_t> channelTable(
|
||||
entryCount * laneCount, metadata.channels.valueAt(0));
|
||||
SmallVector<int64_t> sourceTable(
|
||||
entryCount * laneCount, metadata.sourceCores.valueAt(0));
|
||||
SmallVector<int64_t> targetTable(
|
||||
entryCount * laneCount, metadata.targetCores.valueAt(0));
|
||||
unsigned rank = assembly->resultType.getRank();
|
||||
SmallVector<SmallVector<int64_t>> geometryOffsets(rank, SmallVector<int64_t>(entryCount * laneCount));
|
||||
SmallVector<SmallVector<int64_t>> geometrySizes(rank, SmallVector<int64_t>(entryCount * laneCount, 1));
|
||||
SmallVector<SmallVector<int64_t>> geometryStrides(rank, SmallVector<int64_t>(entryCount * laneCount, 1));
|
||||
for (size_t entryIndex = 0; entryIndex < entryCount; ++entryIndex) {
|
||||
unsigned assemblyEntry = run.assemblyEntries[entryIndex];
|
||||
if (assemblyEntry >= resultPlan->assemblyGeometry.size())
|
||||
return failure();
|
||||
ArrayRef<ScheduledTransferSlice> entrySlices =
|
||||
ArrayRef(run.slices)
|
||||
.slice(run.entryOffsets[entryIndex], run.entryOffsets[entryIndex + 1] - run.entryOffsets[entryIndex]);
|
||||
LogicalTransferMetadataView entry = buildMetadataView(entrySlices);
|
||||
for (size_t index = 0; index < entry.size(); ++index) {
|
||||
size_t tableIndex =
|
||||
entryIndex * laneCount + entry.targetLanes.valueAt(index);
|
||||
channelTable[tableIndex] = entry.channels.valueAt(index);
|
||||
sourceTable[tableIndex] = entry.sourceCores.valueAt(index);
|
||||
targetTable[tableIndex] = entry.targetCores.valueAt(index);
|
||||
unsigned targetLane = entry.targetLanes.valueAt(index);
|
||||
for (unsigned dimension = 0; dimension < rank; ++dimension) {
|
||||
geometryOffsets[dimension][tableIndex] =
|
||||
resultPlan->assemblyGeometry[assemblyEntry].offsets[dimension].valueAt(targetLane);
|
||||
geometrySizes[dimension][tableIndex] =
|
||||
resultPlan->assemblyGeometry[assemblyEntry].sizes[dimension].valueAt(targetLane);
|
||||
geometryStrides[dimension][tableIndex] =
|
||||
resultPlan->assemblyGeometry[assemblyEntry].strides[dimension].valueAt(targetLane);
|
||||
}
|
||||
}
|
||||
}
|
||||
Location loc = exchange->deferred.getLoc();
|
||||
Operation* initial = context.rewriter.clone(*assembly->initialValue);
|
||||
Value assembled = initial->getResult(0);
|
||||
auto emitEntry = [&](Value entryIndex, unsigned staticEntry, Value current) -> FailureOr<Value> {
|
||||
Value tableIndex = entryIndex;
|
||||
if (lane) {
|
||||
Value base = affineMulConst(context.rewriter, loc, entryIndex, laneCount, exchange->deferred);
|
||||
AffineExpr d0 = context.rewriter.getAffineDimExpr(0);
|
||||
AffineExpr d1 = context.rewriter.getAffineDimExpr(1);
|
||||
tableIndex = createOrFoldAffineApply(context.rewriter, loc, d0 + d1, ValueRange {base, lane}, exchange->deferred);
|
||||
}
|
||||
Type fragmentType = run.slices[run.entryOffsets[staticEntry]].family->requirement->publicationFragmentType;
|
||||
auto receive = SpatChannelReceiveOp::create(context.rewriter,
|
||||
loc,
|
||||
fragmentType,
|
||||
lookup(channelTable, tableIndex, exchange->deferred, context, loc),
|
||||
lookup(sourceTable, tableIndex, exchange->deferred, context, loc),
|
||||
lookup(targetTable, tableIndex, exchange->deferred, context, loc));
|
||||
setLogicalTransferMetadata(receive, metadata);
|
||||
auto source =
|
||||
applyAssemblyTransform(receive.getOutput(), assembly->entries[run.assemblyEntries[staticEntry]], context, loc);
|
||||
if (failed(source))
|
||||
return failure();
|
||||
MixedSliceGeometry geometry;
|
||||
for (unsigned dimension = 0; dimension < rank; ++dimension) {
|
||||
geometry.offsets.push_back(
|
||||
lookupGeometry(geometryOffsets[dimension], tableIndex, exchange->deferred, context, loc));
|
||||
geometry.sizes.push_back(lookupGeometry(geometrySizes[dimension], tableIndex, exchange->deferred, context, loc));
|
||||
geometry.strides.push_back(
|
||||
lookupGeometry(geometryStrides[dimension], tableIndex, exchange->deferred, context, loc));
|
||||
}
|
||||
return insertMixedSlice(context.rewriter, loc, *source, current, geometry);
|
||||
};
|
||||
if (entryCount == 1) {
|
||||
auto value = emitEntry(context.constants.getIndex(0), 0, assembled);
|
||||
if (failed(value))
|
||||
return failure();
|
||||
context.assemblies[exchange] = *value;
|
||||
return *value;
|
||||
}
|
||||
const DeferredInsertAssemblyEntryTemplate& firstEntry = assembly->entries[run.assemblyEntries.front()];
|
||||
bool sameTransform = llvm::all_of(ArrayRef(run.assemblyEntries).drop_front(), [&](unsigned entryIndex) {
|
||||
const DeferredInsertAssemblyEntryTemplate& entry = assembly->entries[entryIndex];
|
||||
return entry.sourceTransform == firstEntry.sourceTransform && entry.sourceType == firstEntry.sourceType;
|
||||
});
|
||||
if (!sameTransform)
|
||||
return failure();
|
||||
auto loop = buildNormalizedScfFor(
|
||||
context.rewriter,
|
||||
loc,
|
||||
context.constants.getIndex(0),
|
||||
context.constants.getIndex(entryCount),
|
||||
context.constants.getIndex(1),
|
||||
ValueRange {assembled},
|
||||
[&](OpBuilder&, Location, Value index, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) -> LogicalResult {
|
||||
auto value = emitEntry(index, 0, iterArgs.front());
|
||||
if (failed(value))
|
||||
return failure();
|
||||
yielded.push_back(*value);
|
||||
return success();
|
||||
});
|
||||
if (failed(loop))
|
||||
return failure();
|
||||
context.assemblies[exchange] = loop->results.front();
|
||||
return loop->results.front();
|
||||
}
|
||||
|
||||
static RequirementFamily* getAvailabilityRequirement(const AvailabilityAlternative& alternative) {
|
||||
return std::visit(
|
||||
[](const auto& value) -> RequirementFamily* {
|
||||
using Alternative = std::decay_t<decltype(value)>;
|
||||
if constexpr (std::is_same_v<Alternative, EmitReceiveRun>)
|
||||
return value.slices.front().family->requirement;
|
||||
else
|
||||
return value.families.front()->requirement;
|
||||
},
|
||||
alternative.source);
|
||||
}
|
||||
|
||||
static FailureOr<Value> emitAvailabilityAlternative(const AvailabilityAlternative& alternative,
|
||||
Value lane,
|
||||
unsigned laneCount,
|
||||
DeferredEmissionContext& context) {
|
||||
if (auto receive = std::get_if<EmitReceiveRun>(&alternative.source))
|
||||
return emitReceiveValue(*receive, lane, laneCount, context);
|
||||
return materializeLocalValue(
|
||||
std::get<MaterializeLocalFamily>(alternative.source), lane, laneCount,
|
||||
context);
|
||||
}
|
||||
|
||||
static void recordAvailability(const ResolveAvailability& availability,
|
||||
Value value,
|
||||
DeferredEmissionContext& context) {
|
||||
for (const AvailabilityAlternative& alternative : availability.alternatives) {
|
||||
if (auto receive = std::get_if<EmitReceiveRun>(&alternative.source)) {
|
||||
for (const ScheduledTransferSlice& slice : receive->slices)
|
||||
context.receives[slice.family->requirement] = value;
|
||||
continue;
|
||||
}
|
||||
for (LocalAvailabilityFamily* family :
|
||||
std::get<MaterializeLocalFamily>(alternative.source).families)
|
||||
context.receives[family->requirement] = value;
|
||||
}
|
||||
}
|
||||
|
||||
static FailureOr<SmallVector<Value>> emitSelection(
|
||||
Value branch, ArrayRef<Type> resultTypes, size_t branchCount, Location loc,
|
||||
DeferredEmissionContext& context,
|
||||
llvm::function_ref<LogicalResult(Region&, size_t)> emitRegion) {
|
||||
SmallVector<Value> results;
|
||||
if (branchCount == 2) {
|
||||
Value condition = arith::CmpIOp::create(context.rewriter, loc,
|
||||
arith::CmpIPredicate::eq, branch, context.constants.getIndex(0));
|
||||
auto selection = scf::IfOp::create(
|
||||
context.rewriter, loc, resultTypes, condition, true);
|
||||
if (failed(emitRegion(selection.getThenRegion(), 0))
|
||||
|| failed(emitRegion(selection.getElseRegion(), 1)))
|
||||
return failure();
|
||||
llvm::append_range(results, selection.getResults());
|
||||
return results;
|
||||
}
|
||||
SmallVector<int64_t> cases;
|
||||
for (size_t index = 0; index + 1 < branchCount; ++index)
|
||||
cases.push_back(index);
|
||||
auto selection = scf::IndexSwitchOp::create(
|
||||
context.rewriter, loc, resultTypes, branch, cases, cases.size());
|
||||
for (auto [index, region] : llvm::enumerate(selection.getCaseRegions()))
|
||||
if (failed(emitRegion(region, index)))
|
||||
return failure();
|
||||
if (failed(emitRegion(selection.getDefaultRegion(), branchCount - 1)))
|
||||
return failure();
|
||||
llvm::append_range(results, selection.getResults());
|
||||
return results;
|
||||
}
|
||||
|
||||
static LogicalResult emitAvailability(const ResolveAvailability& availability,
|
||||
Value lane,
|
||||
unsigned laneCount,
|
||||
DeferredEmissionContext& context) {
|
||||
if (availability.alternatives.empty())
|
||||
return failure();
|
||||
if (availability.alternatives.size() == 1) {
|
||||
auto value = emitAvailabilityAlternative(availability.alternatives.front(), lane, laneCount, context);
|
||||
if (failed(value))
|
||||
return failure();
|
||||
recordAvailability(availability, *value, context);
|
||||
return success();
|
||||
}
|
||||
if (!lane)
|
||||
return failure();
|
||||
const AvailabilityAlternative& first = availability.alternatives.front();
|
||||
RequirementFamily* requirement = getAvailabilityRequirement(first);
|
||||
Location loc = requirement->exchange->deferred.getLoc();
|
||||
SmallVector<Type> resultTypes {requirement->publicationFragmentType};
|
||||
auto emitRegion = [&](Region& region, const AvailabilityAlternative& alternative) {
|
||||
OpBuilder::InsertionGuard guard(context.rewriter);
|
||||
Block& block = region.empty() ? *context.rewriter.createBlock(®ion)
|
||||
: region.front();
|
||||
auto yield = block.empty() ? scf::YieldOp() : dyn_cast<scf::YieldOp>(&block.back());
|
||||
if (yield)
|
||||
context.rewriter.setInsertionPoint(yield);
|
||||
else
|
||||
context.rewriter.setInsertionPointToEnd(&block);
|
||||
auto value = emitAvailabilityAlternative(alternative, lane, laneCount, context);
|
||||
if (failed(value))
|
||||
return failure();
|
||||
if (yield)
|
||||
yield->setOperands(ValueRange {*value});
|
||||
else
|
||||
scf::YieldOp::create(context.rewriter, loc, ValueRange {*value});
|
||||
return success();
|
||||
};
|
||||
SmallVector<int64_t> branchByLane(laneCount);
|
||||
SmallVector<bool> assigned(laneCount);
|
||||
for (auto [branch, alternative] : llvm::enumerate(availability.alternatives))
|
||||
for (LaneInterval interval : alternative.lanes.intervals())
|
||||
for (unsigned activeLane = interval.begin; activeLane < interval.end;
|
||||
++activeLane) {
|
||||
if (assigned[activeLane])
|
||||
return failure();
|
||||
assigned[activeLane] = true;
|
||||
branchByLane[activeLane] = branch;
|
||||
}
|
||||
Value branch = emitStaticIntLookup(StaticIntSequence::fromValues(branchByLane),
|
||||
lane,
|
||||
requirement->exchange->deferred,
|
||||
context.constants,
|
||||
context.rewriter,
|
||||
loc);
|
||||
auto selected = emitSelection(branch, resultTypes,
|
||||
availability.alternatives.size(), loc, context,
|
||||
[&](Region& region, size_t index) {
|
||||
return emitRegion(region, availability.alternatives[index]);
|
||||
});
|
||||
if (failed(selected))
|
||||
return failure();
|
||||
Value value = selected->front();
|
||||
recordAvailability(availability, value, context);
|
||||
return success();
|
||||
}
|
||||
|
||||
static FailureOr<SmallVector<Value>> emitInstructions(
|
||||
const BoundaryInstructionList& list, Value lane, unsigned laneCount,
|
||||
ArrayRef<DeferredResultPlan> results, DeferredEmissionContext& context);
|
||||
|
||||
static FailureOr<SmallVector<Value>> emitDispatch(
|
||||
const LaneDispatch& dispatch, Value lane, unsigned laneCount,
|
||||
ArrayRef<DeferredResultPlan> results, DeferredEmissionContext& context) {
|
||||
if (dispatch.branches.empty())
|
||||
return failure();
|
||||
if (dispatch.branches.size() == 1)
|
||||
return emitInstructions(
|
||||
dispatch.branches.front(), lane, laneCount, results, context);
|
||||
if (!lane)
|
||||
return failure();
|
||||
|
||||
SmallVector<Type> resultTypes;
|
||||
for (DeferredExchangePlan* exchange : dispatch.producedExchanges)
|
||||
resultTypes.push_back(exchange->deferred.getOutput().getType());
|
||||
Location loc = dispatch.producedExchanges.empty()
|
||||
? lane.getLoc()
|
||||
: dispatch.producedExchanges.front()->deferred.getLoc();
|
||||
Value branch = emitStaticIntLookup(dispatch.branchByLane,
|
||||
lane,
|
||||
dispatch.producedExchanges.empty()
|
||||
? lane.getDefiningOp()
|
||||
: dispatch.producedExchanges.front()->deferred,
|
||||
context.constants,
|
||||
context.rewriter,
|
||||
loc);
|
||||
auto receives = context.receives;
|
||||
auto assemblies = context.assemblies;
|
||||
auto projectionAssemblies = context.projectionAssemblies;
|
||||
auto emitRegion = [&](Region& region,
|
||||
const BoundaryInstructionList& instructions) {
|
||||
context.receives = receives;
|
||||
context.assemblies = assemblies;
|
||||
context.projectionAssemblies = projectionAssemblies;
|
||||
OpBuilder::InsertionGuard guard(context.rewriter);
|
||||
Block& block = region.empty() ? *context.rewriter.createBlock(®ion)
|
||||
: region.front();
|
||||
auto yield = block.empty() ? scf::YieldOp()
|
||||
: dyn_cast<scf::YieldOp>(&block.back());
|
||||
if (yield)
|
||||
context.rewriter.setInsertionPoint(yield);
|
||||
else
|
||||
context.rewriter.setInsertionPointToEnd(&block);
|
||||
auto values = emitInstructions(
|
||||
instructions, lane, laneCount, results, context);
|
||||
if (failed(values))
|
||||
return failure();
|
||||
if (yield)
|
||||
yield->setOperands(*values);
|
||||
else
|
||||
scf::YieldOp::create(context.rewriter, loc, *values);
|
||||
return success();
|
||||
};
|
||||
|
||||
auto produced = emitSelection(branch, resultTypes, dispatch.branches.size(),
|
||||
loc, context, [&](Region& region, size_t index) {
|
||||
return emitRegion(region, dispatch.branches[index]);
|
||||
});
|
||||
if (failed(produced))
|
||||
return failure();
|
||||
context.receives = std::move(receives);
|
||||
context.assemblies = std::move(assemblies);
|
||||
context.projectionAssemblies = std::move(projectionAssemblies);
|
||||
return *produced;
|
||||
}
|
||||
|
||||
static FailureOr<SmallVector<Value>> emitInstructions(
|
||||
const BoundaryInstructionList& list, Value lane, unsigned laneCount,
|
||||
ArrayRef<DeferredResultPlan> results, DeferredEmissionContext& context) {
|
||||
SmallVector<Value> produced;
|
||||
for (const BoundaryInstruction& instruction : list.instructions) {
|
||||
if (auto send = std::get_if<EmitSendRun>(&instruction)) {
|
||||
if (failed(emitConditionalSendRun(*send, lane, laneCount, context)))
|
||||
return failure();
|
||||
continue;
|
||||
}
|
||||
if (auto availability = std::get_if<ResolveAvailability>(&instruction)) {
|
||||
if (failed(emitAvailability(*availability, lane, laneCount, context)))
|
||||
return failure();
|
||||
continue;
|
||||
}
|
||||
if (auto assembly = std::get_if<EmitReceiveAssemblyRun>(&instruction)) {
|
||||
auto value = assembly->projectionLeaf
|
||||
? emitProjectionAssemblyRun(
|
||||
*assembly, lane, laneCount, context)
|
||||
: emitAssemblyRun(
|
||||
*assembly, lane, laneCount, results, context);
|
||||
if (failed(value))
|
||||
return failure();
|
||||
continue;
|
||||
}
|
||||
if (auto result = std::get_if<ProduceDeferredResult>(&instruction)) {
|
||||
const DeferredResultPlan* plan = findResultPlan(results, result->exchange);
|
||||
if (!plan)
|
||||
return failure();
|
||||
auto value = realizeDeferredResult(*plan, result->lanes, lane, context);
|
||||
if (failed(value))
|
||||
return failure();
|
||||
produced.push_back(*value);
|
||||
continue;
|
||||
}
|
||||
auto values = emitDispatch(
|
||||
**std::get_if<std::unique_ptr<LaneDispatch>>(&instruction), lane,
|
||||
laneCount, results, context);
|
||||
if (failed(values))
|
||||
return failure();
|
||||
llvm::append_range(produced, *values);
|
||||
}
|
||||
return produced;
|
||||
}
|
||||
|
||||
static SmallVector<DeferredExchangePlan*> getProducedExchanges(
|
||||
const BoundaryInstructionList& list) {
|
||||
SmallVector<DeferredExchangePlan*> exchanges;
|
||||
for (const BoundaryInstruction& instruction : list.instructions) {
|
||||
if (auto result = std::get_if<ProduceDeferredResult>(&instruction))
|
||||
exchanges.push_back(result->exchange);
|
||||
else if (auto dispatch =
|
||||
std::get_if<std::unique_ptr<LaneDispatch>>(&instruction))
|
||||
llvm::append_range(exchanges, (*dispatch)->producedExchanges);
|
||||
}
|
||||
return exchanges;
|
||||
}
|
||||
|
||||
static void setInsertionAtBoundary(IRRewriter& rewriter, const BoundaryKey& key) {
|
||||
ScheduledInfo& scheduled = *key.scheduled;
|
||||
if (key.insertionStep < scheduled.stepAnchors.size() && scheduled.stepAnchors[key.insertionStep]->getBlock()) {
|
||||
rewriter.setInsertionPoint(scheduled.stepAnchors[key.insertionStep]);
|
||||
return;
|
||||
}
|
||||
rewriter.setInsertionPoint(scheduled.blocks.front()->getTerminator());
|
||||
}
|
||||
|
||||
static LogicalResult
|
||||
replaceResults(ArrayRef<DeferredExchangePlan*> exchanges, ValueRange replacements, DeferredEraseSet& erase) {
|
||||
if (exchanges.size() != replacements.size())
|
||||
return failure();
|
||||
for (auto [exchange, replacement] : llvm::zip_equal(exchanges, replacements)) {
|
||||
exchange->deferred.getOutput().replaceAllUsesWith(replacement);
|
||||
erase.insert(exchange->deferred);
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
static LogicalResult emitBoundary(const BoundaryProgram& boundary,
|
||||
ArrayRef<DeferredResultPlan> results,
|
||||
DeferredEmissionContext& context,
|
||||
DeferredEraseSet& erase) {
|
||||
setInsertionAtBoundary(context.rewriter, boundary.key);
|
||||
unsigned laneCount = boundary.key.scheduled->cores.size();
|
||||
Value lane;
|
||||
if (auto batch = dyn_cast<SpatScheduledComputeBatch>(boundary.key.scheduled->op))
|
||||
lane = *batch.getLaneArgument();
|
||||
SmallVector<DeferredExchangePlan*> exchanges =
|
||||
getProducedExchanges(boundary.root);
|
||||
auto values = emitInstructions(
|
||||
boundary.root, lane, laneCount, results, context);
|
||||
return failed(values) ? failure()
|
||||
: replaceResults(exchanges, *values, erase);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
LogicalResult realizeDeferredBoundaries(ArrayRef<BoundaryProgram> boundaries,
|
||||
ArrayRef<DeferredResultPlan> results,
|
||||
DeferredEmissionContext& context,
|
||||
DeferredEraseSet& erase) {
|
||||
ScheduledInfo* scheduled = nullptr;
|
||||
for (const BoundaryProgram& boundary : boundaries) {
|
||||
if (scheduled != boundary.key.scheduled) {
|
||||
context.receives.clear();
|
||||
context.assemblies.clear();
|
||||
context.projectionAssemblies.clear();
|
||||
scheduled = boundary.key.scheduled;
|
||||
}
|
||||
if (failed(emitBoundary(boundary, results, context, erase)))
|
||||
return boundary.key.scheduled->op->emitOpError("phase 2 failed to realize a communication boundary");
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
@@ -0,0 +1,30 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/IR/PatternMatch.h"
|
||||
|
||||
#include "DeferredBoundaryPlanning.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
struct DeferredEmissionContext {
|
||||
DeferredEmissionContext(mlir::IRRewriter& rewriter,
|
||||
ConstantPool& constants)
|
||||
: rewriter(rewriter), constants(constants) {}
|
||||
|
||||
mlir::IRRewriter& rewriter;
|
||||
ConstantPool& constants;
|
||||
llvm::DenseMap<RequirementFamily*, mlir::Value> receives;
|
||||
llvm::DenseMap<DeferredExchangePlan*, mlir::Value> assemblies;
|
||||
llvm::DenseMap<std::pair<DeferredExchangePlan*, unsigned>, mlir::Value>
|
||||
projectionAssemblies;
|
||||
};
|
||||
|
||||
using DeferredEraseSet = llvm::SetVector<mlir::Operation*>;
|
||||
|
||||
mlir::LogicalResult realizeDeferredBoundaries(mlir::ArrayRef<BoundaryProgram> boundaries,
|
||||
mlir::ArrayRef<DeferredResultPlan> results,
|
||||
DeferredEmissionContext& context,
|
||||
DeferredEraseSet& erase);
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
+320
-264
@@ -1,190 +1,255 @@
|
||||
#include "DeferredCommunicationDeadlock.hpp"
|
||||
|
||||
#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
|
||||
#include "llvm/ADT/DenseMap.h"
|
||||
#include "llvm/ADT/DenseSet.h"
|
||||
|
||||
#include <map>
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace {
|
||||
|
||||
enum class EventKind { Compute, Send, Receive };
|
||||
|
||||
struct Event {
|
||||
EventKind kind = EventKind::Compute;
|
||||
uint64_t exchangeId = 0;
|
||||
uint64_t channel = 0;
|
||||
};
|
||||
|
||||
static LogicalResult simulate(Operation *anchor,
|
||||
ArrayRef<SmallVector<Event>> streams,
|
||||
StringRef phase) {
|
||||
SmallVector<size_t> cursor(streams.size());
|
||||
DenseMap<uint64_t, unsigned> headSends;
|
||||
DenseMap<uint64_t, unsigned> headReceives;
|
||||
SmallVector<unsigned> readyComputes;
|
||||
SmallVector<uint64_t> readyExchanges;
|
||||
size_t computeCursor = 0;
|
||||
size_t exchangeCursor = 0;
|
||||
unsigned finishedStreams = 0;
|
||||
struct PlannedStreamCursor {
|
||||
unsigned step = 0;
|
||||
size_t slice = 0;
|
||||
size_t offset = 0;
|
||||
};
|
||||
|
||||
static std::optional<Event> getPlannedHead(
|
||||
unsigned stream, PlannedStreamCursor &cursor, unsigned stepCount,
|
||||
const ScheduledCommunicationPlan &plan) {
|
||||
while (cursor.slice < plan.slices.size()) {
|
||||
const ScheduledTransferSlice &slice = plan.slices[cursor.slice];
|
||||
ExternalTransferFamily &family = *slice.family;
|
||||
size_t begin = slice.familyOffset + cursor.offset;
|
||||
size_t length = slice.transferCount - cursor.offset;
|
||||
auto source = family.sourceStreams.find(stream, begin, length);
|
||||
auto target = family.targetStreams.find(stream, begin, length);
|
||||
std::optional<size_t> index = source;
|
||||
EventKind kind = EventKind::Send;
|
||||
unsigned insertionStep = slice.sourceInsertionStep;
|
||||
if (target && (!index || *target < *index)) {
|
||||
index = *target;
|
||||
kind = EventKind::Receive;
|
||||
insertionStep = slice.targetInsertionStep;
|
||||
}
|
||||
if (!index) {
|
||||
++cursor.slice;
|
||||
cursor.offset = 0;
|
||||
continue;
|
||||
}
|
||||
cursor.offset = *index - slice.familyOffset;
|
||||
if (cursor.step < insertionStep)
|
||||
return Event {EventKind::Compute, 0};
|
||||
return Event {kind, static_cast<uint64_t>(
|
||||
family.channelIds.valueAt(*index))};
|
||||
}
|
||||
if (cursor.step < stepCount)
|
||||
return Event {EventKind::Compute, 0};
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
static LogicalResult simulatePlanned(
|
||||
Operation *anchor, ArrayRef<unsigned> stepCounts,
|
||||
const ScheduledCommunicationPlan &plan) {
|
||||
SmallVector<PlannedStreamCursor> cursors(stepCounts.size());
|
||||
DenseMap<uint64_t, unsigned> sends, receives;
|
||||
SmallVector<unsigned> readyComputes;
|
||||
SmallVector<uint64_t> readyChannels;
|
||||
size_t computeCursor = 0, channelCursor = 0;
|
||||
unsigned finished = 0;
|
||||
auto registerHead = [&](unsigned stream) {
|
||||
if (cursor[stream] == streams[stream].size()) {
|
||||
++finishedStreams;
|
||||
auto head = getPlannedHead(
|
||||
stream, cursors[stream], stepCounts[stream], plan);
|
||||
if (!head) {
|
||||
++finished;
|
||||
return;
|
||||
}
|
||||
const Event &event = streams[stream][cursor[stream]];
|
||||
if (event.kind == EventKind::Compute) {
|
||||
if (head->kind == EventKind::Compute) {
|
||||
readyComputes.push_back(stream);
|
||||
return;
|
||||
}
|
||||
DenseMap<uint64_t, unsigned> &heads =
|
||||
event.kind == EventKind::Send ? headSends : headReceives;
|
||||
DenseMap<uint64_t, unsigned> &peers =
|
||||
event.kind == EventKind::Send ? headReceives : headSends;
|
||||
heads[event.exchangeId] = stream;
|
||||
if (peers.contains(event.exchangeId))
|
||||
readyExchanges.push_back(event.exchangeId);
|
||||
auto &own = head->kind == EventKind::Send ? sends : receives;
|
||||
auto &peer = head->kind == EventKind::Send ? receives : sends;
|
||||
own[head->channel] = stream;
|
||||
if (peer.contains(head->channel))
|
||||
readyChannels.push_back(head->channel);
|
||||
};
|
||||
for (unsigned stream = 0; stream < streams.size(); ++stream)
|
||||
for (unsigned stream = 0; stream < stepCounts.size(); ++stream)
|
||||
registerHead(stream);
|
||||
|
||||
while (computeCursor != readyComputes.size()
|
||||
|| exchangeCursor != readyExchanges.size()) {
|
||||
|| channelCursor != readyChannels.size()) {
|
||||
if (computeCursor != readyComputes.size()) {
|
||||
unsigned stream = readyComputes[computeCursor++];
|
||||
++cursor[stream];
|
||||
++cursors[stream].step;
|
||||
registerHead(stream);
|
||||
continue;
|
||||
}
|
||||
|
||||
uint64_t exchange = readyExchanges[exchangeCursor++];
|
||||
auto send = headSends.find(exchange);
|
||||
auto receive = headReceives.find(exchange);
|
||||
if (send == headSends.end() || receive == headReceives.end())
|
||||
uint64_t channel = readyChannels[channelCursor++];
|
||||
auto send = sends.find(channel);
|
||||
auto receive = receives.find(channel);
|
||||
if (send == sends.end() || receive == receives.end())
|
||||
continue;
|
||||
unsigned source = send->second;
|
||||
unsigned target = receive->second;
|
||||
headSends.erase(send);
|
||||
headReceives.erase(receive);
|
||||
++cursor[source];
|
||||
++cursor[target];
|
||||
sends.erase(send);
|
||||
receives.erase(receive);
|
||||
++cursors[source].offset;
|
||||
++cursors[target].offset;
|
||||
registerHead(source);
|
||||
registerHead(target);
|
||||
}
|
||||
|
||||
if (finishedStreams == streams.size())
|
||||
if (finished == stepCounts.size())
|
||||
return success();
|
||||
InFlightDiagnostic diagnostic = anchor->emitError()
|
||||
<< phase << " communication rendezvous simulation made no progress";
|
||||
InFlightDiagnostic diagnostic = anchor->emitError(
|
||||
"planned communication rendezvous simulation made no progress");
|
||||
unsigned reported = 0;
|
||||
for (unsigned stream = 0; stream < streams.size() && reported < 8; ++stream) {
|
||||
if (cursor[stream] == streams[stream].size())
|
||||
for (unsigned stream = 0;
|
||||
stream < stepCounts.size() && reported < 8; ++stream) {
|
||||
auto head = getPlannedHead(
|
||||
stream, cursors[stream], stepCounts[stream], plan);
|
||||
if (!head)
|
||||
continue;
|
||||
const Event &event = streams[stream][cursor[stream]];
|
||||
diagnostic << (reported == 0 ? "; blocked " : ", ") << "stream " << stream
|
||||
<< " at exchange " << event.exchangeId;
|
||||
++reported;
|
||||
diagnostic << (reported++ == 0 ? "; blocked " : ", ")
|
||||
<< "stream " << stream << " at channel " << head->channel;
|
||||
}
|
||||
return failure();
|
||||
}
|
||||
|
||||
static std::optional<int64_t> getI64Attr(Operation *op, StringRef name) {
|
||||
if (auto attr = op->getAttrOfType<IntegerAttr>(name))
|
||||
return attr.getInt();
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
static LogicalResult getI64ArrayAttr(
|
||||
Operation *op, StringRef name,
|
||||
std::optional<SmallVector<int64_t>> &values) {
|
||||
Attribute attr = op->getAttr(name);
|
||||
if (!attr)
|
||||
return success();
|
||||
if (auto array = dyn_cast<DenseI64ArrayAttr>(attr)) {
|
||||
values.emplace(array.asArrayRef());
|
||||
return success();
|
||||
}
|
||||
auto elements = dyn_cast<DenseIntElementsAttr>(attr);
|
||||
auto type = elements ? dyn_cast<RankedTensorType>(elements.getType())
|
||||
: RankedTensorType();
|
||||
if (!elements || !type || type.getRank() != 1
|
||||
|| !type.getElementType().isInteger(64))
|
||||
return op->emitOpError() << "has invalid " << name << " metadata";
|
||||
values.emplace();
|
||||
values->reserve(elements.getNumElements());
|
||||
for (const APInt &value : elements.getValues<APInt>())
|
||||
values->push_back(value.getSExtValue());
|
||||
return success();
|
||||
}
|
||||
|
||||
struct RealizedLogicalTransfer {
|
||||
int64_t channelId = -1;
|
||||
int64_t parentExchangeId = -1;
|
||||
int64_t parentTransferCount = 0;
|
||||
int64_t sourceCore = -1;
|
||||
int64_t targetCore = -1;
|
||||
struct RealizedOperation {
|
||||
Operation *op = nullptr;
|
||||
bool send = false;
|
||||
StaticIntSequence channels;
|
||||
StaticIntSequence parents;
|
||||
StaticIntSequence counts;
|
||||
StaticIntSequence sources;
|
||||
StaticIntSequence targets;
|
||||
};
|
||||
|
||||
static LogicalResult forEachRealizedLogicalTransfer(
|
||||
Operation *op,
|
||||
function_ref<LogicalResult(const RealizedLogicalTransfer &)> callback) {
|
||||
auto scalarChannel = getI64Attr(op, "raptor.channel_id");
|
||||
std::optional<SmallVector<int64_t>> batchChannels;
|
||||
if (failed(getI64ArrayAttr(
|
||||
op, "raptor.batch_channel_ids", batchChannels)))
|
||||
return failure();
|
||||
if (scalarChannel && batchChannels)
|
||||
return op->emitOpError(
|
||||
"mixes scalar and compact logical transfer metadata");
|
||||
|
||||
if (scalarChannel) {
|
||||
auto exchange = getI64Attr(op, "raptor.exchange_id");
|
||||
auto parent = getI64Attr(op, "raptor.parent_exchange_id");
|
||||
auto count = getI64Attr(op, "raptor.parent_transfer_count");
|
||||
auto source = getI64Attr(op, "raptor.source_core");
|
||||
auto target = getI64Attr(op, "raptor.target_core");
|
||||
if (!exchange || !parent || !count || !source || !target)
|
||||
return op->emitOpError(
|
||||
"is missing scalar logical transfer metadata");
|
||||
RealizedLogicalTransfer transfer {
|
||||
*scalarChannel, *parent, *count, *source, *target};
|
||||
if (*exchange != transfer.channelId || transfer.channelId < 0
|
||||
|| transfer.parentExchangeId < 0 || transfer.parentTransferCount <= 0
|
||||
|| transfer.sourceCore < 0 || transfer.targetCore < 0)
|
||||
return op->emitOpError("has invalid scalar logical transfer metadata");
|
||||
return callback(transfer);
|
||||
static FailureOr<size_t> getBatchTransferCount(Operation *op) {
|
||||
if (auto count = op->getAttrOfType<IntegerAttr>(
|
||||
"raptor.batch_transfer_count")) {
|
||||
if (count.getInt() > 0)
|
||||
return count.getInt();
|
||||
return op->emitOpError("has invalid compact transfer count"), failure();
|
||||
}
|
||||
if (op->hasAttr("raptor.batch_channel_ids_encoding"))
|
||||
return op->emitOpError("is missing compact transfer count"), failure();
|
||||
Attribute channels = op->getAttr("raptor.batch_channel_ids");
|
||||
if (auto array = dyn_cast_or_null<DenseI64ArrayAttr>(channels))
|
||||
return array.empty()
|
||||
? FailureOr<size_t>(failure())
|
||||
: FailureOr<size_t>(array.size());
|
||||
if (auto elements = dyn_cast_or_null<DenseIntElementsAttr>(channels);
|
||||
elements && elements.getNumElements() > 0)
|
||||
return elements.getNumElements();
|
||||
return op->emitOpError("has invalid legacy compact transfer metadata"),
|
||||
failure();
|
||||
}
|
||||
|
||||
std::optional<SmallVector<int64_t>> sources, targets, parents, counts;
|
||||
if (failed(getI64ArrayAttr(op, "raptor.batch_source_cores", sources))
|
||||
|| failed(getI64ArrayAttr(op, "raptor.batch_target_cores", targets))
|
||||
|| failed(getI64ArrayAttr(
|
||||
op, "raptor.batch_parent_exchange_ids", parents))
|
||||
|| failed(getI64ArrayAttr(
|
||||
op, "raptor.batch_parent_transfer_counts", counts)))
|
||||
static FailureOr<RealizedOperation> parseRealizedOperation(Operation *op) {
|
||||
bool scalar = op->hasAttr("raptor.channel_id");
|
||||
bool batch = op->hasAttr("raptor.batch_channel_ids");
|
||||
if (scalar == batch) {
|
||||
op->emitOpError(
|
||||
"must have exactly one scalar or compact metadata form");
|
||||
return failure();
|
||||
if (!batchChannels || !sources || !targets || !parents || !counts)
|
||||
return op->emitOpError(
|
||||
"is missing compact logical transfer metadata");
|
||||
size_t size = batchChannels->size();
|
||||
if (size == 0 || sources->size() != size || targets->size() != size
|
||||
|| parents->size() != size || counts->size() != size)
|
||||
return op->emitOpError(
|
||||
"has non-parallel compact logical transfer metadata");
|
||||
for (auto values : llvm::zip_equal(
|
||||
*batchChannels, *parents, *counts, *sources, *targets)) {
|
||||
RealizedLogicalTransfer transfer {
|
||||
std::get<0>(values), std::get<1>(values), std::get<2>(values),
|
||||
std::get<3>(values), std::get<4>(values)};
|
||||
if (transfer.channelId < 0 || transfer.parentExchangeId < 0
|
||||
|| transfer.parentTransferCount <= 0 || transfer.sourceCore < 0
|
||||
|| transfer.targetCore < 0)
|
||||
return op->emitOpError("has invalid compact logical transfer metadata");
|
||||
if (failed(callback(transfer)))
|
||||
}
|
||||
auto batchCount = scalar ? FailureOr<size_t>(1)
|
||||
: getBatchTransferCount(op);
|
||||
if (failed(batchCount))
|
||||
return failure();
|
||||
size_t size = *batchCount;
|
||||
auto channels = getStaticIntSequenceAttr(
|
||||
op, scalar ? "raptor.channel_id" : "raptor.batch_channel_ids", size);
|
||||
auto parents = getStaticIntSequenceAttr(
|
||||
op, scalar ? "raptor.parent_exchange_id"
|
||||
: "raptor.batch_parent_exchange_ids", size);
|
||||
auto counts = getStaticIntSequenceAttr(
|
||||
op, scalar ? "raptor.parent_transfer_count"
|
||||
: "raptor.batch_parent_transfer_counts", size);
|
||||
auto sources = getStaticIntSequenceAttr(
|
||||
op, scalar ? "raptor.source_core" : "raptor.batch_source_cores", size);
|
||||
auto targets = getStaticIntSequenceAttr(
|
||||
op, scalar ? "raptor.target_core" : "raptor.batch_target_cores", size);
|
||||
if (failed(channels) || failed(parents) || failed(counts)
|
||||
|| failed(sources) || failed(targets))
|
||||
return failure();
|
||||
if (scalar) {
|
||||
auto exchange = op->getAttrOfType<IntegerAttr>("raptor.exchange_id");
|
||||
if (!exchange || exchange.getInt() != channels->valueAt(0)) {
|
||||
op->emitOpError("has inconsistent scalar exchange metadata");
|
||||
return failure();
|
||||
}
|
||||
}
|
||||
return RealizedOperation {op, isa<SpatChannelSendOp>(op),
|
||||
*channels, *parents, *counts, *sources, *targets};
|
||||
}
|
||||
|
||||
struct CoreTransferSequences {
|
||||
DenseMap<int64_t, StaticIntSequenceChain> sends;
|
||||
DenseMap<int64_t, StaticIntSequenceChain> receives;
|
||||
};
|
||||
|
||||
struct ExpectedFamily {
|
||||
ExternalTransferFamily *family = nullptr;
|
||||
int64_t firstChannel = 0;
|
||||
int64_t endChannel = 0;
|
||||
};
|
||||
|
||||
static void appendByCore(DenseMap<int64_t, StaticIntSequenceChain> &result,
|
||||
const StaticIntSequence &channels,
|
||||
const StaticIntSequence &cores, size_t begin,
|
||||
size_t count) {
|
||||
size_t end = begin + count;
|
||||
cores.forEachEqualRun(
|
||||
[&](int64_t core, size_t runBegin, size_t runCount) {
|
||||
size_t selectedBegin = std::max(begin, runBegin);
|
||||
size_t selectedEnd = std::min(end, runBegin + runCount);
|
||||
if (selectedBegin < selectedEnd)
|
||||
result[core].append(
|
||||
channels, selectedBegin, selectedEnd - selectedBegin);
|
||||
});
|
||||
}
|
||||
|
||||
static LogicalResult compareSequences(
|
||||
func::FuncOp funcOp,
|
||||
const DenseMap<int64_t, StaticIntSequenceChain> &expected,
|
||||
const DenseMap<int64_t, StaticIntSequenceChain> &actual, StringRef kind) {
|
||||
if (expected.size() != actual.size())
|
||||
return funcOp.emitOpError()
|
||||
<< "realized " << kind << " stream set differs from plan";
|
||||
for (const auto &[core, sequence] : expected) {
|
||||
auto found = actual.find(core);
|
||||
if (found == actual.end())
|
||||
return funcOp.emitOpError() << "realized " << kind
|
||||
<< " stream is missing on core " << core;
|
||||
StaticIntSequenceChainCursor expectedCursor(sequence);
|
||||
StaticIntSequenceChainCursor actualCursor(found->second);
|
||||
uint64_t ordinal = 0;
|
||||
while (!expectedCursor.done() && !actualCursor.done()
|
||||
&& expectedCursor.value() == actualCursor.value()) {
|
||||
expectedCursor.advance();
|
||||
actualCursor.advance();
|
||||
++ordinal;
|
||||
}
|
||||
if (expectedCursor.done() && actualCursor.done())
|
||||
continue;
|
||||
return funcOp.emitOpError()
|
||||
<< "realized " << kind << " logical order differs on core "
|
||||
<< core << " at ordinal " << ordinal << ": expected channel "
|
||||
<< (expectedCursor.done() ? -1 : expectedCursor.value())
|
||||
<< ", actual channel "
|
||||
<< (actualCursor.done() ? -1 : actualCursor.value());
|
||||
}
|
||||
return success();
|
||||
}
|
||||
@@ -192,145 +257,136 @@ static LogicalResult forEachRealizedLogicalTransfer(
|
||||
} // namespace
|
||||
|
||||
LogicalResult verifyPlannedCommunicationDeadlockFree(
|
||||
Operation *anchor,
|
||||
unsigned streamCount,
|
||||
ArrayRef<unsigned> stepCounts,
|
||||
ArrayRef<PlannedCommunicationTransfer> transfers) {
|
||||
if (stepCounts.size() != streamCount)
|
||||
return anchor->emitError("communication plan stream count does not match step counts");
|
||||
|
||||
SmallVector<SmallVector<Event>> streams(streamCount);
|
||||
SmallVector<SmallVector<SmallVector<Event>>> atBoundary(streamCount);
|
||||
for (unsigned stream = 0; stream < streamCount; ++stream)
|
||||
atBoundary[stream].resize(stepCounts[stream] + 1);
|
||||
for (const PlannedCommunicationTransfer &transfer : transfers) {
|
||||
if (transfer.sourceStream >= streamCount || transfer.targetStream >= streamCount
|
||||
|| transfer.producerStep >= stepCounts[transfer.sourceStream]
|
||||
|| transfer.consumerStep >= stepCounts[transfer.targetStream]
|
||||
|| transfer.sourceInsertionStep > stepCounts[transfer.sourceStream]
|
||||
|| transfer.targetInsertionStep > stepCounts[transfer.targetStream]
|
||||
|| transfer.sourceInsertionStep <= transfer.producerStep
|
||||
|| transfer.targetInsertionStep > transfer.consumerStep)
|
||||
return anchor->emitError("communication plan references an invalid stream step");
|
||||
atBoundary[transfer.sourceStream][transfer.sourceInsertionStep].push_back(
|
||||
{EventKind::Send, transfer.exchangeId});
|
||||
atBoundary[transfer.targetStream][transfer.targetInsertionStep].push_back(
|
||||
{EventKind::Receive, transfer.exchangeId});
|
||||
}
|
||||
for (unsigned stream = 0; stream < streamCount; ++stream) {
|
||||
for (unsigned step = 0; step < stepCounts[stream]; ++step) {
|
||||
llvm::append_range(streams[stream], atBoundary[stream][step]);
|
||||
streams[stream].push_back({EventKind::Compute, 0});
|
||||
Operation *anchor, ArrayRef<unsigned> stepCounts,
|
||||
const ScheduledCommunicationPlan &plan) {
|
||||
for (const ScheduledTransferSlice &slice : plan.slices) {
|
||||
ExternalTransferFamily &family = *slice.family;
|
||||
for (size_t offset = 0; offset < slice.transferCount; ++offset) {
|
||||
size_t index = slice.familyOffset + offset;
|
||||
unsigned source = family.sourceStreams.valueAt(index);
|
||||
unsigned target = family.targetStreams.valueAt(index);
|
||||
if (source >= stepCounts.size() || target >= stepCounts.size()
|
||||
|| slice.sourceInsertionStep > stepCounts[source]
|
||||
|| slice.targetInsertionStep > stepCounts[target]
|
||||
|| slice.sourceInsertionStep <= family.requirement->producer->step
|
||||
|| slice.targetInsertionStep
|
||||
> family.requirement->exchange->consumerStep)
|
||||
return anchor->emitError(
|
||||
"communication plan references an invalid stream step");
|
||||
}
|
||||
llvm::append_range(streams[stream], atBoundary[stream].back());
|
||||
}
|
||||
return simulate(anchor, streams, "planned");
|
||||
return simulatePlanned(anchor, stepCounts, plan);
|
||||
}
|
||||
|
||||
LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) {
|
||||
struct LogicalOperation {
|
||||
Operation *op = nullptr;
|
||||
RealizedLogicalTransfer transfer;
|
||||
};
|
||||
DenseMap<int64_t, SmallVector<LogicalOperation, 2>> operationsByExchange;
|
||||
struct ParentExchange {
|
||||
std::optional<int64_t> expectedTransfers;
|
||||
DenseSet<int64_t> channels;
|
||||
};
|
||||
DenseMap<int64_t, ParentExchange> parentExchanges;
|
||||
DenseMap<int64_t, unsigned> streamByCore;
|
||||
SmallVector<int64_t> cores;
|
||||
LogicalResult verifyRealizedCommunicationDeadlockFree(
|
||||
func::FuncOp funcOp, const ScheduledCommunicationPlan &plan) {
|
||||
SmallVector<ExpectedFamily> families;
|
||||
DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
|
||||
for (const ScheduledTransferSlice &slice : plan.slices) {
|
||||
ExternalTransferFamily *family = slice.family;
|
||||
if (familyIndex.count(family))
|
||||
continue;
|
||||
size_t count = family->channelIds.size();
|
||||
int64_t first = family->channelIds.valueAt(0);
|
||||
for (size_t index = 1; index < count; ++index)
|
||||
if (family->channelIds.valueAt(index)
|
||||
!= first + static_cast<int64_t>(index))
|
||||
return funcOp.emitOpError(
|
||||
"planned communication family has non-consecutive channels");
|
||||
familyIndex[family] = families.size();
|
||||
families.push_back({family, first, first + static_cast<int64_t>(count)});
|
||||
}
|
||||
llvm::sort(families, [](const ExpectedFamily &lhs,
|
||||
const ExpectedFamily &rhs) {
|
||||
return lhs.firstChannel < rhs.firstChannel;
|
||||
});
|
||||
familyIndex.clear();
|
||||
std::map<int64_t, unsigned> familyByFirstChannel;
|
||||
int64_t nextChannel = 0;
|
||||
for (auto [index, expected] : llvm::enumerate(families)) {
|
||||
if (expected.firstChannel != nextChannel)
|
||||
return funcOp.emitOpError(
|
||||
"planned communication channels are not exactly contiguous");
|
||||
nextChannel = expected.endChannel;
|
||||
familyIndex[expected.family] = index;
|
||||
familyByFirstChannel.emplace(expected.firstChannel, index);
|
||||
}
|
||||
if (static_cast<uint64_t>(nextChannel) != plan.logicalTransferCount)
|
||||
return funcOp.emitOpError(
|
||||
"planned communication channel count is inconsistent");
|
||||
|
||||
CoreTransferSequences expected;
|
||||
for (const ScheduledTransferSlice &slice : plan.slices) {
|
||||
ExternalTransferFamily &family = *slice.family;
|
||||
appendByCore(expected.sends, family.channelIds, family.sourceCores,
|
||||
slice.familyOffset, slice.transferCount);
|
||||
appendByCore(expected.receives, family.channelIds, family.targetCores,
|
||||
slice.familyOffset, slice.transferCount);
|
||||
}
|
||||
|
||||
CoreTransferSequences actual;
|
||||
SmallVector<std::unique_ptr<StaticIntSequence>> actualChannels;
|
||||
bool invalid = false;
|
||||
funcOp.walk([&](Operation *op) {
|
||||
if (!isa<SpatChannelSendOp, SpatChannelReceiveOp>(op))
|
||||
if (invalid || !isa<SpatChannelSendOp, SpatChannelReceiveOp>(op))
|
||||
return;
|
||||
if (failed(forEachRealizedLogicalTransfer(
|
||||
op, [&](const RealizedLogicalTransfer &transfer) -> LogicalResult {
|
||||
operationsByExchange[transfer.channelId].push_back(
|
||||
{op, transfer});
|
||||
ParentExchange &parent =
|
||||
parentExchanges[transfer.parentExchangeId];
|
||||
if (parent.expectedTransfers
|
||||
&& *parent.expectedTransfers
|
||||
!= transfer.parentTransferCount)
|
||||
return op->emitOpError(
|
||||
"declares an inconsistent parent transfer count");
|
||||
parent.expectedTransfers = transfer.parentTransferCount;
|
||||
parent.channels.insert(transfer.channelId);
|
||||
for (int64_t core : {transfer.sourceCore, transfer.targetCore})
|
||||
if (!llvm::is_contained(cores, core))
|
||||
cores.push_back(core);
|
||||
return success();
|
||||
})))
|
||||
auto realized = parseRealizedOperation(op);
|
||||
if (failed(realized)) {
|
||||
invalid = true;
|
||||
});
|
||||
llvm::sort(cores);
|
||||
for (auto [index, core] : llvm::enumerate(cores))
|
||||
streamByCore[core] = index;
|
||||
|
||||
SmallVector<SmallVector<Event>> streams(cores.size());
|
||||
funcOp.walk([&](Operation *op) {
|
||||
if (!isa<SpatChannelSendOp, SpatChannelReceiveOp>(op))
|
||||
return;
|
||||
if (failed(forEachRealizedLogicalTransfer(
|
||||
op, [&](const RealizedLogicalTransfer &transfer) {
|
||||
unsigned stream = streamByCore.lookup(
|
||||
isa<SpatChannelSendOp>(op) ? transfer.sourceCore
|
||||
: transfer.targetCore);
|
||||
streams[stream].push_back(
|
||||
{isa<SpatChannelSendOp>(op) ? EventKind::Send
|
||||
: EventKind::Receive,
|
||||
static_cast<uint64_t>(transfer.channelId)});
|
||||
return success();
|
||||
})))
|
||||
invalid = true;
|
||||
}
|
||||
Type payloadType = realized->send
|
||||
? cast<SpatChannelSendOp>(op).getInput().getType()
|
||||
: cast<SpatChannelReceiveOp>(op).getOutput().getType();
|
||||
for (size_t index = 0; index < realized->channels.size(); ++index) {
|
||||
int64_t channel = realized->channels.valueAt(index);
|
||||
auto upper = familyByFirstChannel.upper_bound(channel);
|
||||
if (upper == familyByFirstChannel.begin()) {
|
||||
op->emitOpError("references an unknown logical channel");
|
||||
invalid = true;
|
||||
return;
|
||||
}
|
||||
ExpectedFamily &expected = families[std::prev(upper)->second];
|
||||
if (channel >= expected.endChannel) {
|
||||
op->emitOpError("references an unknown logical channel");
|
||||
invalid = true;
|
||||
return;
|
||||
}
|
||||
ExternalTransferFamily &family = *expected.family;
|
||||
size_t familyOffset = channel - expected.firstChannel;
|
||||
RequirementFamily &requirement = *family.requirement;
|
||||
if (realized->parents.valueAt(index)
|
||||
!= static_cast<int64_t>(requirement.exchange->exchangeId)
|
||||
|| realized->counts.valueAt(index)
|
||||
!= requirement.exchange->externalTransferCount
|
||||
|| realized->sources.valueAt(index)
|
||||
!= family.sourceCores.valueAt(familyOffset)
|
||||
|| realized->targets.valueAt(index)
|
||||
!= family.targetCores.valueAt(familyOffset)
|
||||
|| payloadType != requirement.publicationFragmentType) {
|
||||
op->emitOpError(
|
||||
"logical transfer metadata differs from its symbolic family");
|
||||
invalid = true;
|
||||
return;
|
||||
}
|
||||
}
|
||||
if (invalid)
|
||||
return;
|
||||
actualChannels.push_back(
|
||||
std::make_unique<StaticIntSequence>(std::move(realized->channels)));
|
||||
appendByCore(realized->send ? actual.sends : actual.receives,
|
||||
*actualChannels.back(),
|
||||
realized->send ? realized->sources : realized->targets,
|
||||
0, actualChannels.back()->size());
|
||||
});
|
||||
if (invalid)
|
||||
return failure();
|
||||
for (const auto &entry : parentExchanges)
|
||||
if (!entry.second.expectedTransfers
|
||||
|| entry.second.channels.size()
|
||||
!= static_cast<size_t>(*entry.second.expectedTransfers))
|
||||
return funcOp.emitOpError()
|
||||
<< "parent exchange " << entry.first
|
||||
<< " does not contain its declared lane transfer set";
|
||||
|
||||
for (const auto &entry : operationsByExchange) {
|
||||
if (entry.second.size() != 2
|
||||
|| isa<SpatChannelSendOp>(entry.second[0].op)
|
||||
== isa<SpatChannelSendOp>(entry.second[1].op)) {
|
||||
return funcOp.emitOpError()
|
||||
<< "exchange " << entry.first
|
||||
<< " does not have exactly one send and one receive (sends="
|
||||
<< llvm::count_if(entry.second, [](const LogicalOperation &item) {
|
||||
return isa<SpatChannelSendOp>(item.op);
|
||||
})
|
||||
<< ", receives="
|
||||
<< llvm::count_if(entry.second, [](const LogicalOperation &item) {
|
||||
return isa<SpatChannelReceiveOp>(item.op);
|
||||
})
|
||||
<< ")";
|
||||
}
|
||||
const LogicalOperation &first = entry.second[0];
|
||||
const LogicalOperation &second = entry.second[1];
|
||||
const LogicalOperation &sendRecord =
|
||||
isa<SpatChannelSendOp>(first.op) ? first : second;
|
||||
const LogicalOperation &receiveRecord =
|
||||
isa<SpatChannelReceiveOp>(first.op) ? first : second;
|
||||
auto send = cast<SpatChannelSendOp>(sendRecord.op);
|
||||
auto receive = cast<SpatChannelReceiveOp>(receiveRecord.op);
|
||||
if (send.getInput().getType() != receive.getOutput().getType())
|
||||
return send.emitOpError("send and receive payload types do not match");
|
||||
if (receiveRecord.transfer.sourceCore != sendRecord.transfer.sourceCore
|
||||
|| receiveRecord.transfer.targetCore
|
||||
!= sendRecord.transfer.targetCore
|
||||
|| receiveRecord.transfer.parentExchangeId
|
||||
!= sendRecord.transfer.parentExchangeId
|
||||
|| receiveRecord.transfer.parentTransferCount
|
||||
!= sendRecord.transfer.parentTransferCount)
|
||||
return receive.emitOpError("receive core metadata does not match its send");
|
||||
}
|
||||
return simulate(funcOp, streams, "realized");
|
||||
if (failed(compareSequences(
|
||||
funcOp, expected.sends, actual.sends, "send"))
|
||||
|| failed(compareSequences(
|
||||
funcOp, expected.receives, actual.receives, "receive")))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
|
||||
+6
-15
@@ -1,27 +1,18 @@
|
||||
#pragma once
|
||||
|
||||
#include "DeferredCommunicationScheduling.hpp"
|
||||
|
||||
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||||
#include "mlir/Support/LLVM.h"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
struct PlannedCommunicationTransfer {
|
||||
uint64_t exchangeId = 0;
|
||||
uint64_t parentExchangeId = 0;
|
||||
unsigned sourceStream = 0;
|
||||
unsigned targetStream = 0;
|
||||
unsigned producerStep = 0;
|
||||
unsigned consumerStep = 0;
|
||||
unsigned sourceInsertionStep = 0;
|
||||
unsigned targetInsertionStep = 0;
|
||||
};
|
||||
|
||||
mlir::LogicalResult verifyPlannedCommunicationDeadlockFree(
|
||||
mlir::Operation *anchor,
|
||||
unsigned streamCount,
|
||||
mlir::ArrayRef<unsigned> stepCounts,
|
||||
mlir::ArrayRef<PlannedCommunicationTransfer> transfers);
|
||||
const ScheduledCommunicationPlan &plan);
|
||||
|
||||
mlir::LogicalResult verifyRealizedCommunicationDeadlockFree(mlir::func::FuncOp funcOp);
|
||||
mlir::LogicalResult verifyRealizedCommunicationDeadlockFree(
|
||||
mlir::func::FuncOp funcOp,
|
||||
const ScheduledCommunicationPlan &plan);
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/IR/Operation.h"
|
||||
|
||||
#include "llvm/ADT/ArrayRef.h"
|
||||
#include "llvm/ADT/DenseMap.h"
|
||||
#include "llvm/ADT/SetVector.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <variant>
|
||||
|
||||
#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
struct LaneInterval {
|
||||
unsigned begin = 0;
|
||||
unsigned end = 0;
|
||||
|
||||
bool operator==(const LaneInterval& other) const { return begin == other.begin && end == other.end; }
|
||||
};
|
||||
|
||||
class LaneSet {
|
||||
public:
|
||||
static LaneSet all(unsigned laneCount) { return range(0, laneCount); }
|
||||
static LaneSet range(unsigned begin, unsigned end) {
|
||||
LaneSet lanes;
|
||||
if (begin < end)
|
||||
lanes.ranges.push_back({begin, end});
|
||||
return lanes;
|
||||
}
|
||||
|
||||
bool empty() const { return ranges.empty(); }
|
||||
size_t size() const {
|
||||
size_t result = 0;
|
||||
for (LaneInterval range : ranges)
|
||||
result += range.end - range.begin;
|
||||
return result;
|
||||
}
|
||||
bool contains(unsigned lane) const {
|
||||
return llvm::any_of(ranges, [&](LaneInterval range) { return range.begin <= lane && lane < range.end; });
|
||||
}
|
||||
llvm::ArrayRef<LaneInterval> intervals() const { return ranges; }
|
||||
|
||||
LaneSet intersect(const LaneSet& other) const {
|
||||
LaneSet result;
|
||||
for (LaneInterval lhs : ranges)
|
||||
for (LaneInterval rhs : other.ranges) {
|
||||
unsigned begin = std::max(lhs.begin, rhs.begin);
|
||||
unsigned end = std::min(lhs.end, rhs.end);
|
||||
if (begin < end)
|
||||
result.append(begin, end);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
LaneSet subtract(const LaneSet& other) const {
|
||||
LaneSet result;
|
||||
for (LaneInterval source : ranges) {
|
||||
unsigned cursor = source.begin;
|
||||
for (LaneInterval removed : other.ranges) {
|
||||
if (removed.end <= cursor || removed.begin >= source.end)
|
||||
continue;
|
||||
if (cursor < removed.begin)
|
||||
result.append(cursor, std::min(removed.begin, source.end));
|
||||
cursor = std::max(cursor, removed.end);
|
||||
if (cursor >= source.end)
|
||||
break;
|
||||
}
|
||||
if (cursor < source.end)
|
||||
result.append(cursor, source.end);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
LaneSet unite(const LaneSet& other) const {
|
||||
llvm::SmallVector<LaneInterval, 4> combined(ranges.begin(), ranges.end());
|
||||
llvm::append_range(combined, other.ranges);
|
||||
llvm::sort(combined, [](LaneInterval lhs, LaneInterval rhs) { return lhs.begin < rhs.begin; });
|
||||
LaneSet normalized;
|
||||
for (LaneInterval range : combined)
|
||||
normalized.append(range.begin, range.end);
|
||||
return normalized;
|
||||
}
|
||||
|
||||
bool operator==(const LaneSet& other) const { return ranges == other.ranges; }
|
||||
|
||||
private:
|
||||
void append(unsigned begin, unsigned end) {
|
||||
if (begin >= end)
|
||||
return;
|
||||
if (!ranges.empty() && begin <= ranges.back().end) {
|
||||
ranges.back().end = std::max(ranges.back().end, end);
|
||||
return;
|
||||
}
|
||||
ranges.push_back({begin, end});
|
||||
}
|
||||
|
||||
llvm::SmallVector<LaneInterval, 2> ranges;
|
||||
};
|
||||
|
||||
struct RequirementCoordinate {
|
||||
unsigned leafIndex = 0;
|
||||
unsigned selectedPosition = 0;
|
||||
|
||||
bool operator==(const RequirementCoordinate& other) const {
|
||||
return leafIndex == other.leafIndex && selectedPosition == other.selectedPosition;
|
||||
}
|
||||
};
|
||||
|
||||
enum class DeferredLeafForm {
|
||||
DirectSource,
|
||||
GraphBatchProjection
|
||||
};
|
||||
enum class DeferredAssemblySourceTransform {
|
||||
Identity,
|
||||
AddLeadingUnitDimension,
|
||||
RemoveLeadingUnitDimension
|
||||
};
|
||||
|
||||
struct DeferredSliceTemplate {
|
||||
llvm::SmallVector<mlir::OpFoldResult> offsets;
|
||||
llvm::SmallVector<mlir::OpFoldResult> sizes;
|
||||
llvm::SmallVector<mlir::OpFoldResult> strides;
|
||||
};
|
||||
|
||||
struct DeferredProjectionLeafTemplate {
|
||||
DeferredLeafForm form = DeferredLeafForm::DirectSource;
|
||||
mlir::Value sourceRoot;
|
||||
mlir::Value replacementRoot;
|
||||
mlir::tensor::ExtractSliceOp leadingProjection;
|
||||
DeferredSliceTemplate leadingGeometry;
|
||||
DeferredSliceTemplate innerGeometry;
|
||||
mlir::RankedTensorType reconstructedType;
|
||||
};
|
||||
|
||||
struct DeferredInsertAssemblyEntryTemplate {
|
||||
RequirementCoordinate coordinate;
|
||||
DeferredAssemblySourceTransform sourceTransform = DeferredAssemblySourceTransform::Identity;
|
||||
mlir::RankedTensorType sourceType;
|
||||
DeferredSliceTemplate targetGeometry;
|
||||
};
|
||||
|
||||
struct DeferredInsertAssemblyTemplate {
|
||||
mlir::tensor::EmptyOp initialValue;
|
||||
mlir::RankedTensorType resultType;
|
||||
llvm::SmallVector<DeferredInsertAssemblyEntryTemplate> entries;
|
||||
};
|
||||
|
||||
struct DeferredProgramTemplate {
|
||||
SpatDeferredCommunicationOp deferred;
|
||||
mlir::Value scheduledLane;
|
||||
mlir::Value yieldedValue;
|
||||
llvm::SmallVector<DeferredProjectionLeafTemplate, 0> leaves;
|
||||
llvm::SmallVector<mlir::Operation*> residualOps;
|
||||
std::optional<DeferredInsertAssemblyTemplate> insertAssembly;
|
||||
};
|
||||
|
||||
struct ScheduledInfo;
|
||||
|
||||
struct ProducedValue {
|
||||
ScheduledInfo* scheduled = nullptr;
|
||||
unsigned step = 0;
|
||||
unsigned resultIndex = 0;
|
||||
int64_t graphId = -1;
|
||||
int64_t core = -1;
|
||||
int64_t laneStart = 0;
|
||||
int64_t laneCount = 1;
|
||||
unsigned scheduledLane = 0;
|
||||
int64_t publishedSlotStart = 0;
|
||||
int64_t publishedSlotCount = 1;
|
||||
mlir::Value payload;
|
||||
mlir::Value published;
|
||||
};
|
||||
|
||||
struct ScheduledInfo {
|
||||
mlir::Operation* op = nullptr;
|
||||
llvm::SmallVector<mlir::Block*> blocks;
|
||||
llvm::SmallVector<mlir::Operation*> stepAnchors;
|
||||
llvm::SmallVector<int64_t> cores;
|
||||
llvm::SmallVector<int64_t> stepSourceIds;
|
||||
llvm::SmallVector<int64_t> resultOffsets;
|
||||
llvm::SmallVector<int64_t> resultCounts;
|
||||
llvm::SmallVector<ProducedValue*> produced;
|
||||
llvm::SmallVector<unsigned> streamIds;
|
||||
|
||||
bool isBatch() const { return mlir::isa<SpatScheduledComputeBatch>(op); }
|
||||
};
|
||||
|
||||
struct DeferredExchangePlan;
|
||||
|
||||
struct RequirementFamily {
|
||||
DeferredExchangePlan* exchange = nullptr;
|
||||
RequirementCoordinate coordinate;
|
||||
LaneSet targetLanes;
|
||||
ProducedValue* producer = nullptr;
|
||||
mlir::Type publicationFragmentType;
|
||||
std::optional<StaticIntSequence> graphLanes;
|
||||
std::optional<StaticIntSequence> producerLocalOffsets;
|
||||
};
|
||||
|
||||
struct LocalAvailabilityFamily {
|
||||
RequirementFamily* requirement = nullptr;
|
||||
LaneSet targetLanes;
|
||||
};
|
||||
|
||||
struct ExternalTransferFamily {
|
||||
RequirementFamily* requirement = nullptr;
|
||||
LaneSet targetLanes;
|
||||
ScheduledInfo* sourceScheduled = nullptr;
|
||||
ScheduledInfo* targetScheduled = nullptr;
|
||||
StaticIntSequence sourceStreams = StaticIntSequence::uniform(0, 1);
|
||||
StaticIntSequence targetStreams = StaticIntSequence::uniform(0, 1);
|
||||
StaticIntSequence sourceCores = StaticIntSequence::uniform(0, 1);
|
||||
StaticIntSequence targetCores = StaticIntSequence::uniform(0, 1);
|
||||
StaticIntSequence channelIds = StaticIntSequence::uniform(0, 1);
|
||||
};
|
||||
|
||||
struct DeferredExchangePlan {
|
||||
uint64_t exchangeId = 0;
|
||||
SpatDeferredCommunicationOp deferred;
|
||||
ScheduledInfo* target = nullptr;
|
||||
unsigned consumerStep = 0;
|
||||
unsigned targetLaneCount = 1;
|
||||
DeferredProgramTemplate program;
|
||||
llvm::SmallVector<RequirementFamily, 0> requirements;
|
||||
llvm::SmallVector<LocalAvailabilityFamily> local;
|
||||
llvm::SmallVector<ExternalTransferFamily, 0> external;
|
||||
unsigned externalTransferCount = 0;
|
||||
};
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
+157
-3421
File diff suppressed because it is too large
Load Diff
+409
@@ -0,0 +1,409 @@
|
||||
#include "llvm/ADT/MapVector.h"
|
||||
|
||||
#include <limits>
|
||||
#include <queue>
|
||||
|
||||
#include "DeferredCommunicationScheduling.hpp"
|
||||
#include "DeferredTransferPlanning.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
using namespace mlir;
|
||||
namespace {
|
||||
|
||||
struct StreamThreshold {
|
||||
unsigned stream = 0;
|
||||
unsigned completedStep = 0;
|
||||
};
|
||||
|
||||
struct TransferGroup {
|
||||
DeferredExchangePlan* exchange = nullptr;
|
||||
SmallVector<ScheduledTransferSlice> ordered;
|
||||
SmallVector<StreamThreshold> dependencies;
|
||||
unsigned unsatisfied = 0;
|
||||
bool ready = false;
|
||||
bool scheduled = false;
|
||||
BoundaryCost cost;
|
||||
std::tuple<unsigned, unsigned, unsigned, uint64_t> originalPriority;
|
||||
std::optional<TransferEmissionSignature> firstSignature;
|
||||
};
|
||||
|
||||
struct StreamProgress {
|
||||
unsigned completedStep = 0;
|
||||
SmallVector<unsigned> pendingAtStep;
|
||||
SmallVector<SmallVector<unsigned>> unlockedAtStep;
|
||||
bool queued = false;
|
||||
};
|
||||
|
||||
static size_t hashSignature(const TransferEmissionSignature& signature) {
|
||||
return llvm::hash_combine(signature.scheduled,
|
||||
signature.payload.getAsOpaquePointer(),
|
||||
signature.fragmentType.getAsOpaquePointer(),
|
||||
signature.hasGraphLane,
|
||||
signature.sourceIsBatch);
|
||||
}
|
||||
|
||||
static bool betterStaticPriority(const TransferGroup& lhs, const TransferGroup& rhs) {
|
||||
auto left = std::tuple(lhs.cost.instructionCount,
|
||||
lhs.cost.branchRegions,
|
||||
std::numeric_limits<unsigned>::max()
|
||||
- lhs.cost.absorbedTransfers,
|
||||
lhs.cost.lookupEntries,
|
||||
lhs.originalPriority);
|
||||
auto right = std::tuple(rhs.cost.instructionCount,
|
||||
rhs.cost.branchRegions,
|
||||
std::numeric_limits<unsigned>::max()
|
||||
- rhs.cost.absorbedTransfers,
|
||||
rhs.cost.lookupEntries,
|
||||
rhs.originalPriority);
|
||||
return left < right;
|
||||
}
|
||||
|
||||
static bool orderPermutationCycles(SmallVectorImpl<ScheduledTransferSlice>& slices) {
|
||||
if (slices.size() < 2)
|
||||
return false;
|
||||
DenseMap<unsigned, unsigned> edgeBySource;
|
||||
DenseMap<unsigned, unsigned> edgeByTarget;
|
||||
for (auto [index, slice] : llvm::enumerate(slices)) {
|
||||
ExternalTransferFamily& family = *slice.family;
|
||||
if (slice.transferCount != 1 || family.sourceScheduled != family.targetScheduled)
|
||||
return false;
|
||||
unsigned source = family.requirement->producer->scheduledLane;
|
||||
unsigned target = family.targetLanes.intervals().front().begin + slice.familyOffset;
|
||||
if (!edgeBySource.try_emplace(source, index).second || !edgeByTarget.try_emplace(target, index).second)
|
||||
return false;
|
||||
}
|
||||
SmallVector<ScheduledTransferSlice> ordered;
|
||||
SmallVector<bool> emitted(slices.size());
|
||||
auto appendChain = [&](unsigned source) {
|
||||
while (true) {
|
||||
auto edge = edgeBySource.find(source);
|
||||
if (edge == edgeBySource.end() || emitted[edge->second])
|
||||
break;
|
||||
unsigned index = edge->second;
|
||||
emitted[index] = true;
|
||||
ordered.push_back(slices[index]);
|
||||
ExternalTransferFamily& family = *slices[index].family;
|
||||
source = family.targetLanes.intervals().front().begin + slices[index].familyOffset;
|
||||
}
|
||||
};
|
||||
SmallVector<unsigned> pathStarts;
|
||||
for (auto [source, index] : edgeBySource)
|
||||
if (!edgeByTarget.count(source))
|
||||
pathStarts.push_back(source);
|
||||
llvm::sort(pathStarts);
|
||||
for (unsigned source : pathStarts)
|
||||
appendChain(source);
|
||||
while (ordered.size() != slices.size()) {
|
||||
unsigned source = std::numeric_limits<unsigned>::max();
|
||||
for (auto [candidate, index] : edgeBySource)
|
||||
if (!emitted[index])
|
||||
source = std::min(source, candidate);
|
||||
appendChain(source);
|
||||
}
|
||||
slices = std::move(ordered);
|
||||
return true;
|
||||
}
|
||||
|
||||
static void orderGroupSlices(TransferGroup& group) {
|
||||
SmallVector<TransferEmissionSignature> signatures;
|
||||
DenseMap<size_t, SmallVector<unsigned>> signatureIdsByHash;
|
||||
SmallVector<SmallVector<ScheduledTransferSlice>> slicesBySignature;
|
||||
SmallVector<unsigned> order;
|
||||
for (ExternalTransferFamily& family : group.exchange->external) {
|
||||
ScheduledTransferSlice slice {&family, 0, family.targetLanes.size()};
|
||||
TransferEmissionSignature signature = getTransferEmissionSignature(family);
|
||||
size_t hash = hashSignature(signature);
|
||||
std::optional<unsigned> id;
|
||||
for (unsigned candidate : signatureIdsByHash.lookup(hash))
|
||||
if (signatures[candidate] == signature) {
|
||||
id = candidate;
|
||||
break;
|
||||
}
|
||||
if (!id) {
|
||||
id = signatures.size();
|
||||
signatures.push_back(signature);
|
||||
signatureIdsByHash[hash].push_back(*id);
|
||||
slicesBySignature.emplace_back();
|
||||
order.push_back(*id);
|
||||
}
|
||||
slicesBySignature[*id].push_back(slice);
|
||||
}
|
||||
llvm::stable_sort(order, [&](unsigned lhs, unsigned rhs) {
|
||||
auto count = [&](unsigned id) {
|
||||
size_t result = 0;
|
||||
for (const ScheduledTransferSlice& slice : slicesBySignature[id])
|
||||
result += slice.transferCount;
|
||||
return result;
|
||||
};
|
||||
return count(lhs) > count(rhs);
|
||||
});
|
||||
for (unsigned id : order) {
|
||||
if (!orderPermutationCycles(slicesBySignature[id]))
|
||||
llvm::stable_sort(slicesBySignature[id], [](const ScheduledTransferSlice& lhs, const ScheduledTransferSlice& rhs) {
|
||||
ExternalTransferFamily& left = *lhs.family;
|
||||
ExternalTransferFamily& right = *rhs.family;
|
||||
return std::tuple(left.sourceStreams.valueAt(lhs.familyOffset), left.targetLanes.intervals().front().begin)
|
||||
> std::tuple(right.sourceStreams.valueAt(rhs.familyOffset), right.targetLanes.intervals().front().begin);
|
||||
});
|
||||
llvm::append_range(group.ordered, slicesBySignature[id]);
|
||||
}
|
||||
group.cost = estimateCanonicalBoundaryCost({}, group.ordered);
|
||||
if (!group.ordered.empty())
|
||||
group.firstSignature = getTransferEmissionSignature(*group.ordered.front().family);
|
||||
}
|
||||
|
||||
static SmallVector<TransferGroup> buildGroups(DeferredTransferPlan& plan) {
|
||||
SmallVector<TransferGroup> groups;
|
||||
for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges) {
|
||||
if (exchange->external.empty())
|
||||
continue;
|
||||
TransferGroup group;
|
||||
group.exchange = exchange.get();
|
||||
DenseMap<unsigned, unsigned> thresholdByStream;
|
||||
unsigned minSource = std::numeric_limits<unsigned>::max();
|
||||
unsigned minTarget = std::numeric_limits<unsigned>::max();
|
||||
for (ExternalTransferFamily& family : exchange->external) {
|
||||
unsigned source = family.sourceStreams.valueAt(0);
|
||||
thresholdByStream[source] = std::max(thresholdByStream.lookup(source), family.requirement->producer->step + 1);
|
||||
minSource = std::min(minSource, source);
|
||||
for (size_t index = 0; index < family.targetStreams.size(); ++index)
|
||||
minTarget = std::min<unsigned>(minTarget, family.targetStreams.valueAt(index));
|
||||
}
|
||||
for (LocalAvailabilityFamily& local : exchange->local)
|
||||
for (LaneInterval interval : local.targetLanes.intervals())
|
||||
for (unsigned lane = interval.begin; lane < interval.end; ++lane) {
|
||||
unsigned stream = exchange->target->streamIds[lane];
|
||||
thresholdByStream[stream] = std::max(thresholdByStream.lookup(stream), local.requirement->producer->step + 1);
|
||||
}
|
||||
SmallVector<unsigned> streams;
|
||||
for (auto [stream, threshold] : thresholdByStream)
|
||||
streams.push_back(stream);
|
||||
llvm::sort(streams);
|
||||
for (unsigned stream : streams)
|
||||
group.dependencies.push_back({stream, thresholdByStream.lookup(stream)});
|
||||
group.unsatisfied =
|
||||
llvm::count_if(group.dependencies, [](StreamThreshold dependency) { return dependency.completedStep != 0; });
|
||||
group.originalPriority = {exchange->consumerStep, minSource, minTarget, exchange->exchangeId};
|
||||
orderGroupSlices(group);
|
||||
groups.push_back(std::move(group));
|
||||
}
|
||||
return groups;
|
||||
}
|
||||
|
||||
static unsigned tailExtension(const TransferGroup& group,
|
||||
const TransferEmissionSignature& tail,
|
||||
unsigned tailBoundary,
|
||||
ArrayRef<StreamProgress> streams) {
|
||||
unsigned extension = 0;
|
||||
for (const ScheduledTransferSlice& slice : group.ordered) {
|
||||
ExternalTransferFamily& family = *slice.family;
|
||||
if (!(getTransferEmissionSignature(family) == tail))
|
||||
break;
|
||||
for (size_t index = 0; index < slice.transferCount; ++index) {
|
||||
size_t familyIndex = slice.familyOffset + index;
|
||||
if (streams[family.sourceStreams.valueAt(familyIndex)].completedStep
|
||||
!= tailBoundary)
|
||||
return extension;
|
||||
++extension;
|
||||
}
|
||||
}
|
||||
return extension;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
TransferEmissionSignature getTransferEmissionSignature(const ExternalTransferFamily& family) {
|
||||
ProducedValue* producer = family.requirement->producer;
|
||||
return {producer->scheduled,
|
||||
producer->payload,
|
||||
family.requirement->publicationFragmentType,
|
||||
family.requirement->graphLanes.has_value(),
|
||||
producer->scheduled->isBatch()};
|
||||
}
|
||||
|
||||
BoundaryCost estimateCanonicalBoundaryCost(
|
||||
ArrayRef<ScheduledTransferSlice> existingTail,
|
||||
ArrayRef<ScheduledTransferSlice> candidate) {
|
||||
BoundaryCost cost;
|
||||
std::optional<TransferEmissionSignature> previous;
|
||||
RequirementFamily* previousRequirement = nullptr;
|
||||
if (!existingTail.empty()) {
|
||||
previous = getTransferEmissionSignature(*existingTail.back().family);
|
||||
previousRequirement = existingTail.back().family->requirement;
|
||||
}
|
||||
for (const ScheduledTransferSlice& slice : candidate) {
|
||||
TransferEmissionSignature signature = getTransferEmissionSignature(*slice.family);
|
||||
if (!previous || !(*previous == signature))
|
||||
++cost.instructionCount;
|
||||
else
|
||||
cost.absorbedTransfers += slice.transferCount;
|
||||
previous = signature;
|
||||
if (previousRequirement != slice.family->requirement)
|
||||
++cost.instructionCount;
|
||||
previousRequirement = slice.family->requirement;
|
||||
cost.lookupEntries += slice.transferCount;
|
||||
}
|
||||
return cost;
|
||||
}
|
||||
|
||||
FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp funcOp, DeferredTransferPlan& plan) {
|
||||
SmallVector<TransferGroup> groups = buildGroups(plan);
|
||||
SmallVector<StreamProgress> streams(plan.stepCounts.size());
|
||||
for (auto [stream, progress] : llvm::enumerate(streams)) {
|
||||
progress.pendingAtStep.resize(plan.stepCounts[stream]);
|
||||
progress.unlockedAtStep.resize(plan.stepCounts[stream] + 1);
|
||||
}
|
||||
for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges)
|
||||
for (ExternalTransferFamily& family : exchange->external)
|
||||
for (size_t index = 0; index < family.targetStreams.size(); ++index)
|
||||
++streams[family.targetStreams.valueAt(index)].pendingAtStep[exchange->consumerStep];
|
||||
for (auto [groupIndex, group] : llvm::enumerate(groups))
|
||||
for (StreamThreshold dependency : group.dependencies)
|
||||
streams[dependency.stream].unlockedAtStep[dependency.completedStep].push_back(groupIndex);
|
||||
|
||||
auto compare = [&](unsigned lhs, unsigned rhs) { return betterStaticPriority(groups[rhs], groups[lhs]); };
|
||||
using Heap = std::priority_queue<unsigned, std::vector<unsigned>, decltype(compare)>;
|
||||
Heap ready(compare);
|
||||
DenseMap<size_t, std::unique_ptr<Heap>> bySignature;
|
||||
auto addReady = [&](unsigned index) {
|
||||
TransferGroup& group = groups[index];
|
||||
if (group.ready || group.scheduled)
|
||||
return;
|
||||
group.ready = true;
|
||||
ready.push(index);
|
||||
if (group.firstSignature) {
|
||||
auto& bucket = bySignature[hashSignature(*group.firstSignature)];
|
||||
if (!bucket)
|
||||
bucket = std::make_unique<Heap>(compare);
|
||||
bucket->push(index);
|
||||
}
|
||||
};
|
||||
for (auto [index, group] : llvm::enumerate(groups))
|
||||
if (group.unsatisfied == 0)
|
||||
addReady(index);
|
||||
|
||||
std::queue<unsigned> advanceable;
|
||||
auto enqueue = [&](unsigned stream) {
|
||||
StreamProgress& progress = streams[stream];
|
||||
if (!progress.queued && progress.completedStep < plan.stepCounts[stream]
|
||||
&& progress.pendingAtStep[progress.completedStep] == 0) {
|
||||
progress.queued = true;
|
||||
advanceable.push(stream);
|
||||
}
|
||||
};
|
||||
for (unsigned stream = 0; stream < streams.size(); ++stream)
|
||||
enqueue(stream);
|
||||
auto advance = [&] {
|
||||
bool changed = false;
|
||||
while (!advanceable.empty()) {
|
||||
unsigned stream = advanceable.front();
|
||||
advanceable.pop();
|
||||
StreamProgress& progress = streams[stream];
|
||||
progress.queued = false;
|
||||
if (progress.completedStep == plan.stepCounts[stream] || progress.pendingAtStep[progress.completedStep] != 0)
|
||||
continue;
|
||||
++progress.completedStep;
|
||||
changed = true;
|
||||
for (unsigned groupIndex : progress.unlockedAtStep[progress.completedStep])
|
||||
if (--groups[groupIndex].unsatisfied == 0)
|
||||
addReady(groupIndex);
|
||||
enqueue(stream);
|
||||
}
|
||||
return changed;
|
||||
};
|
||||
|
||||
ScheduledCommunicationPlan result;
|
||||
unsigned finishedGroups = 0;
|
||||
while (finishedGroups != groups.size()) {
|
||||
bool progressed = advance();
|
||||
std::optional<unsigned> chosen;
|
||||
unsigned bestExtension = 0;
|
||||
if (!result.slices.empty()) {
|
||||
ScheduledTransferSlice& tailSlice = result.slices.back();
|
||||
ExternalTransferFamily& tailFamily = *tailSlice.family;
|
||||
TransferEmissionSignature tail = getTransferEmissionSignature(tailFamily);
|
||||
size_t hash = hashSignature(tail);
|
||||
auto bucket = bySignature.find(hash);
|
||||
SmallVector<unsigned> inspected;
|
||||
while (bucket != bySignature.end() && !bucket->second->empty()) {
|
||||
unsigned candidate = bucket->second->top();
|
||||
bucket->second->pop();
|
||||
TransferGroup& group = groups[candidate];
|
||||
if (!group.ready || group.scheduled)
|
||||
continue;
|
||||
inspected.push_back(candidate);
|
||||
unsigned extension = tailExtension(
|
||||
group, tail, tailSlice.sourceInsertionStep, streams);
|
||||
if (extension > bestExtension
|
||||
|| (extension == bestExtension && extension != 0
|
||||
&& (!chosen || betterStaticPriority(group, groups[*chosen])))) {
|
||||
chosen = candidate;
|
||||
bestExtension = extension;
|
||||
}
|
||||
}
|
||||
for (unsigned candidate : inspected)
|
||||
if (!chosen || candidate != *chosen)
|
||||
bucket->second->push(candidate);
|
||||
}
|
||||
while (!chosen && !ready.empty()) {
|
||||
unsigned candidate = ready.top();
|
||||
ready.pop();
|
||||
if (groups[candidate].ready && !groups[candidate].scheduled)
|
||||
chosen = candidate;
|
||||
}
|
||||
if (chosen) {
|
||||
TransferGroup& group = groups[*chosen];
|
||||
group.ready = false;
|
||||
group.scheduled = true;
|
||||
++finishedGroups;
|
||||
for (ScheduledTransferSlice slice : group.ordered) {
|
||||
ExternalTransferFamily& family = *slice.family;
|
||||
size_t end = slice.familyOffset + slice.transferCount;
|
||||
for (size_t begin = slice.familyOffset; begin < end;) {
|
||||
unsigned sourceStep = streams[
|
||||
family.sourceStreams.valueAt(begin)].completedStep;
|
||||
unsigned targetStep = streams[
|
||||
family.targetStreams.valueAt(begin)].completedStep;
|
||||
size_t next = begin + 1;
|
||||
while (next < end
|
||||
&& streams[family.sourceStreams.valueAt(next)].completedStep
|
||||
== sourceStep
|
||||
&& streams[family.targetStreams.valueAt(next)].completedStep
|
||||
== targetStep)
|
||||
++next;
|
||||
result.slices.push_back(
|
||||
{&family, begin, next - begin, sourceStep, targetStep});
|
||||
result.logicalTransferCount += next - begin;
|
||||
begin = next;
|
||||
}
|
||||
for (size_t index = slice.familyOffset; index < end; ++index) {
|
||||
unsigned stream = family.targetStreams.valueAt(index);
|
||||
unsigned& pending = streams[stream].pendingAtStep[group.exchange->consumerStep];
|
||||
--pending;
|
||||
if (pending == 0 && streams[stream].completedStep == group.exchange->consumerStep)
|
||||
enqueue(stream);
|
||||
}
|
||||
}
|
||||
progressed = true;
|
||||
}
|
||||
if (progressed)
|
||||
continue;
|
||||
InFlightDiagnostic diagnostic =
|
||||
funcOp.emitOpError("global communication scheduler made no progress before IR mutation");
|
||||
unsigned reported = 0;
|
||||
for (const TransferGroup& group : groups) {
|
||||
if (group.scheduled || reported++ == 8)
|
||||
continue;
|
||||
diagnostic << "; blocked parent " << group.exchange->exchangeId;
|
||||
auto dependency = llvm::find_if(group.dependencies, [&](StreamThreshold item) {
|
||||
return streams[item.stream].completedStep < item.completedStep;
|
||||
});
|
||||
if (dependency != group.dependencies.end())
|
||||
diagnostic << " stream " << dependency->stream << " requires " << dependency->completedStep;
|
||||
}
|
||||
return failure();
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
+55
@@ -0,0 +1,55 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||||
|
||||
#include "DeferredCommunicationModel.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
struct DeferredTransferPlan;
|
||||
|
||||
struct ScheduledTransferSlice {
|
||||
ExternalTransferFamily* family = nullptr;
|
||||
size_t familyOffset = 0;
|
||||
size_t transferCount = 0;
|
||||
unsigned sourceInsertionStep = 0;
|
||||
unsigned targetInsertionStep = 0;
|
||||
};
|
||||
|
||||
struct ScheduledCommunicationPlan {
|
||||
llvm::SmallVector<ScheduledTransferSlice> slices;
|
||||
uint64_t logicalTransferCount = 0;
|
||||
};
|
||||
|
||||
struct TransferEmissionSignature {
|
||||
ScheduledInfo* scheduled = nullptr;
|
||||
mlir::Value payload;
|
||||
mlir::Type fragmentType;
|
||||
bool hasGraphLane = false;
|
||||
bool sourceIsBatch = false;
|
||||
|
||||
bool operator==(const TransferEmissionSignature& other) const {
|
||||
return scheduled == other.scheduled && payload == other.payload
|
||||
&& fragmentType == other.fragmentType
|
||||
&& hasGraphLane == other.hasGraphLane
|
||||
&& sourceIsBatch == other.sourceIsBatch;
|
||||
}
|
||||
};
|
||||
|
||||
struct BoundaryCost {
|
||||
unsigned instructionCount = 0;
|
||||
unsigned branchRegions = 0;
|
||||
unsigned absorbedTransfers = 0;
|
||||
unsigned lookupEntries = 0;
|
||||
};
|
||||
|
||||
TransferEmissionSignature getTransferEmissionSignature(const ExternalTransferFamily& family);
|
||||
|
||||
BoundaryCost estimateCanonicalBoundaryCost(
|
||||
mlir::ArrayRef<ScheduledTransferSlice> existingTail,
|
||||
mlir::ArrayRef<ScheduledTransferSlice> candidate);
|
||||
|
||||
mlir::FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(mlir::func::FuncOp funcOp,
|
||||
DeferredTransferPlan& plan);
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
+356
-431
File diff suppressed because it is too large
Load Diff
+30
-61
@@ -1,8 +1,6 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/IR/Operation.h"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
#include "DeferredCommunicationModel.hpp"
|
||||
|
||||
#include "llvm/ADT/DenseMap.h"
|
||||
|
||||
@@ -17,63 +15,32 @@ mlir::FailureOr<int64_t> evaluateDeferredIndex(
|
||||
mlir::FailureOr<int64_t> evaluateDeferredIndex(
|
||||
mlir::OpFoldResult value, const StaticIndexEnvironment &environment);
|
||||
|
||||
enum class DeferredLeafKind { ScalarSource, GraphBatchProjection, GraphBatchIdentity };
|
||||
|
||||
struct StaticSliceGeometry {
|
||||
llvm::SmallVector<int64_t> offsets;
|
||||
llvm::SmallVector<int64_t> sizes;
|
||||
llvm::SmallVector<int64_t> strides;
|
||||
};
|
||||
|
||||
struct DeferredProjectionLeaf {
|
||||
DeferredLeafKind kind = DeferredLeafKind::ScalarSource;
|
||||
unsigned sourceOperandIndex = 0;
|
||||
mlir::Value replacementRoot;
|
||||
mlir::tensor::ExtractSliceOp leadingProjection;
|
||||
llvm::SmallVector<int64_t> physicalSlots;
|
||||
StaticSliceGeometry innerGeometry;
|
||||
mlir::RankedTensorType reconstructedType;
|
||||
};
|
||||
|
||||
struct DeferredInsertAssemblyEntry {
|
||||
unsigned requirementIndex = 0;
|
||||
unsigned leafIndex = 0;
|
||||
StaticSliceGeometry targetGeometry;
|
||||
};
|
||||
|
||||
struct DeferredInsertAssembly {
|
||||
mlir::tensor::EmptyOp initialValue;
|
||||
mlir::RankedTensorType resultType;
|
||||
llvm::SmallVector<DeferredInsertAssemblyEntry, 0> entries;
|
||||
};
|
||||
|
||||
struct SpecializedDeferredProgram {
|
||||
SpatDeferredCommunicationOp deferred;
|
||||
std::optional<unsigned> targetScheduledLane;
|
||||
mlir::Value scheduledLane;
|
||||
mlir::Value yieldedValue;
|
||||
llvm::SmallVector<DeferredProjectionLeaf, 0> leaves;
|
||||
llvm::SmallVector<mlir::Operation *> residualOps;
|
||||
llvm::DenseMap<mlir::Value, int64_t> staticValues;
|
||||
std::optional<DeferredInsertAssembly> insertAssembly;
|
||||
};
|
||||
|
||||
struct ResolvedDeferredSource {
|
||||
unsigned sourceOperandIndex = 0;
|
||||
mlir::Value selectedValue;
|
||||
};
|
||||
|
||||
mlir::FailureOr<std::optional<ResolvedDeferredSource>> tryResolveDeferredSource(
|
||||
mlir::Value value, SpatDeferredCommunicationOp deferred,
|
||||
const StaticIndexEnvironment &environment);
|
||||
mlir::FailureOr<ResolvedDeferredSource> requireResolvedDeferredSource(
|
||||
mlir::Value value, SpatDeferredCommunicationOp deferred,
|
||||
const StaticIndexEnvironment &environment);
|
||||
mlir::LogicalResult verifyDeferredProgramContract(
|
||||
SpatDeferredCommunicationOp deferred);
|
||||
mlir::FailureOr<SpecializedDeferredProgram> analyzeDeferredProgram(
|
||||
SpatDeferredCommunicationOp deferred,
|
||||
std::optional<unsigned> targetScheduledLane);
|
||||
|
||||
mlir::FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
|
||||
SpatDeferredCommunicationOp deferred);
|
||||
|
||||
class DeferredLaneValueEvaluator {
|
||||
public:
|
||||
DeferredLaneValueEvaluator(const DeferredProgramTemplate &program,
|
||||
unsigned laneCount);
|
||||
|
||||
mlir::FailureOr<StaticIntSequence> evaluate(mlir::Value value);
|
||||
mlir::FailureOr<StaticIntSequence> evaluate(mlir::OpFoldResult value);
|
||||
mlir::FailureOr<StaticIntSequence> resolveSourceOperandIndices(
|
||||
mlir::Value sourceRoot);
|
||||
|
||||
private:
|
||||
const DeferredProgramTemplate &program;
|
||||
unsigned laneCount;
|
||||
llvm::DenseMap<mlir::Value, StaticIntSequence> values;
|
||||
llvm::DenseMap<mlir::Value, StaticIntSequence> sourceOperands;
|
||||
};
|
||||
|
||||
mlir::FailureOr<llvm::SmallVector<unsigned>>
|
||||
getPossibleDeferredSourceOperandIndices(
|
||||
mlir::Value sourceRoot, SpatDeferredCommunicationOp deferred);
|
||||
|
||||
struct GraphBatchPublicationMap {
|
||||
mlir::RankedTensorType physicalResultType;
|
||||
@@ -107,11 +74,13 @@ template <> struct DenseMapInfo<onnx_mlir::spatial::GraphBatchPublicationKey> {
|
||||
static onnx_mlir::spatial::GraphBatchPublicationKey getTombstoneKey() {
|
||||
return {DenseMapInfo<mlir::Operation *>::getTombstoneKey(), 0};
|
||||
}
|
||||
static unsigned getHashValue(const onnx_mlir::spatial::GraphBatchPublicationKey &key) {
|
||||
static unsigned getHashValue(
|
||||
const onnx_mlir::spatial::GraphBatchPublicationKey &key) {
|
||||
return hash_combine(key.graphBatch, key.resultIndex);
|
||||
}
|
||||
static bool isEqual(const onnx_mlir::spatial::GraphBatchPublicationKey &lhs,
|
||||
const onnx_mlir::spatial::GraphBatchPublicationKey &rhs) {
|
||||
static bool isEqual(
|
||||
const onnx_mlir::spatial::GraphBatchPublicationKey &lhs,
|
||||
const onnx_mlir::spatial::GraphBatchPublicationKey &rhs) {
|
||||
return lhs == rhs;
|
||||
}
|
||||
};
|
||||
|
||||
@@ -0,0 +1,306 @@
|
||||
#include "mlir/Dialect/Arith/IR/Arith.h"
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/IR/IRMapping.h"
|
||||
|
||||
#include "DeferredResultRealization.hpp"
|
||||
#include "DeferredBoundaryRealization.hpp"
|
||||
#include "DeferredProjectionAnalysis.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
using namespace mlir;
|
||||
namespace {
|
||||
|
||||
static Value getSequencePosition(
|
||||
const LaneSet& familyLanes, Value lane, Operation* anchor, DeferredEmissionContext& context, Location loc) {
|
||||
unsigned begin = familyLanes.intervals().front().begin;
|
||||
if (!lane)
|
||||
return context.constants.getIndex(0);
|
||||
return affineAddConst(context.rewriter, loc, lane, -static_cast<int64_t>(begin), anchor);
|
||||
}
|
||||
|
||||
static OpFoldResult materializeGeometryValue(
|
||||
const StaticIntSequence& sequence, Value lane, Operation* anchor, DeferredEmissionContext& context, Location loc) {
|
||||
if (sequence.getKind() == StaticIntSequenceKind::Uniform)
|
||||
return context.rewriter.getIndexAttr(sequence.valueAt(0));
|
||||
return emitStaticIntLookup(sequence, lane, anchor, context.constants, context.rewriter, loc);
|
||||
}
|
||||
|
||||
static MixedSliceGeometry materializeGeometry(const DeferredResultPlan::SliceGeometry& geometry,
|
||||
Value lane,
|
||||
Operation* anchor,
|
||||
DeferredEmissionContext& context,
|
||||
Location loc) {
|
||||
MixedSliceGeometry result;
|
||||
for (const StaticIntSequence& value : geometry.offsets)
|
||||
result.offsets.push_back(materializeGeometryValue(value, lane, anchor, context, loc));
|
||||
for (const StaticIntSequence& value : geometry.sizes)
|
||||
result.sizes.push_back(materializeGeometryValue(value, lane, anchor, context, loc));
|
||||
for (const StaticIntSequence& value : geometry.strides)
|
||||
result.strides.push_back(materializeGeometryValue(value, lane, anchor, context, loc));
|
||||
return result;
|
||||
}
|
||||
|
||||
static RequirementFamily*
|
||||
findRequirement(const DeferredResultPlan& plan, RequirementCoordinate coordinate, unsigned representativeLane) {
|
||||
RequirementFamily* match = nullptr;
|
||||
for (RequirementFamily* requirement : plan.requirements)
|
||||
if (requirement->coordinate == coordinate && requirement->targetLanes.contains(representativeLane)) {
|
||||
assert(!match && "requirement coordinate is not unique");
|
||||
match = requirement;
|
||||
}
|
||||
return match;
|
||||
}
|
||||
|
||||
static bool isIdentityGeometry(const DeferredResultPlan::SliceGeometry& geometry, RankedTensorType type) {
|
||||
if (geometry.offsets.size() != static_cast<size_t>(type.getRank()))
|
||||
return false;
|
||||
for (auto [dimension, values] :
|
||||
llvm::enumerate(llvm::zip_equal(geometry.offsets, geometry.sizes, geometry.strides))) {
|
||||
const StaticIntSequence& offset = std::get<0>(values);
|
||||
const StaticIntSequence& size = std::get<1>(values);
|
||||
const StaticIntSequence& stride = std::get<2>(values);
|
||||
if (offset.getKind() != StaticIntSequenceKind::Uniform || size.getKind() != StaticIntSequenceKind::Uniform
|
||||
|| stride.getKind() != StaticIntSequenceKind::Uniform || offset.valueAt(0) != 0
|
||||
|| size.valueAt(0) != type.getDimSize(dimension) || stride.valueAt(0) != 1)
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static FailureOr<Value> applyInnerGeometry(Value fragment,
|
||||
const DeferredProjectionLeafTemplate& leaf,
|
||||
const DeferredResultPlan::SliceGeometry& geometry,
|
||||
Value lane,
|
||||
DeferredEmissionContext& context) {
|
||||
if (leaf.form != DeferredLeafForm::GraphBatchProjection)
|
||||
return fragment;
|
||||
auto fragmentType = dyn_cast<RankedTensorType>(fragment.getType());
|
||||
if (!fragmentType || geometry.offsets.size() != static_cast<size_t>(fragmentType.getRank()))
|
||||
return failure();
|
||||
if (isIdentityGeometry(geometry, fragmentType))
|
||||
return fragment;
|
||||
SmallVector<int64_t> shape(leaf.reconstructedType.getShape().drop_front());
|
||||
auto resultType = RankedTensorType::get(shape, fragmentType.getElementType());
|
||||
Value sourceRoot = leaf.sourceRoot;
|
||||
MixedSliceGeometry mixed =
|
||||
materializeGeometry(geometry, lane, sourceRoot.getParentBlock()->getParentOp(), context, sourceRoot.getLoc());
|
||||
return extractMixedSliceOrIdentity(context.rewriter, leaf.sourceRoot.getLoc(), fragment, resultType, mixed);
|
||||
}
|
||||
|
||||
static FailureOr<Value> reconstructLeaf(const DeferredResultPlan& plan,
|
||||
unsigned leafIndex,
|
||||
const LaneSet& activeLanes,
|
||||
Value lane,
|
||||
DeferredEmissionContext& context) {
|
||||
DeferredExchangePlan& exchange = *plan.exchange;
|
||||
const DeferredProjectionLeafTemplate& leaf = exchange.program.leaves[leafIndex];
|
||||
unsigned representative = activeLanes.intervals().front().begin;
|
||||
if (Value assembled = context.projectionAssemblies.lookup({&exchange, leafIndex})) {
|
||||
for (RequirementFamily* requirement : plan.requirements) {
|
||||
if (requirement->coordinate.leafIndex != leafIndex || !requirement->targetLanes.contains(representative))
|
||||
continue;
|
||||
Value local = context.receives.lookup(requirement);
|
||||
if (!local)
|
||||
continue;
|
||||
auto shaped = applyInnerGeometry(local, leaf, plan.innerGeometry[leafIndex], lane, context);
|
||||
auto source = succeeded(shaped)
|
||||
? addLeadingUnitTensorDimension(context.rewriter, exchange.deferred.getLoc(), *shaped)
|
||||
: FailureOr<Value>(failure());
|
||||
if (failed(source))
|
||||
return failure();
|
||||
auto sourceType = cast<RankedTensorType>(source->getType());
|
||||
MixedSliceGeometry geometry;
|
||||
geometry.offsets.assign(sourceType.getRank(), context.rewriter.getIndexAttr(0));
|
||||
geometry.offsets.front() = context.rewriter.getIndexAttr(requirement->coordinate.selectedPosition);
|
||||
for (int64_t dimension : sourceType.getShape())
|
||||
geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
|
||||
geometry.strides.assign(sourceType.getRank(), context.rewriter.getIndexAttr(1));
|
||||
assembled = insertMixedSlice(context.rewriter, exchange.deferred.getLoc(), *source, assembled, geometry);
|
||||
}
|
||||
return assembled;
|
||||
}
|
||||
unsigned positionCount = 0;
|
||||
for (RequirementFamily* requirement : plan.requirements)
|
||||
if (requirement->coordinate.leafIndex == leafIndex && requirement->targetLanes.contains(representative))
|
||||
positionCount = std::max(positionCount, requirement->coordinate.selectedPosition + 1);
|
||||
if (positionCount == 0)
|
||||
return failure();
|
||||
SmallVector<Value> expanded;
|
||||
for (unsigned position = 0; position < positionCount; ++position) {
|
||||
RequirementFamily* requirement = findRequirement(plan, {leafIndex, position}, representative);
|
||||
if (!requirement)
|
||||
return failure();
|
||||
auto fragment = materializeDeferredRequirement(*requirement, activeLanes, lane, context);
|
||||
if (failed(fragment))
|
||||
return failure();
|
||||
auto shaped = applyInnerGeometry(*fragment, leaf, plan.innerGeometry[leafIndex], lane, context);
|
||||
if (failed(shaped))
|
||||
return failure();
|
||||
if (positionCount == 1 && shaped->getType() == leaf.reconstructedType)
|
||||
return *shaped;
|
||||
auto value = addLeadingUnitTensorDimension(context.rewriter, exchange.deferred.getLoc(), *shaped);
|
||||
if (failed(value))
|
||||
return failure();
|
||||
expanded.push_back(*value);
|
||||
}
|
||||
if (expanded.size() == 1 && expanded.front().getType() == leaf.reconstructedType)
|
||||
return expanded.front();
|
||||
constexpr size_t maxConcatInputs = 64;
|
||||
while (expanded.size() > 1) {
|
||||
SmallVector<Value> next;
|
||||
next.reserve((expanded.size() + maxConcatInputs - 1) / maxConcatInputs);
|
||||
for (size_t index = 0; index < expanded.size(); index += maxConcatInputs) {
|
||||
ValueRange inputs = ValueRange(expanded).slice(index, std::min(maxConcatInputs, expanded.size() - index));
|
||||
if (inputs.size() == 1) {
|
||||
next.push_back(inputs.front());
|
||||
continue;
|
||||
}
|
||||
auto first = cast<RankedTensorType>(inputs.front().getType());
|
||||
SmallVector<int64_t> shape(first.getShape());
|
||||
shape.front() = 0;
|
||||
for (Value input : inputs)
|
||||
shape.front() += cast<RankedTensorType>(input.getType()).getDimSize(0);
|
||||
auto resultType = RankedTensorType::get(shape, first.getElementType());
|
||||
next.push_back(
|
||||
tensor::ConcatOp::create(context.rewriter, exchange.deferred.getLoc(), resultType, 0, inputs).getResult());
|
||||
}
|
||||
expanded = std::move(next);
|
||||
}
|
||||
return expanded.front().getType() == leaf.reconstructedType ? FailureOr<Value>(expanded.front())
|
||||
: FailureOr<Value>(failure());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
FailureOr<DeferredResultPlan>
|
||||
buildDeferredResultPlan(DeferredExchangePlan& exchange) {
|
||||
DeferredResultPlan result;
|
||||
result.exchange = &exchange;
|
||||
for (RequirementFamily& requirement : exchange.requirements)
|
||||
result.requirements.push_back(&requirement);
|
||||
DeferredLaneValueEvaluator evaluator(
|
||||
exchange.program, exchange.targetLaneCount);
|
||||
auto buildGeometry = [&](const DeferredSliceTemplate& source,
|
||||
DeferredResultPlan::SliceGeometry& target) {
|
||||
auto append = [&](ArrayRef<OpFoldResult> values,
|
||||
SmallVectorImpl<StaticIntSequence>& sequences) {
|
||||
for (OpFoldResult value : values) {
|
||||
auto sequence = evaluator.evaluate(value);
|
||||
if (failed(sequence))
|
||||
return failure();
|
||||
sequences.push_back(*sequence);
|
||||
}
|
||||
return success();
|
||||
};
|
||||
return success(succeeded(append(source.offsets, target.offsets))
|
||||
&& succeeded(append(source.sizes, target.sizes))
|
||||
&& succeeded(append(source.strides, target.strides)));
|
||||
};
|
||||
for (const DeferredProjectionLeafTemplate& leaf : exchange.program.leaves) {
|
||||
DeferredResultPlan::SliceGeometry geometry;
|
||||
if (failed(buildGeometry(leaf.innerGeometry, geometry)))
|
||||
return failure();
|
||||
result.innerGeometry.push_back(std::move(geometry));
|
||||
}
|
||||
if (exchange.program.insertAssembly)
|
||||
for (const DeferredInsertAssemblyEntryTemplate& entry :
|
||||
exchange.program.insertAssembly->entries) {
|
||||
DeferredResultPlan::SliceGeometry geometry;
|
||||
if (failed(buildGeometry(entry.targetGeometry, geometry)))
|
||||
return failure();
|
||||
result.assemblyGeometry.push_back(std::move(geometry));
|
||||
}
|
||||
for (Operation* op : exchange.program.residualOps)
|
||||
op->walk([&](Operation* nested) {
|
||||
for (Value operand : nested->getOperands()) {
|
||||
if ((!operand.getType().isIndex()
|
||||
&& !isa<IntegerType>(operand.getType()))
|
||||
|| result.residualValues.count(operand))
|
||||
continue;
|
||||
if (auto sequence = evaluator.evaluate(operand); succeeded(sequence))
|
||||
result.residualValues.try_emplace(operand, *sequence);
|
||||
}
|
||||
});
|
||||
return result;
|
||||
}
|
||||
|
||||
FailureOr<Value> materializeDeferredRequirement(RequirementFamily& requirement,
|
||||
const LaneSet& activeLanes,
|
||||
Value lane,
|
||||
DeferredEmissionContext& context) {
|
||||
if (Value received = context.receives.lookup(&requirement))
|
||||
return received;
|
||||
ProducedValue& producer = *requirement.producer;
|
||||
Value payload = producer.payload;
|
||||
if (payload.getType() == requirement.publicationFragmentType)
|
||||
return payload;
|
||||
auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
|
||||
auto fragmentType = dyn_cast<RankedTensorType>(requirement.publicationFragmentType);
|
||||
if (!payloadType || !fragmentType || !requirement.producerLocalOffsets
|
||||
|| payloadType.getRank() != fragmentType.getRank() + 1)
|
||||
return failure();
|
||||
Location loc = requirement.exchange->deferred.getLoc();
|
||||
Value position = getSequencePosition(requirement.targetLanes, lane, requirement.exchange->deferred, context, loc);
|
||||
Value offset = emitStaticIntLookup(*requirement.producerLocalOffsets,
|
||||
position,
|
||||
requirement.exchange->deferred,
|
||||
context.constants,
|
||||
context.rewriter,
|
||||
loc);
|
||||
MixedSliceGeometry geometry;
|
||||
geometry.offsets.assign(payloadType.getRank(), context.rewriter.getIndexAttr(0));
|
||||
geometry.sizes.push_back(context.rewriter.getIndexAttr(1));
|
||||
geometry.strides.assign(payloadType.getRank(), context.rewriter.getIndexAttr(1));
|
||||
geometry.offsets.front() = offset;
|
||||
for (int64_t dimension : fragmentType.getShape())
|
||||
geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
|
||||
SmallVector<int64_t> unitShape {1};
|
||||
llvm::append_range(unitShape, fragmentType.getShape());
|
||||
auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType());
|
||||
Value unit = extractMixedSliceOrIdentity(context.rewriter, loc, payload, unitType, geometry);
|
||||
return removeLeadingUnitTensorDimension(context.rewriter, loc, unit, fragmentType);
|
||||
}
|
||||
|
||||
FailureOr<Value> realizeDeferredResult(const DeferredResultPlan& plan,
|
||||
const LaneSet& activeLanes,
|
||||
Value lane,
|
||||
DeferredEmissionContext& context) {
|
||||
DeferredExchangePlan& exchange = *plan.exchange;
|
||||
if (Value assembled = context.assemblies.lookup(&exchange))
|
||||
return assembled;
|
||||
IRMapping mapping;
|
||||
if (exchange.program.scheduledLane)
|
||||
mapping.map(exchange.program.scheduledLane, lane);
|
||||
for (auto [value, sequence] : plan.residualValues) {
|
||||
if (mapping.contains(value))
|
||||
continue;
|
||||
Value selected = emitStaticIntLookup(sequence,
|
||||
lane ? lane : context.constants.getIndex(0),
|
||||
exchange.deferred,
|
||||
context.constants,
|
||||
context.rewriter,
|
||||
exchange.deferred.getLoc());
|
||||
if (!value.getType().isIndex())
|
||||
selected = arith::IndexCastOp::create(context.rewriter, exchange.deferred.getLoc(), value.getType(), selected);
|
||||
mapping.map(value, selected);
|
||||
}
|
||||
for (unsigned leaf = 0; leaf < exchange.program.leaves.size(); ++leaf) {
|
||||
auto value = reconstructLeaf(plan, leaf, activeLanes, lane, context);
|
||||
if (failed(value))
|
||||
return failure();
|
||||
mapping.map(exchange.program.leaves[leaf].replacementRoot, *value);
|
||||
}
|
||||
for (Operation* op : exchange.program.residualOps) {
|
||||
Operation* copy = context.rewriter.clone(*op, mapping);
|
||||
for (auto [oldValue, newValue] : llvm::zip(op->getResults(), copy->getResults()))
|
||||
mapping.map(oldValue, newValue);
|
||||
}
|
||||
Value result = mapping.lookupOrDefault(exchange.program.yieldedValue);
|
||||
return result && result.getType() == exchange.deferred.getOutput().getType() ? FailureOr<Value>(result)
|
||||
: FailureOr<Value>(failure());
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
@@ -0,0 +1,35 @@
|
||||
#pragma once
|
||||
|
||||
#include "DeferredCommunicationModel.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
struct DeferredEmissionContext;
|
||||
|
||||
struct DeferredResultPlan {
|
||||
DeferredExchangePlan* exchange = nullptr;
|
||||
llvm::SmallVector<RequirementFamily*> requirements;
|
||||
struct SliceGeometry {
|
||||
llvm::SmallVector<StaticIntSequence> offsets;
|
||||
llvm::SmallVector<StaticIntSequence> sizes;
|
||||
llvm::SmallVector<StaticIntSequence> strides;
|
||||
};
|
||||
llvm::SmallVector<SliceGeometry, 0> innerGeometry;
|
||||
llvm::SmallVector<SliceGeometry, 0> assemblyGeometry;
|
||||
llvm::DenseMap<mlir::Value, StaticIntSequence> residualValues;
|
||||
};
|
||||
|
||||
mlir::FailureOr<DeferredResultPlan>
|
||||
buildDeferredResultPlan(DeferredExchangePlan& exchange);
|
||||
|
||||
mlir::FailureOr<mlir::Value> realizeDeferredResult(const DeferredResultPlan& plan,
|
||||
const LaneSet& activeLanes,
|
||||
mlir::Value lane,
|
||||
DeferredEmissionContext& context);
|
||||
|
||||
mlir::FailureOr<mlir::Value> materializeDeferredRequirement(RequirementFamily& requirement,
|
||||
const LaneSet& activeLanes,
|
||||
mlir::Value lane,
|
||||
DeferredEmissionContext& context);
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
@@ -0,0 +1,600 @@
|
||||
#include "mlir/Dialect/SCF/IR/SCF.h"
|
||||
|
||||
#include "llvm/ADT/MapVector.h"
|
||||
|
||||
#include "DeferredProjectionAnalysis.hpp"
|
||||
#include "DeferredTransferPlanning.hpp"
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
using namespace mlir;
|
||||
namespace {
|
||||
|
||||
static FailureOr<SmallVector<int64_t>> getI64Array(Operation* op, StringRef name) {
|
||||
auto attr = op->getAttrOfType<DenseI64ArrayAttr>(name);
|
||||
if (!attr)
|
||||
return op->emitOpError() << "phase 2 requires '" << name << "' metadata";
|
||||
return SmallVector<int64_t>(attr.asArrayRef());
|
||||
}
|
||||
|
||||
static FailureOr<SmallVector<int64_t>> getLaneTable(Operation* op, StringRef name, size_t expected) {
|
||||
if (auto array = op->getAttrOfType<DenseI64ArrayAttr>(name)) {
|
||||
if (array.size() != static_cast<int64_t>(expected))
|
||||
return op->emitOpError() << "phase 2 metadata '" << name << "' has the wrong size";
|
||||
return SmallVector<int64_t>(array.asArrayRef());
|
||||
}
|
||||
auto elements = op->getAttrOfType<DenseIntElementsAttr>(name);
|
||||
if (!elements || elements.getNumElements() != static_cast<int64_t>(expected))
|
||||
return op->emitOpError() << "phase 2 requires a correctly-sized '" << name << "' lane table";
|
||||
SmallVector<int64_t> values;
|
||||
for (APInt value : elements.getValues<APInt>())
|
||||
values.push_back(value.getSExtValue());
|
||||
return values;
|
||||
}
|
||||
|
||||
static Block* getScheduledBlock(SpatDeferredCommunicationOp deferred, Operation* scheduled) {
|
||||
Block* block = deferred->getBlock();
|
||||
while (block && block->getParentOp() != scheduled) {
|
||||
Operation* parent = block->getParentOp();
|
||||
block = parent ? parent->getBlock() : nullptr;
|
||||
}
|
||||
return block;
|
||||
}
|
||||
|
||||
static FailureOr<unsigned> getStepIndex(ScheduledInfo& info, Block* block) {
|
||||
auto it = llvm::find(info.blocks, block);
|
||||
return it == info.blocks.end() ? FailureOr<unsigned>(failure())
|
||||
: FailureOr<unsigned>(std::distance(info.blocks.begin(), it));
|
||||
}
|
||||
|
||||
static FailureOr<Value>
|
||||
getScalarStepResult(SpatScheduledCompute scheduled, Block& block, unsigned resultIndex, unsigned resultCount) {
|
||||
auto yield = dyn_cast<SpatBlockYieldOp>(block.getTerminator());
|
||||
if (!yield || resultIndex >= resultCount || yield.getOutputs().size() < resultCount)
|
||||
return scheduled.emitOpError("phase 2 cannot recover a scalar scheduled step result"), failure();
|
||||
return yield.getOutputs()[yield.getOutputs().size() - resultCount + resultIndex];
|
||||
}
|
||||
|
||||
struct BatchPublication {
|
||||
Value payload;
|
||||
tensor::ParallelInsertSliceOp insertion;
|
||||
};
|
||||
|
||||
struct ProducedValueGeometry {
|
||||
int64_t laneStart = 0;
|
||||
int64_t laneCount = 1;
|
||||
int64_t publishedSlotStart = 0;
|
||||
int64_t publishedSlotCount = 1;
|
||||
Value payload;
|
||||
Value published;
|
||||
};
|
||||
|
||||
static FailureOr<BatchPublication>
|
||||
getBatchStepResult(SpatScheduledComputeBatch scheduled, Block& block, unsigned globalResultIndex) {
|
||||
auto parallel = dyn_cast<SpatInParallelOp>(block.getTerminator());
|
||||
if (!parallel)
|
||||
return scheduled.emitOpError("phase 2 cannot recover a batched scheduled step result"), failure();
|
||||
unsigned resultBase = 1 + scheduled.getWeights().size() + scheduled.getInputs().size();
|
||||
for (Operation& op : parallel.getRegion().front()) {
|
||||
auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(op);
|
||||
auto destination = insert ? dyn_cast<BlockArgument>(insert.getDest()) : BlockArgument();
|
||||
if (destination && destination.getOwner() == &block && destination.getArgNumber() == resultBase + globalResultIndex)
|
||||
return BatchPublication {insert.getSource(), insert};
|
||||
}
|
||||
return scheduled.emitOpError("phase 2 cannot find the batched scheduled result insertion"), failure();
|
||||
}
|
||||
|
||||
static FailureOr<ProducedValueGeometry> verifyScheduledPublicationGeometry(
|
||||
ScheduledInfo& info, unsigned step, unsigned globalResult, unsigned lane,
|
||||
int64_t laneStart, int64_t laneCount, Value payload,
|
||||
tensor::ParallelInsertSliceOp insertion = {}) {
|
||||
ProducedValueGeometry result;
|
||||
result.laneStart = laneStart;
|
||||
result.laneCount = info.isBatch() ? laneCount : std::max<int64_t>(laneCount, 1);
|
||||
result.payload = payload;
|
||||
result.published = info.op->getResult(globalResult);
|
||||
auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
|
||||
auto publishedType = dyn_cast<RankedTensorType>(result.published.getType());
|
||||
if (laneStart < 0 || result.laneCount <= 0)
|
||||
return info.op->emitOpError("phase 2 scheduled publication has an invalid source-lane range"), failure();
|
||||
|
||||
if (!info.isBatch()) {
|
||||
result.publishedSlotCount = result.laneCount;
|
||||
if (payload.getType() != result.published.getType()
|
||||
|| (result.laneCount > 1
|
||||
&& (!publishedType || !publishedType.hasStaticShape()
|
||||
|| publishedType.getRank() == 0
|
||||
|| publishedType.getDimSize(0) != result.laneCount)))
|
||||
return info.op->emitOpError("phase 2 scalar scheduled publication is incompatible with its result"), failure();
|
||||
return result;
|
||||
}
|
||||
|
||||
if (!insertion || !payloadType || !payloadType.hasStaticShape()
|
||||
|| !publishedType || !publishedType.hasStaticShape()
|
||||
|| insertion.getDest() != info.blocks[step]->getArgument(
|
||||
1 + info.op->getNumOperands() + globalResult))
|
||||
return info.op->emitOpError("phase 2 batched publication is malformed"), failure();
|
||||
unsigned rank = publishedType.getRank();
|
||||
if (rank == 0 || payloadType.getRank() != rank
|
||||
|| insertion.getMixedOffsets().size() != rank
|
||||
|| insertion.getMixedSizes().size() != rank
|
||||
|| insertion.getMixedStrides().size() != rank)
|
||||
return info.op->emitOpError("phase 2 batched publication geometry rank is invalid"), failure();
|
||||
|
||||
StaticIndexEnvironment environment;
|
||||
environment.bindings[info.blocks[step]->getArgument(0)] = lane;
|
||||
SmallVector<int64_t> offsets, sizes, strides;
|
||||
for (OpFoldResult value : insertion.getMixedOffsets()) {
|
||||
auto evaluated = evaluateDeferredIndex(value, environment);
|
||||
if (failed(evaluated))
|
||||
return info.op->emitOpError("phase 2 batched publication offset is not static"), failure();
|
||||
offsets.push_back(*evaluated);
|
||||
}
|
||||
for (OpFoldResult value : insertion.getMixedSizes()) {
|
||||
auto evaluated = evaluateDeferredIndex(value, environment);
|
||||
if (failed(evaluated))
|
||||
return info.op->emitOpError("phase 2 batched publication size is not static"), failure();
|
||||
sizes.push_back(*evaluated);
|
||||
}
|
||||
for (OpFoldResult value : insertion.getMixedStrides()) {
|
||||
auto evaluated = evaluateDeferredIndex(value, environment);
|
||||
if (failed(evaluated))
|
||||
return info.op->emitOpError("phase 2 batched publication stride is not static"), failure();
|
||||
strides.push_back(*evaluated);
|
||||
}
|
||||
if (offsets.front() < 0 || sizes.front() <= 0 || sizes.front() != result.laneCount
|
||||
|| strides.front() != 1
|
||||
|| offsets.front() + sizes.front() > publishedType.getDimSize(0))
|
||||
return info.op->emitOpError("phase 2 batched publication leading geometry is invalid"), failure();
|
||||
for (unsigned dimension = 1; dimension < rank; ++dimension)
|
||||
if (offsets[dimension] != 0 || strides[dimension] != 1)
|
||||
return info.op->emitOpError("phase 2 batched publication inner geometry is invalid"), failure();
|
||||
for (unsigned dimension = 0; dimension < rank; ++dimension)
|
||||
if (sizes[dimension] != payloadType.getDimSize(dimension))
|
||||
return info.op->emitOpError("phase 2 batched publication payload shape is incompatible"), failure();
|
||||
result.publishedSlotStart = offsets.front();
|
||||
result.publishedSlotCount = sizes.front();
|
||||
return result;
|
||||
}
|
||||
|
||||
static LogicalResult collectScheduledOperations(func::FuncOp funcOp, DeferredTransferPlan& plan) {
|
||||
unsigned nextStream = 0;
|
||||
for (Operation& op : funcOp.getOps()) {
|
||||
if (!isa<SpatScheduledCompute, SpatScheduledComputeBatch>(op))
|
||||
continue;
|
||||
ScheduledInfo info;
|
||||
info.op = &op;
|
||||
Region& body = isa<SpatScheduledCompute>(op) ? cast<SpatScheduledCompute>(op).getBody()
|
||||
: cast<SpatScheduledComputeBatch>(op).getBody();
|
||||
for (Block& block : body) {
|
||||
info.blocks.push_back(&block);
|
||||
info.stepAnchors.push_back(&block.front());
|
||||
}
|
||||
auto sourceIds = getI64Array(&op, "scheduled.step_source_ids");
|
||||
auto offsets = getI64Array(&op, "scheduled.step_result_offsets");
|
||||
auto counts = getI64Array(&op, "scheduled.step_result_counts");
|
||||
if (failed(sourceIds) || failed(offsets) || failed(counts))
|
||||
return failure();
|
||||
info.stepSourceIds = std::move(*sourceIds);
|
||||
info.resultOffsets = std::move(*offsets);
|
||||
info.resultCounts = std::move(*counts);
|
||||
if (info.blocks.size() != info.stepSourceIds.size() || info.blocks.size() != info.resultOffsets.size()
|
||||
|| info.blocks.size() != info.resultCounts.size())
|
||||
return op.emitOpError("phase 2 scheduled metadata does not match its block count");
|
||||
if (auto scalar = dyn_cast<SpatScheduledCompute>(op)) {
|
||||
auto core = scalar->getAttrOfType<IntegerAttr>(kCoreIdAttrName);
|
||||
if (!core)
|
||||
return scalar.emitOpError("phase 2 requires scalar coreId metadata");
|
||||
info.cores.push_back(core.getInt());
|
||||
}
|
||||
else {
|
||||
auto cores = op.getAttrOfType<DenseI32ArrayAttr>(kCoreIdsAttrName);
|
||||
if (!cores)
|
||||
return op.emitOpError("phase 2 requires batch coreIds metadata");
|
||||
for (int32_t core : cores.asArrayRef())
|
||||
info.cores.push_back(core);
|
||||
}
|
||||
for (size_t lane = 0; lane < info.cores.size(); ++lane)
|
||||
info.streamIds.push_back(nextStream++);
|
||||
plan.scheduled.push_back(std::move(info));
|
||||
}
|
||||
plan.stepCounts.resize(nextStream);
|
||||
for (ScheduledInfo& info : plan.scheduled)
|
||||
for (unsigned stream : info.streamIds)
|
||||
plan.stepCounts[stream] = info.blocks.size();
|
||||
return success();
|
||||
}
|
||||
|
||||
static LogicalResult collectProducedValues(DeferredTransferPlan& plan) {
|
||||
for (ScheduledInfo& info : plan.scheduled) {
|
||||
SmallVector<int64_t> laneStarts, laneCounts;
|
||||
size_t tableSize = info.blocks.size() * info.cores.size();
|
||||
auto starts = info.isBatch() ? getLaneTable(info.op, "scheduled.source_lane_starts", tableSize)
|
||||
: getI64Array(info.op, "scheduled.source_lane_starts");
|
||||
auto counts = info.isBatch() ? getLaneTable(info.op, "scheduled.source_lane_counts", tableSize)
|
||||
: getI64Array(info.op, "scheduled.source_lane_counts");
|
||||
if (failed(starts) || failed(counts))
|
||||
return failure();
|
||||
laneStarts = std::move(*starts);
|
||||
laneCounts = std::move(*counts);
|
||||
for (unsigned step = 0; step < info.blocks.size(); ++step) {
|
||||
if (info.resultOffsets[step] < 0 || info.resultCounts[step] < 0)
|
||||
return info.op->emitOpError("phase 2 scheduled result metadata must be non-negative");
|
||||
for (unsigned result = 0; result < static_cast<unsigned>(info.resultCounts[step]); ++result) {
|
||||
unsigned globalResult = info.resultOffsets[step] + result;
|
||||
if (!info.isBatch()) {
|
||||
auto payload = getScalarStepResult(
|
||||
cast<SpatScheduledCompute>(info.op), *info.blocks[step], result, info.resultCounts[step]);
|
||||
if (failed(payload))
|
||||
return failure();
|
||||
auto geometry = verifyScheduledPublicationGeometry(
|
||||
info, step, globalResult, 0, laneStarts[step], laneCounts[step], *payload);
|
||||
if (failed(geometry))
|
||||
return failure();
|
||||
auto produced = std::make_unique<ProducedValue>(ProducedValue {
|
||||
&info, step, result, info.stepSourceIds[step], info.cores.front(),
|
||||
geometry->laneStart, geometry->laneCount, 0,
|
||||
geometry->publishedSlotStart, geometry->publishedSlotCount,
|
||||
geometry->payload, geometry->published});
|
||||
info.produced.push_back(produced.get());
|
||||
plan.producedByGraph[produced->graphId].push_back(produced.get());
|
||||
plan.producedStorage.push_back(std::move(produced));
|
||||
continue;
|
||||
}
|
||||
auto publication =
|
||||
getBatchStepResult(cast<SpatScheduledComputeBatch>(info.op), *info.blocks[step], globalResult);
|
||||
if (failed(publication))
|
||||
return failure();
|
||||
for (unsigned lane = 0; lane < info.cores.size(); ++lane) {
|
||||
size_t laneIndex = step * info.cores.size() + lane;
|
||||
auto geometry = verifyScheduledPublicationGeometry(
|
||||
info, step, globalResult, lane, laneStarts[laneIndex],
|
||||
laneCounts[laneIndex], publication->payload, publication->insertion);
|
||||
if (failed(geometry))
|
||||
return failure();
|
||||
auto produced = std::make_unique<ProducedValue>(ProducedValue {
|
||||
&info, step, result, info.stepSourceIds[step], info.cores[lane],
|
||||
geometry->laneStart, geometry->laneCount, lane,
|
||||
geometry->publishedSlotStart, geometry->publishedSlotCount,
|
||||
geometry->payload, geometry->published});
|
||||
info.produced.push_back(produced.get());
|
||||
plan.producedByGraph[produced->graphId].push_back(produced.get());
|
||||
plan.producedStorage.push_back(std::move(produced));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
static FailureOr<ProducedValue*> findProducer(DeferredTransferPlan& plan,
|
||||
Operation* diagnosticOwner,
|
||||
int64_t graphId,
|
||||
unsigned resultIndex,
|
||||
std::optional<int64_t> graphLane) {
|
||||
ProducedValue* match = nullptr;
|
||||
for (ProducedValue* produced : plan.producedByGraph.lookup(graphId)) {
|
||||
if (produced->resultIndex != resultIndex
|
||||
|| (graphLane && (*graphLane < produced->laneStart || *graphLane >= produced->laneStart + produced->laneCount)))
|
||||
continue;
|
||||
if (match)
|
||||
return diagnosticOwner->emitOpError("phase 2 cannot uniquely resolve graph publication ownership"), failure();
|
||||
match = produced;
|
||||
}
|
||||
if (!match)
|
||||
return diagnosticOwner->emitOpError("phase 2 cannot map a graph lane to a scheduled producer"), failure();
|
||||
if (graphLane) {
|
||||
int64_t relative = *graphLane - match->laneStart;
|
||||
if (relative < 0 || relative >= match->laneCount
|
||||
|| relative >= match->publishedSlotCount)
|
||||
return diagnosticOwner->emitOpError(
|
||||
"phase 2 graph lane is outside its scheduled publication window"), failure();
|
||||
}
|
||||
return match;
|
||||
}
|
||||
|
||||
struct RequirementPoint {
|
||||
ProducedValue* producer = nullptr;
|
||||
Type fragmentType;
|
||||
std::optional<int64_t> graphLane;
|
||||
std::optional<int64_t> localOffset;
|
||||
|
||||
bool sameFamily(const RequirementPoint& other) const {
|
||||
return producer == other.producer && fragmentType == other.fragmentType;
|
||||
}
|
||||
};
|
||||
|
||||
static FailureOr<std::optional<RequirementPoint>>
|
||||
resolveRequirementPoint(DeferredTransferPlan& plan,
|
||||
DeferredExchangePlan& exchange,
|
||||
const DeferredProjectionLeafTemplate& leaf,
|
||||
DeferredLaneValueEvaluator& evaluator,
|
||||
unsigned lane,
|
||||
unsigned position,
|
||||
GraphBatchPublicationCache& publicationCache) {
|
||||
auto sourceIndices = evaluator.resolveSourceOperandIndices(leaf.sourceRoot);
|
||||
if (failed(sourceIndices))
|
||||
return failure();
|
||||
unsigned sourceIndex = sourceIndices->valueAt(lane);
|
||||
auto source = dyn_cast<OpResult>(exchange.deferred.getSources()[sourceIndex]);
|
||||
if (!source)
|
||||
return exchange.deferred.emitOpError("phase 2 requires graph-result deferred sources"), failure();
|
||||
auto graphId = source.getOwner()->getAttrOfType<IntegerAttr>("scheduled.graph_id");
|
||||
if (!graphId)
|
||||
return exchange.deferred.emitOpError("phase 2 cannot identify graph producer"), failure();
|
||||
|
||||
RequirementPoint point;
|
||||
if (auto batch = dyn_cast<SpatGraphComputeBatch>(source.getOwner())) {
|
||||
auto publication = getGraphBatchPublicationMap(batch, source.getResultNumber(), publicationCache);
|
||||
if (failed(publication))
|
||||
return failure();
|
||||
int64_t physicalSlot = position;
|
||||
if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
|
||||
auto offset = evaluator.evaluate(leaf.leadingGeometry.offsets.front());
|
||||
auto size = evaluator.evaluate(leaf.leadingGeometry.sizes.front());
|
||||
auto stride = evaluator.evaluate(leaf.leadingGeometry.strides.front());
|
||||
if (failed(offset) || failed(size) || failed(stride))
|
||||
return failure();
|
||||
if (position >= static_cast<unsigned>(size->valueAt(lane)))
|
||||
return std::optional<RequirementPoint>();
|
||||
physicalSlot = offset->valueAt(lane) + static_cast<int64_t>(position) * stride->valueAt(lane);
|
||||
}
|
||||
else if (position >= (*publication)->physicalSlotToGraphLane.size()) {
|
||||
return std::optional<RequirementPoint>();
|
||||
}
|
||||
if (physicalSlot < 0 || physicalSlot >= static_cast<int64_t>((*publication)->physicalSlotToGraphLane.size()))
|
||||
return exchange.deferred.emitOpError("projection physical slot is outside publication map"), failure();
|
||||
point.graphLane = (*publication)->physicalSlotToGraphLane[physicalSlot];
|
||||
point.fragmentType = (*publication)->publicationFragmentType;
|
||||
auto producer = findProducer(plan, exchange.deferred, graphId.getInt(), source.getResultNumber(), point.graphLane);
|
||||
if (failed(producer))
|
||||
return failure();
|
||||
point.producer = *producer;
|
||||
point.localOffset = *point.graphLane - point.producer->laneStart;
|
||||
}
|
||||
else {
|
||||
if (position != 0 || !isa<SpatGraphCompute>(source.getOwner()))
|
||||
return std::optional<RequirementPoint>();
|
||||
point.fragmentType = source.getType();
|
||||
auto producer = findProducer(plan, exchange.deferred, graphId.getInt(), source.getResultNumber(), std::nullopt);
|
||||
if (failed(producer))
|
||||
return failure();
|
||||
point.producer = *producer;
|
||||
}
|
||||
return std::optional<RequirementPoint>(point);
|
||||
}
|
||||
|
||||
static void appendRequirementFamily(DeferredExchangePlan& exchange,
|
||||
RequirementCoordinate coordinate,
|
||||
unsigned begin,
|
||||
ArrayRef<RequirementPoint> points) {
|
||||
RequirementFamily family;
|
||||
family.exchange = &exchange;
|
||||
family.coordinate = coordinate;
|
||||
family.targetLanes = LaneSet::range(begin, begin + points.size());
|
||||
family.producer = points.front().producer;
|
||||
family.publicationFragmentType = points.front().fragmentType;
|
||||
auto sequence = [&](auto member) -> std::optional<StaticIntSequence> {
|
||||
if (!(points.front().*member))
|
||||
return std::nullopt;
|
||||
SmallVector<int64_t> values;
|
||||
for (const RequirementPoint& point : points)
|
||||
values.push_back(*(point.*member));
|
||||
return StaticIntSequence::fromValues(values);
|
||||
};
|
||||
family.graphLanes = sequence(&RequirementPoint::graphLane);
|
||||
family.producerLocalOffsets = sequence(&RequirementPoint::localOffset);
|
||||
exchange.requirements.push_back(std::move(family));
|
||||
}
|
||||
|
||||
static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan,
|
||||
DeferredExchangePlan& exchange,
|
||||
GraphBatchPublicationCache& publicationCache) {
|
||||
DeferredLaneValueEvaluator evaluator(exchange.program, exchange.targetLaneCount);
|
||||
for (auto leafItem : llvm::enumerate(exchange.program.leaves)) {
|
||||
unsigned leafIndex = leafItem.index();
|
||||
const DeferredProjectionLeafTemplate& leaf = leafItem.value();
|
||||
unsigned positionCount = 1;
|
||||
if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
|
||||
auto sizes = evaluator.evaluate(leaf.leadingGeometry.sizes.front());
|
||||
if (failed(sizes))
|
||||
return failure();
|
||||
for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane)
|
||||
positionCount = std::max<unsigned>(positionCount, sizes->valueAt(lane));
|
||||
}
|
||||
else {
|
||||
auto sources = evaluator.resolveSourceOperandIndices(leaf.sourceRoot);
|
||||
if (failed(sources))
|
||||
return failure();
|
||||
for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane) {
|
||||
auto source = dyn_cast<OpResult>(exchange.deferred.getSources()[sources->valueAt(lane)]);
|
||||
auto type = source ? dyn_cast<RankedTensorType>(source.getType()) : RankedTensorType();
|
||||
if (source && isa<SpatGraphComputeBatch>(source.getOwner()) && type)
|
||||
positionCount = std::max<unsigned>(positionCount, type.getDimSize(0));
|
||||
}
|
||||
}
|
||||
for (unsigned position = 0; position < positionCount; ++position) {
|
||||
unsigned runBegin = 0;
|
||||
SmallVector<RequirementPoint> run;
|
||||
auto flush = [&] {
|
||||
if (!run.empty())
|
||||
appendRequirementFamily(exchange, {static_cast<unsigned>(leafIndex), position}, runBegin, run);
|
||||
run.clear();
|
||||
};
|
||||
for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane) {
|
||||
auto point = resolveRequirementPoint(plan, exchange, leaf, evaluator, lane, position, publicationCache);
|
||||
if (failed(point))
|
||||
return failure();
|
||||
if (!*point) {
|
||||
flush();
|
||||
continue;
|
||||
}
|
||||
if (!run.empty() && !run.front().sameFamily(**point))
|
||||
flush();
|
||||
if (run.empty())
|
||||
runBegin = lane;
|
||||
run.push_back(**point);
|
||||
}
|
||||
flush();
|
||||
}
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
static void buildAvailabilityFamilies(DeferredExchangePlan& exchange, uint64_t& nextChannel) {
|
||||
for (RequirementFamily& requirement : exchange.requirements) {
|
||||
for (LaneInterval interval : requirement.targetLanes.intervals()) {
|
||||
unsigned runBegin = interval.begin;
|
||||
bool runLocal = false;
|
||||
bool haveRun = false;
|
||||
auto flush = [&](unsigned end) {
|
||||
if (!haveRun || runBegin == end)
|
||||
return;
|
||||
LaneSet lanes = LaneSet::range(runBegin, end);
|
||||
if (runLocal) {
|
||||
exchange.local.push_back({&requirement, lanes});
|
||||
}
|
||||
else {
|
||||
size_t count = end - runBegin;
|
||||
unsigned sourceStream = requirement.producer->scheduled->streamIds[requirement.producer->scheduledLane];
|
||||
SmallVector<int64_t> targetStreams, targetCores;
|
||||
for (unsigned lane = runBegin; lane < end; ++lane) {
|
||||
targetStreams.push_back(exchange.target->streamIds[lane]);
|
||||
targetCores.push_back(exchange.target->cores[lane]);
|
||||
}
|
||||
ExternalTransferFamily family;
|
||||
family.requirement = &requirement;
|
||||
family.targetLanes = lanes;
|
||||
family.sourceScheduled = requirement.producer->scheduled;
|
||||
family.targetScheduled = exchange.target;
|
||||
family.sourceStreams = StaticIntSequence::uniform(sourceStream, count);
|
||||
family.targetStreams = StaticIntSequence::fromValues(targetStreams);
|
||||
family.sourceCores = StaticIntSequence::uniform(requirement.producer->core, count);
|
||||
family.targetCores = StaticIntSequence::fromValues(targetCores);
|
||||
family.channelIds = StaticIntSequence::affine(nextChannel, 1, count);
|
||||
nextChannel += count;
|
||||
exchange.externalTransferCount += count;
|
||||
exchange.external.push_back(std::move(family));
|
||||
}
|
||||
};
|
||||
for (unsigned lane = interval.begin; lane < interval.end; ++lane) {
|
||||
unsigned sourceStream = requirement.producer->scheduled->streamIds[requirement.producer->scheduledLane];
|
||||
bool local =
|
||||
sourceStream == exchange.target->streamIds[lane] && requirement.producer->step < exchange.consumerStep;
|
||||
if (haveRun && local != runLocal) {
|
||||
flush(lane);
|
||||
runBegin = lane;
|
||||
}
|
||||
runLocal = local;
|
||||
haveRun = true;
|
||||
}
|
||||
flush(interval.end);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& plan) {
|
||||
DenseMap<Operation*, ScheduledInfo*> scheduledByOp;
|
||||
for (ScheduledInfo& scheduled : plan.scheduled)
|
||||
scheduledByOp[scheduled.op] = &scheduled;
|
||||
SmallVector<SpatDeferredCommunicationOp> deferredOps;
|
||||
funcOp.walk([&](SpatDeferredCommunicationOp op) { deferredOps.push_back(op); });
|
||||
GraphBatchPublicationCache publicationCache;
|
||||
uint64_t nextChannel = 0;
|
||||
for (SpatDeferredCommunicationOp deferred : deferredOps) {
|
||||
Operation* targetOp = deferred->getParentOfType<SpatScheduledCompute>();
|
||||
if (!targetOp)
|
||||
targetOp = deferred->getParentOfType<SpatScheduledComputeBatch>();
|
||||
ScheduledInfo* target = scheduledByOp.lookup(targetOp);
|
||||
auto step = target ? getStepIndex(*target, getScheduledBlock(deferred, targetOp)) : FailureOr<unsigned>(failure());
|
||||
auto program = analyzeDeferredProgramTemplate(deferred);
|
||||
if (!target || failed(step) || failed(program))
|
||||
return deferred.emitOpError("phase 2 cannot normalize deferred communication");
|
||||
auto exchange = std::make_unique<DeferredExchangePlan>();
|
||||
exchange->exchangeId = plan.exchanges.size();
|
||||
exchange->deferred = deferred;
|
||||
exchange->target = target;
|
||||
exchange->consumerStep = *step;
|
||||
exchange->targetLaneCount = target->cores.size();
|
||||
exchange->program = std::move(*program);
|
||||
if (failed(buildRequirementFamilies(plan, *exchange, publicationCache)))
|
||||
return failure();
|
||||
buildAvailabilityFamilies(*exchange, nextChannel);
|
||||
plan.exchanges.push_back(std::move(exchange));
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
static LogicalResult
|
||||
retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBatchPublicationCache& publicationCache) {
|
||||
if (blueprint.getMode() != "fragment_assembly")
|
||||
return success();
|
||||
auto operandIndices = blueprint.getFragmentOperandIndices();
|
||||
auto sourceSlots = blueprint.getFragmentSourceSlots();
|
||||
if (!operandIndices || !sourceSlots)
|
||||
return blueprint.emitOpError("phase 2 requires explicit Blueprint fragment ownership");
|
||||
SmallVector<Value> oldOperands {blueprint.getInput()};
|
||||
llvm::append_range(oldOperands, blueprint.getFragments());
|
||||
SmallVector<Value> publications;
|
||||
SmallVector<int64_t> newOperands, newSlots;
|
||||
for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
|
||||
Value source = oldOperands[operandIndex];
|
||||
auto result = dyn_cast<OpResult>(source);
|
||||
auto batch = result ? dyn_cast<SpatGraphComputeBatch>(result.getOwner()) : SpatGraphComputeBatch();
|
||||
int64_t slot = (*sourceSlots)[fragmentIndex];
|
||||
if (batch) {
|
||||
auto graphId = batch->getAttrOfType<IntegerAttr>("scheduled.graph_id");
|
||||
auto map = getGraphBatchPublicationMap(batch, result.getResultNumber(), publicationCache);
|
||||
if (!graphId || failed(map) || slot < 0 || slot >= static_cast<int64_t>((*map)->physicalSlotToGraphLane.size()))
|
||||
return blueprint.emitOpError("phase 2 cannot resolve Blueprint fragment ownership");
|
||||
int64_t graphLane = (*map)->physicalSlotToGraphLane[slot];
|
||||
auto producer = findProducer(plan, blueprint, graphId.getInt(), result.getResultNumber(), graphLane);
|
||||
if (failed(producer))
|
||||
return failure();
|
||||
source = (*producer)->published;
|
||||
slot = (*producer)->publishedSlotStart + graphLane - (*producer)->laneStart;
|
||||
if (slot < (*producer)->publishedSlotStart
|
||||
|| slot >= (*producer)->publishedSlotStart
|
||||
+ (*producer)->publishedSlotCount)
|
||||
return blueprint.emitOpError(
|
||||
"phase 2 Blueprint slot is outside its scheduled publication window"), failure();
|
||||
}
|
||||
if (Operation* producer = source.getDefiningOp();
|
||||
producer && blueprint->getBlock() == producer->getBlock() && blueprint->isBeforeInBlock(producer))
|
||||
blueprint->moveAfter(producer);
|
||||
auto it = llvm::find(publications, source);
|
||||
if (it == publications.end()) {
|
||||
publications.push_back(source);
|
||||
it = std::prev(publications.end());
|
||||
}
|
||||
newOperands.push_back(std::distance(publications.begin(), it));
|
||||
newSlots.push_back(slot);
|
||||
}
|
||||
blueprint->setOperands(publications);
|
||||
OpBuilder builder(blueprint);
|
||||
blueprint->setAttr("fragmentOperandIndices", builder.getDenseI64ArrayAttr(newOperands));
|
||||
blueprint->setAttr("fragmentSourceSlots", builder.getDenseI64ArrayAttr(newSlots));
|
||||
return success();
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(func::FuncOp funcOp) {
|
||||
DeferredTransferPlan plan;
|
||||
if (failed(collectScheduledOperations(funcOp, plan)) || failed(collectProducedValues(plan))
|
||||
|| failed(buildExchanges(funcOp, plan)))
|
||||
return failure();
|
||||
return std::move(plan);
|
||||
}
|
||||
|
||||
LogicalResult retargetDeferredPublications(func::FuncOp funcOp, DeferredTransferPlan& plan) {
|
||||
GraphBatchPublicationCache publicationCache;
|
||||
SmallVector<SpatBlueprintOp> blueprints;
|
||||
funcOp.walk([&](SpatBlueprintOp blueprint) { blueprints.push_back(blueprint); });
|
||||
for (SpatBlueprintOp blueprint : blueprints)
|
||||
if (failed(retargetBlueprint(plan, blueprint, publicationCache)))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
@@ -0,0 +1,21 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||||
|
||||
#include "DeferredCommunicationModel.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
struct DeferredTransferPlan {
|
||||
llvm::SmallVector<ScheduledInfo, 0> scheduled;
|
||||
llvm::SmallVector<std::unique_ptr<ProducedValue>> producedStorage;
|
||||
llvm::DenseMap<int64_t, llvm::SmallVector<ProducedValue*>> producedByGraph;
|
||||
llvm::SmallVector<std::unique_ptr<DeferredExchangePlan>> exchanges;
|
||||
llvm::SmallVector<unsigned> stepCounts;
|
||||
};
|
||||
|
||||
mlir::FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(mlir::func::FuncOp funcOp);
|
||||
|
||||
mlir::LogicalResult retargetDeferredPublications(mlir::func::FuncOp funcOp, DeferredTransferPlan& plan);
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
@@ -3,7 +3,6 @@
|
||||
#include "ScheduledComputeVerification.hpp"
|
||||
#include "SpatialDataflowCsvExporter.hpp"
|
||||
#include "DeferredCommunicationRealization.hpp"
|
||||
#include "DeferredCommunicationDeadlock.hpp"
|
||||
|
||||
#include "mlir/Pass/Pass.h"
|
||||
|
||||
@@ -101,11 +100,6 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
|
||||
return;
|
||||
}
|
||||
dumpModule(moduleOp, "spatial3_scheduled", /*assumeVerified=*/true);
|
||||
if (failed(verifyRealizedCommunicationDeadlockFree(funcOp))) {
|
||||
moduleOp.emitError("MergeComputeNodes final communication verification failed");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
if (failed(verifyScheduledSpatialInvariants(funcOp))) {
|
||||
moduleOp.emitError("scheduled Spatial phase 2 verification failed");
|
||||
signalPassFailure();
|
||||
|
||||
+8
-19
@@ -105,25 +105,14 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
|
||||
diagnostics.report(transfer.getOperation(), [&](Operation *) {});
|
||||
return;
|
||||
}
|
||||
if (auto scheduled = transfer->getParentOfType<SpatScheduledComputeBatch>()) {
|
||||
for (unsigned lane = 0; lane < static_cast<unsigned>(scheduled.getLaneCount()); ++lane) {
|
||||
auto program = analyzeDeferredProgram(transfer, lane);
|
||||
if (failed(program)) {
|
||||
diagnostics.report(transfer.getOperation(), [&](Operation *) {});
|
||||
continue;
|
||||
}
|
||||
for (const DeferredProjectionLeaf &leaf : program->leaves) {
|
||||
if (leaf.kind == DeferredLeafKind::ScalarSource)
|
||||
continue;
|
||||
auto source = dyn_cast<OpResult>(transfer.getSources()[leaf.sourceOperandIndex]);
|
||||
auto graph = source ? dyn_cast<SpatGraphComputeBatch>(source.getOwner()) : SpatGraphComputeBatch();
|
||||
if (!graph) continue;
|
||||
if (failed(getGraphBatchPublicationMap(graph, source.getResultNumber(), publicationCache)))
|
||||
diagnostics.report(transfer.getOperation(), [&](Operation *) {});
|
||||
}
|
||||
}
|
||||
} else if (failed(analyzeDeferredProgram(transfer, std::nullopt))) {
|
||||
diagnostics.report(transfer.getOperation(), [&](Operation *) {});
|
||||
for (Value source : transfer.getSources()) {
|
||||
auto result = dyn_cast<OpResult>(source);
|
||||
auto batch = result
|
||||
? dyn_cast<SpatGraphComputeBatch>(result.getOwner())
|
||||
: SpatGraphComputeBatch();
|
||||
if (batch && failed(getGraphBatchPublicationMap(
|
||||
batch, result.getResultNumber(), publicationCache)))
|
||||
diagnostics.report(transfer.getOperation(), [&](Operation *) {});
|
||||
}
|
||||
});
|
||||
|
||||
|
||||
Reference in New Issue
Block a user