From c744f388dcac13fe6b47c168c4a5d99eee575183 Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Wed, 15 Jul 2026 10:40:14 +0200 Subject: [PATCH] Better implementation of MergeComputeNodesPass --- src/PIM/Common/CMakeLists.txt | 1 + src/PIM/Common/IR/ConstantUtils.cpp | 26 + src/PIM/Common/IR/ConstantUtils.hpp | 16 + src/PIM/Common/IR/StaticIntSequence.cpp | 539 +++ src/PIM/Common/IR/StaticIntSequence.hpp | 117 + src/PIM/Common/IR/TensorSliceUtils.cpp | 17 + src/PIM/Common/IR/TensorSliceUtils.hpp | 18 + src/PIM/Dialect/Spatial/CMakeLists.txt | 5 + .../DeferredBoundaryPlanning.cpp | 658 +++ .../DeferredBoundaryPlanning.hpp | 102 + .../DeferredBoundaryRealization.cpp | 824 ++++ .../DeferredBoundaryRealization.hpp | 30 + .../DeferredCommunicationDeadlock.cpp | 584 +-- .../DeferredCommunicationDeadlock.hpp | 21 +- .../DeferredCommunicationModel.hpp | 236 ++ .../DeferredCommunicationRealization.cpp | 3578 +---------------- .../DeferredCommunicationScheduling.cpp | 409 ++ .../DeferredCommunicationScheduling.hpp | 55 + .../DeferredProjectionAnalysis.cpp | 787 ++-- .../DeferredProjectionAnalysis.hpp | 91 +- .../DeferredResultRealization.cpp | 306 ++ .../DeferredResultRealization.hpp | 35 + .../DeferredTransferPlanning.cpp | 600 +++ .../DeferredTransferPlanning.hpp | 21 + .../MergeComputeNodesPass.cpp | 6 - .../ScheduledComputeVerification.cpp | 27 +- 26 files changed, 4892 insertions(+), 4217 deletions(-) create mode 100644 src/PIM/Common/IR/StaticIntSequence.cpp create mode 100644 src/PIM/Common/IR/StaticIntSequence.hpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.hpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.hpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp diff --git a/src/PIM/Common/CMakeLists.txt b/src/PIM/Common/CMakeLists.txt index 2056a13..be85ed9 100644 --- a/src/PIM/Common/CMakeLists.txt +++ b/src/PIM/Common/CMakeLists.txt @@ -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 diff --git a/src/PIM/Common/IR/ConstantUtils.cpp b/src/PIM/Common/IR/ConstantUtils.cpp index 8f7c5b7..c1dc1eb 100644 --- a/src/PIM/Common/IR/ConstantUtils.cpp +++ b/src/PIM/Common/IR/ConstantUtils.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(&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(value)).getResult(); + cache.try_emplace(key, constant); + return constant; +} + static std::optional getIndexConstantValue(arith::ConstantOp constantOp) { if (!constantOp.getType().isIndex()) return std::nullopt; diff --git a/src/PIM/Common/IR/ConstantUtils.hpp b/src/PIM/Common/IR/ConstantUtils.hpp index 595738d..119ce2d 100644 --- a/src/PIM/Common/IR/ConstantUtils.hpp +++ b/src/PIM/Common/IR/ConstantUtils.hpp @@ -6,10 +6,26 @@ #include "mlir/IR/Value.h" #include "mlir/Transforms/FoldUtils.h" +#include "llvm/ADT/DenseMap.h" + #include 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, mlir::Value> cache; +}; + mlir::Block* getConstantInsertionBlock(mlir::Operation* anchorOp); mlir::Value diff --git a/src/PIM/Common/IR/StaticIntSequence.cpp b/src/PIM/Common/IR/StaticIntSequence.cpp new file mode 100644 index 0000000..a67884f --- /dev/null +++ b/src/PIM/Common/IR/StaticIntSequence.cpp @@ -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 + +#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(std::numeric_limits::max())) + return false; + int64_t scaled; + return !llvm::MulOverflow(static_cast(index), step, scaled) + && !llvm::AddOverflow(base, scaled, value); +} + +static FailureOr> 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(attr)) + return SmallVector {scalar.getInt()}; + if (auto array = dyn_cast(attr)) + return SmallVector(array.asArrayRef()); + auto elements = dyn_cast(attr); + auto type = elements ? dyn_cast(elements.getType()) + : RankedTensorType(); + if (!elements || !type || type.getRank() != 1 + || !type.getElementType().isInteger(64)) + return op->emitOpError() << "has invalid " << name << " metadata", + failure(); + SmallVector values; + values.reserve(elements.getNumElements()); + for (APInt value : elements.getValues()) + 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 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 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 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(data[run + 1]); + if (index < length) + return data[run]; + index -= length; + } + llvm_unreachable("malformed run-length encoded sequence"); +} + +std::optional 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(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(begin) + && index < static_cast(end) + ? std::optional(index) + : std::nullopt; + } + if (kind == StaticIntSequenceKind::Dense) { + ArrayRef selected = ArrayRef(data).slice(begin, length); + auto found = llvm::find(selected, value); + return found == selected.end() + ? std::nullopt + : std::optional(begin + (found - selected.begin())); + } + size_t runBegin = 0; + for (size_t run = 0; run < data.size(); run += 2) { + size_t runEnd = runBegin + static_cast(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 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 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 indices) const { + assert(!indices.empty() && "empty static integer sequence remap"); + SmallVector 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 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(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::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(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 callback) const { + for (const StaticIntSequenceSlice &slice : slices) + callback(*slice.sequence, slice.begin, slice.count); +} + +void StaticIntSequenceChain::forEachEqualRun( + llvm::function_ref callback) const { + std::optional 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 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 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 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(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 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((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((*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 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(values.size())}, builder.getI64Type()); + Value table = constants.get(type, + DenseElementsAttr::get(type, ArrayRef(values))); + Value selected = tensor::ExtractOp::create( + builder, loc, table, ValueRange {position}); + return arith::IndexCastOp::create( + builder, loc, builder.getIndexType(), selected); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Common/IR/StaticIntSequence.hpp b/src/PIM/Common/IR/StaticIntSequence.hpp new file mode 100644 index 0000000..5fa03da --- /dev/null +++ b/src/PIM/Common/IR/StaticIntSequence.hpp @@ -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 +#include +#include +#include + +namespace onnx_mlir { + +class ConstantPool; + +enum class StaticIntSequenceKind { + Uniform, + Affine, + RunLengthEncoded, + Dense +}; + +class StaticIntSequence { +public: + static StaticIntSequence fromValues(llvm::ArrayRef 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 find(int64_t value, size_t begin, size_t length) const; + StaticIntSequence slice(size_t begin, size_t count) const; + StaticIntSequence remap(llvm::ArrayRef indices) const; + StaticIntSequenceKind getKind() const { return kind; } + + bool operator==(const StaticIntSequence& other) const; + llvm::hash_code hash() const; + + void forEachEqualRun( + llvm::function_ref callback) const; + +private: + friend class StaticIntSequenceChain; + friend void setStaticIntSequenceAttr(mlir::Operation *, llvm::StringRef, + const StaticIntSequence &, size_t); + friend mlir::FailureOr + getStaticIntSequenceAttr(mlir::Operation *, llvm::StringRef, size_t); + + static StaticIntSequence runLengthEncoded( + llvm::ArrayRef runs, size_t count); + StaticIntSequenceKind kind = StaticIntSequenceKind::Dense; + size_t count = 0; + int64_t base = 0; + int64_t step = 0; + llvm::SmallVector 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 callback) const; + void forEachEqualRun( + llvm::function_ref callback) const; + StaticIntSequence canonicalize() const; + +private: + friend class StaticIntSequenceChainCursor; + llvm::SmallVector slices; + llvm::SmallVector> 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 +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 diff --git a/src/PIM/Common/IR/TensorSliceUtils.cpp b/src/PIM/Common/IR/TensorSliceUtils.cpp index dfdbc99..c8991c3 100644 --- a/src/PIM/Common/IR/TensorSliceUtils.cpp +++ b/src/PIM/Common/IR/TensorSliceUtils.cpp @@ -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 addLeadingUnitTensorDimension(OpBuilder& builder, Location loc, Value value) { auto type = dyn_cast(value.getType()); if (!type || !type.hasStaticShape()) diff --git a/src/PIM/Common/IR/TensorSliceUtils.hpp b/src/PIM/Common/IR/TensorSliceUtils.hpp index 92b5dd0..bc3088c 100644 --- a/src/PIM/Common/IR/TensorSliceUtils.hpp +++ b/src/PIM/Common/IR/TensorSliceUtils.hpp @@ -8,6 +8,12 @@ namespace onnx_mlir { +struct MixedSliceGeometry { + llvm::SmallVector offsets; + llvm::SmallVector sizes; + llvm::SmallVector 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 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 addLeadingUnitTensorDimension(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value); diff --git a/src/PIM/Dialect/Spatial/CMakeLists.txt b/src/PIM/Dialect/Spatial/CMakeLists.txt index 1c2b303..a23ed46 100644 --- a/src/PIM/Dialect/Spatial/CMakeLists.txt +++ b/src/PIM/Dialect/Spatial/CMakeLists.txt @@ -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 diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp new file mode 100644 index 0000000..f0a2fac --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.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 value; +}; + +struct SequenceCursor { + LaneSet lanes; + unsigned position = 0; +}; + +struct CanonicalAction { + SemanticKey key; + SmallVector slices; + SmallVector locals; + LaneSet instructionLanes; + LaneSet receiveLanes; + LaneSet localLanes; + DeferredExchangePlan* result = nullptr; +}; + +struct BoundaryWork { + BoundaryProgram program; + SmallVector events; +}; + +static unsigned getBoundaryIndex(SmallVectorImpl& boundaries, + DenseMap& 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 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& keys, + DenseMap>& 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 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 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& states) { + DenseMap byPosition; + SmallVector 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 events) { + if (auto eventId = std::get_if(&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(&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(&token.value); + return success(); +} + +static FailureOr> collectCanonicalActions( + ArrayRef sequence, ArrayRef tokens, + ArrayRef semantics, ArrayRef events, + const LaneSet& lanes) { + SmallVector actions; + for (unsigned semantic : sequence) { + CanonicalAction action; + action.key = semantics[semantic]; + actions.push_back(std::move(action)); + } + SmallVector states { + {lanes, 0} + }; + for (const PendingToken& token : tokens) { + SmallVector 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 actions, size_t begin, SmallVectorImpl& 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 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 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 actions, + size_t begin, + const LaneSet& lanes, + unsigned& leafIndex, + SmallVectorImpl& 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(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 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 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 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 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& classes, + SmallVectorImpl& nodes, + DenseMap, unsigned>& interned) { + SmallVector 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 sequence; +}; + +static SmallVector +buildSequenceClasses(unsigned laneCount, ArrayRef tokens) { + SmallVector classes { + {LaneSet::all(laneCount), 0} + }; + SmallVector nodes { + {0, 0} + }; + DenseMap, unsigned> interned; + for (const PendingToken& token : tokens) + addTokenToClasses(token.lanes, token.semantic, classes, nodes, interned); + + SmallVector result; + DenseMap classBySequence; + SmallVector 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 +getProducedExchanges(const BoundaryInstructionList& list) { + SmallVector result; + for (const BoundaryInstruction& instruction : list.instructions) + if (auto produced = std::get_if(&instruction)) + result.push_back(produced->exchange); + return result; +} + +static LogicalResult buildCanonicalBoundary( + BoundaryProgram& boundary, ArrayRef events, + ArrayRef tokens, ArrayRef semantics) { + unsigned laneCount = boundary.key.scheduled->cores.size(); + SmallVector 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(); + SmallVector 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 buildDeferredBoundaryPlan(DeferredTransferPlan& transfers, + const ScheduledCommunicationPlan& schedule) { + DeferredBoundaryPlan result; + SmallVector boundaries; + DenseMap 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> locals; + DenseMap> exchanges; + for (const std::unique_ptr& 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 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 tokens; + SmallVector semantics; + DenseMap> semanticsByHash; + DenseMap> localsBySemantic; + SmallPtrSet emittedLocals; + for (LocalAvailabilityFamily* local : locals[boundary.key]) { + unsigned semantic = + internSemantic(getLocalKey(*local), semantics, semanticsByHash); + localsBySemantic[semantic].push_back(local); + } + SmallVector eventSemantics; + DenseMap> 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 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(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 diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp new file mode 100644 index 0000000..773567d --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp @@ -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 slices; + LaneSet lanes; +}; +struct EmitReceiveRun { + llvm::SmallVector slices; + llvm::SmallVector entryOffsets; + LaneSet lanes; +}; +struct EmitReceiveAssemblyRun { + llvm::SmallVector slices; + llvm::SmallVector entryOffsets; + llvm::SmallVector assemblyEntries; + std::optional projectionLeaf; + LaneSet lanes; +}; +struct MaterializeLocalFamily { + llvm::SmallVector families; + LaneSet lanes; +}; +using AvailabilitySource = + std::variant; +struct AvailabilityAlternative { + LaneSet lanes; + AvailabilitySource source; +}; +struct ResolveAvailability { + DeferredExchangePlan* exchange = nullptr; + RequirementCoordinate coordinate; + mlir::Type fragmentType; + llvm::SmallVector alternatives; +}; +struct ProduceDeferredResult { + DeferredExchangePlan* exchange = nullptr; + LaneSet lanes; +}; + +struct BoundaryInstructionList; +struct LaneDispatch; +using BoundaryInstruction = std::variant>; +struct BoundaryInstructionList { + llvm::SmallVector instructions; +}; +struct LaneDispatch { + StaticIntSequence branchByLane = StaticIntSequence::uniform(0, 1); + llvm::SmallVector branches; + llvm::SmallVector producedExchanges; +}; +struct BoundaryProgram { + BoundaryKey key; + BoundaryInstructionList root; +}; + +struct DeferredBoundaryPlan { + llvm::SmallVector boundaries; + llvm::SmallVector results; +}; + +mlir::FailureOr buildDeferredBoundaryPlan(DeferredTransferPlan& transfers, + const ScheduledCommunicationPlan& schedule); + +} // namespace onnx_mlir::spatial + +namespace llvm { +template <> +struct DenseMapInfo { + static onnx_mlir::spatial::BoundaryKey getEmptyKey() { + return {DenseMapInfo::getEmptyKey(), 0}; + } + static onnx_mlir::spatial::BoundaryKey getTombstoneKey() { + return {DenseMapInfo::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 diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp new file mode 100644 index 0000000..236ccdc --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp @@ -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 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 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 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 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(payload.getType()); + auto fragmentType = dyn_cast(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 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 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 channelTable( + actionCount * laneCount, logical.channels.valueAt(0)); + SmallVector sourceTable( + actionCount * laneCount, logical.sourceCores.valueAt(0)); + SmallVector targetTable( + actionCount * laneCount, logical.targetCores.valueAt(0)); + SmallVector offsetTable( + actionCount * laneCount, logical.localOffsets.valueAt(0)); + SmallVector 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 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 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 +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 channelTable(laneCount, metadata.channels.valueAt(0)); + SmallVector sourceTable(laneCount, metadata.sourceCores.valueAt(0)); + SmallVector 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 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(block.getTerminator()); + OpBuilder::InsertionGuard guard(context.rewriter); + context.rewriter.setInsertionPoint(yield); + return emitSendRun(run, lane, laneCount, context); +} + +static FailureOr 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 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 results, DeferredExchangePlan* exchange) { + auto it = llvm::find_if(results, [&](const DeferredResultPlan& result) { return result.exchange == exchange; }); + return it == results.end() ? nullptr : &*it; +} + +static FailureOr applyAssemblyTransform(Value source, + const DeferredInsertAssemblyEntryTemplate& entry, + DeferredEmissionContext& context, + Location loc) { + auto sourceType = dyn_cast(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 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 channelTable( + entryCount * laneCount, metadata.channels.valueAt(0)); + SmallVector sourceTable( + entryCount * laneCount, metadata.sourceCores.valueAt(0)); + SmallVector targetTable( + entryCount * laneCount, metadata.targetCores.valueAt(0)); + SmallVector positions(entryCount * laneCount); + for (size_t entry = 0; entry < entryCount; ++entry) { + ArrayRef 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 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(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& 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 emitAssemblyRun(const EmitReceiveAssemblyRun& run, + Value lane, + unsigned laneCount, + ArrayRef 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 channelTable( + entryCount * laneCount, metadata.channels.valueAt(0)); + SmallVector sourceTable( + entryCount * laneCount, metadata.sourceCores.valueAt(0)); + SmallVector targetTable( + entryCount * laneCount, metadata.targetCores.valueAt(0)); + unsigned rank = assembly->resultType.getRank(); + SmallVector> geometryOffsets(rank, SmallVector(entryCount * laneCount)); + SmallVector> geometrySizes(rank, SmallVector(entryCount * laneCount, 1)); + SmallVector> geometryStrides(rank, SmallVector(entryCount * laneCount, 1)); + for (size_t entryIndex = 0; entryIndex < entryCount; ++entryIndex) { + unsigned assemblyEntry = run.assemblyEntries[entryIndex]; + if (assemblyEntry >= resultPlan->assemblyGeometry.size()) + return failure(); + ArrayRef 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 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& 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; + if constexpr (std::is_same_v) + return value.slices.front().family->requirement; + else + return value.families.front()->requirement; + }, + alternative.source); +} + +static FailureOr emitAvailabilityAlternative(const AvailabilityAlternative& alternative, + Value lane, + unsigned laneCount, + DeferredEmissionContext& context) { + if (auto receive = std::get_if(&alternative.source)) + return emitReceiveValue(*receive, lane, laneCount, context); + return materializeLocalValue( + std::get(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(&alternative.source)) { + for (const ScheduledTransferSlice& slice : receive->slices) + context.receives[slice.family->requirement] = value; + continue; + } + for (LocalAvailabilityFamily* family : + std::get(alternative.source).families) + context.receives[family->requirement] = value; + } +} + +static FailureOr> emitSelection( + Value branch, ArrayRef resultTypes, size_t branchCount, Location loc, + DeferredEmissionContext& context, + llvm::function_ref emitRegion) { + SmallVector 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 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 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(&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 branchByLane(laneCount); + SmallVector 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> emitInstructions( + const BoundaryInstructionList& list, Value lane, unsigned laneCount, + ArrayRef results, DeferredEmissionContext& context); + +static FailureOr> emitDispatch( + const LaneDispatch& dispatch, Value lane, unsigned laneCount, + ArrayRef 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 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(&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> emitInstructions( + const BoundaryInstructionList& list, Value lane, unsigned laneCount, + ArrayRef results, DeferredEmissionContext& context) { + SmallVector produced; + for (const BoundaryInstruction& instruction : list.instructions) { + if (auto send = std::get_if(&instruction)) { + if (failed(emitConditionalSendRun(*send, lane, laneCount, context))) + return failure(); + continue; + } + if (auto availability = std::get_if(&instruction)) { + if (failed(emitAvailability(*availability, lane, laneCount, context))) + return failure(); + continue; + } + if (auto assembly = std::get_if(&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(&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>(&instruction), lane, + laneCount, results, context); + if (failed(values)) + return failure(); + llvm::append_range(produced, *values); + } + return produced; +} + +static SmallVector getProducedExchanges( + const BoundaryInstructionList& list) { + SmallVector exchanges; + for (const BoundaryInstruction& instruction : list.instructions) { + if (auto result = std::get_if(&instruction)) + exchanges.push_back(result->exchange); + else if (auto dispatch = + std::get_if>(&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 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 results, + DeferredEmissionContext& context, + DeferredEraseSet& erase) { + setInsertionAtBoundary(context.rewriter, boundary.key); + unsigned laneCount = boundary.key.scheduled->cores.size(); + Value lane; + if (auto batch = dyn_cast(boundary.key.scheduled->op)) + lane = *batch.getLaneArgument(); + SmallVector 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 boundaries, + ArrayRef 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 diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.hpp new file mode 100644 index 0000000..d48e64d --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.hpp @@ -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 receives; + llvm::DenseMap assemblies; + llvm::DenseMap, mlir::Value> + projectionAssemblies; +}; + +using DeferredEraseSet = llvm::SetVector; + +mlir::LogicalResult realizeDeferredBoundaries(mlir::ArrayRef boundaries, + mlir::ArrayRef results, + DeferredEmissionContext& context, + DeferredEraseSet& erase); + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp index 4216e18..e0129a6 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp @@ -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 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> streams, - StringRef phase) { - SmallVector cursor(streams.size()); - DenseMap headSends; - DenseMap headReceives; - SmallVector readyComputes; - SmallVector 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 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 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( + family.channelIds.valueAt(*index))}; + } + if (cursor.step < stepCount) + return Event {EventKind::Compute, 0}; + return std::nullopt; +} + +static LogicalResult simulatePlanned( + Operation *anchor, ArrayRef stepCounts, + const ScheduledCommunicationPlan &plan) { + SmallVector cursors(stepCounts.size()); + DenseMap sends, receives; + SmallVector readyComputes; + SmallVector 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 &heads = - event.kind == EventKind::Send ? headSends : headReceives; - DenseMap &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 getI64Attr(Operation *op, StringRef name) { - if (auto attr = op->getAttrOfType(name)) - return attr.getInt(); - return std::nullopt; -} - -static LogicalResult getI64ArrayAttr( - Operation *op, StringRef name, - std::optional> &values) { - Attribute attr = op->getAttr(name); - if (!attr) - return success(); - if (auto array = dyn_cast(attr)) { - values.emplace(array.asArrayRef()); - return success(); - } - auto elements = dyn_cast(attr); - auto type = elements ? dyn_cast(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()) - 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 callback) { - auto scalarChannel = getI64Attr(op, "raptor.channel_id"); - std::optional> 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 getBatchTransferCount(Operation *op) { + if (auto count = op->getAttrOfType( + "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(channels)) + return array.empty() + ? FailureOr(failure()) + : FailureOr(array.size()); + if (auto elements = dyn_cast_or_null(channels); + elements && elements.getNumElements() > 0) + return elements.getNumElements(); + return op->emitOpError("has invalid legacy compact transfer metadata"), + failure(); +} - std::optional> 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 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(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("raptor.exchange_id"); + if (!exchange || exchange.getInt() != channels->valueAt(0)) { + op->emitOpError("has inconsistent scalar exchange metadata"); return failure(); + } + } + return RealizedOperation {op, isa(op), + *channels, *parents, *counts, *sources, *targets}; +} + +struct CoreTransferSequences { + DenseMap sends; + DenseMap receives; +}; + +struct ExpectedFamily { + ExternalTransferFamily *family = nullptr; + int64_t firstChannel = 0; + int64_t endChannel = 0; +}; + +static void appendByCore(DenseMap &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 &expected, + const DenseMap &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 stepCounts, - ArrayRef transfers) { - if (stepCounts.size() != streamCount) - return anchor->emitError("communication plan stream count does not match step counts"); - - SmallVector> streams(streamCount); - SmallVector>> 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 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> operationsByExchange; - struct ParentExchange { - std::optional expectedTransfers; - DenseSet channels; - }; - DenseMap parentExchanges; - DenseMap streamByCore; - SmallVector cores; +LogicalResult verifyRealizedCommunicationDeadlockFree( + func::FuncOp funcOp, const ScheduledCommunicationPlan &plan) { + SmallVector families; + DenseMap 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(index)) + return funcOp.emitOpError( + "planned communication family has non-consecutive channels"); + familyIndex[family] = families.size(); + families.push_back({family, first, first + static_cast(count)}); + } + llvm::sort(families, [](const ExpectedFamily &lhs, + const ExpectedFamily &rhs) { + return lhs.firstChannel < rhs.firstChannel; + }); + familyIndex.clear(); + std::map 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(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> actualChannels; bool invalid = false; funcOp.walk([&](Operation *op) { - if (!isa(op)) + if (invalid || !isa(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> streams(cores.size()); - funcOp.walk([&](Operation *op) { - if (!isa(op)) return; - if (failed(forEachRealizedLogicalTransfer( - op, [&](const RealizedLogicalTransfer &transfer) { - unsigned stream = streamByCore.lookup( - isa(op) ? transfer.sourceCore - : transfer.targetCore); - streams[stream].push_back( - {isa(op) ? EventKind::Send - : EventKind::Receive, - static_cast(transfer.channelId)}); - return success(); - }))) - invalid = true; + } + Type payloadType = realized->send + ? cast(op).getInput().getType() + : cast(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(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(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(*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(entry.second[0].op) - == isa(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(item.op); - }) - << ", receives=" - << llvm::count_if(entry.second, [](const LogicalOperation &item) { - return isa(item.op); - }) - << ")"; - } - const LogicalOperation &first = entry.second[0]; - const LogicalOperation &second = entry.second[1]; - const LogicalOperation &sendRecord = - isa(first.op) ? first : second; - const LogicalOperation &receiveRecord = - isa(first.op) ? first : second; - auto send = cast(sendRecord.op); - auto receive = cast(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 diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp index 062af6b..a377255 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp @@ -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 stepCounts, - mlir::ArrayRef 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 diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp new file mode 100644 index 0000000..e94fd91 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp @@ -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 +#include +#include + +#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 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 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 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 offsets; + llvm::SmallVector sizes; + llvm::SmallVector 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 entries; +}; + +struct DeferredProgramTemplate { + SpatDeferredCommunicationOp deferred; + mlir::Value scheduledLane; + mlir::Value yieldedValue; + llvm::SmallVector leaves; + llvm::SmallVector residualOps; + std::optional 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 blocks; + llvm::SmallVector stepAnchors; + llvm::SmallVector cores; + llvm::SmallVector stepSourceIds; + llvm::SmallVector resultOffsets; + llvm::SmallVector resultCounts; + llvm::SmallVector produced; + llvm::SmallVector streamIds; + + bool isBatch() const { return mlir::isa(op); } +}; + +struct DeferredExchangePlan; + +struct RequirementFamily { + DeferredExchangePlan* exchange = nullptr; + RequirementCoordinate coordinate; + LaneSet targetLanes; + ProducedValue* producer = nullptr; + mlir::Type publicationFragmentType; + std::optional graphLanes; + std::optional 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 requirements; + llvm::SmallVector local; + llvm::SmallVector external; + unsigned externalTransferCount = 0; +}; + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp index d53b1c9..35f8b97 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp @@ -1,3491 +1,200 @@ -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/Dominance.h" -#include "mlir/IR/IRMapping.h" -#include "mlir/IR/Matchers.h" -#include "mlir/IR/PatternMatch.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/MapVector.h" -#include "llvm/ADT/SetVector.h" -#include "llvm/ADT/SmallPtrSet.h" -#include -#include - -#include "DeferredCommunicationDeadlock.hpp" -#include "DeferredProjectionAnalysis.hpp" #include "DeferredCommunicationRealization.hpp" -#include "ScheduledComputePlan.hpp" -#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp" -#include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +#include "DeferredBoundaryPlanning.hpp" +#include "DeferredBoundaryRealization.hpp" +#include "DeferredCommunicationDeadlock.hpp" +#include "DeferredCommunicationScheduling.hpp" +#include "DeferredTransferPlanning.hpp" + +#include "mlir/IR/Dominance.h" namespace onnx_mlir::spatial { - using namespace mlir; - namespace { -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; - Value payload; - Value published; -}; - -struct ScheduledInfo { - Operation* op = nullptr; - SmallVector blocks; - SmallVector stepAnchors; - SmallVector cores; - SmallVector stepSourceIds; - SmallVector resultOffsets; - SmallVector resultCounts; - SmallVector produced; - SmallVector streamIds; - - bool isBatch() const { return isa(op); } -}; - -struct FragmentRequirement { - unsigned leafIndex = 0; - unsigned selectedPosition = 0; - unsigned sourceOperandIndex = 0; - std::optional physicalSlot; - std::optional graphLane; - ProducedValue *producer = nullptr; - Type publicationFragmentType; - std::optional targetScheduledLane; -}; - -struct FragmentTransfer { - unsigned requirementIndex = 0; - unsigned sourceStream = 0; - unsigned targetStream = 0; - int64_t sourceCore = -1; - int64_t targetCore = -1; - uint64_t channelId = 0; - unsigned globalOrder = 0; - unsigned sourceInsertionStep = 0; - unsigned targetInsertionStep = 0; - bool local = false; -}; - -struct DeferredSpecialization { - std::optional targetScheduledLane; - SpecializedDeferredProgram program; - SmallVector requirements; - SmallVector transfers; -}; - -struct NormalizedDeferredExchange { - uint64_t exchangeId = 0; - SpatDeferredCommunicationOp deferred; - ScheduledInfo *target = nullptr; - unsigned consumerStep = 0; - SmallVector specializations; - unsigned externalTransferCount = 0; -}; - -struct FragmentTransferRef { - NormalizedDeferredExchange *exchange = nullptr; - DeferredSpecialization *specialization = nullptr; - FragmentTransfer *transfer = nullptr; -}; - -struct LogicalTransferMetadata { - uint64_t channelId = 0; - uint64_t parentExchangeId = 0; - unsigned parentTransferCount = 0; - int64_t sourceCore = -1; - int64_t targetCore = -1; -}; - -static LogicalTransferMetadata getLogicalTransferMetadata( - const FragmentTransferRef &ref) { - return {ref.transfer->channelId, ref.exchange->exchangeId, - ref.exchange->externalTransferCount, ref.transfer->sourceCore, - ref.transfer->targetCore}; -} - -struct SendEmissionSignature { - ScheduledInfo *scheduled = nullptr; - Value payload; - Type fragmentType; - bool hasGraphLane = false; - bool sourceIsBatch = false; -}; - -static SendEmissionSignature getSendEmissionSignature( - const FragmentTransferRef &ref) { - const FragmentRequirement &requirement = - ref.specialization->requirements[ref.transfer->requirementIndex]; - return {requirement.producer->scheduled, - requirement.producer->payload, - requirement.publicationFragmentType, - requirement.graphLane.has_value(), - requirement.producer->scheduled->isBatch()}; -} - -static bool hasRepresentableSendOffset(const FragmentTransferRef &ref) { - const FragmentRequirement &requirement = - ref.specialization->requirements[ref.transfer->requirementIndex]; - if (!requirement.graphLane) - return true; - int64_t offset = - *requirement.graphLane - requirement.producer->laneStart; - return offset >= 0 && offset < requirement.producer->laneCount; -} - -static bool haveCompatibleSendEmission( - const FragmentTransferRef &lhs, const FragmentTransferRef &rhs, - unsigned lhsInsertionStep, unsigned rhsInsertionStep) { - SendEmissionSignature left = getSendEmissionSignature(lhs); - SendEmissionSignature right = getSendEmissionSignature(rhs); - return lhsInsertionStep == rhsInsertionStep - && left.scheduled == right.scheduled && left.payload == right.payload - && left.fragmentType == right.fragmentType - && left.hasGraphLane == right.hasGraphLane - && left.sourceIsBatch == right.sourceIsBatch - && hasRepresentableSendOffset(lhs) && hasRepresentableSendOffset(rhs); -} - -struct RealizationPlan { - SmallVector scheduled; - SmallVector> producedStorage; - DenseMap> producedByGraph; - GraphBatchPublicationCache publicationCache; - SmallVector exchanges; - SmallVector external; - DenseMap transferByChannel; - SmallVector stepCounts; -}; - -enum class BoundaryActionKind { Send, Receive }; - -struct BoundaryAction { - BoundaryActionKind kind = BoundaryActionKind::Send; - FragmentTransferRef ref; - unsigned globalOrder = 0; -}; - -struct ScheduledBoundaryPlan { - ScheduledInfo *scheduled = nullptr; - unsigned insertionStep = 0; - SmallVector, 0> actionsByLane; - SmallVector targetExchanges; -}; - -struct BatchBoundaryClass { - SmallVector lanes; -}; - -using ReceiveCache = DenseMap; -using AssemblyCache = DenseMap; -using IndexConstantCache = DenseMap; -using DenseConstantCache = DenseMap; -using DeferredEraseSet = llvm::SetVector; - -static FailureOr> getI64Array(Operation* op, StringRef name) { - auto attr = op->getAttrOfType(name); - if (!attr) - return op->emitOpError() << "phase 2 requires '" << name << "' metadata"; - return SmallVector(attr.asArrayRef().begin(), attr.asArrayRef().end()); -} - -static FailureOr> getLaneTable(Operation* op, StringRef name, size_t expected) { - if (auto array = op->getAttrOfType(name)) { - if (array.size() != static_cast(expected)) - return op->emitOpError() << "phase 2 metadata '" << name << "' has the wrong size"; - return SmallVector(array.asArrayRef().begin(), array.asArrayRef().end()); - } - auto elements = op->getAttrOfType(name); - if (!elements || elements.getNumElements() != static_cast(expected)) - return op->emitOpError() << "phase 2 requires a correctly-sized '" << name << "' lane table"; - SmallVector values; - for (APInt value : elements.getValues()) - 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 getStepIndex(ScheduledInfo& info, Block* block) { - auto it = llvm::find(info.blocks, block); - if (it == info.blocks.end()) - return failure(); - return static_cast(std::distance(info.blocks.begin(), it)); -} - -static FailureOr -getScalarStepResult(SpatScheduledCompute scheduled, Block& block, unsigned resultIndex, unsigned resultCount) { - auto yield = dyn_cast(block.getTerminator()); - if (!yield || resultIndex >= resultCount || yield.getOutputs().size() < resultCount) - return scheduled.emitOpError("phase 2 cannot recover a scalar scheduled step result"); - return yield.getOutputs()[yield.getOutputs().size() - resultCount + resultIndex]; -} - -struct BatchStepResult { - Value payload; - tensor::ParallelInsertSliceOp publication; -}; - -static FailureOr -getBatchStepResult(SpatScheduledComputeBatch scheduled, Block& block, unsigned globalResultIndex) { - auto inParallel = dyn_cast(block.getTerminator()); - if (!inParallel) - return scheduled.emitOpError("phase 2 cannot recover a batched scheduled step result"); - unsigned resultBase = 1 + scheduled.getWeights().size() + scheduled.getInputs().size(); - for (Operation& op : inParallel.getRegion().front()) { - auto insert = dyn_cast(op); - if (!insert) - continue; - auto destination = dyn_cast(insert.getDest()); - if (destination && destination.getOwner() == &block && destination.getArgNumber() == resultBase + globalResultIndex) - return BatchStepResult {insert.getSource(), insert}; - } - return scheduled.emitOpError("phase 2 cannot find the batched result insertion for a scheduled step"); -} - -static LogicalResult collectScheduled(func::FuncOp funcOp, RealizationPlan& plan) { - unsigned nextStream = 0; - for (Operation& op : funcOp.getOps()) { - if (!isa(op)) - continue; - ScheduledInfo info; - info.op = &op; - Region& body = isa(op) ? cast(op).getBody() - : cast(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 step metadata does not match the region block count"); - - if (auto scalar = dyn_cast(op)) { - auto core = scalar->getAttrOfType(kCoreIdAttrName); - if (!core) - return scalar.emitOpError("phase 2 requires scalar coreId metadata"); - info.cores.push_back(core.getInt()); - } - else { - auto cores = op.getAttrOfType(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(); - - for (ScheduledInfo& info : plan.scheduled) { - SmallVector laneStarts; - SmallVector laneCounts; - if (info.isBatch()) { - auto starts = getLaneTable(info.op, "scheduled.source_lane_starts", info.blocks.size() * info.cores.size()); - auto counts = getLaneTable(info.op, "scheduled.source_lane_counts", info.blocks.size() * info.cores.size()); - if (failed(starts) || failed(counts)) - return failure(); - laneStarts = std::move(*starts); - laneCounts = std::move(*counts); - } - else { - auto starts = getI64Array(info.op, "scheduled.source_lane_starts"); - auto counts = 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(info.resultCounts[step]); ++result) { - unsigned globalResult = info.resultOffsets[step] + result; - if (globalResult >= info.op->getNumResults()) - return info.op->emitOpError("phase 2 scheduled result metadata is out of range"); - if (!info.isBatch()) { - auto payload = getScalarStepResult( - cast(info.op), *info.blocks[step], result, info.resultCounts[step]); - if (failed(payload)) - return failure(); - int64_t publishedSlotCount = 1; - if (laneCounts[step] > 1) { - auto publishedType = dyn_cast(info.op->getResult(globalResult).getType()); - if (!publishedType || !publishedType.hasStaticShape() || publishedType.getRank() == 0 - || publishedType.getDimSize(0) != laneCounts[step]) - return info.op->emitOpError("phase 2 scalar scheduled publication geometry is inconsistent"); - publishedSlotCount = laneCounts[step]; - } - auto produced = std::make_unique(ProducedValue { - &info, - step, - result, - info.stepSourceIds[step], - info.cores.front(), - laneStarts[step], - std::max(laneCounts[step], 1), - 0, - 0, - publishedSlotCount, - *payload, - info.op->getResult(globalResult), - }); - info.produced.push_back(produced.get()); - plan.producedByGraph[produced->graphId].push_back(produced.get()); - plan.producedStorage.push_back(std::move(produced)); - continue; - } - auto stepResult = getBatchStepResult(cast(info.op), *info.blocks[step], globalResult); - if (failed(stepResult)) - return failure(); - auto publishedType = dyn_cast(info.op->getResult(globalResult).getType()); - if (!publishedType || !publishedType.hasStaticShape() || publishedType.getRank() == 0) - return info.op->emitOpError("phase 2 batched scheduled publication must have a static leading slot dimension"); - for (unsigned lane = 0; lane < info.cores.size(); ++lane) { - size_t laneIndex = step * info.cores.size() + lane; - StaticIndexEnvironment environment; - environment.bindings[info.blocks[step]->getArgument(0)] = lane; - auto slotStart = evaluateDeferredIndex( - stepResult->publication.getMixedOffsets().front(), environment); - auto slotCount = evaluateDeferredIndex( - stepResult->publication.getMixedSizes().front(), environment); - auto slotStride = evaluateDeferredIndex( - stepResult->publication.getMixedStrides().front(), environment); - if (failed(slotStart) || failed(slotCount) || failed(slotStride) - || *slotStart < 0 || *slotCount <= 0 || *slotStride != 1 - || *slotStart > publishedType.getDimSize(0) - *slotCount) - return info.op->emitOpError("phase 2 batched scheduled publication geometry is invalid"); - auto produced = std::make_unique(ProducedValue { - &info, - step, - result, - info.stepSourceIds[step], - info.cores[lane], - laneStarts[laneIndex], - std::max(laneCounts[laneIndex], 1), - lane, - *slotStart, - *slotCount, - stepResult->payload, - info.op->getResult(globalResult), - }); - 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 -findProducer(RealizationPlan& plan, Operation *diagnosticOwner, int64_t graphId, unsigned resultIndex, - std::optional graphLane) { - ProducedValue* match = nullptr; - for (ProducedValue* produced : plan.producedByGraph.lookup(graphId)) { - if (produced->resultIndex != resultIndex) - continue; - if (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(); - return match; -} - -static LogicalResult retargetBlueprintToScheduledPublications( - RealizationPlan &plan, SpatBlueprintOp blueprint) { - 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 oldOperands {blueprint.getInput()}; - llvm::append_range(oldOperands, blueprint.getFragments()); - SmallVector publications; - SmallVector scheduledOperandIndices; - SmallVector scheduledSourceSlots; - for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) { - if (operandIndex < 0 || operandIndex >= static_cast(oldOperands.size())) - return blueprint.emitOpError("phase 2 Blueprint operand index is out of range"); - auto sourceResult = dyn_cast(oldOperands[operandIndex]); - auto graphBatch = sourceResult - ? dyn_cast(sourceResult.getOwner()) - : SpatGraphComputeBatch(); - if (!graphBatch) { - Value existing = oldOperands[operandIndex]; - auto it = llvm::find(publications, existing); - if (it == publications.end()) { - publications.push_back(existing); - it = std::prev(publications.end()); - } - scheduledOperandIndices.push_back(std::distance(publications.begin(), it)); - scheduledSourceSlots.push_back((*sourceSlots)[fragmentIndex]); - continue; - } - auto graphId = graphBatch->getAttrOfType("scheduled.graph_id"); - auto publication = getGraphBatchPublicationMap( - graphBatch, sourceResult.getResultNumber(), plan.publicationCache); - int64_t physicalSlot = (*sourceSlots)[fragmentIndex]; - if (!graphId || failed(publication) || physicalSlot < 0 || - physicalSlot >= static_cast((*publication)->physicalSlotToGraphLane.size())) - return blueprint.emitOpError("phase 2 cannot resolve Blueprint physical fragment ownership"); - int64_t graphLane = (*publication)->physicalSlotToGraphLane[physicalSlot]; - auto producer = findProducer(plan, blueprint.getOperation(), graphId.getInt(), - sourceResult.getResultNumber(), graphLane); - if (failed(producer)) - return failure(); - auto it = llvm::find(publications, (*producer)->published); - if (it == publications.end()) { - publications.push_back((*producer)->published); - it = std::prev(publications.end()); - } - scheduledOperandIndices.push_back(std::distance(publications.begin(), it)); - int64_t localGraphLane = graphLane - (*producer)->laneStart; - if (localGraphLane < 0 || localGraphLane >= (*producer)->laneCount - || localGraphLane >= (*producer)->publishedSlotCount) - return blueprint.emitOpError("phase 2 Blueprint graph lane is outside its scheduled publication"); - int64_t scheduledSourceSlot = (*producer)->publishedSlotStart + localGraphLane; - auto publishedType = dyn_cast((*producer)->published.getType()); - if (!publishedType || !publishedType.hasStaticShape() || publishedType.getRank() == 0 - || scheduledSourceSlot < 0 || scheduledSourceSlot >= publishedType.getDimSize(0)) - return blueprint.emitOpError("phase 2 Blueprint scheduled publication slot is out of range"); - scheduledSourceSlots.push_back(scheduledSourceSlot); - Operation *scheduled = (*producer)->published.getDefiningOp(); - if (blueprint->getBlock() == scheduled->getBlock() && - blueprint->isBeforeInBlock(scheduled)) - blueprint->moveAfter(scheduled); - } - blueprint->setOperands(publications); - blueprint->setAttr("fragmentOperandIndices", - OpBuilder(blueprint).getDenseI64ArrayAttr(scheduledOperandIndices)); - blueprint->setAttr("fragmentSourceSlots", - OpBuilder(blueprint).getDenseI64ArrayAttr(scheduledSourceSlots)); - return success(); -} - -static LogicalResult configureFragmentTransfer(NormalizedDeferredExchange &exchange, - DeferredSpecialization &specialization, - unsigned requirementIndex, - FragmentTransfer &transfer) { - FragmentRequirement &requirement = specialization.requirements[requirementIndex]; - ProducedValue *producer = requirement.producer; - transfer.requirementIndex = requirementIndex; - transfer.sourceStream = producer->scheduled->streamIds[producer->scheduledLane]; - transfer.sourceCore = producer->core; - unsigned targetLane = specialization.targetScheduledLane.value_or(0); - transfer.targetStream = exchange.target->streamIds[targetLane]; - transfer.targetCore = exchange.target->cores[targetLane]; - transfer.local = transfer.sourceStream == transfer.targetStream && - producer->step < exchange.consumerStep; - return success(); -} - -static LogicalResult buildDeferredSpecialization(RealizationPlan &plan, - NormalizedDeferredExchange &exchange, - std::optional targetScheduledLane) { - auto program = analyzeDeferredProgram(exchange.deferred, targetScheduledLane); - if (failed(program)) - return exchange.deferred.emitOpError("phase 2 deferred projection analysis failed"); - DeferredSpecialization specialization; - specialization.targetScheduledLane = targetScheduledLane; - specialization.program = std::move(*program); - for (auto [leafIndex, leaf] : llvm::enumerate(specialization.program.leaves)) { - Value sourceValue = exchange.deferred.getSources()[leaf.sourceOperandIndex]; - auto sourceResult = dyn_cast(sourceValue); - if (!sourceResult) - return exchange.deferred.emitOpError("phase 2 requires graph-result deferred sources"); - Operation *owner = sourceResult.getOwner(); - auto graphId = owner->getAttrOfType("scheduled.graph_id"); - if (!graphId) - return exchange.deferred.emitOpError("phase 2 cannot identify graph producer"); - unsigned requirementLeafIndex = leafIndex; - unsigned sourceOperandIndex = leaf.sourceOperandIndex; - auto append = [&](std::optional physicalSlot, std::optional graphLane, - Type fragmentType, unsigned selectedPosition) -> LogicalResult { - auto producer = findProducer(plan, exchange.deferred.getOperation(), graphId.getInt(), - sourceResult.getResultNumber(), graphLane); - if (failed(producer)) return failure(); - FragmentRequirement requirement {requirementLeafIndex, selectedPosition, - sourceOperandIndex, physicalSlot, graphLane, *producer, - fragmentType, targetScheduledLane}; - specialization.requirements.push_back(std::move(requirement)); - FragmentTransfer transfer; - if (failed(configureFragmentTransfer(exchange, specialization, - specialization.requirements.size() - 1, transfer))) - return failure(); - specialization.transfers.push_back(std::move(transfer)); - return success(); - }; - if (leaf.kind == DeferredLeafKind::ScalarSource) { - if (!isa(owner)) - return exchange.deferred.emitOpError("scalar deferred leaf requires graph_compute"); - if (failed(append(std::nullopt, std::nullopt, sourceResult.getType(), 0))) return failure(); - continue; - } - auto graphBatch = dyn_cast(owner); - if (!graphBatch) - return exchange.deferred.emitOpError("projection deferred leaf requires graph_compute_batch"); - auto publication = getGraphBatchPublicationMap(graphBatch, sourceResult.getResultNumber(), - plan.publicationCache); - if (failed(publication)) - return exchange.deferred.emitOpError("phase 2 graph-batch publication analysis failed"); - for (auto [position, physicalSlot] : llvm::enumerate(leaf.physicalSlots)) { - if (physicalSlot < 0 || physicalSlot >= static_cast((*publication)->physicalSlotToGraphLane.size())) - return exchange.deferred.emitOpError("projection physical slot is outside publication map"); - int64_t graphLane = (*publication)->physicalSlotToGraphLane[physicalSlot]; - if (failed(append(physicalSlot, graphLane, (*publication)->publicationFragmentType, position))) return failure(); - } - } - if (specialization.program.insertAssembly) - for (DeferredInsertAssemblyEntry &entry : - specialization.program.insertAssembly->entries) { - std::optional requirementIndex; - for (auto [index, requirement] : - llvm::enumerate(specialization.requirements)) - if (requirement.leafIndex == entry.leafIndex - && requirement.selectedPosition == 0) { - if (requirementIndex) - return exchange.deferred.emitOpError( - "phase 2 insert assembly requirement is ambiguous"); - requirementIndex = index; - } - if (!requirementIndex) - return exchange.deferred.emitOpError( - "phase 2 insert assembly requirement is missing"); - entry.requirementIndex = *requirementIndex; - } - exchange.specializations.push_back(std::move(specialization)); - return success(); -} - -static LogicalResult normalizeDeferred(func::FuncOp funcOp, RealizationPlan& plan) { - DenseMap infoByOp; - for (ScheduledInfo& info : plan.scheduled) - infoByOp[info.op] = &info; - - SmallVector deferredOps; - funcOp.walk([&](SpatDeferredCommunicationOp deferred) { deferredOps.push_back(deferred); }); - uint64_t nextExchangeId = 0; - for (SpatDeferredCommunicationOp deferred : deferredOps) { - Operation* targetOp = deferred->getParentOfType(); - if (!targetOp) - targetOp = deferred->getParentOfType(); - ScheduledInfo* target = infoByOp.lookup(targetOp); - Block* scheduledBlock = target ? getScheduledBlock(deferred, targetOp) : nullptr; - auto step = target && scheduledBlock ? getStepIndex(*target, scheduledBlock) : FailureOr(failure()); - if (!target || failed(step)) - return deferred.emitOpError("phase 2 cannot locate the deferred consumer scheduled step"); - NormalizedDeferredExchange exchange; - exchange.exchangeId = nextExchangeId++; - exchange.deferred = deferred; - exchange.target = target; - exchange.consumerStep = *step; - if (!target->isBatch()) { - if (failed(buildDeferredSpecialization(plan, exchange, std::nullopt))) - return deferred.emitOpError("phase 2 failed to build scalar deferred specialization"); - } else { - for (unsigned lane = 0; lane < target->cores.size(); ++lane) - if (failed(buildDeferredSpecialization(plan, exchange, lane))) - return deferred.emitOpError() - << "phase 2 failed to build deferred specialization for scheduled lane " << lane; - } - plan.exchanges.push_back(std::move(exchange)); - } - return success(); -} - -struct GroupDependency { - unsigned stream = 0; - unsigned requiredCompletedStep = 0; -}; - -struct StaticEmissionCost { - unsigned logicalTransferCount = 0; - unsigned sendRunCount = 0; - unsigned transfersCoveredBySendRuns = 0; - unsigned largestSendRun = 0; - unsigned assemblyCoveredReceiveCount = 0; - unsigned estimatedScalarReceiveCount = 0; - unsigned estimatedBatchBehaviorClassCount = 0; -}; - -struct ParentTransferGroup { - uint64_t parentExchangeId = 0; - SmallVector originalTransferIndices; - SmallVector orderedTransferIndices; - SmallVector dependencies; - unsigned unsatisfiedDependencies = 0; - bool ready = false; - bool scheduled = false; - StaticEmissionCost cost; - std::tuple deterministicPriority; - std::optional firstSendSignature; -}; - -struct SchedulerStreamState { - unsigned completedStep = 0; - SmallVector pendingIncomingAtConsumerStep; - SmallVector> groupsUnlockedAtStep; - bool queued = false; -}; - -static bool haveSameSendEmissionSignature( - const SendEmissionSignature &lhs, const SendEmissionSignature &rhs) { - return lhs.scheduled == rhs.scheduled && lhs.payload == rhs.payload - && lhs.fragmentType == rhs.fragmentType - && lhs.hasGraphLane == rhs.hasGraphLane - && lhs.sourceIsBatch == rhs.sourceIsBatch; -} - -static size_t hashSendEmissionSignature( - const SendEmissionSignature &signature) { - return static_cast(llvm::hash_combine( - signature.scheduled, signature.payload.getAsOpaquePointer(), - signature.fragmentType.getAsOpaquePointer(), signature.hasGraphLane, - signature.sourceIsBatch)); -} - -static bool hasBetterStaticPriority(const ParentTransferGroup &lhs, - const ParentTransferGroup &rhs) { - auto addedOps = [](const ParentTransferGroup &group) { - return group.cost.sendRunCount + group.cost.estimatedScalarReceiveCount; - }; - auto lhsScore = std::tuple( - addedOps(lhs), lhs.cost.estimatedBatchBehaviorClassCount, - std::numeric_limits::max() - - lhs.cost.assemblyCoveredReceiveCount, - std::numeric_limits::max() - - lhs.cost.transfersCoveredBySendRuns, - std::numeric_limits::max() - lhs.cost.largestSendRun, - lhs.cost.sendRunCount, lhs.deterministicPriority); - auto rhsScore = std::tuple( - addedOps(rhs), rhs.cost.estimatedBatchBehaviorClassCount, - std::numeric_limits::max() - - rhs.cost.assemblyCoveredReceiveCount, - std::numeric_limits::max() - - rhs.cost.transfersCoveredBySendRuns, - std::numeric_limits::max() - rhs.cost.largestSendRun, - rhs.cost.sendRunCount, rhs.deterministicPriority); - return lhsScore < rhsScore; -} - -static void buildGroupTransferOrder(RealizationPlan &plan, - ParentTransferGroup &group) { - SmallVector semanticOrder; - llvm::SmallDenseSet orderedIndices; - DenseMap indexByChannel; - for (unsigned index : group.originalTransferIndices) - indexByChannel[plan.external[index].exchangeId] = index; - FragmentTransferRef groupRef = plan.transferByChannel.lookup( - plan.external[group.originalTransferIndices.front()].exchangeId); - for (DeferredSpecialization &specialization : - groupRef.exchange->specializations) { - if (!specialization.program.insertAssembly) - continue; - for (const DeferredInsertAssemblyEntry &entry : - specialization.program.insertAssembly->entries) { - if (entry.requirementIndex >= specialization.transfers.size()) - continue; - FragmentTransfer &transfer = - specialization.transfers[entry.requirementIndex]; - if (transfer.local) - continue; - auto it = indexByChannel.find(transfer.channelId); - if (it != indexByChannel.end() - && orderedIndices.insert(it->second).second) - semanticOrder.push_back(it->second); - } - } - for (unsigned index : group.originalTransferIndices) - if (orderedIndices.insert(index).second) - semanticOrder.push_back(index); - - SmallVector> classes; - DenseMap> classesBySignature; - for (unsigned index : semanticOrder) { - FragmentTransferRef ref = - plan.transferByChannel.lookup(plan.external[index].exchangeId); - if (!hasRepresentableSendOffset(ref)) { - classes.push_back({index}); - continue; - } - SendEmissionSignature signature = getSendEmissionSignature(ref); - size_t signatureHash = hashSendEmissionSignature(signature); - SmallVector &candidates = classesBySignature[signatureHash]; - auto candidateIt = llvm::find_if(candidates, [&](unsigned classIndex) { - ArrayRef candidate = classes[classIndex]; - FragmentTransferRef first = plan.transferByChannel.lookup( - plan.external[candidate.front()].exchangeId); - return haveSameSendEmissionSignature( - getSendEmissionSignature(first), signature); - }); - if (candidateIt == candidates.end()) { - candidates.push_back(classes.size()); - classes.push_back({index}); - } else { - classes[*candidateIt].push_back(index); - } - } - llvm::stable_sort(classes, [&](ArrayRef lhs, - ArrayRef rhs) { - if (lhs.size() != rhs.size()) - return lhs.size() > rhs.size(); - const PlannedCommunicationTransfer &left = plan.external[lhs.front()]; - const PlannedCommunicationTransfer &right = plan.external[rhs.front()]; - return std::tie(left.sourceStream, left.targetStream, left.exchangeId) - < std::tie(right.sourceStream, right.targetStream, right.exchangeId); - }); - group.cost.logicalTransferCount = group.originalTransferIndices.size(); - group.cost.sendRunCount = classes.size(); - for (ArrayRef compactClass : classes) { - llvm::append_range(group.orderedTransferIndices, compactClass); - group.cost.largestSendRun = - std::max(group.cost.largestSendRun, compactClass.size()); - if (compactClass.size() >= 2) - group.cost.transfersCoveredBySendRuns += compactClass.size(); - } - if (!group.orderedTransferIndices.empty()) { - FragmentTransferRef first = plan.transferByChannel.lookup( - plan.external[group.orderedTransferIndices.front()].exchangeId); - group.firstSendSignature = getSendEmissionSignature(first); - } -} - -static void buildGroupStaticCost(RealizationPlan &plan, - ParentTransferGroup &group) { - if (group.originalTransferIndices.empty()) - return; - FragmentTransferRef ref = plan.transferByChannel.lookup( - plan.external[group.originalTransferIndices.front()].exchangeId); - unsigned assemblyCoverage = 0; - for (DeferredSpecialization &specialization : ref.exchange->specializations) - if (specialization.program.insertAssembly) { - unsigned externalEntries = llvm::count_if( - specialization.transfers, - [](const FragmentTransfer &transfer) { return !transfer.local; }); - assemblyCoverage = std::max(assemblyCoverage, externalEntries); - } - group.cost.assemblyCoveredReceiveCount = assemblyCoverage; - group.cost.estimatedScalarReceiveCount = - group.cost.logicalTransferCount - - std::min(group.cost.logicalTransferCount, assemblyCoverage); - group.cost.estimatedBatchBehaviorClassCount = - ref.exchange->target->isBatch() ? ref.exchange->specializations.size() : 1; -} - -static SmallVector buildParentTransferGroups( - RealizationPlan &plan) { - llvm::MapVector groupByParent; - SmallVector groups; - for (auto [index, transfer] : llvm::enumerate(plan.external)) { - auto inserted = groupByParent.insert({transfer.parentExchangeId, - groups.size()}); - if (inserted.second) { - ParentTransferGroup group; - group.parentExchangeId = transfer.parentExchangeId; - groups.push_back(std::move(group)); - } - groups[inserted.first->second].originalTransferIndices.push_back(index); - } - - for (ParentTransferGroup &group : groups) { - DenseMap thresholdByStream; - unsigned minConsumer = std::numeric_limits::max(); - unsigned minSource = std::numeric_limits::max(); - unsigned minTarget = std::numeric_limits::max(); - for (unsigned index : group.originalTransferIndices) { - const PlannedCommunicationTransfer &transfer = plan.external[index]; - thresholdByStream[transfer.sourceStream] = std::max( - thresholdByStream.lookup(transfer.sourceStream), - transfer.producerStep + 1); - minConsumer = std::min(minConsumer, transfer.consumerStep); - minSource = std::min(minSource, transfer.sourceStream); - minTarget = std::min(minTarget, transfer.targetStream); - } - FragmentTransferRef ref = plan.transferByChannel.lookup( - plan.external[group.originalTransferIndices.front()].exchangeId); - for (DeferredSpecialization &specialization : ref.exchange->specializations) - for (FragmentTransfer &transfer : specialization.transfers) - if (transfer.local) { - const FragmentRequirement &requirement = - specialization.requirements[transfer.requirementIndex]; - thresholdByStream[transfer.targetStream] = std::max( - thresholdByStream.lookup(transfer.targetStream), - requirement.producer->step + 1); - } - SmallVector dependencyStreams; - dependencyStreams.reserve(thresholdByStream.size()); - for (auto entry : thresholdByStream) - dependencyStreams.push_back(entry.first); - llvm::sort(dependencyStreams); - for (unsigned stream : dependencyStreams) - group.dependencies.push_back( - {stream, thresholdByStream.lookup(stream)}); - group.unsatisfiedDependencies = llvm::count_if( - group.dependencies, [](const GroupDependency &dependency) { - return dependency.requiredCompletedStep != 0; - }); - group.deterministicPriority = - {minConsumer, minSource, minTarget, group.parentExchangeId}; - buildGroupTransferOrder(plan, group); - buildGroupStaticCost(plan, group); - } - return groups; -} - -static LogicalResult scheduleCommunication(func::FuncOp funcOp, - RealizationPlan &plan) { - SmallVector groups = - buildParentTransferGroups(plan); - SmallVector streams(plan.stepCounts.size()); - for (auto [stream, state] : llvm::enumerate(streams)) { - state.pendingIncomingAtConsumerStep.resize(plan.stepCounts[stream]); - state.groupsUnlockedAtStep.resize(plan.stepCounts[stream] + 1); - } - for (const PlannedCommunicationTransfer &transfer : plan.external) - ++streams[transfer.targetStream] - .pendingIncomingAtConsumerStep[transfer.consumerStep]; - for (auto [groupIndex, group] : llvm::enumerate(groups)) - for (const GroupDependency &dependency : group.dependencies) - streams[dependency.stream] - .groupsUnlockedAtStep[dependency.requiredCompletedStep] - .push_back(groupIndex); - - std::function heapCompare = - [&](unsigned lhs, unsigned rhs) { - return hasBetterStaticPriority(groups[rhs], groups[lhs]); - }; - std::priority_queue, - std::function> - readyByStaticPriority(heapCompare); - using ReadyHeap = std::priority_queue< - unsigned, std::vector, - std::function>; - DenseMap> - readyByFirstSendSignature; - auto addReady = [&](unsigned groupIndex) { - ParentTransferGroup &group = groups[groupIndex]; - if (group.ready || group.scheduled) - return; - group.ready = true; - readyByStaticPriority.push(groupIndex); - if (group.firstSendSignature) { - std::unique_ptr &bucket = readyByFirstSendSignature[ - hashSendEmissionSignature(*group.firstSendSignature)]; - if (!bucket) - bucket = std::make_unique(heapCompare); - bucket->push(groupIndex); - } - }; - for (auto [groupIndex, group] : llvm::enumerate(groups)) - if (group.unsatisfiedDependencies == 0) - addReady(groupIndex); - - std::queue advanceable; - auto enqueueIfAdvanceable = [&](unsigned stream) { - SchedulerStreamState &state = streams[stream]; - if (!state.queued && state.completedStep < plan.stepCounts[stream] - && state.pendingIncomingAtConsumerStep[state.completedStep] == 0) { - state.queued = true; - advanceable.push(stream); - } - }; - for (unsigned stream = 0; stream < streams.size(); ++stream) - enqueueIfAdvanceable(stream); - auto advanceStreams = [&] { - bool advanced = false; - while (!advanceable.empty()) { - unsigned stream = advanceable.front(); - advanceable.pop(); - SchedulerStreamState &state = streams[stream]; - state.queued = false; - if (state.completedStep == plan.stepCounts[stream] - || state.pendingIncomingAtConsumerStep[state.completedStep] != 0) - continue; - ++state.completedStep; - advanced = true; - for (unsigned groupIndex : - state.groupsUnlockedAtStep[state.completedStep]) { - ParentTransferGroup &group = groups[groupIndex]; - assert(group.unsatisfiedDependencies != 0); - if (--group.unsatisfiedDependencies == 0) - addReady(groupIndex); - } - enqueueIfAdvanceable(stream); - } - return advanced; - }; - - SmallVector scheduled; - scheduled.reserve(plan.external.size()); - unsigned scheduledGroups = 0; - while (scheduledGroups != groups.size()) { - bool progressed = advanceStreams(); - std::optional chosen; - unsigned chosenTailExtension = 0; - if (!scheduled.empty()) { - const PlannedCommunicationTransfer &tail = scheduled.back(); - FragmentTransferRef tailRef = - plan.transferByChannel.lookup(tail.exchangeId); - size_t signatureHash = - hashSendEmissionSignature(getSendEmissionSignature(tailRef)); - auto bucketIt = readyByFirstSendSignature.find(signatureHash); - SmallVector inspected; - ReadyHeap *bucket = bucketIt == readyByFirstSendSignature.end() - ? nullptr : bucketIt->second.get(); - while (bucket && !bucket->empty() && inspected.size() != 32) { - unsigned groupIndex = bucket->top(); - bucket->pop(); - ParentTransferGroup &group = groups[groupIndex]; - if (!group.ready || group.scheduled) - continue; - inspected.push_back(groupIndex); - if (!group.firstSendSignature - || !haveSameSendEmissionSignature( - getSendEmissionSignature(tailRef), - *group.firstSendSignature)) - continue; - unsigned extension = 0; - for (unsigned transferIndex : group.orderedTransferIndices) { - const PlannedCommunicationTransfer &transfer = - plan.external[transferIndex]; - FragmentTransferRef ref = - plan.transferByChannel.lookup(transfer.exchangeId); - if (!haveCompatibleSendEmission( - tailRef, ref, tail.sourceInsertionStep, - streams[transfer.sourceStream].completedStep)) - break; - ++extension; - } - if (extension == 0) - continue; - if (!chosen || extension > chosenTailExtension - || (extension == chosenTailExtension - && hasBetterStaticPriority(group, groups[*chosen]))) { - chosen = groupIndex; - chosenTailExtension = extension; - } - } - for (unsigned groupIndex : inspected) - if (!chosen || groupIndex != *chosen) - bucket->push(groupIndex); - } - while (!chosen && !readyByStaticPriority.empty()) { - unsigned candidate = readyByStaticPriority.top(); - readyByStaticPriority.pop(); - if (groups[candidate].ready && !groups[candidate].scheduled) - chosen = candidate; - } - if (chosen) { - ParentTransferGroup &group = groups[*chosen]; - group.ready = false; - group.scheduled = true; - ++scheduledGroups; - for (unsigned index : group.orderedTransferIndices) { - PlannedCommunicationTransfer transfer = plan.external[index]; - transfer.sourceInsertionStep = - streams[transfer.sourceStream].completedStep; - transfer.targetInsertionStep = - streams[transfer.targetStream].completedStep; - scheduled.push_back(transfer); - unsigned &pending = streams[transfer.targetStream] - .pendingIncomingAtConsumerStep[transfer.consumerStep]; - assert(pending != 0); - --pending; - if (pending == 0 - && streams[transfer.targetStream].completedStep - == transfer.consumerStep) - enqueueIfAdvanceable(transfer.targetStream); - } - progressed = true; - } - if (progressed) - continue; - - InFlightDiagnostic diagnostic = funcOp.emitOpError( - "global communication scheduler made no progress before IR mutation"); - unsigned reported = 0; - for (const ParentTransferGroup &group : groups) { - if (group.scheduled || reported == 8) - continue; - diagnostic << (reported++ == 0 ? "; blocked " : ", ") - << "parent " << group.parentExchangeId; - auto dependency = llvm::find_if( - group.dependencies, [&](const GroupDependency &candidate) { - return streams[candidate.stream].completedStep - < candidate.requiredCompletedStep; - }); - if (dependency != group.dependencies.end()) - diagnostic << " stream " << dependency->stream << " requires " - << dependency->requiredCompletedStep << " has " - << streams[dependency->stream].completedStep; - const PlannedCommunicationTransfer &first = - plan.external[group.originalTransferIndices.front()]; - diagnostic << " consumer stream " << first.targetStream << " step " - << first.consumerStep; - } - return failure(); - } - plan.external = std::move(scheduled); - return success(); -} - -static FailureOr getSpecializationTargetInsertionStep( - NormalizedDeferredExchange &exchange, DeferredSpecialization &specialization) { - std::optional targetStep; - for (FragmentTransfer &transfer : specialization.transfers) { - if (transfer.local) - continue; - if (targetStep && *targetStep != transfer.targetInsertionStep) - return exchange.deferred.emitOpError() - << "exchange " << exchange.exchangeId << " scheduled lane " - << specialization.targetScheduledLane.value_or(0) << " channel " - << transfer.channelId << " has target insertion step " - << transfer.targetInsertionStep << " instead of " << *targetStep, - failure(); - targetStep = transfer.targetInsertionStep; - } - return targetStep.value_or(exchange.consumerStep); -} - -static LogicalResult verifyExchangeTargetInsertionSteps(RealizationPlan &plan) { - for (NormalizedDeferredExchange &exchange : plan.exchanges) { - std::optional sharedStep; - for (DeferredSpecialization &specialization : exchange.specializations) { - auto step = getSpecializationTargetInsertionStep(exchange, specialization); - if (failed(step)) - return failure(); - if (exchange.target->isBatch() && sharedStep && *sharedStep != *step) - return exchange.deferred.emitOpError() - << "exchange " << exchange.exchangeId << " scheduled lane " - << specialization.targetScheduledLane.value_or(0) - << " has target insertion step " << *step << " instead of shared step " - << *sharedStep; - sharedStep = *step; - } - } - return success(); -} - -static LogicalResult buildAndVerifyPlan(func::FuncOp funcOp, RealizationPlan& plan) { - uint64_t nextChannel = 0; - for (NormalizedDeferredExchange &exchange : plan.exchanges) - for (DeferredSpecialization &specialization : exchange.specializations) - for (FragmentTransfer &transfer : specialization.transfers) { - if (transfer.local) - continue; - transfer.channelId = nextChannel++; - ++exchange.externalTransferCount; - if (!plan.transferByChannel.try_emplace( - transfer.channelId, FragmentTransferRef {&exchange, &specialization, &transfer}).second) - return exchange.deferred.emitOpError("phase 2 assigned a duplicate communication channel"); - const FragmentRequirement &requirement = specialization.requirements[transfer.requirementIndex]; - plan.external.push_back({transfer.channelId, exchange.exchangeId, - transfer.sourceStream, transfer.targetStream, - requirement.producer->step, exchange.consumerStep}); - } - if (failed(scheduleCommunication(funcOp, plan))) - return failure(); - for (auto [order, transfer] : llvm::enumerate(plan.external)) { - FragmentTransferRef ref = plan.transferByChannel.lookup(transfer.exchangeId); - if (!ref.exchange || !ref.specialization || !ref.transfer) - return funcOp.emitOpError("phase 2 cannot resolve a planned communication channel"); - ref.transfer->globalOrder = order; - ref.transfer->sourceInsertionStep = transfer.sourceInsertionStep; - ref.transfer->targetInsertionStep = transfer.targetInsertionStep; - } - if (failed(verifyExchangeTargetInsertionSteps(plan))) - return failure(); - return verifyPlannedCommunicationDeadlockFree(funcOp, plan.stepCounts.size(), plan.stepCounts, plan.external); -} - -static FailureOr> -buildScheduledBoundaryPlans(RealizationPlan &plan) { - SmallVector boundaries; - DenseMap, unsigned> boundaryIndex; - SmallVector> actionCounts; - auto getBoundary = [&](ScheduledInfo *scheduled, - unsigned insertionStep) -> ScheduledBoundaryPlan & { - auto key = std::make_pair(scheduled, insertionStep); - if (auto it = boundaryIndex.find(key); it != boundaryIndex.end()) - return boundaries[it->second]; - ScheduledBoundaryPlan boundary; - boundary.scheduled = scheduled; - boundary.insertionStep = insertionStep; - boundary.actionsByLane.resize(scheduled->cores.size()); - boundaryIndex[key] = boundaries.size(); - boundaries.push_back(std::move(boundary)); - actionCounts.emplace_back(scheduled->cores.size()); - return boundaries.back(); - }; - - for (const PlannedCommunicationTransfer &transfer : plan.external) { - FragmentTransferRef ref = - plan.transferByChannel.lookup(transfer.exchangeId); - if (!ref.exchange || !ref.specialization || !ref.transfer) - return failure(); - FragmentRequirement &requirement = - ref.specialization->requirements[ref.transfer->requirementIndex]; - getBoundary(requirement.producer->scheduled, - ref.transfer->sourceInsertionStep); - getBoundary(ref.exchange->target, ref.transfer->targetInsertionStep); - } - for (NormalizedDeferredExchange &exchange : plan.exchanges) { - unsigned targetStep = exchange.consumerStep; - if (exchange.externalTransferCount != 0) { - auto step = getSpecializationTargetInsertionStep( - exchange, exchange.specializations.front()); - if (failed(step)) - return failure(); - targetStep = *step; - } - getBoundary(exchange.target, targetStep); - } - - for (const PlannedCommunicationTransfer &transfer : plan.external) { - FragmentTransferRef ref = - plan.transferByChannel.lookup(transfer.exchangeId); - FragmentRequirement &requirement = - ref.specialization->requirements[ref.transfer->requirementIndex]; - unsigned sourceBoundary = boundaryIndex.lookup( - {requirement.producer->scheduled, - ref.transfer->sourceInsertionStep}); - unsigned targetBoundary = boundaryIndex.lookup( - {ref.exchange->target, ref.transfer->targetInsertionStep}); - ++actionCounts[sourceBoundary][requirement.producer->scheduledLane]; - ++actionCounts[targetBoundary][ - ref.specialization->targetScheduledLane.value_or(0)]; - } - for (auto [boundary, counts] : llvm::zip(boundaries, actionCounts)) - for (auto [actions, count] : llvm::zip(boundary.actionsByLane, counts)) - actions.reserve(count); - - for (auto [globalOrder, transfer] : llvm::enumerate(plan.external)) { - FragmentTransferRef ref = - plan.transferByChannel.lookup(transfer.exchangeId); - if (!ref.exchange || !ref.specialization || !ref.transfer) - return failure(); - FragmentRequirement &requirement = - ref.specialization->requirements[ref.transfer->requirementIndex]; - getBoundary(requirement.producer->scheduled, - ref.transfer->sourceInsertionStep) - .actionsByLane[requirement.producer->scheduledLane] - .push_back({BoundaryActionKind::Send, ref, - static_cast(globalOrder)}); - unsigned targetLane = - ref.specialization->targetScheduledLane.value_or(0); - getBoundary(ref.exchange->target, ref.transfer->targetInsertionStep) - .actionsByLane[targetLane] - .push_back({BoundaryActionKind::Receive, ref, - static_cast(globalOrder)}); - } - - for (NormalizedDeferredExchange &exchange : plan.exchanges) { - unsigned targetStep = exchange.consumerStep; - if (exchange.externalTransferCount != 0) { - auto step = getSpecializationTargetInsertionStep( - exchange, exchange.specializations.front()); - if (failed(step)) - return failure(); - targetStep = *step; - } - getBoundary(exchange.target, targetStep).targetExchanges.push_back(&exchange); - } - - DenseMap scheduledOrder; - for (auto [index, scheduled] : llvm::enumerate(plan.scheduled)) - scheduledOrder[&scheduled] = index; - llvm::stable_sort(boundaries, [&](const ScheduledBoundaryPlan &lhs, - const ScheduledBoundaryPlan &rhs) { - return std::tie(scheduledOrder[lhs.scheduled], lhs.insertionStep) - < std::tie(scheduledOrder[rhs.scheduled], rhs.insertionStep); - }); - for (ScheduledBoundaryPlan &boundary : boundaries) { - for (ArrayRef actions : boundary.actionsByLane) { - if (actions.size() < 2) - continue; - for (auto pair : llvm::zip(actions.drop_back(), actions.drop_front())) - if (std::get<0>(pair).globalOrder >= std::get<1>(pair).globalOrder) - return boundary.scheduled->op->emitOpError( - "phase 2 boundary actions are not in committed order"), - failure(); - } - } - return boundaries; -} - -static LogicalResult validateScalarLinearization(ScheduledInfo& info) { +static LogicalResult validateScalarLinearization(ScheduledInfo &info) { auto scheduled = cast(info.op); for (unsigned index = 1; index < info.blocks.size(); ++index) { - auto previousYield = dyn_cast(info.blocks[index - 1]->getTerminator()); - if (!previousYield || previousYield.getOutputs().size() != info.blocks[index]->getNumArguments()) - return scheduled.emitOpError("phase 2 cannot linearize malformed scalar scheduled blocks"); - for (auto [argument, value] : llvm::zip(info.blocks[index]->getArguments(), previousYield.getOutputs())) { + auto previous = dyn_cast( + info.blocks[index - 1]->getTerminator()); + if (!previous + || previous.getOutputs().size() + != info.blocks[index]->getNumArguments()) + return scheduled.emitOpError( + "phase 2 cannot linearize malformed scalar scheduled blocks"); + for (auto [argument, value] : llvm::zip( + info.blocks[index]->getArguments(), previous.getOutputs())) { if (argument.getType() == value.getType()) continue; - for (Operation* user : argument.getUsers()) + for (Operation *user : argument.getUsers()) if (!isa(user)) return scheduled.emitOpError( - "phase 2 cannot linearize a mismatched carried value used by scheduled computation"); + "phase 2 cannot linearize a live mismatched carried value"); } } return success(); } -static LogicalResult linearizeScalar(ScheduledInfo& info, IRRewriter& rewriter) { +static LogicalResult linearizeScalar(ScheduledInfo &info, + IRRewriter &rewriter) { auto scheduled = cast(info.op); if (failed(validateScalarLinearization(info))) return failure(); - Block* first = info.blocks.front(); + Block *first = info.blocks.front(); SmallVector> incoming(info.blocks.size()); for (unsigned index = 1; index < info.blocks.size(); ++index) { - auto previousYield = cast(info.blocks[index - 1]->getTerminator()); - incoming[index].assign(previousYield.getOutputs().begin(), previousYield.getOutputs().end()); + auto previous = cast( + info.blocks[index - 1]->getTerminator()); + incoming[index].assign(previous.getOutputs().begin(), + previous.getOutputs().end()); } - IRMapping carriedValues; + IRMapping carried; for (unsigned index = 1; index < info.blocks.size(); ++index) - for (auto [argument, value] : llvm::zip(info.blocks[index]->getArguments(), incoming[index])) { - Value resolved = carriedValues.lookupOrDefault(value); + for (auto [argument, value] : llvm::zip( + info.blocks[index]->getArguments(), incoming[index])) { + Value resolved = carried.lookupOrDefault(value); if (argument.getType() == resolved.getType()) { - carriedValues.map(argument, resolved); + carried.map(argument, resolved); argument.replaceAllUsesWith(resolved); } } - - for (unsigned index = 1; index < info.blocks.size(); ++index) { - Block* block = info.blocks[index]; - for (Operation& op : llvm::make_early_inc_range(block->without_terminator())) + for (unsigned index = 1; index < info.blocks.size(); ++index) + for (Operation &op : llvm::make_early_inc_range( + info.blocks[index]->without_terminator())) op.moveBefore(first->getTerminator()); - } - for (Block* block : info.blocks) + for (Block *block : info.blocks) cast(block->getTerminator()).erase(); - SmallVector outputs(scheduled.getNumResults()); - for (ProducedValue* produced : info.produced) { - unsigned globalResult = info.resultOffsets[produced->step] + produced->resultIndex; - outputs[globalResult] = produced->payload; + for (ProducedValue *produced : info.produced) { + unsigned result = info.resultOffsets[produced->step] + + produced->resultIndex; + outputs[result] = produced->payload; } - if (llvm::any_of(outputs, [](Value value) { return !value; })) - return scheduled.emitOpError("phase 2 cannot recover every scheduled scalar result"); + if (llvm::any_of(outputs, [](Value output) { return !output; })) + return scheduled.emitOpError( + "phase 2 cannot recover every scheduled scalar result"); rewriter.setInsertionPointToEnd(first); SpatYieldOp::create(rewriter, scheduled.getLoc(), outputs); scheduled->setAttr("scheduled.realized", rewriter.getBoolAttr(true)); for (unsigned index = 1; index < info.blocks.size(); ++index) { - for (auto [argumentIndex, argument] : llvm::enumerate(info.blocks[index]->getArguments())) - if (!argument.use_empty()) { - InFlightDiagnostic diagnostic = - scheduled.emitOpError("phase 2 scalar linearization left a live carried block argument"); - diagnostic << " at block " << index << ", argument " << argumentIndex << ", first user " - << (*argument.getUsers().begin())->getName(); - return failure(); - } + if (llvm::any_of(info.blocks[index]->getArguments(), + [](BlockArgument argument) { + return !argument.use_empty(); + })) + return scheduled.emitOpError( + "phase 2 scalar linearization left a live block argument"); info.blocks[index]->erase(); } info.blocks.assign(1, first); return success(); } -static LogicalResult linearizeBatch(ScheduledInfo& info, IRRewriter& rewriter) { +static LogicalResult linearizeBatch(ScheduledInfo &info, + IRRewriter &rewriter) { auto scheduled = cast(info.op); - Block* first = info.blocks.front(); + Block *first = info.blocks.front(); SmallVector terminators; - for (Block* block : info.blocks) { - auto inParallel = dyn_cast(block->getTerminator()); - if (!inParallel) - return scheduled.emitOpError("phase 2 cannot linearize a batch block without spat.in_parallel"); - terminators.push_back(inParallel); + for (Block *block : info.blocks) { + auto parallel = dyn_cast(block->getTerminator()); + if (!parallel) + return scheduled.emitOpError( + "phase 2 cannot linearize a batch block without spat.in_parallel"); + terminators.push_back(parallel); } for (unsigned index = 1; index < info.blocks.size(); ++index) - for (auto [argument, firstArgument] : llvm::zip(info.blocks[index]->getArguments(), first->getArguments())) { + for (auto [argument, firstArgument] : llvm::zip( + info.blocks[index]->getArguments(), first->getArguments())) { if (argument.getType() != firstArgument.getType()) - return scheduled.emitOpError("phase 2 cannot linearize incompatible batch block arguments"); + return scheduled.emitOpError( + "phase 2 cannot linearize incompatible batch block arguments"); argument.replaceAllUsesWith(firstArgument); } - for (unsigned index = 1; index < info.blocks.size(); ++index) - for (Operation& op : llvm::make_early_inc_range(info.blocks[index]->without_terminator())) + for (Operation &op : llvm::make_early_inc_range( + info.blocks[index]->without_terminator())) op.moveBefore(first->getTerminator()); rewriter.setInsertionPoint(first->getTerminator()); auto combined = SpatInParallelOp::create(rewriter, scheduled.getLoc()); - Block& combinedBlock = combined.getRegion().front(); - for (SpatInParallelOp terminator : terminators) - for (Operation& op : llvm::make_early_inc_range(terminator.getRegion().front())) + Block &combinedBlock = combined.getRegion().front(); + for (SpatInParallelOp parallel : terminators) + for (Operation &op : llvm::make_early_inc_range( + parallel.getRegion().front())) op.moveBefore(&combinedBlock, combinedBlock.end()); - for (SpatInParallelOp terminator : terminators) - terminator.erase(); + for (SpatInParallelOp parallel : terminators) + parallel.erase(); scheduled->setAttr("scheduled.realized", rewriter.getBoolAttr(true)); for (unsigned index = 1; index < info.blocks.size(); ++index) { - for (auto [argumentIndex, argument] : llvm::enumerate(info.blocks[index]->getArguments())) - if (!argument.use_empty()) { - InFlightDiagnostic diagnostic = - scheduled.emitOpError("phase 2 batch linearization left a live carried block argument"); - diagnostic << " at block " << index << ", argument " << argumentIndex << ", first user " - << (*argument.getUsers().begin())->getName(); - return failure(); - } + if (llvm::any_of(info.blocks[index]->getArguments(), + [](BlockArgument argument) { + return !argument.use_empty(); + })) + return scheduled.emitOpError( + "phase 2 batch linearization left a live block argument"); info.blocks[index]->erase(); } info.blocks.assign(1, first); return success(); } -static void -setTransferAttrs(Operation* op, uint64_t exchangeId, uint64_t channelId, int64_t sourceCore, int64_t targetCore) { - OpBuilder builder(op); - op->setAttr("raptor.exchange_id", builder.getI64IntegerAttr(exchangeId)); - op->setAttr("raptor.channel_id", builder.getI64IntegerAttr(channelId)); - op->setAttr("raptor.source_core", builder.getI64IntegerAttr(sourceCore)); - op->setAttr("raptor.target_core", builder.getI64IntegerAttr(targetCore)); -} - -static void replaceDeferred(SpatDeferredCommunicationOp deferred, - Value replacement, - DeferredEraseSet &erase) { - deferred.getOutput().replaceAllUsesWith(replacement); - erase.insert(deferred); -} - -static Value getOrCreateCachedIndexConstant( - IndexConstantCache &cache, IRRewriter &rewriter, Operation *anchor, - int64_t value) { - if (Value constant = cache.lookup(value)) - return constant; - Value constant = createConstantAtHostBlockStart( - rewriter, anchor, rewriter.getIndexAttr(value)); - cache[value] = constant; - return constant; -} - -static Value getOrCreateCachedTypedConstant( - DenseConstantCache &cache, IRRewriter &rewriter, Operation *anchor, - TypedAttr value) { - if (Value constant = cache.lookup(value)) - return constant; - Value constant = createConstantAtHostBlockStart( - rewriter, anchor, value); - cache[value] = constant; - return constant; -} - -static void setInsertionAtBoundary(IRRewriter& rewriter, ScheduledInfo& scheduled, unsigned step) { - if (step < scheduled.stepAnchors.size() && scheduled.stepAnchors[step]->getBlock()) { - rewriter.setInsertionPoint(scheduled.stepAnchors[step]); - return; - } - rewriter.setInsertionPoint(scheduled.blocks.front()->getTerminator()); -} - -static FailureOr materializeProducedFragment(const FragmentRequirement &requirement, - IRRewriter &rewriter, Location loc) { - Value payload = requirement.producer->payload; - if (!requirement.graphLane) { - if (payload.getType() != requirement.publicationFragmentType) - return failure(); - return payload; - } - if (payload.getType() == requirement.publicationFragmentType) - return payload; - auto payloadType = dyn_cast(payload.getType()); - auto fragmentType = dyn_cast(requirement.publicationFragmentType); - if (!payloadType || !fragmentType || payloadType.getRank() != fragmentType.getRank() + 1 || - payloadType.getDimSize(0) != requirement.producer->laneCount || - !llvm::equal(payloadType.getShape().drop_front(), fragmentType.getShape())) - return failure(); - int64_t localOffset = *requirement.graphLane - requirement.producer->laneStart; - if (localOffset < 0 || localOffset >= requirement.producer->laneCount) - return failure(); - SmallVector offsets(payloadType.getRank(), rewriter.getIndexAttr(0)); - SmallVector sizes, strides(payloadType.getRank(), rewriter.getIndexAttr(1)); - offsets[0] = rewriter.getIndexAttr(localOffset); - sizes.push_back(rewriter.getIndexAttr(1)); - for (int64_t dim : fragmentType.getShape()) sizes.push_back(rewriter.getIndexAttr(dim)); - SmallVector unitShape {1}; llvm::append_range(unitShape, fragmentType.getShape()); - auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType()); - Value unit = tensor::ExtractSliceOp::create(rewriter, loc, unitType, payload, offsets, sizes, strides); - return removeLeadingUnitTensorDimension(rewriter, loc, unit, fragmentType); -} - -static FailureOr materializeProducedFragment( - const FragmentRequirement &requirement, Value localOffset, - IRRewriter &rewriter, Location loc) { - Value payload = requirement.producer->payload; - if (payload.getType() == requirement.publicationFragmentType) - return payload; - auto payloadType = dyn_cast(payload.getType()); - auto fragmentType = dyn_cast(requirement.publicationFragmentType); - if (!requirement.graphLane || !payloadType || !fragmentType - || payloadType.getRank() != fragmentType.getRank() + 1 - || payloadType.getDimSize(0) != requirement.producer->laneCount - || !llvm::equal(payloadType.getShape().drop_front(), fragmentType.getShape())) - return failure(); - SmallVector offsets(payloadType.getRank(), rewriter.getIndexAttr(0)); - SmallVector sizes, strides(payloadType.getRank(), rewriter.getIndexAttr(1)); - offsets[0] = localOffset; - sizes.push_back(rewriter.getIndexAttr(1)); - for (int64_t dim : fragmentType.getShape()) - sizes.push_back(rewriter.getIndexAttr(dim)); - SmallVector unitShape {1}; - llvm::append_range(unitShape, fragmentType.getShape()); - auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType()); - Value unit = tensor::ExtractSliceOp::create( - rewriter, loc, unitType, payload, offsets, sizes, strides); - return removeLeadingUnitTensorDimension(rewriter, loc, unit, fragmentType); -} - -static FailureOr applyInnerFragmentGeometry(Value fragment, const DeferredProjectionLeaf &leaf, - IRRewriter &rewriter, Location loc) { - if (leaf.kind != DeferredLeafKind::GraphBatchProjection) - return fragment; - auto type = dyn_cast(fragment.getType()); - if (!type || leaf.innerGeometry.offsets.size() != static_cast(type.getRank())) return failure(); - bool identity = true; - for (unsigned i = 0; i < type.getRank(); ++i) - identity &= leaf.innerGeometry.offsets[i] == 0 && leaf.innerGeometry.sizes[i] == type.getDimSize(i) && - leaf.innerGeometry.strides[i] == 1; - if (identity) return fragment; - if (leaf.reconstructedType.getRank() != type.getRank() + 1) return failure(); - SmallVector shape(leaf.reconstructedType.getShape().drop_front()); - auto innerType = RankedTensorType::get(shape, type.getElementType()); - SmallVector offsets, sizes, strides; - for (unsigned i = 0; i < type.getRank(); ++i) { - offsets.push_back(rewriter.getIndexAttr(leaf.innerGeometry.offsets[i])); - sizes.push_back(rewriter.getIndexAttr(leaf.innerGeometry.sizes[i])); - strides.push_back(rewriter.getIndexAttr(leaf.innerGeometry.strides[i])); - } - return tensor::ExtractSliceOp::create(rewriter, loc, innerType, fragment, offsets, sizes, strides).getResult(); -} - -static FailureOr reconstructProjectionLeaf(const DeferredProjectionLeaf &leaf, unsigned leafIndex, - const DeferredSpecialization &specialization, - ArrayRef available, IRRewriter &rewriter, Location loc) { - SmallVector selected(leaf.kind == DeferredLeafKind::ScalarSource ? 1 : leaf.physicalSlots.size()); - for (auto [index, requirement] : llvm::enumerate(specialization.requirements)) - if (requirement.leafIndex == leafIndex && requirement.selectedPosition < selected.size()) { - if (selected[requirement.selectedPosition]) return failure(); - selected[requirement.selectedPosition] = available[index]; - } - if (llvm::any_of(selected, [](Value value) { return !value; })) return failure(); - if (leaf.kind == DeferredLeafKind::ScalarSource) return selected.front(); - SmallVector shaped; - shaped.reserve(selected.size()); - for (Value fragment : selected) { - auto fragmentResult = applyInnerFragmentGeometry(fragment, leaf, rewriter, loc); - if (failed(fragmentResult)) - return failure(); - shaped.push_back(*fragmentResult); - } - if (shaped.size() == 1 && shaped.front().getType() == leaf.reconstructedType) - return shaped.front(); - SmallVector expanded; - expanded.reserve(shaped.size()); - for (Value fragment : shaped) { - auto unit = addLeadingUnitTensorDimension(rewriter, loc, fragment); - if (failed(unit)) - return failure(); - expanded.push_back(*unit); - } - if (expanded.size() == 1 - && expanded.front().getType() == leaf.reconstructedType) - return expanded.front(); - Value result = tensor::ConcatOp::create( - rewriter, loc, leaf.reconstructedType, 0, expanded).getResult(); - return result.getType() == leaf.reconstructedType ? FailureOr(result) - : FailureOr(failure()); -} - -static FailureOr materializeDeferredSpecializationResult(DeferredSpecialization &specialization, - ArrayRef available, - IndexConstantCache &indexConstants, - DenseConstantCache &typedConstants, - IRRewriter &rewriter, Location loc) { - IRMapping mapping; - for (auto [value, staticValue] : specialization.program.staticValues) { - if (mapping.contains(value)) - continue; - Type type = value.getType(); - if (!type.isIndex() && !isa(type)) - return failure(); - Value constant = type.isIndex() - ? getOrCreateCachedIndexConstant( - indexConstants, rewriter, specialization.program.deferred, - staticValue) - : getOrCreateCachedTypedConstant( - typedConstants, rewriter, specialization.program.deferred, - rewriter.getIntegerAttr(type, staticValue)); - mapping.map(value, constant); - } - for (auto [index, leaf] : llvm::enumerate(specialization.program.leaves)) { - auto reconstructed = reconstructProjectionLeaf(leaf, index, specialization, available, rewriter, loc); - if (failed(reconstructed)) return failure(); - mapping.map(leaf.replacementRoot, *reconstructed); - } - if (specialization.targetScheduledLane && !mapping.contains(specialization.program.scheduledLane)) { - mapping.map(specialization.program.scheduledLane, - getOrCreateCachedIndexConstant( - indexConstants, rewriter, - specialization.program.deferred, - *specialization.targetScheduledLane)); - } - for (Operation *op : specialization.program.residualOps) { - Operation *copy = rewriter.clone(*op, mapping); - for (auto pair : llvm::zip(op->getResults(), copy->getResults())) mapping.map(std::get<0>(pair), std::get<1>(pair)); - } - Value result = mapping.lookupOrDefault(specialization.program.yieldedValue); - return result && result.getType() == specialization.program.deferred.getOutput().getType() - ? FailureOr(result) : FailureOr(failure()); -} - -static LogicalResult emitFragmentSendHere( - NormalizedDeferredExchange &exchange, - DeferredSpecialization &specialization, FragmentTransfer &transfer, - IndexConstantCache &indexConstants, IRRewriter &rewriter) { - Location loc = exchange.deferred.getLoc(); - auto payload = materializeProducedFragment( - specialization.requirements[transfer.requirementIndex], rewriter, loc); - if (failed(payload)) return failure(); - auto send = SpatChannelSendOp::create( - rewriter, loc, - getOrCreateCachedIndexConstant(indexConstants, rewriter, - exchange.deferred, transfer.channelId), - getOrCreateCachedIndexConstant(indexConstants, rewriter, - exchange.deferred, transfer.sourceCore), - getOrCreateCachedIndexConstant(indexConstants, rewriter, - exchange.deferred, transfer.targetCore), - *payload); - setTransferAttrs(send, transfer.channelId, transfer.channelId, - transfer.sourceCore, transfer.targetCore); - send->setAttr("raptor.parent_exchange_id", - rewriter.getI64IntegerAttr(exchange.exchangeId)); - send->setAttr("raptor.parent_transfer_count", - rewriter.getI64IntegerAttr(exchange.externalTransferCount)); - return success(); -} - -static FailureOr emitFragmentReceiveHere( - const FragmentTransferRef &ref, ReceiveCache &receiveCache, - IndexConstantCache &indexConstants, IRRewriter &rewriter) { - FragmentTransfer &transfer = *ref.transfer; - FragmentRequirement &requirement = - ref.specialization->requirements[transfer.requirementIndex]; - Location loc = ref.exchange->deferred.getLoc(); - auto receive = SpatChannelReceiveOp::create( - rewriter, loc, requirement.publicationFragmentType, - getOrCreateCachedIndexConstant(indexConstants, rewriter, - ref.exchange->deferred, - transfer.channelId), - getOrCreateCachedIndexConstant(indexConstants, rewriter, - ref.exchange->deferred, - transfer.sourceCore), - getOrCreateCachedIndexConstant(indexConstants, rewriter, - ref.exchange->deferred, - transfer.targetCore)); - setTransferAttrs(receive, transfer.channelId, transfer.channelId, - transfer.sourceCore, transfer.targetCore); - receive->setAttr("raptor.parent_exchange_id", - rewriter.getI64IntegerAttr(ref.exchange->exchangeId)); - receive->setAttr("raptor.parent_transfer_count", - rewriter.getI64IntegerAttr( - ref.exchange->externalTransferCount)); - receiveCache[&transfer] = receive.getOutput(); - return receive.getOutput(); -} - -static FailureOr receiveFragment(NormalizedDeferredExchange &exchange, - DeferredSpecialization &specialization, - FragmentTransfer &transfer, - ReceiveCache &receiveCache, - IRRewriter &rewriter) { - if (transfer.local) - return materializeProducedFragment(specialization.requirements[transfer.requirementIndex], rewriter, - exchange.deferred.getLoc()); - if (Value received = receiveCache.lookup(&transfer)) - return received; - return exchange.deferred.emitOpError( - "phase 2 boundary did not materialize a planned receive"), - failure(); -} - -static FailureOr> tryRealizeAnalyzedInsertAssembly( - NormalizedDeferredExchange &exchange, - DeferredSpecialization &specialization, IRRewriter &rewriter); - -static FailureOr realizeSpecialization(NormalizedDeferredExchange &exchange, - DeferredSpecialization &specialization, - ReceiveCache &receiveCache, - AssemblyCache &assemblyCache, - IndexConstantCache &indexConstants, - DenseConstantCache &typedConstants, - IRRewriter &rewriter) { - if (Value assembled = assemblyCache.lookup(&specialization)) - return assembled; - SmallVector available(specialization.requirements.size()); - for (FragmentTransfer &transfer : specialization.transfers) - if (transfer.local) { - auto value = receiveFragment( - exchange, specialization, transfer, receiveCache, rewriter); - if (failed(value)) return failure(); - available[transfer.requirementIndex] = *value; - } - SmallVector ordered; - for (FragmentTransfer &transfer : specialization.transfers) - if (!transfer.local) - ordered.push_back(&transfer); - llvm::stable_sort(ordered, [](FragmentTransfer *lhs, FragmentTransfer *rhs) { - return lhs->globalOrder < rhs->globalOrder; - }); - for (FragmentTransfer *transfer : ordered) { - auto value = receiveFragment( - exchange, specialization, *transfer, receiveCache, rewriter); - if (failed(value)) return failure(); - available[transfer->requirementIndex] = *value; - } - return materializeDeferredSpecializationResult( - specialization, available, indexConstants, typedConstants, rewriter, - exchange.deferred.getLoc()); -} - -static Value buildLaneLookup(ArrayRef values, Value lane, - Operation *anchor, IRRewriter &rewriter, - Location loc) { - Value table = createI64LookupTableConstant(rewriter, anchor, values); - Value selected = tensor::ExtractOp::create( - rewriter, loc, table, ValueRange {lane}); - return arith::IndexCastOp::create( - rewriter, loc, rewriter.getIndexType(), selected); -} - -static Value buildLaneLookupCached( - ArrayRef values, Value lane, Operation *anchor, - DenseConstantCache &denseConstants, IRRewriter &rewriter, Location loc) { - auto type = RankedTensorType::get( - {static_cast(values.size())}, rewriter.getI64Type()); - Value table = getOrCreateCachedTypedConstant( - denseConstants, rewriter, anchor, - DenseElementsAttr::get(type, values)); - Value selected = tensor::ExtractOp::create( - rewriter, loc, table, ValueRange {lane}); - return arith::IndexCastOp::create( - rewriter, loc, rewriter.getIndexType(), selected); -} - -static void setBatchTransferAttrs( - Operation *op, ArrayRef refs, - RewriterBase &rewriter) { - SmallVector channels, parents, counts, sources, targets; - for (const FragmentTransferRef &ref : refs) { - LogicalTransferMetadata metadata = getLogicalTransferMetadata(ref); - channels.push_back(metadata.channelId); - parents.push_back(metadata.parentExchangeId); - counts.push_back(metadata.parentTransferCount); - sources.push_back(metadata.sourceCore); - targets.push_back(metadata.targetCore); - } - auto dense = [&](ArrayRef values) -> DenseIntElementsAttr { - auto type = RankedTensorType::get( - {static_cast(values.size())}, rewriter.getI64Type()); - if (llvm::all_equal(values)) - values = values.take_front(); - return DenseIntElementsAttr::get(type, values); - }; - op->setAttr("raptor.batch_channel_ids", dense(channels)); - op->setAttr("raptor.batch_source_cores", dense(sources)); - op->setAttr("raptor.batch_target_cores", dense(targets)); - op->setAttr("raptor.batch_parent_exchange_ids", dense(parents)); - op->setAttr("raptor.batch_parent_transfer_counts", dense(counts)); -} - -static FailureOr> tryRealizeAnalyzedInsertAssembly( - NormalizedDeferredExchange &exchange, - DeferredSpecialization &specialization, IRRewriter &rewriter) { - const std::optional &assembly = - specialization.program.insertAssembly; - if (!assembly || assembly->entries.empty()) - return std::optional(); - - SmallVector ordered; - DenseMap entryByRequirement; - ordered.reserve(assembly->entries.size()); - for (const DeferredInsertAssemblyEntry &entry : assembly->entries) { - if (entry.requirementIndex >= specialization.transfers.size()) - return std::optional(); - FragmentTransfer &transfer = - specialization.transfers[entry.requirementIndex]; - if (transfer.local - || transfer.requirementIndex != entry.requirementIndex) - return std::optional(); - ordered.push_back(&transfer); - entryByRequirement[entry.requirementIndex] = &entry; - } - llvm::stable_sort(ordered, [](FragmentTransfer *lhs, - FragmentTransfer *rhs) { - return lhs->globalOrder < rhs->globalOrder; - }); - FragmentRequirement &reference = specialization.requirements[ - ordered.front()->requirementIndex]; - for (auto [index, transfer] : llvm::enumerate(ordered)) { - const FragmentRequirement &requirement = - specialization.requirements[transfer->requirementIndex]; - if (requirement.publicationFragmentType - != reference.publicationFragmentType - || transfer->globalOrder != ordered.front()->globalOrder + index - || transfer->targetInsertionStep - != ordered.front()->targetInsertionStep) - return std::optional(); - } - if (ordered.size() != specialization.requirements.size()) - return std::optional(); - - Location loc = exchange.deferred.getLoc(); - Operation *clonedInitial = rewriter.clone(*assembly->initialValue); - Value assembled = clonedInitial->getResult(0); - SmallVector channels, sources, targets; - SmallVector metadataRefs; - for (FragmentTransfer *transfer : ordered) { - channels.push_back(transfer->channelId); - sources.push_back(transfer->sourceCore); - targets.push_back(transfer->targetCore); - metadataRefs.push_back({&exchange, &specialization, transfer}); - } - auto insertReceived = [&](Value index, Value current) -> FailureOr { - Value channel = buildLaneLookup( - channels, index, exchange.deferred, rewriter, loc); - Value source = buildLaneLookup( - sources, index, exchange.deferred, rewriter, loc); - Value target = buildLaneLookup( - targets, index, exchange.deferred, rewriter, loc); - auto receive = SpatChannelReceiveOp::create( - rewriter, loc, reference.publicationFragmentType, - channel, source, target); - setBatchTransferAttrs(receive, metadataRefs, rewriter); - - auto buildGeometry = [&](auto member, - unsigned dimension) -> OpFoldResult { - SmallVector values; - for (FragmentTransfer *transfer : ordered) { - const DeferredInsertAssemblyEntry *entry = - entryByRequirement.lookup(transfer->requirementIndex); - values.push_back((entry->targetGeometry.*member)[dimension]); - } - if (llvm::all_equal(values)) - return rewriter.getIndexAttr(values.front()); - return buildLaneLookup( - values, index, exchange.deferred, rewriter, loc); - }; - unsigned rank = assembly->resultType.getRank(); - SmallVector offsets, sizes, strides; - for (unsigned dimension = 0; dimension < rank; ++dimension) { - offsets.push_back(buildGeometry( - &StaticSliceGeometry::offsets, dimension)); - sizes.push_back(buildGeometry( - &StaticSliceGeometry::sizes, dimension)); - strides.push_back(buildGeometry( - &StaticSliceGeometry::strides, dimension)); - } - return tensor::InsertSliceOp::create( - rewriter, loc, receive.getOutput(), current, offsets, sizes, strides) - .getResult(); - }; - if (ordered.size() == 1) { - Value zero = getOrCreateIndexConstant( - rewriter, exchange.deferred.getOperation(), 0); - auto result = insertReceived(zero, assembled); +static LogicalResult linearizeScheduled( + DeferredTransferPlan &plan, IRRewriter &rewriter) { + for (ScheduledInfo &info : plan.scheduled) { + LogicalResult result = info.isBatch() + ? linearizeBatch(info, rewriter) : linearizeScalar(info, rewriter); if (failed(result)) return failure(); - return std::optional(*result); - } - Value lower = getOrCreateIndexConstant( - rewriter, exchange.deferred.getOperation(), 0); - Value upper = getOrCreateIndexConstant( - rewriter, exchange.deferred.getOperation(), ordered.size()); - Value step = getOrCreateIndexConstant( - rewriter, exchange.deferred.getOperation(), 1); - auto loop = buildNormalizedScfFor( - rewriter, loc, lower, upper, step, ValueRange {assembled}, - [&](OpBuilder &, Location, Value index, ValueRange iterArgs, - SmallVectorImpl &yielded) -> LogicalResult { - auto result = insertReceived(index, iterArgs.front()); - if (failed(result)) - return failure(); - yielded.push_back(*result); - return success(); - }); - if (failed(loop)) - return failure(); - return std::optional(loop->results.front()); -} - -static FailureOr materializeBatchLocalFragment( - ArrayRef requirements, Value lane, - SpatDeferredCommunicationOp deferred, IRRewriter &rewriter, - ArrayRef transfers = {}) { - FragmentRequirement *reference = nullptr; - SmallVector offsets; - offsets.reserve(requirements.size()); - bool commonPayload = true; - for (auto [index, requirement] : llvm::enumerate(requirements)) { - bool active = transfers.empty() || transfers[index]->local; - if (!active) { - offsets.push_back(0); - continue; - } - if (!reference) - reference = requirement; - if (requirement->producer->payload != reference->producer->payload - || requirement->publicationFragmentType - != reference->publicationFragmentType - || requirement->graphLane.has_value() != reference->graphLane.has_value()) - commonPayload = false; - offsets.push_back(requirement->graphLane - ? *requirement->graphLane - - requirement->producer->laneStart - : 0); - } - if (!reference) - return failure(); - if (commonPayload) { - if (!reference->graphLane) - return materializeProducedFragment(*reference, rewriter, - deferred.getLoc()); - int64_t placeholder = *reference->graphLane - reference->producer->laneStart; - for (auto [index, transfer] : llvm::enumerate(transfers)) - if (!transfer->local) - offsets[index] = placeholder; - Value offset = buildLaneLookup(offsets, lane, deferred.getOperation(), - rewriter, deferred.getLoc()); - return materializeProducedFragment(*reference, offset, rewriter, - deferred.getLoc()); - } - - struct LocalClass { - FragmentRequirement *representative = nullptr; - SmallVector lanes; - }; - SmallVector classes; - for (auto [laneIndex, requirementValue] : llvm::enumerate(requirements)) { - FragmentRequirement *requirement = requirementValue; - if (!transfers.empty() && !transfers[laneIndex]->local) - continue; - auto sameBehavior = [&](const LocalClass &candidate) { - FragmentRequirement *other = candidate.representative; - return requirement->producer->payload == other->producer->payload - && requirement->publicationFragmentType - == other->publicationFragmentType - && requirement->graphLane.has_value() - == other->graphLane.has_value(); - }; - auto classIt = llvm::find_if(classes, sameBehavior); - if (classIt == classes.end()) { - classes.push_back({requirement, {static_cast(laneIndex)}}); - } else { - classIt->lanes.push_back(laneIndex); - } - } - if (classes.empty() - || llvm::any_of(classes, [&](const LocalClass &localClass) { - return localClass.representative->publicationFragmentType - != reference->publicationFragmentType; - })) - return failure(); - - auto materializeClass = [&](const LocalClass &localClass) - -> FailureOr { - FragmentRequirement *representative = localClass.representative; - if (!representative->graphLane) - return materializeProducedFragment( - *representative, rewriter, deferred.getLoc()); - int64_t defaultOffset = - *representative->graphLane - representative->producer->laneStart; - SmallVector classOffsets(requirements.size(), defaultOffset); - for (unsigned laneIndex : localClass.lanes) { - FragmentRequirement *requirement = requirements[laneIndex]; - int64_t offset = - *requirement->graphLane - requirement->producer->laneStart; - if (offset < 0 || offset >= requirement->producer->laneCount) - return failure(); - classOffsets[laneIndex] = offset; - } - Value offset = buildLaneLookup( - classOffsets, lane, deferred.getOperation(), rewriter, - deferred.getLoc()); - return materializeProducedFragment( - *representative, offset, rewriter, deferred.getLoc()); - }; - if (classes.size() == 1) - return materializeClass(classes.front()); - - auto buildMembership = [&](ArrayRef members) -> Value { - if (members.size() == 1) { - Value member = getOrCreateIndexConstant( - rewriter, deferred.getOperation(), members.front()); - return arith::CmpIOp::create( - rewriter, deferred.getLoc(), arith::CmpIPredicate::eq, lane, - member); - } - SmallVector membership(requirements.size()); - for (unsigned member : members) - membership[member] = 1; - Value selected = buildLaneLookup( - membership, lane, deferred.getOperation(), rewriter, - deferred.getLoc()); - Value zero = getOrCreateIndexConstant( - rewriter, deferred.getOperation(), 0); - return arith::CmpIOp::create( - rewriter, deferred.getLoc(), arith::CmpIPredicate::ne, selected, - zero); - }; - if (classes.size() == 2) { - unsigned thenIndex = classes[0].lanes.size() <= classes[1].lanes.size() - ? 0 : 1; - unsigned elseIndex = 1 - thenIndex; - Value condition = buildMembership(classes[thenIndex].lanes); - auto selection = scf::IfOp::create( - rewriter, deferred.getLoc(), - TypeRange {reference->publicationFragmentType}, condition, true); - auto buildBranch = [&](Region ®ion, const LocalClass &localClass) - -> LogicalResult { - OpBuilder::InsertionGuard guard(rewriter); - auto existingYield = dyn_cast_or_null( - region.front().empty() ? nullptr : ®ion.front().back()); - if (existingYield) - rewriter.setInsertionPoint(existingYield); - else - rewriter.setInsertionPointToStart(®ion.front()); - auto value = materializeClass(localClass); - if (failed(value)) - return failure(); - if (existingYield) - existingYield->setOperands(*value); - else - scf::YieldOp::create(rewriter, deferred.getLoc(), *value); - return success(); - }; - if (failed(buildBranch(selection.getThenRegion(), classes[thenIndex])) - || failed(buildBranch(selection.getElseRegion(), classes[elseIndex]))) - return failure(); - return selection.getResult(0); - } - - llvm::stable_sort(classes, [](const LocalClass &lhs, - const LocalClass &rhs) { - return lhs.lanes.front() < rhs.lanes.front(); - }); - SmallVector classByLane(requirements.size()); - for (auto [classIndex, localClass] : llvm::enumerate(classes)) - for (unsigned laneIndex : localClass.lanes) - classByLane[laneIndex] = classIndex; - Value classId = buildLaneLookup(classByLane, lane, deferred.getOperation(), - rewriter, deferred.getLoc()); - SmallVector cases; - for (int64_t classIndex = 0; - classIndex < static_cast(classes.size()) - 1; - ++classIndex) - cases.push_back(classIndex); - auto selection = scf::IndexSwitchOp::create( - rewriter, deferred.getLoc(), - TypeRange {reference->publicationFragmentType}, classId, cases, - cases.size()); - auto buildRegion = [&](Region ®ion, const LocalClass &localClass) - -> LogicalResult { - OpBuilder::InsertionGuard guard(rewriter); - Block *block = rewriter.createBlock(®ion); - rewriter.setInsertionPointToEnd(block); - auto value = materializeClass(localClass); - if (failed(value)) - return failure(); - scf::YieldOp::create(rewriter, deferred.getLoc(), *value); - return success(); - }; - for (auto [region, localClass] : - llvm::zip(selection.getCaseRegions(), ArrayRef(classes).drop_back())) - if (failed(buildRegion(region, localClass))) - return failure(); - if (failed(buildRegion(selection.getDefaultRegion(), classes.back()))) - return failure(); - return selection.getResult(0); -} - -struct DynamicSliceGeometry { - SmallVector offsets; - SmallVector sizes; - SmallVector strides; - RankedTensorType resultType; - bool identity = false; -}; - -static FailureOr buildBatchInnerGeometry( - ArrayRef specializations, unsigned leafIndex, - ArrayRef classLanes, Value lane, Operation *anchor, - IRRewriter &rewriter, Location loc) { - if (classLanes.empty()) - return failure(); - const DeferredProjectionLeaf &reference = - specializations[classLanes.front()].program.leaves[leafIndex]; - unsigned rank = reference.innerGeometry.offsets.size(); - if (reference.innerGeometry.sizes.size() != rank - || reference.innerGeometry.strides.size() != rank) - return failure(); - - DynamicSliceGeometry geometry; - geometry.identity = true; - auto buildDimension = [&](auto member, unsigned dimension) - -> FailureOr { - SmallVector values(specializations.size(), - (reference.innerGeometry.*member)[dimension]); - SmallVector activeValues; - for (unsigned classLane : classLanes) { - if (classLane >= specializations.size() - || leafIndex - >= specializations[classLane].program.leaves.size()) - return failure(); - const DeferredProjectionLeaf &leaf = - specializations[classLane].program.leaves[leafIndex]; - if ((leaf.innerGeometry.*member).size() != rank - || leaf.reconstructedType != reference.reconstructedType) - return failure(); - values[classLane] = (leaf.innerGeometry.*member)[dimension]; - activeValues.push_back(values[classLane]); - } - if (llvm::all_equal(activeValues)) - return OpFoldResult(rewriter.getIndexAttr(activeValues.front())); - return OpFoldResult(buildLaneLookup(values, lane, anchor, rewriter, loc)); - }; - for (unsigned dimension = 0; dimension < rank; ++dimension) { - auto offset = buildDimension(&StaticSliceGeometry::offsets, dimension); - auto size = buildDimension(&StaticSliceGeometry::sizes, dimension); - auto stride = buildDimension(&StaticSliceGeometry::strides, dimension); - if (failed(offset) || failed(size) || failed(stride)) - return failure(); - geometry.offsets.push_back(*offset); - geometry.sizes.push_back(*size); - geometry.strides.push_back(*stride); - } - if (reference.kind != DeferredLeafKind::GraphBatchProjection) { - geometry.resultType = reference.reconstructedType; - return geometry; - } - if (reference.reconstructedType.getRank() != rank + 1) - return failure(); - geometry.resultType = RankedTensorType::get( - reference.reconstructedType.getShape().drop_front(), - reference.reconstructedType.getElementType()); - for (unsigned classLane : classLanes) { - const DeferredProjectionLeaf &leaf = - specializations[classLane].program.leaves[leafIndex]; - for (unsigned dimension = 0; dimension < rank; ++dimension) - geometry.identity &= leaf.innerGeometry.offsets[dimension] == 0 - && leaf.innerGeometry.sizes[dimension] - == geometry.resultType.getDimSize(dimension) - && leaf.innerGeometry.strides[dimension] == 1; - } - return geometry; -} - -static FailureOr reconstructBatchProjectionLeaf( - NormalizedDeferredExchange &exchange, unsigned leafIndex, - const BatchBoundaryClass &group, ArrayRef available, Value lane, - IRRewriter &rewriter) { - unsigned referenceLane = group.lanes.front(); - DeferredSpecialization &reference = - exchange.specializations[referenceLane]; - const DeferredProjectionLeaf &leaf = reference.program.leaves[leafIndex]; - SmallVector selected( - leaf.kind == DeferredLeafKind::ScalarSource ? 1 - : leaf.physicalSlots.size()); - for (auto [requirementIndex, requirement] : - llvm::enumerate(reference.requirements)) - if (requirement.leafIndex == leafIndex - && requirement.selectedPosition < selected.size()) { - if (selected[requirement.selectedPosition]) - return failure(); - selected[requirement.selectedPosition] = available[requirementIndex]; - } - if (llvm::any_of(selected, [](Value value) { return !value; })) - return failure(); - if (leaf.kind == DeferredLeafKind::ScalarSource) - return selected.front(); - - auto geometry = buildBatchInnerGeometry( - exchange.specializations, leafIndex, group.lanes, lane, - exchange.deferred, rewriter, exchange.deferred.getLoc()); - if (failed(geometry)) - return failure(); - SmallVector shaped; - for (Value fragment : selected) { - if (geometry->identity) { - shaped.push_back(fragment); - continue; - } - shaped.push_back(tensor::ExtractSliceOp::create( - rewriter, exchange.deferred.getLoc(), geometry->resultType, fragment, - geometry->offsets, geometry->sizes, geometry->strides)); - } - if (shaped.size() == 1 && shaped.front().getType() == leaf.reconstructedType) - return shaped.front(); - SmallVector expanded; - for (Value fragment : shaped) { - auto unit = addLeadingUnitTensorDimension( - rewriter, exchange.deferred.getLoc(), fragment); - if (failed(unit)) - return failure(); - expanded.push_back(*unit); - } - if (expanded.size() == 1 - && expanded.front().getType() == leaf.reconstructedType) - return expanded.front(); - return tensor::ConcatOp::create( - rewriter, exchange.deferred.getLoc(), leaf.reconstructedType, 0, - expanded).getResult(); -} - -static FailureOr materializeBatchResult( - NormalizedDeferredExchange &exchange, ArrayRef available, - const BatchBoundaryClass &group, Value lane, IRRewriter &rewriter) { - unsigned referenceLane = group.lanes.front(); - DeferredSpecialization &reference = exchange.specializations[referenceLane]; - IRMapping mapping; - if (reference.program.scheduledLane) - mapping.map(reference.program.scheduledLane, lane); - for (auto [value, staticValue] : reference.program.staticValues) { - if (mapping.contains(value)) - continue; - SmallVector values; - values.reserve(exchange.specializations.size()); - for (DeferredSpecialization &specialization : exchange.specializations) { - auto it = specialization.program.staticValues.find(value); - if (it == specialization.program.staticValues.end()) - return failure(); - values.push_back(it->second); - } - Value selected = buildLaneLookup(values, lane, exchange.deferred, - rewriter, exchange.deferred.getLoc()); - if (!value.getType().isIndex()) - selected = arith::IndexCastOp::create( - rewriter, exchange.deferred.getLoc(), value.getType(), selected); - mapping.map(value, selected); - } - for (auto [index, leaf] : llvm::enumerate(reference.program.leaves)) { - auto reconstructed = reconstructBatchProjectionLeaf( - exchange, index, group, available, lane, rewriter); - if (failed(reconstructed)) - return failure(); - mapping.map(leaf.replacementRoot, *reconstructed); - } - for (Operation *op : reference.program.residualOps) { - Operation *copy = rewriter.clone(*op, mapping); - for (auto [original, result] : llvm::zip(op->getResults(), - copy->getResults())) - mapping.map(original, result); - } - Value result = mapping.lookupOrDefault(reference.program.yieldedValue); - return result && result.getType() == exchange.deferred.getOutput().getType() - ? FailureOr(result) - : FailureOr(failure()); -} - -static FailureOr tryEmitLoopedScalarSends( - ArrayRef sends, IRRewriter &rewriter) { - if (sends.size() < 2) - return false; - FragmentTransferRef first = sends.front(); - FragmentRequirement &reference = first.specialization->requirements[ - first.transfer->requirementIndex]; - ProducedValue *producer = reference.producer; - if (producer->scheduled->isBatch()) - return false; - - SmallVector channels, sources, targets, offsets; - for (FragmentTransferRef ref : sends) { - FragmentRequirement &requirement = ref.specialization->requirements[ - ref.transfer->requirementIndex]; - if (requirement.producer->scheduled != producer->scheduled - || ref.transfer->sourceInsertionStep - != first.transfer->sourceInsertionStep - || requirement.producer->payload != producer->payload - || requirement.publicationFragmentType - != reference.publicationFragmentType - || requirement.graphLane.has_value() != reference.graphLane.has_value()) - return false; - int64_t offset = requirement.graphLane - ? *requirement.graphLane - - requirement.producer->laneStart - : 0; - if (offset < 0 || offset >= requirement.producer->laneCount) - return failure(); - channels.push_back(ref.transfer->channelId); - sources.push_back(ref.transfer->sourceCore); - targets.push_back(ref.transfer->targetCore); - offsets.push_back(offset); - } - - Location loc = first.exchange->deferred.getLoc(); - Value lower = getOrCreateIndexConstant( - rewriter, first.exchange->deferred.getOperation(), 0); - Value upper = getOrCreateIndexConstant( - rewriter, first.exchange->deferred.getOperation(), sends.size()); - Value step = getOrCreateIndexConstant( - rewriter, first.exchange->deferred.getOperation(), 1); - auto loop = buildNormalizedScfFor( - rewriter, loc, lower, upper, step, ValueRange {}, - [&](OpBuilder &, Location, Value index, ValueRange, - SmallVectorImpl &) -> LogicalResult { - Value channel = buildLaneLookup( - channels, index, first.exchange->deferred, rewriter, loc); - Value source = buildLaneLookup( - sources, index, first.exchange->deferred, rewriter, loc); - Value target = buildLaneLookup( - targets, index, first.exchange->deferred, rewriter, loc); - Value payload; - if (reference.graphLane) { - Value offset = buildLaneLookup( - offsets, index, first.exchange->deferred, rewriter, loc); - auto materialized = materializeProducedFragment( - reference, offset, rewriter, loc); - if (failed(materialized)) - return failure(); - payload = *materialized; - } - else { - auto materialized = materializeProducedFragment( - reference, rewriter, loc); - if (failed(materialized)) - return failure(); - payload = *materialized; - } - auto send = SpatChannelSendOp::create( - rewriter, loc, channel, source, target, payload); - setBatchTransferAttrs(send, sends, rewriter); - return success(); - }); - if (failed(loop)) - return failure(); - return true; -} - -static FailureOr tryEmitConsecutiveAnalyzedAssembly( - ArrayRef actions, size_t start, - AssemblyCache &assemblyCache, IRRewriter &rewriter) { - const BoundaryAction &first = actions[start]; - DeferredSpecialization *specialization = first.ref.specialization; - if (first.kind != BoundaryActionKind::Receive) - return 0; - size_t transferCount = llvm::count_if( - specialization->transfers, - [](const FragmentTransfer &transfer) { return !transfer.local; }); - if (transferCount == 0 || start + transferCount > actions.size()) - return 0; - llvm::SmallDenseSet expected; - for (FragmentTransfer &transfer : specialization->transfers) - if (!transfer.local) - expected.insert(&transfer); - for (const BoundaryAction &action : actions.slice(start, transferCount)) - if (action.kind != BoundaryActionKind::Receive - || action.ref.specialization != specialization - || !expected.erase(action.ref.transfer)) - return 0; - if (!expected.empty()) - return 0; - FailureOr> assembled = - tryRealizeAnalyzedInsertAssembly( - *first.ref.exchange, *specialization, rewriter); - if (failed(assembled)) - return failure(); - if (!*assembled) - return 0; - assemblyCache[specialization] = **assembled; - return transferCount; -} - -static LogicalResult realizeScalarBoundary( - ScheduledBoundaryPlan &boundary, ReceiveCache &receiveCache, - AssemblyCache &assemblyCache, IndexConstantCache &indexConstants, - DenseConstantCache &typedConstants, IRRewriter &rewriter, - DeferredEraseSet &erase) { - setInsertionAtBoundary( - rewriter, *boundary.scheduled, boundary.insertionStep); - ArrayRef actions = boundary.actionsByLane.front(); - for (size_t index = 0; index < actions.size();) { - const BoundaryAction &action = actions[index]; - if (action.kind == BoundaryActionKind::Receive) { - auto compacted = tryEmitConsecutiveAnalyzedAssembly( - actions, index, assemblyCache, rewriter); - if (failed(compacted)) - return failure(); - if (*compacted != 0) { - index += *compacted; - continue; - } - if (failed(emitFragmentReceiveHere( - action.ref, receiveCache, indexConstants, rewriter))) - return failure(); - ++index; - continue; - } - - size_t end = index + 1; - while (end < actions.size() - && actions[end].kind == BoundaryActionKind::Send - && haveCompatibleSendEmission( - action.ref, actions[end].ref, boundary.insertionStep, - boundary.insertionStep)) - ++end; - SmallVector sends; - for (const BoundaryAction &send : actions.slice(index, end - index)) - sends.push_back(send.ref); - bool compacted = false; - if (sends.size() >= 2) { - auto result = tryEmitLoopedScalarSends(sends, rewriter); - if (failed(result)) - return failure(); - compacted = *result; - } - if (!compacted) - for (FragmentTransferRef ref : sends) - if (failed(emitFragmentSendHere( - *ref.exchange, *ref.specialization, *ref.transfer, - indexConstants, rewriter))) - return failure(); - index = end; - } - - for (NormalizedDeferredExchange *exchange : boundary.targetExchanges) { - auto result = realizeSpecialization( - *exchange, exchange->specializations.front(), receiveCache, - assemblyCache, indexConstants, typedConstants, rewriter); - if (failed(result)) - return failure(); - replaceDeferred(exchange->deferred, *result, erase); } return success(); } -struct BatchActionSemanticKey { - BoundaryActionKind kind = BoundaryActionKind::Send; - uint64_t parentExchangeId = 0; - unsigned requirementIndex = 0; - unsigned leafIndex = 0; - unsigned selectedPosition = 0; - Type publicationFragmentType; - std::optional sendSignature; - - bool operator==(const BatchActionSemanticKey &other) const { - if (kind != other.kind - || publicationFragmentType != other.publicationFragmentType) - return false; - if (kind == BoundaryActionKind::Send) - return sendSignature && other.sendSignature - && haveSameSendEmissionSignature(*sendSignature, - *other.sendSignature); - return parentExchangeId == other.parentExchangeId - && requirementIndex == other.requirementIndex - && leafIndex == other.leafIndex - && selectedPosition == other.selectedPosition - && !sendSignature && !other.sendSignature; - } -}; - -struct DeferredLeafStructuralKey { - DeferredLeafKind kind = DeferredLeafKind::ScalarSource; - RankedTensorType reconstructedType; - unsigned selectedPositionCount = 0; - unsigned innerGeometryRank = 0; - bool identityGeometry = false; - - bool operator==(const DeferredLeafStructuralKey &other) const { - return kind == other.kind && reconstructedType == other.reconstructedType - && selectedPositionCount == other.selectedPositionCount - && innerGeometryRank == other.innerGeometryRank - && identityGeometry == other.identityGeometry; - } -}; - -struct DeferredInsertAssemblyStructuralKey { - RankedTensorType resultType; - SmallVector requirementIndices; - SmallVector geometryRanks; - - bool operator==(const DeferredInsertAssemblyStructuralKey &other) const { - return resultType == other.resultType - && requirementIndices == other.requirementIndices - && geometryRanks == other.geometryRanks; - } -}; - -struct BatchExchangeSemanticKey { - uint64_t exchangeId = 0; - SmallVector residualOps; - Value yieldedValue; - SmallVector leaves; - SmallVector localByRequirement; - std::optional assembly; - - bool operator==(const BatchExchangeSemanticKey &other) const { - return exchangeId == other.exchangeId && residualOps == other.residualOps - && yieldedValue == other.yieldedValue && leaves == other.leaves - && localByRequirement == other.localByRequirement - && assembly == other.assembly; - } -}; - -struct BatchLaneSemanticKey { - SmallVector actions; - SmallVector exchanges; - - bool operator==(const BatchLaneSemanticKey &other) const { - return actions == other.actions && exchanges == other.exchanges; - } -}; - -static BatchLaneSemanticKey buildBatchLaneSemanticKey( - const ScheduledBoundaryPlan &boundary, unsigned lane) { - BatchLaneSemanticKey key; - for (const BoundaryAction &action : boundary.actionsByLane[lane]) { - const FragmentRequirement &requirement = - action.ref.specialization->requirements[ - action.ref.transfer->requirementIndex]; - BatchActionSemanticKey actionKey { - action.kind, action.ref.exchange->exchangeId, - action.ref.transfer->requirementIndex, requirement.leafIndex, - requirement.selectedPosition, requirement.publicationFragmentType, - std::nullopt}; - if (action.kind == BoundaryActionKind::Send) - actionKey.sendSignature = getSendEmissionSignature(action.ref); - key.actions.push_back(std::move(actionKey)); - } - for (NormalizedDeferredExchange *exchange : boundary.targetExchanges) { - DeferredSpecialization &specialization = exchange->specializations[lane]; - BatchExchangeSemanticKey exchangeKey; - exchangeKey.exchangeId = exchange->exchangeId; - exchangeKey.residualOps = specialization.program.residualOps; - exchangeKey.yieldedValue = specialization.program.yieldedValue; - for (const DeferredProjectionLeaf &leaf : specialization.program.leaves) { - bool identity = leaf.kind != DeferredLeafKind::GraphBatchProjection; - if (!identity) { - identity = true; - for (auto [offset, stride] : llvm::zip( - leaf.innerGeometry.offsets, leaf.innerGeometry.strides)) - identity &= offset == 0 && stride == 1; - } - exchangeKey.leaves.push_back( - {leaf.kind, leaf.reconstructedType, - static_cast(leaf.kind == DeferredLeafKind::ScalarSource - ? 1 - : leaf.physicalSlots.size()), - static_cast(leaf.innerGeometry.offsets.size()), identity}); - } - for (const FragmentTransfer &transfer : specialization.transfers) - exchangeKey.localByRequirement.push_back(transfer.local); - if (specialization.program.insertAssembly) { - DeferredInsertAssemblyStructuralKey assemblyKey; - assemblyKey.resultType = - specialization.program.insertAssembly->resultType; - for (const DeferredInsertAssemblyEntry &entry : - specialization.program.insertAssembly->entries) { - assemblyKey.requirementIndices.push_back(entry.requirementIndex); - assemblyKey.geometryRanks.push_back(entry.targetGeometry.offsets.size()); - } - exchangeKey.assembly = std::move(assemblyKey); - } - key.exchanges.push_back(std::move(exchangeKey)); - } - return key; -} - -static size_t hashBatchLaneSemanticKey(const BatchLaneSemanticKey &key) { - llvm::hash_code hash = llvm::hash_value(key.actions.size()); - for (const BatchActionSemanticKey &action : key.actions) { - hash = llvm::hash_combine( - hash, static_cast(action.kind), - action.publicationFragmentType.getAsOpaquePointer()); - if (action.sendSignature) - hash = llvm::hash_combine( - hash, hashSendEmissionSignature(*action.sendSignature)); - else - hash = llvm::hash_combine( - hash, action.parentExchangeId, action.requirementIndex, - action.leafIndex, action.selectedPosition); - } - for (const BatchExchangeSemanticKey &exchange : key.exchanges) { - hash = llvm::hash_combine( - hash, exchange.exchangeId, exchange.yieldedValue.getAsOpaquePointer(), - exchange.residualOps.size(), exchange.leaves.size(), - exchange.localByRequirement.size(), exchange.assembly.has_value()); - for (Operation *op : exchange.residualOps) - hash = llvm::hash_combine(hash, op); - for (const DeferredLeafStructuralKey &leaf : exchange.leaves) - hash = llvm::hash_combine( - hash, static_cast(leaf.kind), - leaf.reconstructedType.getAsOpaquePointer(), - leaf.selectedPositionCount, leaf.innerGeometryRank, - leaf.identityGeometry); - for (bool local : exchange.localByRequirement) - hash = llvm::hash_combine(hash, local); - if (exchange.assembly) { - hash = llvm::hash_combine( - hash, exchange.assembly->resultType.getAsOpaquePointer()); - for (auto values : llvm::zip_equal( - exchange.assembly->requirementIndices, - exchange.assembly->geometryRanks)) - hash = llvm::hash_combine( - hash, std::get<0>(values), std::get<1>(values)); - } - } - return static_cast(hash); -} - -static SmallVector classifyBatchBoundary( - const ScheduledBoundaryPlan &boundary) { - SmallVector classes; - SmallVector classKeys; - DenseMap> classesByHash; - for (unsigned lane = 0; lane < boundary.actionsByLane.size(); ++lane) { - BatchLaneSemanticKey key = buildBatchLaneSemanticKey(boundary, lane); - size_t hash = hashBatchLaneSemanticKey(key); - std::optional matchingClass; - for (unsigned classIndex : classesByHash.lookup(hash)) - if (classKeys[classIndex] == key) { - matchingClass = classIndex; - break; - } - if (!matchingClass) { - classesByHash[hash].push_back(classes.size()); - classes.push_back({{lane}}); - classKeys.push_back(std::move(key)); - } else { - classes[*matchingClass].lanes.push_back(lane); - } - } - return classes; -} - -static Value buildBatchActionTableIndex( - Value lane, Value actionIndex, unsigned laneCount, Operation *anchor, - IRRewriter &rewriter, Location loc) { - // Runtime tables are action-major over every lane; compact metadata is - // action-major over active class lanes only. - Value width = getOrCreateIndexConstant(rewriter, anchor, laneCount); - Value actionBase = arith::MulIOp::create( - rewriter, loc, actionIndex, width); - return arith::AddIOp::create(rewriter, loc, actionBase, lane); -} - -static LogicalResult emitBatchSendRun( - const ScheduledBoundaryPlan &boundary, const BatchBoundaryClass &group, - size_t start, size_t runLength, Value lane, - DenseConstantCache &denseConstants, IRRewriter &rewriter) { - unsigned laneCount = boundary.actionsByLane.size(); - const BoundaryAction &firstAction = - boundary.actionsByLane[group.lanes.front()][start]; - FragmentRequirement &reference = - firstAction.ref.specialization->requirements[ - firstAction.ref.transfer->requirementIndex]; - size_t tableSize = static_cast(laneCount) * runLength; - SmallVector channels(tableSize, - firstAction.ref.transfer->channelId); - SmallVector sources(tableSize, - firstAction.ref.transfer->sourceCore); - SmallVector targets(tableSize, - firstAction.ref.transfer->targetCore); - int64_t defaultOffset = reference.graphLane - ? *reference.graphLane - reference.producer->laneStart : 0; - SmallVector offsets(tableSize, defaultOffset); - SmallVector metadataRefs; - for (size_t actionOffset = 0; actionOffset < runLength; - ++actionOffset) - for (unsigned classLane : group.lanes) { - const BoundaryAction &action = - boundary.actionsByLane[classLane][start + actionOffset]; - FragmentRequirement &requirement = - action.ref.specialization->requirements[ - action.ref.transfer->requirementIndex]; - if (action.kind != BoundaryActionKind::Send - || !haveCompatibleSendEmission( - firstAction.ref, action.ref, boundary.insertionStep, - boundary.insertionStep)) - return failure(); - size_t tableIndex = actionOffset * laneCount + classLane; - channels[tableIndex] = action.ref.transfer->channelId; - sources[tableIndex] = action.ref.transfer->sourceCore; - targets[tableIndex] = action.ref.transfer->targetCore; - if (requirement.graphLane) - offsets[tableIndex] = - *requirement.graphLane - requirement.producer->laneStart; - metadataRefs.push_back(action.ref); - } - - auto emitOne = [&](Value tableIndex) -> LogicalResult { - Location loc = firstAction.ref.exchange->deferred.getLoc(); - Value payload; - if (reference.graphLane) { - Value offset = buildLaneLookupCached( - offsets, tableIndex, firstAction.ref.exchange->deferred, - denseConstants, rewriter, loc); - auto materialized = materializeProducedFragment( - reference, offset, rewriter, loc); - if (failed(materialized)) - return failure(); - payload = *materialized; - } else { - auto materialized = materializeProducedFragment( - reference, rewriter, loc); - if (failed(materialized)) - return failure(); - payload = *materialized; - } - auto send = SpatChannelSendOp::create( - rewriter, loc, - buildLaneLookupCached(channels, tableIndex, - firstAction.ref.exchange->deferred, - denseConstants, rewriter, loc), - buildLaneLookupCached(sources, tableIndex, - firstAction.ref.exchange->deferred, - denseConstants, rewriter, loc), - buildLaneLookupCached(targets, tableIndex, - firstAction.ref.exchange->deferred, - denseConstants, rewriter, loc), - payload); - setBatchTransferAttrs(send, metadataRefs, rewriter); - return success(); - }; - if (runLength == 1) - return emitOne(lane); - Location loc = firstAction.ref.exchange->deferred.getLoc(); - Value lower = getOrCreateIndexConstant( - rewriter, firstAction.ref.exchange->deferred, 0); - Value upper = getOrCreateIndexConstant( - rewriter, firstAction.ref.exchange->deferred, runLength); - Value step = getOrCreateIndexConstant( - rewriter, firstAction.ref.exchange->deferred, 1); - auto loop = buildNormalizedScfFor( - rewriter, loc, lower, upper, step, ValueRange {}, - [&](OpBuilder &, Location, Value actionIndex, ValueRange, - SmallVectorImpl &) -> LogicalResult { - Value tableIndex = buildBatchActionTableIndex( - lane, actionIndex, laneCount, - firstAction.ref.exchange->deferred, rewriter, loc); - return emitOne(tableIndex); - }); - return success(succeeded(loop)); -} - -static LogicalResult emitBatchReceiveAction( - const ScheduledBoundaryPlan &boundary, const BatchBoundaryClass &group, - size_t actionIndex, Value lane, ReceiveCache &receiveCache, - DenseConstantCache &denseConstants, IRRewriter &rewriter) { - unsigned laneCount = boundary.actionsByLane.size(); - const BoundaryAction &first = - boundary.actionsByLane[group.lanes.front()][actionIndex]; - FragmentRequirement &reference = first.ref.specialization->requirements[ - first.ref.transfer->requirementIndex]; - SmallVector channels(laneCount, first.ref.transfer->channelId); - SmallVector sources(laneCount, first.ref.transfer->sourceCore); - SmallVector targets(laneCount, first.ref.transfer->targetCore); - SmallVector refs; - for (unsigned classLane : group.lanes) { - const BoundaryAction &action = - boundary.actionsByLane[classLane][actionIndex]; - FragmentRequirement &requirement = - action.ref.specialization->requirements[ - action.ref.transfer->requirementIndex]; - if (action.kind != BoundaryActionKind::Receive - || action.ref.exchange != first.ref.exchange - || requirement.publicationFragmentType - != reference.publicationFragmentType) - return failure(); - channels[classLane] = action.ref.transfer->channelId; - sources[classLane] = action.ref.transfer->sourceCore; - targets[classLane] = action.ref.transfer->targetCore; - refs.push_back(action.ref); - } - Location loc = first.ref.exchange->deferred.getLoc(); - auto receive = SpatChannelReceiveOp::create( - rewriter, loc, reference.publicationFragmentType, - buildLaneLookupCached(channels, lane, first.ref.exchange->deferred, - denseConstants, rewriter, loc), - buildLaneLookupCached(sources, lane, first.ref.exchange->deferred, - denseConstants, rewriter, loc), - buildLaneLookupCached(targets, lane, first.ref.exchange->deferred, - denseConstants, rewriter, loc)); - setBatchTransferAttrs(receive, refs, rewriter); - for (const FragmentTransferRef &ref : refs) - receiveCache[ref.transfer] = receive.getOutput(); - return success(); -} - -static FailureOr> tryEmitBatchAnalyzedAssembly( - ScheduledBoundaryPlan &boundary, const BatchBoundaryClass &group, - NormalizedDeferredExchange &exchange, size_t startAction, Value lane, - ReceiveCache &receiveCache, DenseConstantCache &denseConstants, - IRRewriter &rewriter) { - (void)receiveCache; - unsigned referenceLane = group.lanes.front(); - DeferredSpecialization &reference = - exchange.specializations[referenceLane]; - const std::optional &assembly = - reference.program.insertAssembly; - if (!assembly || assembly->entries.empty()) - return std::optional(); - size_t entryCount = assembly->entries.size(); - if (startAction + entryCount - > boundary.actionsByLane[referenceLane].size()) - return std::optional(); - - Type fragmentType; - for (unsigned classLane : group.lanes) { - DeferredSpecialization &specialization = - exchange.specializations[classLane]; - if (!specialization.program.insertAssembly - || specialization.program.insertAssembly->entries.size() - != entryCount) - return std::optional(); - for (auto [entryIndex, entry] : llvm::enumerate( - specialization.program.insertAssembly->entries)) { - if (entry.requirementIndex >= specialization.requirements.size()) - return std::optional(); - const BoundaryAction &action = - boundary.actionsByLane[classLane][startAction + entryIndex]; - const FragmentRequirement &requirement = - specialization.requirements[entry.requirementIndex]; - if (action.kind != BoundaryActionKind::Receive - || action.ref.exchange != &exchange - || action.ref.transfer->requirementIndex != entry.requirementIndex - || action.ref.specialization != &specialization - || (fragmentType - && requirement.publicationFragmentType != fragmentType)) - return std::optional(); - fragmentType = requirement.publicationFragmentType; - } - } - if (!fragmentType) - return std::optional(); - - unsigned laneCount = boundary.actionsByLane.size(); - size_t tableSize = entryCount * laneCount; - const BoundaryAction &first = - boundary.actionsByLane[referenceLane][startAction]; - SmallVector channels(tableSize, first.ref.transfer->channelId); - SmallVector sources(tableSize, first.ref.transfer->sourceCore); - SmallVector targets(tableSize, first.ref.transfer->targetCore); - unsigned rank = assembly->resultType.getRank(); - SmallVector> offsets( - rank, SmallVector(tableSize)); - SmallVector> sizes( - rank, SmallVector(tableSize, 1)); - SmallVector> strides( - rank, SmallVector(tableSize, 1)); - SmallVector metadataRefs; - metadataRefs.reserve(entryCount * group.lanes.size()); - for (size_t entryIndex = 0; entryIndex < entryCount; ++entryIndex) - for (unsigned classLane : group.lanes) { - DeferredSpecialization &specialization = - exchange.specializations[classLane]; - const DeferredInsertAssemblyEntry &entry = - specialization.program.insertAssembly->entries[entryIndex]; - const BoundaryAction &action = - boundary.actionsByLane[classLane][startAction + entryIndex]; - size_t tableIndex = entryIndex * laneCount + classLane; - channels[tableIndex] = action.ref.transfer->channelId; - sources[tableIndex] = action.ref.transfer->sourceCore; - targets[tableIndex] = action.ref.transfer->targetCore; - for (unsigned dimension = 0; dimension < rank; ++dimension) { - offsets[dimension][tableIndex] = entry.targetGeometry.offsets[dimension]; - sizes[dimension][tableIndex] = entry.targetGeometry.sizes[dimension]; - strides[dimension][tableIndex] = entry.targetGeometry.strides[dimension]; - } - metadataRefs.push_back(action.ref); - } - - Operation *initial = rewriter.clone(*assembly->initialValue); - Value assembled = initial->getResult(0); - Location loc = exchange.deferred.getLoc(); - auto emitEntry = [&](Value tableIndex, Value current) -> FailureOr { - auto receive = SpatChannelReceiveOp::create( - rewriter, loc, fragmentType, - buildLaneLookupCached(channels, tableIndex, exchange.deferred, - denseConstants, rewriter, loc), - buildLaneLookupCached(sources, tableIndex, exchange.deferred, - denseConstants, rewriter, loc), - buildLaneLookupCached(targets, tableIndex, exchange.deferred, - denseConstants, rewriter, loc)); - setBatchTransferAttrs(receive, metadataRefs, rewriter); - SmallVector targetOffsets, targetSizes, targetStrides; - for (unsigned dimension = 0; dimension < rank; ++dimension) { - targetOffsets.push_back(buildLaneLookupCached( - offsets[dimension], tableIndex, exchange.deferred, denseConstants, - rewriter, loc)); - targetSizes.push_back(buildLaneLookupCached( - sizes[dimension], tableIndex, exchange.deferred, denseConstants, - rewriter, loc)); - targetStrides.push_back(buildLaneLookupCached( - strides[dimension], tableIndex, exchange.deferred, denseConstants, - rewriter, loc)); - } - return tensor::InsertSliceOp::create( - rewriter, loc, receive.getOutput(), current, targetOffsets, - targetSizes, targetStrides).getResult(); - }; - if (entryCount == 1) { - auto result = emitEntry(lane, assembled); - if (failed(result)) - return failure(); - return std::optional(*result); - } - - Value lower = getOrCreateIndexConstant(rewriter, exchange.deferred, 0); - Value upper = getOrCreateIndexConstant( - rewriter, exchange.deferred, entryCount); - Value step = getOrCreateIndexConstant(rewriter, exchange.deferred, 1); - auto loop = buildNormalizedScfFor( - rewriter, loc, lower, upper, step, ValueRange {assembled}, - [&](OpBuilder &, Location, Value actionIndex, ValueRange iterArgs, - SmallVectorImpl &yielded) -> LogicalResult { - Value tableIndex = buildBatchActionTableIndex( - lane, actionIndex, laneCount, exchange.deferred, rewriter, loc); - auto result = emitEntry(tableIndex, iterArgs.front()); - if (failed(result)) - return failure(); - yielded.push_back(*result); - return success(); - }); - if (failed(loop)) - return failure(); - return std::optional(loop->results.front()); -} - -static FailureOr> emitBatchBoundaryClass( - ScheduledBoundaryPlan &boundary, const BatchBoundaryClass &group, - Value lane, DenseConstantCache &denseConstants, IRRewriter &rewriter) { - ReceiveCache receiveCache; - AssemblyCache assemblyCache; - ArrayRef representative = - boundary.actionsByLane[group.lanes.front()]; - for (size_t index = 0; index < representative.size();) { - if (representative[index].kind == BoundaryActionKind::Receive) { - NormalizedDeferredExchange &exchange = - *representative[index].ref.exchange; - auto assembled = tryEmitBatchAnalyzedAssembly( - boundary, group, exchange, index, lane, receiveCache, - denseConstants, rewriter); - if (failed(assembled)) - return failure(); - if (*assembled) { - DeferredSpecialization &specialization = - exchange.specializations[group.lanes.front()]; - assemblyCache[&specialization] = **assembled; - index += specialization.program.insertAssembly->entries.size(); - continue; - } - if (failed(emitBatchReceiveAction( - boundary, group, index, lane, receiveCache, denseConstants, - rewriter))) - return failure(); - ++index; - continue; - } - size_t end = index + 1; - while (end < representative.size() - && representative[end].kind == BoundaryActionKind::Send - && haveCompatibleSendEmission( - representative[index].ref, representative[end].ref, - boundary.insertionStep, boundary.insertionStep)) - ++end; - if (failed(emitBatchSendRun( - boundary, group, index, end - index, lane, denseConstants, - rewriter))) - return failure(); - index = end; - } - - SmallVector results; - for (NormalizedDeferredExchange *exchange : boundary.targetExchanges) { - unsigned referenceLane = group.lanes.front(); - DeferredSpecialization &specialization = - exchange->specializations[referenceLane]; - if (Value assembled = assemblyCache.lookup(&specialization)) { - results.push_back(assembled); - continue; - } - SmallVector available; - available.reserve(specialization.requirements.size()); - for (unsigned requirementIndex = 0; - requirementIndex < specialization.requirements.size(); - ++requirementIndex) { - FragmentTransfer &transfer = - specialization.transfers[requirementIndex]; - if (!transfer.local) { - Value received = receiveCache.lookup(&transfer); - if (!received) - return failure(); - available.push_back(received); - continue; - } - SmallVector requirements; - SmallVector transfers; - for (DeferredSpecialization &laneSpecialization : - exchange->specializations) { - requirements.push_back( - &laneSpecialization.requirements[requirementIndex]); - transfers.push_back( - &laneSpecialization.transfers[requirementIndex]); - } - auto local = materializeBatchLocalFragment( - requirements, lane, exchange->deferred, rewriter, transfers); - if (failed(local)) - return failure(); - available.push_back(*local); - } - auto result = materializeBatchResult( - *exchange, available, group, lane, rewriter); - if (failed(result)) - return failure(); - results.push_back(*result); - } - return results; -} - -static Value buildClassMembershipCondition( - ArrayRef members, unsigned laneCount, Value lane, - Operation *anchor, DenseConstantCache &denseConstants, - IRRewriter &rewriter, Location loc) { - if (members.size() == 1) { - Value selected = getOrCreateIndexConstant( - rewriter, anchor, members.front()); - return arith::CmpIOp::create( - rewriter, loc, arith::CmpIPredicate::eq, lane, selected); - } - SmallVector membership(laneCount); - for (unsigned member : members) - membership[member] = 1; - Value selected = buildLaneLookupCached( - membership, lane, anchor, denseConstants, rewriter, loc); - Value zero = getOrCreateIndexConstant(rewriter, anchor, 0); - return arith::CmpIOp::create( - rewriter, loc, arith::CmpIPredicate::ne, selected, zero); -} - -static FailureOr> tryEmitSimpleParentBoundary( - ScheduledBoundaryPlan &boundary, - ArrayRef classes, Value lane, - DenseConstantCache &denseConstants, IRRewriter &rewriter) { - if (classes.size() <= 2 || boundary.targetExchanges.size() != 1) - return std::optional(); - NormalizedDeferredExchange *exchange = boundary.targetExchanges.front(); - if (llvm::any_of(exchange->specializations, - [](const DeferredSpecialization &specialization) { - return specialization.transfers.size() != 1; - })) - return std::optional(); - - SmallVector remoteLanes, localLanes; - SmallVector preSendLanes, postSendLanes, localSendLanes; - const BoundaryAction *referenceReceive = nullptr; - for (unsigned classLane = 0; - classLane < boundary.actionsByLane.size(); ++classLane) { - ArrayRef actions = boundary.actionsByLane[classLane]; - if (actions.size() > 2) - return std::optional(); - const BoundaryAction *send = nullptr; - const BoundaryAction *receive = nullptr; - for (const BoundaryAction &action : actions) { - if (action.ref.exchange != exchange) - return std::optional(); - const BoundaryAction **slot = action.kind == BoundaryActionKind::Send - ? &send : &receive; - if (*slot) - return std::optional(); - *slot = &action; - } - FragmentTransfer &targetTransfer = - exchange->specializations[classLane].transfers.front(); - if ((receive == nullptr) != targetTransfer.local) - return std::optional(); - if (receive) { - const FragmentRequirement &requirement = - receive->ref.specialization->requirements[ - receive->ref.transfer->requirementIndex]; - if (referenceReceive) { - const FragmentRequirement &reference = - referenceReceive->ref.specialization->requirements[ - referenceReceive->ref.transfer->requirementIndex]; - if (receive->ref.transfer->requirementIndex - != referenceReceive->ref.transfer->requirementIndex - || requirement.leafIndex != reference.leafIndex - || requirement.selectedPosition != reference.selectedPosition - || requirement.publicationFragmentType - != reference.publicationFragmentType) - return std::optional(); - } else { - referenceReceive = receive; - } - } - if (!receive) { - localLanes.push_back(classLane); - if (send) - localSendLanes.push_back(classLane); - continue; - } - remoteLanes.push_back(classLane); - if (send) - (actions.front().kind == BoundaryActionKind::Send - ? preSendLanes : postSendLanes).push_back(classLane); - } - if (remoteLanes.empty()) - return std::optional(); - - auto filteredBoundary = [&](ArrayRef lanes, - std::optional kind, - bool withTarget) { - ScheduledBoundaryPlan filtered; - filtered.scheduled = boundary.scheduled; - filtered.insertionStep = boundary.insertionStep; - filtered.actionsByLane.resize(boundary.actionsByLane.size()); - for (unsigned classLane : lanes) - for (const BoundaryAction &action : - boundary.actionsByLane[classLane]) - if (!kind || action.kind == *kind) - filtered.actionsByLane[classLane].push_back(action); - if (withTarget) - filtered.targetExchanges.push_back(exchange); - return filtered; - }; - auto emitClass = [&](ScheduledBoundaryPlan &filtered, - ArrayRef lanes) -> FailureOr> { - BatchBoundaryClass group; - llvm::append_range(group.lanes, lanes); - return emitBatchBoundaryClass( - filtered, group, lane, denseConstants, rewriter); - }; - auto emitOptionalSends = [&](ArrayRef executingLanes, - ArrayRef sendLanes) -> LogicalResult { - if (sendLanes.empty()) - return success(); - ScheduledBoundaryPlan sends = filteredBoundary( - sendLanes, BoundaryActionKind::Send, false); - auto emit = [&]() -> LogicalResult { - auto results = emitClass(sends, sendLanes); - return failed(results) ? failure() : success(); - }; - if (sendLanes.size() == executingLanes.size()) - return emit(); - Value condition = buildClassMembershipCondition( - sendLanes, boundary.actionsByLane.size(), lane, - boundary.scheduled->op, denseConstants, rewriter, - boundary.scheduled->op->getLoc()); - auto selection = scf::IfOp::create( - rewriter, boundary.scheduled->op->getLoc(), TypeRange {}, condition, - false); - OpBuilder::InsertionGuard guard(rewriter); - Region ®ion = selection.getThenRegion(); - auto yield = dyn_cast_or_null( - region.front().empty() ? nullptr : ®ion.front().back()); - if (yield) - rewriter.setInsertionPoint(yield); - else - rewriter.setInsertionPointToStart(®ion.front()); - if (failed(emit())) - return failure(); - if (!yield) - scf::YieldOp::create(rewriter, boundary.scheduled->op->getLoc()); - return success(); - }; - auto emitRemote = [&]() -> FailureOr { - if (failed(emitOptionalSends(remoteLanes, preSendLanes))) - return failure(); - ScheduledBoundaryPlan receives = filteredBoundary( - remoteLanes, BoundaryActionKind::Receive, true); - auto results = emitClass(receives, remoteLanes); - if (failed(results) || results->size() != 1) - return failure(); - if (failed(emitOptionalSends(remoteLanes, postSendLanes))) - return failure(); - return results->front(); - }; - auto emitLocal = [&]() -> FailureOr { - if (failed(emitOptionalSends(localLanes, localSendLanes))) - return failure(); - ScheduledBoundaryPlan local = filteredBoundary( - localLanes, BoundaryActionKind::Receive, true); - auto results = emitClass(local, localLanes); - if (failed(results) || results->size() != 1) - return failure(); - return results->front(); - }; - if (localLanes.empty()) { - auto result = emitRemote(); - if (failed(result)) - return failure(); - return std::optional(*result); - } - - bool remoteIsThen = remoteLanes.size() <= localLanes.size(); - ArrayRef thenLanes = remoteIsThen - ? ArrayRef(remoteLanes) : ArrayRef(localLanes); - Value condition = buildClassMembershipCondition( - thenLanes, boundary.actionsByLane.size(), lane, - boundary.scheduled->op, denseConstants, rewriter, - boundary.scheduled->op->getLoc()); - auto selection = scf::IfOp::create( - rewriter, boundary.scheduled->op->getLoc(), - TypeRange {exchange->deferred.getOutput().getType()}, condition, true); - auto buildBranch = [&](Region ®ion, auto &&emit) -> LogicalResult { - OpBuilder::InsertionGuard guard(rewriter); - auto yield = dyn_cast_or_null( - region.front().empty() ? nullptr : ®ion.front().back()); - if (yield) - rewriter.setInsertionPoint(yield); - else - rewriter.setInsertionPointToStart(®ion.front()); - auto result = emit(); - if (failed(result)) - return failure(); - if (yield) - yield->setOperands(ValueRange {*result}); - else - scf::YieldOp::create( - rewriter, boundary.scheduled->op->getLoc(), ValueRange {*result}); - return success(); - }; - if (failed(buildBranch(selection.getThenRegion(), [&] { - return remoteIsThen ? emitRemote() : emitLocal(); - })) - || failed(buildBranch(selection.getElseRegion(), [&] { - return remoteIsThen ? emitLocal() : emitRemote(); - }))) - return failure(); - return std::optional(selection.getResult(0)); -} - -static LogicalResult realizeBatchBoundaryPart( - ScheduledBoundaryPlan &boundary, IRRewriter &rewriter, - DenseConstantCache &denseConstants, - DeferredEraseSet &erase) { - auto scheduled = cast(boundary.scheduled->op); - Value lane = *scheduled.getLaneArgument(); - SmallVector classes = classifyBatchBoundary(boundary); - auto simple = tryEmitSimpleParentBoundary( - boundary, classes, lane, denseConstants, rewriter); - if (failed(simple)) - return failure(); - if (*simple) { - replaceDeferred( - boundary.targetExchanges.front()->deferred, **simple, erase); - return success(); - } - SmallVector resultTypes; - for (NormalizedDeferredExchange *exchange : boundary.targetExchanges) - resultTypes.push_back(exchange->deferred.getOutput().getType()); - SmallVector replacements; - if (classes.size() == 1) { - auto results = emitBatchBoundaryClass( - boundary, classes.front(), lane, denseConstants, rewriter); - if (failed(results)) - return failure(); - replacements = std::move(*results); - } else if (classes.size() == 2) { - unsigned thenIndex = classes[0].lanes.size() <= classes[1].lanes.size() - ? 0 : 1; - unsigned elseIndex = 1 - thenIndex; - Value condition; - if (classes[thenIndex].lanes.size() == 1) { - Value selectedLane = getOrCreateIndexConstant( - rewriter, scheduled, classes[thenIndex].lanes.front()); - condition = arith::CmpIOp::create( - rewriter, scheduled.getLoc(), arith::CmpIPredicate::eq, lane, - selectedLane); - } else { - SmallVector membership(boundary.actionsByLane.size()); - for (unsigned member : classes[thenIndex].lanes) - membership[member] = 1; - Value selected = buildLaneLookupCached( - membership, lane, scheduled, denseConstants, rewriter, - scheduled.getLoc()); - Value zero = getOrCreateIndexConstant(rewriter, scheduled, 0); - condition = arith::CmpIOp::create( - rewriter, scheduled.getLoc(), arith::CmpIPredicate::ne, selected, - zero); - } - auto selection = scf::IfOp::create( - rewriter, scheduled.getLoc(), resultTypes, condition, true); - auto buildBranch = [&](Region ®ion, const BatchBoundaryClass &group) - -> LogicalResult { - OpBuilder::InsertionGuard guard(rewriter); - auto existingYield = dyn_cast_or_null( - region.front().empty() ? nullptr : ®ion.front().back()); - if (existingYield) - rewriter.setInsertionPoint(existingYield); - else - rewriter.setInsertionPointToStart(®ion.front()); - auto results = emitBatchBoundaryClass( - boundary, group, lane, denseConstants, rewriter); - if (failed(results)) - return failure(); - if (existingYield) - existingYield->setOperands(*results); - else - scf::YieldOp::create(rewriter, scheduled.getLoc(), *results); - return success(); - }; - if (failed(buildBranch(selection.getThenRegion(), classes[thenIndex])) - || failed(buildBranch(selection.getElseRegion(), classes[elseIndex]))) - return failure(); - replacements.assign(selection.getResults().begin(), - selection.getResults().end()); - } else { - llvm::stable_sort(classes, [](const BatchBoundaryClass &lhs, - const BatchBoundaryClass &rhs) { - return lhs.lanes.front() < rhs.lanes.front(); - }); - SmallVector classByLane(boundary.actionsByLane.size()); - for (auto [classIndex, group] : llvm::enumerate(classes)) - for (unsigned member : group.lanes) - classByLane[member] = classIndex; - Value classId = buildLaneLookupCached( - classByLane, lane, scheduled, denseConstants, rewriter, - scheduled.getLoc()); - SmallVector cases; - for (int64_t classIndex = 0; - classIndex < static_cast(classes.size()) - 1; - ++classIndex) - cases.push_back(classIndex); - auto selection = scf::IndexSwitchOp::create( - rewriter, scheduled.getLoc(), resultTypes, classId, cases, - cases.size()); - auto buildRegion = [&](Region ®ion, const BatchBoundaryClass &group) - -> LogicalResult { - OpBuilder::InsertionGuard guard(rewriter); - Block *block = rewriter.createBlock(®ion); - rewriter.setInsertionPointToEnd(block); - auto results = emitBatchBoundaryClass( - boundary, group, lane, denseConstants, rewriter); - if (failed(results)) - return failure(); - scf::YieldOp::create(rewriter, scheduled.getLoc(), *results); - return success(); - }; - for (auto [region, group] : llvm::zip( - selection.getCaseRegions(), ArrayRef(classes).drop_back())) - if (failed(buildRegion(region, group))) - return failure(); - if (failed(buildRegion(selection.getDefaultRegion(), classes.back()))) - return failure(); - replacements.assign(selection.getResults().begin(), - selection.getResults().end()); - } - - if (replacements.size() != boundary.targetExchanges.size()) - return failure(); - for (auto [exchange, replacement] : - llvm::zip_equal(boundary.targetExchanges, replacements)) - replaceDeferred(exchange->deferred, replacement, erase); - return success(); -} - -static LogicalResult realizeBatchBoundary( - ScheduledBoundaryPlan &boundary, IRRewriter &rewriter, - DenseConstantCache &denseConstants, - DeferredEraseSet &erase) { - setInsertionAtBoundary( - rewriter, *boundary.scheduled, boundary.insertionStep); - if (classifyBatchBoundary(boundary).size() <= 2) - return realizeBatchBoundaryPart( - boundary, rewriter, denseConstants, erase); - - SmallVector> orderedParents; - llvm::SmallDenseSet seenParents; - for (ArrayRef actions : boundary.actionsByLane) - for (const BoundaryAction &action : actions) - if (seenParents.insert(action.ref.exchange->exchangeId).second) - orderedParents.push_back( - {action.globalOrder, action.ref.exchange->exchangeId}); - llvm::sort(orderedParents); - - SmallVector parts; - DenseMap partByParent; - for (auto [globalOrder, parent] : orderedParents) { - (void)globalOrder; - ScheduledBoundaryPlan part; - part.scheduled = boundary.scheduled; - part.insertionStep = boundary.insertionStep; - part.actionsByLane.resize(boundary.actionsByLane.size()); - partByParent[parent] = parts.size(); - parts.push_back(std::move(part)); - } - for (auto [lane, actions] : llvm::enumerate(boundary.actionsByLane)) - for (const BoundaryAction &action : actions) - parts[partByParent.lookup(action.ref.exchange->exchangeId)] - .actionsByLane[lane] - .push_back(action); - for (NormalizedDeferredExchange *exchange : boundary.targetExchanges) { - auto it = partByParent.find(exchange->exchangeId); - if (it == partByParent.end()) { - ScheduledBoundaryPlan part; - part.scheduled = boundary.scheduled; - part.insertionStep = boundary.insertionStep; - part.actionsByLane.resize(boundary.actionsByLane.size()); - partByParent[exchange->exchangeId] = parts.size(); - parts.push_back(std::move(part)); - it = partByParent.find(exchange->exchangeId); - } - parts[it->second].targetExchanges.push_back(exchange); - } - - for (ScheduledBoundaryPlan &part : parts) - if (failed(realizeBatchBoundaryPart( - part, rewriter, denseConstants, erase))) - return failure(); - return success(); -} - -static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { - IRRewriter rewriter(funcOp.getContext()); - auto boundaries = buildScheduledBoundaryPlans(plan); - if (failed(boundaries)) - return funcOp.emitOpError("phase 2 failed to build scheduled boundary plans"); - for (ScheduledInfo& info : plan.scheduled) { - if (info.isBatch()) { - if (failed(linearizeBatch(info, rewriter))) - return failure(); - } - else if (failed(linearizeScalar(info, rewriter))) { - return failure(); - } - } - - DeferredEraseSet erase; - ReceiveCache receiveCache; - AssemblyCache assemblyCache; - IndexConstantCache indexConstants; - DenseConstantCache denseConstants; - for (Operation &op : *getConstantInsertionBlock(funcOp)) { - auto constant = dyn_cast(&op); - if (!constant) - continue; - if (std::optional value = - matchConstantIndexValue(constant.getResult())) - indexConstants.try_emplace(*value, constant.getResult()); - if (auto typed = dyn_cast(constant.getValue())) - denseConstants.try_emplace(typed, constant.getResult()); - } - for (ScheduledBoundaryPlan &boundary : *boundaries) { - if (boundary.scheduled->isBatch()) { - if (failed(realizeBatchBoundary( - boundary, rewriter, denseConstants, erase))) - return boundary.scheduled->op->emitOpError( - "phase 2 failed to realize a batch communication boundary"); - } else if (failed(realizeScalarBoundary( - boundary, receiveCache, assemblyCache, indexConstants, - denseConstants, rewriter, erase))) { - return boundary.scheduled->op->emitOpError( - "phase 2 failed to realize a scalar communication boundary"); - } - } - for (Operation* op : erase) - if (op && op->getBlock() && isa(op)) { - if (!op->use_empty()) - return op->emitOpError("phase 2 cannot erase deferred communication with live uses"); - rewriter.eraseOp(op); - } - - SmallVector blueprints; - funcOp.walk([&](SpatBlueprintOp blueprint) { blueprints.push_back(blueprint); }); - for (SpatBlueprintOp blueprint : blueprints) - if (failed(retargetBlueprintToScheduledPublications(plan, blueprint))) - return failure(); - - for (Operation& op : funcOp.getOps()) { +static LogicalResult replaceFinalGraphPublications( + func::FuncOp funcOp, DeferredTransferPlan &plan) { + for (Operation &op : funcOp.getOps()) { if (!isa(op)) continue; auto graphId = op.getAttrOfType("scheduled.graph_id"); if (!graphId) continue; for (auto [resultIndex, result] : llvm::enumerate(op.getResults())) { - SmallVector externalUses; - for (OpOperand& use : result.getUses()) + SmallVector externalUses; + for (OpOperand &use : result.getUses()) if (!isa(use.getOwner())) externalUses.push_back(&use); if (externalUses.empty()) continue; SmallVector exact; - for (ProducedValue* produced : plan.producedByGraph.lookup(graphId.getInt())) - if (produced->resultIndex == resultIndex && produced->published.getType() == result.getType() + for (ProducedValue *produced : + plan.producedByGraph.lookup(graphId.getInt())) + if (produced->resultIndex == resultIndex + && produced->published.getType() == result.getType() && !llvm::is_contained(exact, produced->published)) exact.push_back(produced->published); - if (exact.size() == 1) { - for (OpOperand* use : externalUses) { - Operation *consumer = use->getOwner(); - Operation *producer = exact.front().getDefiningOp(); - if (consumer->getBlock() == producer->getBlock() && - consumer->isBeforeInBlock(producer)) - consumer->moveAfter(producer); - use->set(exact.front()); - } - continue; + if (exact.size() != 1) + return op.emitOpError( + "phase 2 final publication ownership changed after planning"); + for (OpOperand *use : externalUses) { + Operation *consumer = use->getOwner(); + Operation *producer = exact.front().getDefiningOp(); + if (consumer->getBlock() == producer->getBlock() + && consumer->isBeforeInBlock(producer)) + consumer->moveAfter(producer); + use->set(exact.front()); } - return op.emitOpError("phase 2 final publication ownership changed after planning"); } } + return success(); +} - SmallVector graphOps; - for (Operation& op : funcOp.getOps()) +static LogicalResult eraseOldGraph(func::FuncOp funcOp, + IRRewriter &rewriter) { + SmallVector graphOps; + for (Operation &op : funcOp.getOps()) if (isa(op)) graphOps.push_back(&op); - for (Operation* op : llvm::reverse(graphOps)) { - if (!op->use_empty()) { - InFlightDiagnostic diagnostic = - op->emitOpError("phase 2 cannot erase an old graph compute with an unreconstructed live result"); - for (auto [resultIndex, result] : llvm::enumerate(op->getResults())) - if (!result.use_empty()) { - diagnostic << "; result " << resultIndex << " first used by " << (*result.getUsers().begin())->getName(); - break; - } - return failure(); - } + for (Operation *op : llvm::reverse(graphOps)) { + if (!op->use_empty()) + return op->emitOpError( + "phase 2 cannot erase an old graph compute with live results"); rewriter.eraseOp(op); } SmallVector deadBlueprints; @@ -3495,8 +204,12 @@ static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { }); for (SpatBlueprintOp blueprint : deadBlueprints) rewriter.eraseOp(blueprint); + return success(); +} + +static LogicalResult verifyDominance(func::FuncOp funcOp) { DominanceInfo dominance(funcOp); - WalkResult dominanceResult = funcOp.walk([&](Operation *op) { + WalkResult result = funcOp.walk([&](Operation *op) { for (auto [index, operand] : llvm::enumerate(op->getOperands())) if (!dominance.dominates(operand, op)) { op->emitOpError() << "phase 2 produced non-dominating operand " @@ -3505,32 +218,55 @@ static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { } return WalkResult::advance(); }); - if (dominanceResult.wasInterrupted()) - return failure(); - return success(); + return success(!result.wasInterrupted()); } } // namespace LogicalResult realizeDeferredCommunication(func::FuncOp funcOp) { - RealizationPlan plan; - if (failed(collectScheduled(funcOp, plan))) - return funcOp.emitOpError("phase 2 failed to collect scheduled producers"); - if (failed(normalizeDeferred(funcOp, plan))) - return funcOp.emitOpError("phase 2 failed to normalize deferred projections"); - if (failed(buildAndVerifyPlan(funcOp, plan))) - return funcOp.emitOpError("phase 2 failed to plan fragment communication"); - if (failed(realizePlan(funcOp, plan))) - return failure(); + auto transfers = buildDeferredTransferPlan(funcOp); + if (failed(transfers)) + return funcOp.emitOpError( + "phase 2 failed to build symbolic transfer families"); + auto schedule = scheduleDeferredCommunication(funcOp, *transfers); + if (failed(schedule) + || failed(verifyPlannedCommunicationDeadlockFree( + funcOp, transfers->stepCounts, *schedule))) + return funcOp.emitOpError( + "phase 2 failed to schedule symbolic communication"); + auto boundaries = buildDeferredBoundaryPlan(*transfers, *schedule); + if (failed(boundaries)) + return funcOp.emitOpError( + "phase 2 failed to build sparse boundary programs"); - bool hasDeferred = false; - funcOp.walk([&](SpatDeferredCommunicationOp deferred) { - deferred.emitOpError("phase 2 left an unrealized deferred communication"); - hasDeferred = true; - }); - if (hasDeferred) + IRRewriter rewriter(funcOp.getContext()); + if (failed(linearizeScheduled(*transfers, rewriter))) return failure(); - return success(); + ConstantPool constants(funcOp, rewriter); + DeferredEmissionContext context(rewriter, constants); + DeferredEraseSet erase; + if (failed(realizeDeferredBoundaries( + boundaries->boundaries, boundaries->results, context, erase))) + return failure(); + for (Operation *op : erase) { + if (!op->use_empty()) + return op->emitOpError( + "phase 2 cannot erase deferred communication with live uses"); + rewriter.eraseOp(op); + } + if (failed(retargetDeferredPublications(funcOp, *transfers)) + || failed(replaceFinalGraphPublications(funcOp, *transfers)) + || failed(eraseOldGraph(funcOp, rewriter)) + || failed(verifyDominance(funcOp)) + || failed(verifyRealizedCommunicationDeadlockFree(funcOp, *schedule))) + return failure(); + bool deferredRemains = false; + funcOp.walk([&](SpatDeferredCommunicationOp deferred) { + deferred.emitOpError( + "phase 2 left an unrealized deferred communication"); + deferredRemains = true; + }); + return success(!deferredRemains); } } // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp new file mode 100644 index 0000000..77c6509 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp @@ -0,0 +1,409 @@ +#include "llvm/ADT/MapVector.h" + +#include +#include + +#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 ordered; + SmallVector dependencies; + unsigned unsatisfied = 0; + bool ready = false; + bool scheduled = false; + BoundaryCost cost; + std::tuple originalPriority; + std::optional firstSignature; +}; + +struct StreamProgress { + unsigned completedStep = 0; + SmallVector pendingAtStep; + SmallVector> 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::max() + - lhs.cost.absorbedTransfers, + lhs.cost.lookupEntries, + lhs.originalPriority); + auto right = std::tuple(rhs.cost.instructionCount, + rhs.cost.branchRegions, + std::numeric_limits::max() + - rhs.cost.absorbedTransfers, + rhs.cost.lookupEntries, + rhs.originalPriority); + return left < right; +} + +static bool orderPermutationCycles(SmallVectorImpl& slices) { + if (slices.size() < 2) + return false; + DenseMap edgeBySource; + DenseMap 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 ordered; + SmallVector 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 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::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 signatures; + DenseMap> signatureIdsByHash; + SmallVector> slicesBySignature; + SmallVector 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 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 buildGroups(DeferredTransferPlan& plan) { + SmallVector groups; + for (const std::unique_ptr& exchange : plan.exchanges) { + if (exchange->external.empty()) + continue; + TransferGroup group; + group.exchange = exchange.get(); + DenseMap thresholdByStream; + unsigned minSource = std::numeric_limits::max(); + unsigned minTarget = std::numeric_limits::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(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 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 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 existingTail, + ArrayRef candidate) { + BoundaryCost cost; + std::optional 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 scheduleDeferredCommunication(func::FuncOp funcOp, DeferredTransferPlan& plan) { + SmallVector groups = buildGroups(plan); + SmallVector 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& 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, decltype(compare)>; + Heap ready(compare); + DenseMap> 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(compare); + bucket->push(index); + } + }; + for (auto [index, group] : llvm::enumerate(groups)) + if (group.unsatisfied == 0) + addReady(index); + + std::queue 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 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 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 diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp new file mode 100644 index 0000000..13f4e64 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp @@ -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 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 existingTail, + mlir::ArrayRef candidate); + +mlir::FailureOr scheduleDeferredCommunication(mlir::func::FuncOp funcOp, + DeferredTransferPlan& plan); + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp index 1122882..42c836e 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp @@ -3,30 +3,36 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/Matchers.h" + #include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallSet.h" + #include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" -#include - namespace onnx_mlir::spatial { using namespace mlir; namespace { static FailureOr getSignedInt64(IntegerAttr value) { - return value.getValue().isSignedIntN(64) ? FailureOr(value.getValue().getSExtValue()) - : FailureOr(failure()); + return value.getValue().isSignedIntN(64) + ? FailureOr(value.getValue().getSExtValue()) + : FailureOr(failure()); } -static FailureOr evaluate(Value value, const StaticIndexEnvironment &environment, - llvm::SmallDenseSet &visiting); +static FailureOr evaluate( + Value value, const StaticIndexEnvironment &environment, + llvm::SmallDenseSet &visiting); -static FailureOr evaluateDenseExtract(tensor::ExtractOp extract, - const StaticIndexEnvironment &environment, - llvm::SmallDenseSet &visiting) { +static FailureOr evaluateDenseExtract( + tensor::ExtractOp extract, const StaticIndexEnvironment &environment, + llvm::SmallDenseSet &visiting) { auto constant = extract.getTensor().getDefiningOp(); - auto elements = constant ? dyn_cast(constant.getValue()) : DenseIntElementsAttr(); - auto type = elements ? dyn_cast(elements.getType()) : RankedTensorType(); + auto elements = constant + ? dyn_cast(constant.getValue()) + : DenseIntElementsAttr(); + auto type = elements + ? dyn_cast(elements.getType()) : RankedTensorType(); if (!elements || !type || !type.hasStaticShape() || extract.getIndices().size() != static_cast(type.getRank())) return failure(); @@ -34,104 +40,126 @@ static FailureOr evaluateDenseExtract(tensor::ExtractOp extract, for (auto [index, dim] : llvm::zip(extract.getIndices(), type.getShape())) { auto folded = evaluate(index, environment, visiting); if (failed(folded) || *folded < 0 || *folded >= dim - || llvm::MulOverflow(linear, dim, linear) || llvm::AddOverflow(linear, *folded, linear)) + || llvm::MulOverflow(linear, dim, linear) + || llvm::AddOverflow(linear, *folded, linear)) return failure(); } APInt value = elements.getValues()[linear]; - return value.isSignedIntN(64) ? FailureOr(value.getSExtValue()) - : FailureOr(failure()); + return value.isSignedIntN(64) + ? FailureOr(value.getSExtValue()) + : FailureOr(failure()); } -static FailureOr evaluate(Value value, const StaticIndexEnvironment &environment, - llvm::SmallDenseSet &visiting) { - if (auto it = environment.bindings.find(value); it != environment.bindings.end()) +static FailureOr evaluate( + Value value, const StaticIndexEnvironment &environment, + llvm::SmallDenseSet &visiting) { + if (auto it = environment.bindings.find(value); + it != environment.bindings.end()) return it->second; Attribute constant; if (matchPattern(value, m_Constant(&constant))) if (auto integer = dyn_cast_or_null(constant)) return getSignedInt64(integer); - if (isa(value) || !value.getDefiningOp()) - return failure(); - if (!visiting.insert(value).second) + if (isa(value) || !value.getDefiningOp() + || !visiting.insert(value).second) return failure(); if (auto extract = value.getDefiningOp()) { auto result = evaluateDenseExtract(extract, environment, visiting); visiting.erase(value); return result; } - - Operation *definingOp = value.getDefiningOp(); - if (definingOp->getNumRegions() != 0 || !isPureIndexComputationOp(definingOp)) { + Operation *op = value.getDefiningOp(); + if (op->getNumRegions() != 0 || !isPureIndexComputationOp(op)) { visiting.erase(value); return failure(); } - SmallVector operandConstants; - operandConstants.reserve(definingOp->getNumOperands()); - Builder builder(definingOp->getContext()); - for (Value operand : definingOp->getOperands()) { + SmallVector operands; + Builder builder(op->getContext()); + for (Value operand : op->getOperands()) { auto folded = evaluate(operand, environment, visiting); if (failed(folded)) { visiting.erase(value); return failure(); } - operandConstants.push_back(builder.getIntegerAttr(operand.getType(), *folded)); + operands.push_back(builder.getIntegerAttr(operand.getType(), *folded)); } SmallVector results; - if (failed(definingOp->fold(operandConstants, results)) || results.size() != 1) { + if (failed(op->fold(operands, results)) || results.size() != 1) { visiting.erase(value); return failure(); } FailureOr result = failure(); - if (auto integer = dyn_cast(results.front())) { - if (auto attr = dyn_cast(integer)) - result = getSignedInt64(attr); - } else if (auto foldedValue = dyn_cast(results.front())) { - result = evaluate(foldedValue, environment, visiting); + if (auto attr = dyn_cast(results.front())) { + if (auto integer = dyn_cast(attr)) + result = getSignedInt64(integer); + } else if (auto folded = dyn_cast(results.front())) { + result = evaluate(folded, environment, visiting); } visiting.erase(value); return result; } -static FailureOr> sourceArgument(Value value, SpatDeferredCommunicationOp deferred, - const StaticIndexEnvironment &environment) { - while (auto cast = value.getDefiningOp()) value = cast.getSource(); +static FailureOr> sourceArgument( + Value value, SpatDeferredCommunicationOp deferred, + const StaticIndexEnvironment &environment) { + while (auto cast = value.getDefiningOp()) + value = cast.getSource(); if (auto argument = dyn_cast(value); argument && argument.getOwner() == &deferred.getBody().front() && argument.getArgNumber() < deferred.getSources().size()) return std::optional(argument.getArgNumber()); auto result = dyn_cast(value); - auto selection = result ? dyn_cast(result.getOwner()) : scf::IndexSwitchOp(); - if (!selection || result.getResultNumber() != 0 || selection.getNumResults() != 1) + auto selection = result + ? dyn_cast(result.getOwner()) : scf::IndexSwitchOp(); + if (!selection || result.getResultNumber() != 0 + || selection.getNumResults() != 1) return std::optional(); auto selector = evaluateDeferredIndex(selection.getArg(), environment); if (failed(selector)) return failure(); - Region *selectedRegion = &selection.getDefaultRegion(); - for (auto [caseValue, region] : llvm::zip(selection.getCases(), selection.getCaseRegions())) + Region *selected = &selection.getDefaultRegion(); + for (auto [caseValue, region] : + llvm::zip(selection.getCases(), selection.getCaseRegions())) if (caseValue == *selector) { - selectedRegion = ®ion; + selected = ®ion; break; } - if (!selectedRegion->hasOneBlock()) - return failure(); - Block &block = selectedRegion->front(); - auto yield = dyn_cast(block.getTerminator()); - if (!yield || yield.getResults().size() != 1) - return failure(); - for (Operation &op : block.without_terminator()) - if (!isa(op)) - return failure(); - return sourceArgument(yield.getResults().front(), deferred, environment); + auto yield = selected->hasOneBlock() + ? dyn_cast(selected->front().getTerminator()) + : scf::YieldOp(); + return yield && yield.getResults().size() == 1 + ? sourceArgument(yield.getResults().front(), deferred, environment) + : FailureOr>(failure()); } -static SpatGraphComputeBatch graphBatchOwner(Value value) { - if (auto result = dyn_cast(value)) - return dyn_cast(result.getOwner()); - return {}; +static void collectSourceArguments(Value value, + SpatDeferredCommunicationOp deferred, + llvm::SmallSet &indices) { + while (auto cast = value.getDefiningOp()) + value = cast.getSource(); + if (auto argument = dyn_cast(value)) { + if (argument.getOwner() == &deferred.getBody().front() + && argument.getArgNumber() < deferred.getSources().size()) + indices.insert(argument.getArgNumber()); + return; + } + auto result = dyn_cast(value); + auto selection = result + ? dyn_cast(result.getOwner()) : scf::IndexSwitchOp(); + if (!selection || result.getResultNumber() != 0) + return; + for (Region ®ion : selection.getCaseRegions()) + collectSourceArguments( + cast(region.front().getTerminator()).getResults().front(), + deferred, indices); + collectSourceArguments( + cast(selection.getDefaultRegion().front().getTerminator()) + .getResults().front(), deferred, indices); } -static Value getEnclosingScheduledLane(SpatDeferredCommunicationOp deferred, - SpatScheduledComputeBatch scheduled) { +static Value getEnclosingScheduledLane( + SpatDeferredCommunicationOp deferred, + SpatScheduledComputeBatch scheduled) { Block *block = deferred->getBlock(); while (block && block->getParentOp() != scheduled) { Operation *parent = block->getParentOp(); @@ -140,15 +168,6 @@ static Value getEnclosingScheduledLane(SpatDeferredCommunicationOp deferred, return block && !block->empty() ? block->getArgument(0) : Value(); } -static bool isInsideDeferredLoop(Operation *op, - SpatDeferredCommunicationOp deferred) { - for (Operation *parent = op->getParentOp(); parent && parent != deferred; - parent = parent->getParentOp()) - if (isa(parent)) - return true; - return false; -} - static bool isAllowedStaticIndexExpression( Value value, Value scheduledLane, llvm::SmallDenseSet &visiting) { @@ -164,14 +183,15 @@ static bool isAllowedStaticIndexExpression( bool allowed = op->getNumRegions() == 0 && (isPureIndexComputationOp(op) || isCompileTimeOp(op)) && llvm::all_of(op->getOperands(), [&](Value operand) { - return isAllowedStaticIndexExpression(operand, scheduledLane, - visiting); + return isAllowedStaticIndexExpression( + operand, scheduledLane, visiting); }); visiting.erase(value); return allowed; } -static bool isAllowedStaticIndexExpression(Value value, Value scheduledLane) { +static bool isAllowedStaticIndexExpression(Value value, + Value scheduledLane) { llvm::SmallDenseSet visiting; return isAllowedStaticIndexExpression(value, scheduledLane, visiting); } @@ -185,13 +205,10 @@ static bool originatesFromDeferredSource( return argument.getOwner() == &deferred.getBody().front() && argument.getArgNumber() < deferred.getSources().size(); Operation *op = value.getDefiningOp(); - if (!op) - return false; - if (isa(op) && op->getBlock() == &deferred.getBody().front()) - return true; - return llvm::any_of(op->getOperands(), [&](Value operand) { - return originatesFromDeferredSource(operand, deferred, visited); - }); + return op && (isa(op) + || llvm::any_of(op->getOperands(), [&](Value operand) { + return originatesFromDeferredSource(operand, deferred, visited); + })); } static bool originatesFromDeferredSource( @@ -207,43 +224,29 @@ static LogicalResult verifyCanonicalSourceSelector( || selection.getNumResults() != 1 || !selection.getArg().getType().isIndex() || !isAllowedStaticIndexExpression(selection.getArg(), scheduledLane)) - return selection.emitOpError("is not a canonical deferred source selector"); + return selection.emitOpError( + "is not a canonical deferred source selector"); if (laneCount < 2 - || selection.getCases().size() != static_cast(laneCount - 1) - || selection.getCaseRegions().size() != selection.getCases().size()) + || selection.getCases().size() != static_cast(laneCount - 1)) return selection.emitOpError( "must cover every non-default scheduled lane"); for (auto [index, caseValue] : llvm::enumerate(selection.getCases())) if (caseValue != static_cast(index)) return selection.emitOpError( "must use consecutive scheduled-lane cases starting at zero"); - auto verifyRegion = [&](Region ®ion) -> LogicalResult { - if (!region.hasOneBlock()) - return selection.emitOpError("source-selector region must have one block"); - Block &block = region.front(); - auto yield = dyn_cast(block.getTerminator()); - if (!yield || yield.getResults().size() != 1 - || yield.getResults().front().getType() != selection.getResult(0).getType()) + auto yield = region.hasOneBlock() + ? dyn_cast(region.front().getTerminator()) : scf::YieldOp(); + if (!yield || yield.getResults().size() != 1) return selection.emitOpError( - "source-selector region must yield one exact result type"); - for (Operation &op : block.without_terminator()) + "source-selector region must yield one result"); + for (Operation &op : region.front().without_terminator()) if (!isa(op)) return selection.emitOpError( - "source-selector regions may contain only tensor.cast before scf.yield"); - Value source = yield.getResults().front(); - while (auto cast = source.getDefiningOp()) { - if (cast->getBlock() != &block) - return selection.emitOpError( - "source-selector casts must be local to their region"); - source = cast.getSource(); - } - auto argument = dyn_cast(source); - if (!argument || argument.getOwner() != &deferred.getBody().front() - || argument.getArgNumber() >= deferred.getSources().size()) - return selection.emitOpError( - "source-selector branch must resolve to a deferred source argument"); - return success(); + "source-selector regions may contain only tensor.cast"); + llvm::SmallSet indices; + collectSourceArguments(yield.getResults().front(), deferred, indices); + return success(indices.size() == 1); }; for (Region ®ion : selection.getCaseRegions()) if (failed(verifyRegion(region))) @@ -251,165 +254,93 @@ static LogicalResult verifyCanonicalSourceSelector( return verifyRegion(selection.getDefaultRegion()); } -static bool valueDependsOnAny( - Value value, const llvm::SmallDenseSet &dependencies, - llvm::SmallDenseSet &visited) { - if (dependencies.contains(value)) - return true; - if (!visited.insert(value).second) - return false; - Operation *op = value.getDefiningOp(); - return op && llvm::any_of(op->getOperands(), [&](Value operand) { - return valueDependsOnAny(operand, dependencies, visited); - }); +static bool isInsideDeferredLoop(Operation *op, + SpatDeferredCommunicationOp deferred) { + for (Operation *parent = op->getParentOp(); parent && parent != deferred; + parent = parent->getParentOp()) + if (isa(parent)) + return true; + return false; } -static bool valueDependsOnAny( - Value value, const llvm::SmallDenseSet &dependencies) { - llvm::SmallDenseSet visited; - return valueDependsOnAny(value, dependencies, visited); +static bool haveLeadingUnitDifference(RankedTensorType larger, + RankedTensorType smaller) { + return larger && smaller && larger.hasStaticShape() && smaller.hasStaticShape() + && larger.getRank() == smaller.getRank() + 1 + && larger.getDimSize(0) == 1 + && larger.getShape().drop_front() == smaller.getShape(); } -static LogicalResult specializeDeferredLoopStaticValues( - scf::ForOp loop, SpatDeferredCommunicationOp deferred, - const StaticIndexEnvironment &environment, - SpecializedDeferredProgram &program, - llvm::SmallDenseSet dynamicIndices = {}) { - auto record = [&](Value value, StringRef diagnostic) -> LogicalResult { - auto folded = evaluateDeferredIndex(value, environment); - if (failed(folded)) - return deferred.emitOpError(diagnostic); - program.staticValues.try_emplace(value, *folded); - return success(); - }; - auto step = evaluateDeferredIndex(loop.getStep(), environment); - if (failed(record(loop.getLowerBound(), - "deferred shaping loop lower bound did not specialize")) - || failed(record(loop.getUpperBound(), - "deferred shaping loop upper bound did not specialize")) - || failed(step) || *step <= 0) - return deferred.emitOpError( - "deferred shaping loop requires specialized bounds and a positive step"); - program.staticValues.try_emplace(loop.getStep(), *step); - - dynamicIndices.insert(loop.getInductionVar()); - for (Operation &nested : loop.getBody()->without_terminator()) { - if (auto nestedLoop = dyn_cast(nested)) { - if (failed(specializeDeferredLoopStaticValues( - nestedLoop, deferred, environment, program, dynamicIndices))) - return failure(); - continue; - } - for (Value operand : nested.getOperands()) { - if ((!operand.getType().isIndex() - && !isa(operand.getType())) - || valueDependsOnAny(operand, dynamicIndices)) - continue; - if (failed(record(operand, - "deferred shaping loop captured an index that did not specialize"))) - return failure(); - } - } - return success(); +static std::optional getSourceTransform( + RankedTensorType publication, RankedTensorType inserted) { + if (publication == inserted) + return DeferredAssemblySourceTransform::Identity; + if (haveLeadingUnitDifference(inserted, publication)) + return DeferredAssemblySourceTransform::AddLeadingUnitDimension; + if (haveLeadingUnitDifference(publication, inserted)) + return DeferredAssemblySourceTransform::RemoveLeadingUnitDimension; + return std::nullopt; } -static FailureOr> -analyzeDeferredInsertAssembly( - const SpecializedDeferredProgram &program, - const StaticIndexEnvironment &environment) { +static FailureOr> +analyzeInsertAssembly(const DeferredProgramTemplate &program) { auto finalInsert = program.yieldedValue.getDefiningOp(); if (!finalInsert || program.leaves.empty()) - return std::optional(); - + return std::optional(); DenseMap leafByRoot; - for (auto [leafIndex, leaf] : llvm::enumerate(program.leaves)) { - if (leaf.physicalSlots.size() != 1 - || !leafByRoot.try_emplace(leaf.replacementRoot, leafIndex).second) - return std::optional(); - } + for (auto [index, leaf] : llvm::enumerate(program.leaves)) + if (!leafByRoot.try_emplace(leaf.replacementRoot, index).second) + return std::optional(); SmallPtrSet consumed; - SmallVector reverseEntries; + SmallVector reverseEntries; llvm::SmallDenseSet insertedLeaves; Value current = program.yieldedValue; while (auto insert = current.getDefiningOp()) { Value source = insert.getSource(); - SmallVector rankShaping; - while (Operation *shape = source.getDefiningOp()) { - if (auto collapse = dyn_cast(shape)) - source = collapse.getSrc(); - else if (auto expand = dyn_cast(shape)) - source = expand.getSrc(); - else - break; - rankShaping.push_back(shape); - } - auto leafIt = leafByRoot.find(source); - if (leafIt == leafByRoot.end() - || !insertedLeaves.insert(leafIt->second).second) - return std::optional(); - - DeferredInsertAssemblyEntry entry; - entry.leafIndex = leafIt->second; - entry.requirementIndex = leafIt->second; - auto foldGeometry = [&](ArrayRef values, - SmallVectorImpl &folded) { - for (OpFoldResult value : values) { - auto result = evaluateDeferredIndex(value, environment); - if (failed(result)) - return failure(); - folded.push_back(*result); - } - return success(); - }; - if (failed(foldGeometry(insert.getMixedOffsets(), - entry.targetGeometry.offsets)) - || failed(foldGeometry(insert.getMixedSizes(), - entry.targetGeometry.sizes)) - || failed(foldGeometry(insert.getMixedStrides(), - entry.targetGeometry.strides))) - return std::optional(); + Operation *shape = source.getDefiningOp(); + if (isa_and_nonnull(shape)) + source = shape->getOperand(0); + if (source.getDefiningOp() + || source.getDefiningOp()) + return std::optional(); + auto leaf = leafByRoot.find(source); + if (leaf == leafByRoot.end() + || !insertedLeaves.insert(leaf->second).second) + return std::optional(); + auto publicationType = dyn_cast(source.getType()); + auto insertedType = dyn_cast(insert.getSourceType()); + auto transform = getSourceTransform(publicationType, insertedType); + if (!transform) + return std::optional(); + DeferredInsertAssemblyEntryTemplate entry; + entry.coordinate = {leaf->second, 0}; + entry.sourceTransform = *transform; + entry.sourceType = insertedType; + entry.targetGeometry = {SmallVector(insert.getMixedOffsets()), + SmallVector(insert.getMixedSizes()), + SmallVector(insert.getMixedStrides())}; reverseEntries.push_back(std::move(entry)); consumed.insert(insert); - consumed.insert(rankShaping.begin(), rankShaping.end()); + if (shape) + consumed.insert(shape); current = insert.getDest(); } - auto initial = current.getDefiningOp(); auto resultType = dyn_cast(program.yieldedValue.getType()); if (!initial || !initial.getDynamicSizes().empty() || !resultType - || !resultType.hasStaticShape() || initial.getType() != resultType + || initial.getType() != resultType || insertedLeaves.size() != program.leaves.size()) - return std::optional(); + return std::optional(); consumed.insert(initial); - - for (const DeferredInsertAssemblyEntry &entry : reverseEntries) { - const StaticSliceGeometry &geometry = entry.targetGeometry; - if (geometry.offsets.size() != static_cast(resultType.getRank()) - || geometry.sizes.size() != geometry.offsets.size() - || geometry.strides.size() != geometry.offsets.size()) - return std::optional(); - for (auto [offset, size, stride, dim] : - llvm::zip_equal(geometry.offsets, geometry.sizes, - geometry.strides, resultType.getShape())) { - int64_t span = 0; - int64_t last = 0; - if (offset < 0 || size <= 0 || stride <= 0 - || llvm::MulOverflow(size - 1, stride, span) - || llvm::AddOverflow(offset, span, last) || last >= dim) - return std::optional(); - } - } if (llvm::any_of(program.residualOps, - [&](Operation *op) { return !consumed.contains(op); }) - || consumed.size() != program.residualOps.size()) - return std::optional(); - - DeferredInsertAssembly assembly; + [&](Operation *op) { return !consumed.contains(op); })) + return std::optional(); + DeferredInsertAssemblyTemplate assembly; assembly.initialValue = initial; assembly.resultType = resultType; assembly.entries.assign(reverseEntries.rbegin(), reverseEntries.rend()); - return std::optional(std::move(assembly)); + return std::optional(std::move(assembly)); } } // namespace @@ -417,72 +348,38 @@ analyzeDeferredInsertAssembly( LogicalResult verifyDeferredProgramContract( SpatDeferredCommunicationOp deferred) { if (!deferred.getBody().hasOneBlock()) - return deferred.emitOpError("deferred program must have exactly one body block"); + return deferred.emitOpError( + "deferred program must have exactly one body block"); Block &body = deferred.getBody().front(); auto terminator = dyn_cast(body.getTerminator()); if (!terminator || terminator.getOutputs().size() != 1) return deferred.emitOpError( "deferred program must have exactly one yielded value"); - Value scheduledLane; int64_t laneCount = 1; if (auto scheduled = deferred->getParentOfType()) { scheduledLane = getEnclosingScheduledLane(deferred, scheduled); laneCount = scheduled.getLaneCount(); - if (!scheduledLane || laneCount <= 0) - return deferred.emitOpError( - "deferred program cannot locate a valid enclosing scheduled lane"); } - bool invalid = false; - WalkResult result = deferred.getBody().walk([&](Operation *op) { - if (invalid) - return WalkResult::interrupt(); + deferred.getBody().walk([&](Operation *op) { auto reject = [&](StringRef message) { op->emitOpError(message); invalid = true; - return WalkResult::interrupt(); }; - - if (auto yield = dyn_cast(op)) { - if (op != body.getTerminator()) - return reject("must be the unique top-level deferred terminator"); - return WalkResult::advance(); - } - if (auto yield = dyn_cast(op)) { - Operation *parent = op->getParentOp(); - if (!parent || !isa(parent) - || op != op->getBlock()->getTerminator()) - return reject("is not a valid deferred control-flow terminator"); - return WalkResult::advance(); - } + if (invalid || isa(op)) + return; if (auto selection = dyn_cast(op)) { if (failed(verifyCanonicalSourceSelector( - selection, deferred, scheduledLane, laneCount))) { + selection, deferred, scheduledLane, laneCount))) invalid = true; - return WalkResult::interrupt(); - } - return WalkResult::advance(); + return; } if (auto loop = dyn_cast(op)) { - auto isStaticRankedTensor = [](Type type) { - auto tensor = dyn_cast(type); - return tensor && tensor.hasStaticShape(); - }; - if (!llvm::all_of(loop.getInitArgs(), [&](Value value) { - return isStaticRankedTensor(value.getType()); - }) - || !llvm::all_of(loop.getResultTypes(), isStaticRankedTensor)) - return reject( - "requires static ranked tensor iter arguments and results"); - auto yield = dyn_cast(loop.getBody()->getTerminator()); - if (!yield || yield.getResults().getTypes() != loop.getResultTypes()) - return reject("yield types must exactly match loop result types"); for (Value bound : {loop.getLowerBound(), loop.getUpperBound(), loop.getStep()}) if (!isAllowedStaticIndexExpression(bound, scheduledLane)) - return reject( - "bounds must use only constants, the scheduled lane, and pure index computation"); + return reject("has a non-static deferred loop bound"); for (int64_t lane = 0; lane < laneCount; ++lane) { StaticIndexEnvironment environment; if (scheduledLane) @@ -491,222 +388,250 @@ LogicalResult verifyDeferredProgramContract( auto upper = evaluateDeferredIndex(loop.getUpperBound(), environment); auto step = evaluateDeferredIndex(loop.getStep(), environment); if (failed(lower) || failed(upper) || failed(step) || *step <= 0) - return reject( - "bounds must specialize for every scheduled lane with a positive step"); + return reject("has a loop bound that does not specialize"); } - return WalkResult::advance(); + return; } if (op->getNumRegions() != 0) - return reject("unsupported region operation in deferred program"); + return reject("contains an unsupported region operation"); if (!isShapingOnlyOp(op) && !isCompileTimeOp(op) && !isPureIndexComputationOp(op)) - return reject( - "deferred program permits only shaping, compile-time, or pure index operations"); - - for (Value operand : op->getOperands()) { - if (auto argument = dyn_cast(operand)) { - Operation *owner = argument.getOwner()->getParentOp(); - bool local = owner == deferred - || (owner && deferred->isAncestor(owner)); - if (!local && operand != scheduledLane) - return reject("captures an unsupported external block argument"); - } + return reject("contains an unsupported deferred operation"); + for (Value operand : op->getOperands()) if (isInsideDeferredLoop(op, deferred) && originatesFromDeferredSource(operand, deferred)) - return reject( - "deferred source projection must remain outside residual loops"); - } - return WalkResult::advance(); + return reject("projects a deferred source inside a residual loop"); }); - return success(!invalid && !result.wasInterrupted()); + return success(!invalid); } -FailureOr evaluateDeferredIndex(Value value, const StaticIndexEnvironment &environment) { +FailureOr evaluateDeferredIndex( + Value value, const StaticIndexEnvironment &environment) { llvm::SmallDenseSet visiting; return evaluate(value, environment, visiting); } -FailureOr evaluateDeferredIndex(OpFoldResult value, const StaticIndexEnvironment &environment) { + +FailureOr evaluateDeferredIndex( + OpFoldResult value, const StaticIndexEnvironment &environment) { if (auto attr = dyn_cast(value)) - if (auto integer = dyn_cast(attr)) return getSignedInt64(integer); - if (auto dynamic = dyn_cast(value)) return evaluateDeferredIndex(dynamic, environment); + if (auto integer = dyn_cast(attr)) + return getSignedInt64(integer); + if (auto dynamic = dyn_cast(value)) + return evaluateDeferredIndex(dynamic, environment); return failure(); } -FailureOr> tryResolveDeferredSource(Value value, SpatDeferredCommunicationOp deferred, - const StaticIndexEnvironment &environment) { - auto index = sourceArgument(value, deferred, environment); - if (failed(index)) +FailureOr> getPossibleDeferredSourceOperandIndices( + Value sourceRoot, SpatDeferredCommunicationOp deferred) { + llvm::SmallSet indices; + collectSourceArguments(sourceRoot, deferred, indices); + if (indices.empty()) return failure(); - if (!*index) - return std::optional(); - return std::optional(ResolvedDeferredSource {**index, deferred.getSources()[**index]}); + return SmallVector(indices.begin(), indices.end()); } -FailureOr requireResolvedDeferredSource(Value value, SpatDeferredCommunicationOp deferred, - const StaticIndexEnvironment &environment) { - auto source = tryResolveDeferredSource(value, deferred, environment); - if (failed(source) || !*source) - return deferred.emitOpError("cannot statically resolve deferred source selection"), failure(); - return **source; +DeferredLaneValueEvaluator::DeferredLaneValueEvaluator( + const DeferredProgramTemplate &program, unsigned laneCount) + : program(program), laneCount(laneCount) {} + +FailureOr DeferredLaneValueEvaluator::evaluate(Value value) { + if (auto it = values.find(value); it != values.end()) + return it->second; + if (value == program.scheduledLane) { + StaticIntSequence result = StaticIntSequence::affine(0, 1, laneCount); + values.try_emplace(value, result); + return result; + } + SmallVector evaluated; + evaluated.reserve(laneCount); + for (unsigned lane = 0; lane < laneCount; ++lane) { + StaticIndexEnvironment environment; + if (program.scheduledLane) + environment.bindings[program.scheduledLane] = lane; + auto result = evaluateDeferredIndex(value, environment); + if (failed(result)) + return failure(); + evaluated.push_back(*result); + } + StaticIntSequence result = StaticIntSequence::fromValues(evaluated); + values.try_emplace(value, result); + return result; } -FailureOr analyzeDeferredProgram(SpatDeferredCommunicationOp deferred, - std::optional targetScheduledLane) { - // The Phase 1 verifier must establish the deferred-program contract once; - // this analysis only specializes lane-dependent static semantics. +FailureOr DeferredLaneValueEvaluator::evaluate( + OpFoldResult value) { + if (auto attr = dyn_cast(value)) { + auto integer = dyn_cast(attr); + if (!integer) + return failure(); + auto number = getSignedInt64(integer); + return succeeded(number) + ? FailureOr( + StaticIntSequence::uniform(*number, laneCount)) + : FailureOr(failure()); + } + return evaluate(cast(value)); +} + +FailureOr +DeferredLaneValueEvaluator::resolveSourceOperandIndices(Value sourceRoot) { + if (auto it = sourceOperands.find(sourceRoot); it != sourceOperands.end()) + return it->second; + SmallVector indices; + indices.reserve(laneCount); + for (unsigned lane = 0; lane < laneCount; ++lane) { + StaticIndexEnvironment environment; + if (program.scheduledLane) + environment.bindings[program.scheduledLane] = lane; + auto index = sourceArgument(sourceRoot, program.deferred, environment); + if (failed(index) || !*index) + return failure(); + indices.push_back(**index); + } + StaticIntSequence result = StaticIntSequence::fromValues(indices); + sourceOperands.try_emplace(sourceRoot, result); + return result; +} + +FailureOr analyzeDeferredProgramTemplate( + SpatDeferredCommunicationOp deferred) { Block &body = deferred.getBody().front(); auto yield = dyn_cast(body.getTerminator()); if (!yield || yield.getOutputs().size() != 1) - return deferred.emitOpError("requires exactly one deferred yielded value"), failure(); - StaticIndexEnvironment environment; - if (auto scheduled = deferred->getParentOfType()) { - if (!targetScheduledLane) return deferred.emitOpError("scheduled-batch deferred program requires lane specialization"), failure(); - Value scheduledLane = getEnclosingScheduledLane(deferred, scheduled); - if (!scheduledLane) - return deferred.emitOpError("cannot locate the enclosing scheduled lane"), failure(); - environment.bindings[scheduledLane] = *targetScheduledLane; - } else if (targetScheduledLane) return deferred.emitOpError("scalar deferred program cannot have a target lane"), failure(); - SpecializedDeferredProgram program; + return deferred.emitOpError( + "requires one deferred yielded value"), failure(); + DeferredProgramTemplate program; program.deferred = deferred; - program.targetScheduledLane = targetScheduledLane; + program.yieldedValue = yield.getOutputs().front(); if (auto scheduled = deferred->getParentOfType()) program.scheduledLane = getEnclosingScheduledLane(deferred, scheduled); - program.yieldedValue = yield.getOutputs().front(); + llvm::SmallDenseSet visited; std::function visit = [&](Value value) -> LogicalResult { - if (!visited.insert(value).second) return success(); - // A graph-batch projection is semantically different from the canonical - // source-selector scaffold even though both contain extract/collapse ops. + if (!visited.insert(value).second) + return success(); if (auto slice = value.getDefiningOp()) { - auto source = tryResolveDeferredSource(slice.getSource(), deferred, environment); - if (failed(source)) return failure(); - if (*source && graphBatchOwner((*source)->selectedValue)) { - auto type = dyn_cast((*source)->selectedValue.getType()); - auto result = dyn_cast(slice.getResult().getType()); - if (!type || !result || type.getRank() == 0) return deferred.emitOpError("graph projection requires ranked tensors"); - auto offset = evaluateDeferredIndex(slice.getMixedOffsets().front(), environment); - if (failed(offset)) return deferred.emitOpError("graph projection leading offset is not statically resolvable"); - auto size = evaluateDeferredIndex(slice.getMixedSizes().front(), environment); - if (failed(size)) return deferred.emitOpError("graph projection leading size is not statically resolvable"); - auto stride = evaluateDeferredIndex(slice.getMixedStrides().front(), environment); - if (failed(stride)) return deferred.emitOpError("graph projection leading stride is not statically resolvable"); - if (*offset < 0) return deferred.emitOpError("graph projection leading offset is negative"); - if (*size <= 0) return deferred.emitOpError("graph projection leading size must be positive"); - if (*stride <= 0) return deferred.emitOpError("graph projection leading stride must be positive"); - DeferredProjectionLeaf leaf; - leaf.kind = DeferredLeafKind::GraphBatchProjection; leaf.sourceOperandIndex = (*source)->sourceOperandIndex; - leaf.replacementRoot = value; leaf.leadingProjection = slice; leaf.reconstructedType = result; - for (int64_t i = 0; i < *size; ++i) { - int64_t slot; - if (llvm::MulOverflow(i, *stride, slot) || llvm::AddOverflow(*offset, slot, slot) - || slot >= type.getDimSize(0)) - return deferred.emitOpError("graph projection selects a physical slot outside the batch"); - leaf.physicalSlots.push_back(slot); - } - for (unsigned i = 1; i < type.getRank(); ++i) { - auto innerOffset = evaluateDeferredIndex(slice.getMixedOffsets()[i], environment); - auto innerSize = evaluateDeferredIndex(slice.getMixedSizes()[i], environment); - auto innerStride = evaluateDeferredIndex(slice.getMixedStrides()[i], environment); - if (failed(innerOffset) || failed(innerSize) || failed(innerStride)) - return deferred.emitOpError("graph projection has unresolved inner geometry"); - leaf.innerGeometry.offsets.push_back(*innerOffset); leaf.innerGeometry.sizes.push_back(*innerSize); leaf.innerGeometry.strides.push_back(*innerStride); - } - program.leaves.push_back(std::move(leaf)); return success(); + auto sources = getPossibleDeferredSourceOperandIndices( + slice.getSource(), deferred); + bool graphProjection = succeeded(sources) + && llvm::all_of(*sources, [&](unsigned index) { + auto result = dyn_cast(deferred.getSources()[index]); + return result && isa(result.getOwner()); + }); + if (graphProjection) { + DeferredProjectionLeafTemplate leaf; + leaf.form = DeferredLeafForm::GraphBatchProjection; + leaf.sourceRoot = slice.getSource(); + leaf.replacementRoot = value; + leaf.leadingProjection = slice; + leaf.leadingGeometry = { + SmallVector(slice.getMixedOffsets()), + SmallVector(slice.getMixedSizes()), + SmallVector(slice.getMixedStrides())}; + leaf.innerGeometry = { + SmallVector( + ArrayRef(slice.getMixedOffsets()).drop_front()), + SmallVector( + ArrayRef(slice.getMixedSizes()).drop_front()), + SmallVector( + ArrayRef(slice.getMixedStrides()).drop_front())}; + leaf.reconstructedType = cast(value.getType()); + program.leaves.push_back(std::move(leaf)); + return success(); } } - - auto source = tryResolveDeferredSource(value, deferred, environment); - if (failed(source)) return failure(); - if (*source) { - auto type = dyn_cast((*source)->selectedValue.getType()); - if (!type) return deferred.emitOpError("deferred source is not a ranked tensor"); - DeferredProjectionLeaf leaf; - leaf.sourceOperandIndex = (*source)->sourceOperandIndex; - leaf.replacementRoot = value; - leaf.reconstructedType = type; - if (graphBatchOwner((*source)->selectedValue)) { - leaf.kind = DeferredLeafKind::GraphBatchIdentity; - for (int64_t slot = 0; slot < type.getDimSize(0); ++slot) leaf.physicalSlots.push_back(slot); - } else leaf.kind = DeferredLeafKind::ScalarSource; - program.leaves.push_back(std::move(leaf)); + if (succeeded(getPossibleDeferredSourceOperandIndices(value, deferred))) { + auto type = dyn_cast(value.getType()); + if (!type) + return deferred.emitOpError( + "deferred source is not a ranked tensor"); + program.leaves.push_back({DeferredLeafForm::DirectSource, value, value, + {}, {}, {}, type}); return success(); } - if (value.getType().isIndex() || isa(value.getType())) { - auto folded = evaluateDeferredIndex(value, environment); - if (failed(folded)) - return deferred.emitOpError("deferred index expression is not statically evaluable after specialization"); - program.staticValues.try_emplace(value, *folded); + if (value.getType().isIndex() || isa(value.getType())) return success(); - } Operation *op = value.getDefiningOp(); - if (!op || op->getBlock() != &body) - return deferred.emitOpError("deferred residual escapes its verified body"); + if (!op || (op->getBlock() != &body && !isa(op))) + return deferred.emitOpError( + "deferred residual escapes its verified body"); if (auto loop = dyn_cast(op)) { - if (failed(specializeDeferredLoopStaticValues( - loop, deferred, environment, program))) - return failure(); for (Value init : loop.getInitArgs()) if (failed(visit(init))) return failure(); - program.residualOps.push_back(op); - return success(); + } else { + for (Value operand : op->getOperands()) + if (failed(visit(operand))) + return failure(); } - for (Value operand : op->getOperands()) if (failed(visit(operand))) return failure(); - program.residualOps.push_back(op); return success(); + program.residualOps.push_back(op); + return success(); }; - if (failed(visit(program.yieldedValue))) return failure(); - StaticIndexEnvironment assemblyEnvironment = environment; - for (auto [value, staticValue] : program.staticValues) - assemblyEnvironment.bindings.try_emplace(value, staticValue); - auto assembly = analyzeDeferredInsertAssembly(program, assemblyEnvironment); + if (failed(visit(program.yieldedValue))) + return failure(); + auto assembly = analyzeInsertAssembly(program); if (failed(assembly)) return failure(); program.insertAssembly = std::move(*assembly); - return std::move(program); + return program; } FailureOr getGraphBatchPublicationMap( - SpatGraphComputeBatch graphBatch, unsigned resultIndex, GraphBatchPublicationCache &cache) { + SpatGraphComputeBatch graphBatch, unsigned resultIndex, + GraphBatchPublicationCache &cache) { GraphBatchPublicationKey key {graphBatch.getOperation(), resultIndex}; - if (auto it = cache.find(key); it != cache.end()) return &it->second; - auto resultType = dyn_cast(graphBatch.getResult(resultIndex).getType()); + if (auto it = cache.find(key); it != cache.end()) + return &it->second; + auto resultType = dyn_cast( + graphBatch.getResult(resultIndex).getType()); auto output = graphBatch.getOutputArgument(resultIndex); auto lane = graphBatch.getLaneArgument(); if (!resultType || !output || !lane || resultType.getRank() == 0) - return graphBatch.emitOpError("graph batch publication is malformed"), failure(); + return graphBatch.emitOpError( + "graph batch publication is malformed"), failure(); tensor::ParallelInsertSliceOp publication; - auto parallel = dyn_cast(graphBatch.getBody().front().getTerminator()); - if (!parallel) return graphBatch.emitOpError("graph batch lacks publication region"), failure(); - for (Operation &op : parallel.getRegion().front()) if (auto insert = dyn_cast(op); insert && insert.getDest() == *output) { - if (publication) return graphBatch.emitOpError("graph result has multiple publications"), failure(); - publication = insert; - } - if (!publication) return graphBatch.emitOpError("graph result lacks publication"), failure(); - auto fragment = dyn_cast(publication.getSource().getType()); - if (!fragment || resultType.getRank() != fragment.getRank() + 1) return graphBatch.emitOpError("graph publication fragment type is invalid"), failure(); + auto parallel = dyn_cast( + graphBatch.getBody().front().getTerminator()); + if (!parallel) + return graphBatch.emitOpError( + "graph batch lacks publication region"), failure(); + for (Operation &op : parallel.getRegion().front()) + if (auto insert = dyn_cast(op); + insert && insert.getDest() == *output) + publication = insert; + auto fragment = publication + ? dyn_cast(publication.getSource().getType()) + : RankedTensorType(); + if (!publication || !fragment + || resultType.getRank() != fragment.getRank() + 1) + return graphBatch.emitOpError( + "graph publication fragment type is invalid"), failure(); GraphBatchPublicationMap map; map.physicalResultType = resultType; map.publicationFragmentType = fragment; map.graphLaneToPhysicalSlot.resize(graphBatch.getLaneCount(), -1); map.physicalSlotToGraphLane.resize(resultType.getDimSize(0), -1); for (int64_t index = 0; index < graphBatch.getLaneCount(); ++index) { - StaticIndexEnvironment environment; environment.bindings[*lane] = index; - auto slot = evaluateDeferredIndex(publication.getMixedOffsets().front(), environment); - auto size = evaluateDeferredIndex(publication.getMixedSizes().front(), environment); - auto stride = evaluateDeferredIndex(publication.getMixedStrides().front(), environment); - if (failed(slot) || failed(size) || failed(stride) || *size != 1 || *stride != 1 || *slot < 0 || *slot >= resultType.getDimSize(0)) - return graphBatch.emitOpError("graph publication leading geometry is invalid"), failure(); - for (unsigned dim = 1; dim < resultType.getRank(); ++dim) { - auto offset = evaluateDeferredIndex(publication.getMixedOffsets()[dim], environment); - auto extent = evaluateDeferredIndex(publication.getMixedSizes()[dim], environment); - auto step = evaluateDeferredIndex(publication.getMixedStrides()[dim], environment); - if (failed(offset) || failed(extent) || failed(step) || *offset != 0 || *extent != fragment.getDimSize(dim - 1) || *step != 1) - return graphBatch.emitOpError("graph publication inner geometry is invalid"), failure(); - } - if (map.physicalSlotToGraphLane[*slot] != -1) return graphBatch.emitOpError("graph publication has duplicate physical slot"), failure(); - map.graphLaneToPhysicalSlot[index] = *slot; map.physicalSlotToGraphLane[*slot] = index; + StaticIndexEnvironment environment; + environment.bindings[*lane] = index; + auto slot = evaluateDeferredIndex( + publication.getMixedOffsets().front(), environment); + auto size = evaluateDeferredIndex( + publication.getMixedSizes().front(), environment); + auto stride = evaluateDeferredIndex( + publication.getMixedStrides().front(), environment); + if (failed(slot) || failed(size) || failed(stride) + || *size != 1 || *stride != 1 || *slot < 0 + || *slot >= resultType.getDimSize(0) + || map.physicalSlotToGraphLane[*slot] != -1) + return graphBatch.emitOpError( + "graph publication leading geometry is invalid"), failure(); + map.graphLaneToPhysicalSlot[index] = *slot; + map.physicalSlotToGraphLane[*slot] = index; } - if (llvm::is_contained(map.physicalSlotToGraphLane, -1)) return graphBatch.emitOpError("graph publication has missing physical slot"), failure(); + if (llvm::is_contained(map.physicalSlotToGraphLane, -1)) + return graphBatch.emitOpError( + "graph publication has a missing physical slot"), failure(); return &cache.try_emplace(key, std::move(map)).first->second; } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp index 0fc418b..0513771 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp @@ -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 evaluateDeferredIndex( mlir::FailureOr evaluateDeferredIndex( mlir::OpFoldResult value, const StaticIndexEnvironment &environment); -enum class DeferredLeafKind { ScalarSource, GraphBatchProjection, GraphBatchIdentity }; - -struct StaticSliceGeometry { - llvm::SmallVector offsets; - llvm::SmallVector sizes; - llvm::SmallVector strides; -}; - -struct DeferredProjectionLeaf { - DeferredLeafKind kind = DeferredLeafKind::ScalarSource; - unsigned sourceOperandIndex = 0; - mlir::Value replacementRoot; - mlir::tensor::ExtractSliceOp leadingProjection; - llvm::SmallVector 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 entries; -}; - -struct SpecializedDeferredProgram { - SpatDeferredCommunicationOp deferred; - std::optional targetScheduledLane; - mlir::Value scheduledLane; - mlir::Value yieldedValue; - llvm::SmallVector leaves; - llvm::SmallVector residualOps; - llvm::DenseMap staticValues; - std::optional insertAssembly; -}; - -struct ResolvedDeferredSource { - unsigned sourceOperandIndex = 0; - mlir::Value selectedValue; -}; - -mlir::FailureOr> tryResolveDeferredSource( - mlir::Value value, SpatDeferredCommunicationOp deferred, - const StaticIndexEnvironment &environment); -mlir::FailureOr requireResolvedDeferredSource( - mlir::Value value, SpatDeferredCommunicationOp deferred, - const StaticIndexEnvironment &environment); mlir::LogicalResult verifyDeferredProgramContract( SpatDeferredCommunicationOp deferred); -mlir::FailureOr analyzeDeferredProgram( - SpatDeferredCommunicationOp deferred, - std::optional targetScheduledLane); + +mlir::FailureOr analyzeDeferredProgramTemplate( + SpatDeferredCommunicationOp deferred); + +class DeferredLaneValueEvaluator { +public: + DeferredLaneValueEvaluator(const DeferredProgramTemplate &program, + unsigned laneCount); + + mlir::FailureOr evaluate(mlir::Value value); + mlir::FailureOr evaluate(mlir::OpFoldResult value); + mlir::FailureOr resolveSourceOperandIndices( + mlir::Value sourceRoot); + +private: + const DeferredProgramTemplate &program; + unsigned laneCount; + llvm::DenseMap values; + llvm::DenseMap sourceOperands; +}; + +mlir::FailureOr> +getPossibleDeferredSourceOperandIndices( + mlir::Value sourceRoot, SpatDeferredCommunicationOp deferred); struct GraphBatchPublicationMap { mlir::RankedTensorType physicalResultType; @@ -107,11 +74,13 @@ template <> struct DenseMapInfo { static onnx_mlir::spatial::GraphBatchPublicationKey getTombstoneKey() { return {DenseMapInfo::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; } }; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp new file mode 100644 index 0000000..e51b4f5 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp @@ -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(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(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 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(fragment.getType()); + if (!fragmentType || geometry.offsets.size() != static_cast(fragmentType.getRank())) + return failure(); + if (isIdentityGeometry(geometry, fragmentType)) + return fragment; + SmallVector 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 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(failure()); + if (failed(source)) + return failure(); + auto sourceType = cast(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 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 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(inputs.front().getType()); + SmallVector shape(first.getShape()); + shape.front() = 0; + for (Value input : inputs) + shape.front() += cast(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(expanded.front()) + : FailureOr(failure()); +} + +} // namespace + +FailureOr +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 values, + SmallVectorImpl& 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(operand.getType())) + || result.residualValues.count(operand)) + continue; + if (auto sequence = evaluator.evaluate(operand); succeeded(sequence)) + result.residualValues.try_emplace(operand, *sequence); + } + }); + return result; +} + +FailureOr 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(payload.getType()); + auto fragmentType = dyn_cast(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 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 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(result) + : FailureOr(failure()); +} + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.hpp new file mode 100644 index 0000000..8946cf2 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include "DeferredCommunicationModel.hpp" + +namespace onnx_mlir::spatial { + +struct DeferredEmissionContext; + +struct DeferredResultPlan { + DeferredExchangePlan* exchange = nullptr; + llvm::SmallVector requirements; + struct SliceGeometry { + llvm::SmallVector offsets; + llvm::SmallVector sizes; + llvm::SmallVector strides; + }; + llvm::SmallVector innerGeometry; + llvm::SmallVector assemblyGeometry; + llvm::DenseMap residualValues; +}; + +mlir::FailureOr +buildDeferredResultPlan(DeferredExchangePlan& exchange); + +mlir::FailureOr realizeDeferredResult(const DeferredResultPlan& plan, + const LaneSet& activeLanes, + mlir::Value lane, + DeferredEmissionContext& context); + +mlir::FailureOr materializeDeferredRequirement(RequirementFamily& requirement, + const LaneSet& activeLanes, + mlir::Value lane, + DeferredEmissionContext& context); + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp new file mode 100644 index 0000000..fc59f26 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp @@ -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> getI64Array(Operation* op, StringRef name) { + auto attr = op->getAttrOfType(name); + if (!attr) + return op->emitOpError() << "phase 2 requires '" << name << "' metadata"; + return SmallVector(attr.asArrayRef()); +} + +static FailureOr> getLaneTable(Operation* op, StringRef name, size_t expected) { + if (auto array = op->getAttrOfType(name)) { + if (array.size() != static_cast(expected)) + return op->emitOpError() << "phase 2 metadata '" << name << "' has the wrong size"; + return SmallVector(array.asArrayRef()); + } + auto elements = op->getAttrOfType(name); + if (!elements || elements.getNumElements() != static_cast(expected)) + return op->emitOpError() << "phase 2 requires a correctly-sized '" << name << "' lane table"; + SmallVector values; + for (APInt value : elements.getValues()) + 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 getStepIndex(ScheduledInfo& info, Block* block) { + auto it = llvm::find(info.blocks, block); + return it == info.blocks.end() ? FailureOr(failure()) + : FailureOr(std::distance(info.blocks.begin(), it)); +} + +static FailureOr +getScalarStepResult(SpatScheduledCompute scheduled, Block& block, unsigned resultIndex, unsigned resultCount) { + auto yield = dyn_cast(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 +getBatchStepResult(SpatScheduledComputeBatch scheduled, Block& block, unsigned globalResultIndex) { + auto parallel = dyn_cast(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(op); + auto destination = insert ? dyn_cast(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 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(laneCount, 1); + result.payload = payload; + result.published = info.op->getResult(globalResult); + auto payloadType = dyn_cast(payload.getType()); + auto publishedType = dyn_cast(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 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(op)) + continue; + ScheduledInfo info; + info.op = &op; + Region& body = isa(op) ? cast(op).getBody() + : cast(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(op)) { + auto core = scalar->getAttrOfType(kCoreIdAttrName); + if (!core) + return scalar.emitOpError("phase 2 requires scalar coreId metadata"); + info.cores.push_back(core.getInt()); + } + else { + auto cores = op.getAttrOfType(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 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(info.resultCounts[step]); ++result) { + unsigned globalResult = info.resultOffsets[step] + result; + if (!info.isBatch()) { + auto payload = getScalarStepResult( + cast(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 { + &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(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 { + &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 findProducer(DeferredTransferPlan& plan, + Operation* diagnosticOwner, + int64_t graphId, + unsigned resultIndex, + std::optional 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 graphLane; + std::optional localOffset; + + bool sameFamily(const RequirementPoint& other) const { + return producer == other.producer && fragmentType == other.fragmentType; + } +}; + +static FailureOr> +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(exchange.deferred.getSources()[sourceIndex]); + if (!source) + return exchange.deferred.emitOpError("phase 2 requires graph-result deferred sources"), failure(); + auto graphId = source.getOwner()->getAttrOfType("scheduled.graph_id"); + if (!graphId) + return exchange.deferred.emitOpError("phase 2 cannot identify graph producer"), failure(); + + RequirementPoint point; + if (auto batch = dyn_cast(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(size->valueAt(lane))) + return std::optional(); + physicalSlot = offset->valueAt(lane) + static_cast(position) * stride->valueAt(lane); + } + else if (position >= (*publication)->physicalSlotToGraphLane.size()) { + return std::optional(); + } + if (physicalSlot < 0 || physicalSlot >= static_cast((*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(source.getOwner())) + return std::optional(); + 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(point); +} + +static void appendRequirementFamily(DeferredExchangePlan& exchange, + RequirementCoordinate coordinate, + unsigned begin, + ArrayRef 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 { + if (!(points.front().*member)) + return std::nullopt; + SmallVector 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(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(exchange.deferred.getSources()[sources->valueAt(lane)]); + auto type = source ? dyn_cast(source.getType()) : RankedTensorType(); + if (source && isa(source.getOwner()) && type) + positionCount = std::max(positionCount, type.getDimSize(0)); + } + } + for (unsigned position = 0; position < positionCount; ++position) { + unsigned runBegin = 0; + SmallVector run; + auto flush = [&] { + if (!run.empty()) + appendRequirementFamily(exchange, {static_cast(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 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 scheduledByOp; + for (ScheduledInfo& scheduled : plan.scheduled) + scheduledByOp[scheduled.op] = &scheduled; + SmallVector deferredOps; + funcOp.walk([&](SpatDeferredCommunicationOp op) { deferredOps.push_back(op); }); + GraphBatchPublicationCache publicationCache; + uint64_t nextChannel = 0; + for (SpatDeferredCommunicationOp deferred : deferredOps) { + Operation* targetOp = deferred->getParentOfType(); + if (!targetOp) + targetOp = deferred->getParentOfType(); + ScheduledInfo* target = scheduledByOp.lookup(targetOp); + auto step = target ? getStepIndex(*target, getScheduledBlock(deferred, targetOp)) : FailureOr(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(); + 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 oldOperands {blueprint.getInput()}; + llvm::append_range(oldOperands, blueprint.getFragments()); + SmallVector publications; + SmallVector newOperands, newSlots; + for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) { + Value source = oldOperands[operandIndex]; + auto result = dyn_cast(source); + auto batch = result ? dyn_cast(result.getOwner()) : SpatGraphComputeBatch(); + int64_t slot = (*sourceSlots)[fragmentIndex]; + if (batch) { + auto graphId = batch->getAttrOfType("scheduled.graph_id"); + auto map = getGraphBatchPublicationMap(batch, result.getResultNumber(), publicationCache); + if (!graphId || failed(map) || slot < 0 || slot >= static_cast((*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 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 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 diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp new file mode 100644 index 0000000..44f44dc --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include "mlir/Dialect/Func/IR/FuncOps.h" + +#include "DeferredCommunicationModel.hpp" + +namespace onnx_mlir::spatial { + +struct DeferredTransferPlan { + llvm::SmallVector scheduled; + llvm::SmallVector> producedStorage; + llvm::DenseMap> producedByGraph; + llvm::SmallVector> exchanges; + llvm::SmallVector stepCounts; +}; + +mlir::FailureOr buildDeferredTransferPlan(mlir::func::FuncOp funcOp); + +mlir::LogicalResult retargetDeferredPublications(mlir::func::FuncOp funcOp, DeferredTransferPlan& plan); + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp index 53f2488..0de2b17 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp @@ -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 : PassWrappergetParentOfType()) { - for (unsigned lane = 0; lane < static_cast(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(transfer.getSources()[leaf.sourceOperandIndex]); - auto graph = source ? dyn_cast(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(source); + auto batch = result + ? dyn_cast(result.getOwner()) + : SpatGraphComputeBatch(); + if (batch && failed(getGraphBatchPublicationMap( + batch, result.getResultNumber(), publicationCache))) + diagnostics.report(transfer.getOperation(), [&](Operation *) {}); } });