fix vgg16 diamonds
Validate Operations / validate-operations (push) Waiting to run

updat ops validations
This commit is contained in:
NiccoloN
2026-07-22 17:54:09 +02:00
parent b491ff77b1
commit 45578ef4c4
29 changed files with 734 additions and 421 deletions
@@ -32,6 +32,18 @@ schedule semantics, scheduling granularity, and available parallelism.
An optimization is acceptable only when its static runtime proxies are equal or
better.
## Asymptotic cost
For each new or materially changed compiler algorithm, identify the relevant
input size and target linear or sublinear time and space. Avoid repeated full-IR
walks, nested scans, and per-operation recomputation when indexing, caching, or
a single traversal can express the same behavior.
When linear-or-better complexity is not possible, use the lowest justified
complexity and report the actual time and space Big-O, the input variable, and
why a lower bound is not practical. Include that cost in the final report; do
not hide a superlinear path behind small current test sizes.
## Forbidden optimization trades
Do not:
@@ -9,6 +9,53 @@ using namespace mlir;
namespace onnx_mlir {
FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
RankedTensorType logicalType,
ArrayRef<FragmentAssemblyEntry> entries,
StringRef physicalLayout,
StringRef indexMap,
PatternRewriter& rewriter,
Location loc) {
auto physicalType = dyn_cast<RankedTensorType>(physicalBatch.getType());
if (!physicalType || !physicalType.hasStaticShape() || !logicalType || !logicalType.hasStaticShape()
|| physicalType.getRank() != logicalType.getRank() + 1 || entries.empty())
return emitError(loc, "invalid static physical batch for fragment assembly"), failure();
const int64_t rank = logicalType.getRank();
const int64_t laneCount = physicalType.getDimSize(0);
if (laneCount <= 0)
return emitError(loc, "fragment assembly requires at least one physical source slot"), failure();
const int64_t fragmentElements = physicalType.getNumElements() / laneCount;
SmallVector<int64_t> operandIndices(entries.size(), 0), sourceSlots, sourceOffsets, offsets, sizes,
strides(entries.size() * rank, 1);
for (const FragmentAssemblyEntry& entry : entries) {
if (entry.sourceSlot < 0 || entry.sourceSlot >= laneCount || entry.sourceOffset < 0
|| entry.destinationOffsets.size() != static_cast<size_t>(rank)
|| entry.sizes.size() != static_cast<size_t>(rank))
return emitError(loc, "invalid fragment assembly entry"), failure();
int64_t entryElements = 1;
for (int64_t dim = 0; dim < rank; ++dim) {
if (entry.destinationOffsets[dim] < 0 || entry.sizes[dim] <= 0
|| entry.destinationOffsets[dim] + entry.sizes[dim] > logicalType.getDimSize(dim))
return emitError(loc, "fragment assembly entry exceeds the logical tensor"), failure();
entryElements *= entry.sizes[dim];
}
if (entry.sourceOffset + entryElements > fragmentElements)
return emitError(loc, "fragment assembly entry exceeds its physical source slot"), failure();
sourceSlots.push_back(entry.sourceSlot);
sourceOffsets.push_back(entry.sourceOffset);
llvm::append_range(offsets, entry.destinationOffsets);
llvm::append_range(sizes, entry.sizes);
}
return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
rewriter.getStringAttr("nchw"), rewriter.getStringAttr(physicalLayout),
rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
rewriter.getStringAttr(indexMap), rewriter.getStringAttr("fragment_assembly"),
rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
}
Value sumTensors(ArrayRef<Value> tensors, PatternRewriter& rewriter) {
if (tensors.size() == 1)
return tensors[0];
@@ -19,6 +19,13 @@
namespace onnx_mlir {
struct FragmentAssemblyEntry {
int64_t sourceSlot;
int64_t sourceOffset;
llvm::SmallVector<int64_t, 4> destinationOffsets;
llvm::SmallVector<int64_t, 4> sizes;
};
namespace detail {
inline mlir::ValueRange getBlockArgs(mlir::Block* block) { return mlir::ValueRange(block->getArguments()); }
@@ -407,4 +414,12 @@ mlir::Value materializeOrComputeUnary(mlir::Value input,
mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::PatternRewriter& rewriter);
mlir::FailureOr<mlir::Value> createFragmentAssemblyBlueprint(mlir::Value physicalBatch,
mlir::RankedTensorType logicalType,
llvm::ArrayRef<FragmentAssemblyEntry> entries,
llvm::StringRef physicalLayout,
llvm::StringRef indexMap,
mlir::PatternRewriter& rewriter,
mlir::Location loc);
} // namespace onnx_mlir
@@ -4,6 +4,7 @@
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -12,6 +13,22 @@ using namespace mlir;
namespace onnx_mlir {
FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, RankedTensorType logicalType) {
auto storageType = dyn_cast<RankedTensorType>(storage.getType());
if (!storageType || !storageType.hasStaticShape() || !logicalType || !logicalType.hasStaticShape()
|| storageType.getRank() != 5 || logicalType.getRank() != 4 || logicalType.getDimSize(0) != 1
|| storageType.getElementType() != logicalType.getElementType()
|| storageType.getDimSize(1) != 1 || storageType.getDimSize(2) <= 0
|| storageType.getDimSize(3) != 1 || storageType.getDimSize(4) != logicalType.getDimSize(3))
return failure();
const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(2));
if (storageType.getDimSize(0) != logicalType.getDimSize(2) * tilesPerRow)
return failure();
return RowStripPhysicalValue {storage, logicalType,
RankedTensorType::get(storageType.getShape().drop_front(), storageType.getElementType(), storageType.getEncoding()),
tilesPerRow};
}
RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
logicalType.getElementType(),
@@ -123,42 +140,39 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
return batchOp->getResult(0);
}
FailureOr<Value>
createRowStripAssemblyBlueprint(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
auto storageType = dyn_cast<RankedTensorType>(storage.getType());
if (!storageType || storageType != getRowStripStorageType(logicalType))
return failure();
auto [offsets, sizes] = buildRowStripMetadata(logicalType);
int64_t height = logicalType.getDimSize(2);
SmallVector<int64_t> operandIndices(height, 0), sourceSlots, sourceOffsets(height, 0), strides(height * 4, 1);
for (int64_t row = 0; row < height; ++row)
sourceSlots.push_back(row);
return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, storage, ValueRange {},
rewriter.getStringAttr("nchw"), rewriter.getStringAttr("nchw_row_strip"),
rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
rewriter.getStringAttr("nchw_row_strip_fragments"), rewriter.getStringAttr("fragment_assembly"),
rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
PatternRewriter& rewriter,
Location loc) {
SmallVector<FragmentAssemblyEntry> entries;
const int64_t tileChannels = value.fragmentType.getDimSize(1);
for (int64_t row = 0; row < value.logicalType.getDimSize(2); ++row)
for (int64_t tile = 0; tile < value.tilesPerRow; ++tile) {
const int64_t channelOffset = tile * tileChannels;
entries.push_back({row * value.tilesPerRow + tile, 0, {0, channelOffset, row, 0},
{1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
value.logicalType.getDimSize(3)}});
}
return createFragmentAssemblyBlueprint(value.storage, value.logicalType, entries, "nchw_row_strip",
kRowStripIndexMap, rewriter, loc);
}
FailureOr<Value>
applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
auto fragmentType = getRowStripFragmentType(logicalType);
auto storageType = getRowStripStorageType(logicalType);
FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
auto storageType = cast<RankedTensorType>(value.storage.getType());
const int64_t laneCount = storageType.getDimSize(0);
auto batchOp = createSpatComputeBatch(rewriter,
loc,
TypeRange {storageType},
logicalType.getDimSize(2),
laneCount,
{},
ValueRange {storage},
ValueRange {value.storage},
[&](detail::SpatComputeBatchBodyArgs args) {
Value fragment =
extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
insertRowStripFragment(
fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
if (failed(fragment)) return failure();
Value relu = spatial::SpatReluOp::create(
rewriter, loc, value.fragmentType, *fragment).getResult();
publishGraphBatchPhysicalFragment(
rewriter, loc, relu, args.outputs.front(), args.lane);
return success();
});
if (failed(batchOp))
@@ -166,41 +180,47 @@ applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter&
return batchOp->getResult(0);
}
FailureOr<Value>
applyRowStripBiasAdd(Value storage, RankedTensorType logicalType, Value bias, PatternRewriter& rewriter, Location loc) {
FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
Value bias,
PatternRewriter& rewriter,
Location loc) {
DenseElementsAttr denseAttr;
if (!isSupportedBiasAddValue(bias, logicalType, &denseAttr))
if (!isSupportedBiasAddValue(bias, value.logicalType, &denseAttr))
return failure();
auto fragmentType = getRowStripFragmentType(logicalType);
auto storageType = getRowStripStorageType(logicalType);
FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, value.logicalType);
if (failed(channelValues)) return failure();
auto storageType = cast<RankedTensorType>(value.storage.getType());
auto biasStorageType = spatial::getGraphBatchPhysicalResultType(value.tilesPerRow, value.fragmentType);
SmallVector<Attribute> biasValues(
biasStorageType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(value.fragmentType.getElementType())));
const int64_t tileChannels = value.fragmentType.getDimSize(1);
const int64_t width = value.fragmentType.getDimSize(3);
for (int64_t channel = 0; channel < value.logicalType.getDimSize(1); ++channel)
for (int64_t w = 0; w < width; ++w)
biasValues[((channel / tileChannels) * tileChannels + channel % tileChannels) * width + w] =
(*channelValues)[channel];
Value biasStorage = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
DenseElementsAttr::get(biasStorageType, biasValues), biasStorageType);
const int64_t laneCount = storageType.getDimSize(0);
auto batchOp = createSpatComputeBatch(rewriter,
loc,
TypeRange {storageType},
logicalType.getDimSize(2),
laneCount,
{},
ValueRange {storage},
ValueRange {value.storage, biasStorage},
[&](detail::SpatComputeBatchBodyArgs args) {
Value fragment =
extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
Value constant;
if (denseAttr.isSplat()) {
constant = getOrCreateConstant(
rewriter,
rewriter.getInsertionBlock()->getParentOp(),
DenseElementsAttr::get(fragmentType, denseAttr.getSplatValue<Attribute>()),
fragmentType);
}
else {
FailureOr<Value> perChannel =
createPerChannelConstantFragment(denseAttr, fragmentType, rewriter);
if (failed(perChannel))
return failure();
constant = *perChannel;
}
fragment =
spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
insertRowStripFragment(
fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
rewriter, loc, args.inputs[0], args.lane, value.fragmentType);
Value tile = affineModConst(rewriter, loc, args.lane,
value.tilesPerRow, anchorOp);
FailureOr<Value> constant = extractGraphBatchPhysicalFragment(
rewriter, loc, args.inputs[1], tile, value.fragmentType);
if (failed(fragment) || failed(constant)) return failure();
Value added = spatial::SpatVAddOp::create(
rewriter, loc, value.fragmentType, *fragment, *constant).getResult();
publishGraphBatchPhysicalFragment(
rewriter, loc, added, args.outputs.front(), args.lane);
return success();
});
if (failed(batchOp))
@@ -11,10 +11,13 @@ inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragmen
struct RowStripPhysicalValue {
mlir::Value storage;
mlir::RankedTensorType logicalType;
llvm::SmallVector<int64_t, 16> fragmentOffsets;
llvm::SmallVector<int64_t, 16> fragmentSizes;
mlir::RankedTensorType fragmentType;
int64_t tilesPerRow;
};
mlir::FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(mlir::Value storage,
mlir::RankedTensorType logicalType);
std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
buildRowStripMetadata(mlir::RankedTensorType type);
@@ -50,18 +53,15 @@ mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
mlir::PatternRewriter& rewriter,
mlir::Location loc);
mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(mlir::Value storage,
mlir::RankedTensorType logicalType,
mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
mlir::PatternRewriter& rewriter,
mlir::Location loc);
mlir::FailureOr<mlir::Value> applyRowStripRelu(mlir::Value storage,
mlir::RankedTensorType logicalType,
mlir::FailureOr<mlir::Value> applyRowStripRelu(const RowStripPhysicalValue& value,
mlir::PatternRewriter& rewriter,
mlir::Location loc);
mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(mlir::Value storage,
mlir::RankedTensorType logicalType,
mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
mlir::Value bias,
mlir::PatternRewriter& rewriter,
mlir::Location loc);
@@ -45,38 +45,30 @@ static FailureOr<RowStripPhysicalValue> buildRowStripValue(spatial::SpatBlueprin
auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
if (!logicalType)
return blueprint.emitOpError("requires ranked logical output type"), failure();
RowStripPhysicalValue value;
value.storage = storage;
value.logicalType = logicalType;
value.fragmentOffsets.append(blueprint.getFragmentOffsets().begin(), blueprint.getFragmentOffsets().end());
value.fragmentSizes.append(blueprint.getFragmentSizes().begin(), blueprint.getFragmentSizes().end());
if (blueprint.getIndexMap() != kRowStripIndexMap)
return blueprint.emitOpError("requires the canonical row-strip index map"), failure();
auto storageType = dyn_cast<RankedTensorType>(storage.getType());
if (!storageType || storageType != getRowStripStorageType(logicalType))
FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
if (failed(value))
return blueprint.emitOpError("requires physical row-strip fragment storage"), failure();
return value;
return *value;
}
static FailureOr<Value>
lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
return applyRowStripRelu(input.storage, input.logicalType, rewriter, planOp.getLoc());
return applyRowStripRelu(input, rewriter, planOp.getLoc());
}
static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
spatial::SpatBiasAddPlanOp planOp,
PatternRewriter& rewriter) {
return applyRowStripBiasAdd(input.storage, input.logicalType, planOp.getBias(), rewriter, planOp.getLoc());
return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc());
}
static FailureOr<Value>
materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
return failure();
auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType);
if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes))
return failure();
return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc);
}
static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
@@ -168,9 +160,12 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
});
if (rowStripBlueprint != planOp.getResult().getUsers().end()) {
rewriter.setInsertionPoint(planOp);
std::optional<Value> physicalInput;
if (succeeded(rowStripInput))
physicalInput = rowStripInput->storage;
FailureOr<Value> lowered = lowerSelectedConv2DPlan(
planOp,
succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->storage} : std::nullopt,
physicalInput,
/*emitRowStripLayout=*/true,
rewriter);
if (failed(lowered)) {
@@ -255,8 +250,23 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
rewriter.setInsertionPoint(planOp);
std::optional<Value> physicalInput;
if (succeeded(input)) {
if (input->tilesPerRow == 1) {
physicalInput = input->storage;
}
else {
FailureOr<Value> denseInput = materializeRowStripToDense(*input, planOp.getLoc(), rewriter);
if (failed(denseInput)) {
planOp.emitOpError("failed to materialize tiled row-strip input for MaxPool");
signalPassFailure();
return;
}
planOp.getInputMutable().assign(*denseInput);
}
}
FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
planOp, succeeded(input) ? std::optional<Value> {input->storage} : std::nullopt, rewriter);
planOp, physicalInput, rewriter);
if (failed(lowered)) {
planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
signalPassFailure();
@@ -2326,54 +2326,25 @@ static Value maybeUnpackChunkRows(Value gemmRows,
return unpackCompute.getResult(0);
}
static Value createChunkedConvRows(const ConvLoweringState& state,
const PreparedConvInput& preparedInput,
Value weightMatrix,
Value biasMatrix,
DenseElementsAttr wDenseAttr,
DenseElementsAttr biasDenseAttr,
int64_t forcedPackFactor,
uint64_t chunkPositions,
PatternRewriter& rewriter,
Location loc) {
SmallVector<Value> chunkRows;
static Value createStreamedConvRows(const ConvLoweringState& state,
const PreparedConvInput& preparedInput,
Value weightMatrix,
Value biasMatrix,
DenseElementsAttr wDenseAttr,
DenseElementsAttr biasDenseAttr,
int64_t forcedPackFactor,
PatternRewriter& rewriter,
Location loc) {
const int64_t totalPatches = state.batchSize * state.outHeight * state.outWidth;
for (int64_t chunkStart = 0; chunkStart < totalPatches; chunkStart += static_cast<int64_t>(chunkPositions)) {
const int64_t chunkNumPatches = std::min<int64_t>(static_cast<int64_t>(chunkPositions), totalPatches - chunkStart);
ConvGemmPlan chunkPlan = buildConvGemmPlan(state,
static_cast<bool>(wDenseAttr),
!state.hasBias || static_cast<bool>(biasDenseAttr),
chunkStart,
chunkNumPatches,
forcedPackFactor);
Value chunkInputRows = createIm2colRows(state, preparedInput, chunkPlan, rewriter, loc);
Value chunkB = buildPackedWeights(wDenseAttr, weightMatrix, state, chunkPlan, rewriter, loc);
Value gemmBias = createZeroGemmBias(chunkPlan.gemmOutputRowsType, rewriter);
if (state.hasBias)
gemmBias = state.b;
Value chunkC = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, chunkPlan, rewriter, loc);
Value chunkGemmRows = ONNXGemmOp::create(rewriter,
loc,
chunkPlan.gemmOutputRowsType,
chunkInputRows,
chunkB,
chunkC,
APFloat(1.0f),
APFloat(1.0f),
/*transA=*/0,
/*transB=*/0)
.getY();
chunkRows.push_back(maybeUnpackChunkRows(chunkGemmRows, chunkPlan, rewriter, loc));
}
if (chunkRows.size() == 1)
return chunkRows.front();
auto rowType = RankedTensorType::get({totalPatches, state.numChannelsOut}, state.outType.getElementType());
auto collectRows = createSpatCompute(rewriter, loc, TypeRange {rowType}, {}, chunkRows, [&](ValueRange rows) {
spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, /*axis=*/0, rows));
});
return collectRows.getResult(0);
ConvGemmPlan plan = buildConvGemmPlan(state, static_cast<bool>(wDenseAttr),
!state.hasBias || static_cast<bool>(biasDenseAttr), 0, totalPatches, forcedPackFactor);
Value inputRows = createIm2colRows(state, preparedInput, plan, rewriter, loc);
Value packedWeights = buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc);
Value gemmBias = state.hasBias ? state.b : createZeroGemmBias(plan.gemmOutputRowsType, rewriter);
Value packedBias = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc);
Value gemmRows = ONNXGemmOp::create(rewriter, loc, plan.gemmOutputRowsType, inputRows,
packedWeights, packedBias, APFloat(1.0f), APFloat(1.0f), 0, 0).getY();
return maybeUnpackChunkRows(gemmRows, plan, rewriter, loc);
}
static Value rewritePackedIm2ColConv(const ConvLoweringState& state,
@@ -2444,16 +2415,13 @@ static Value rewriteStreamedConv(const ConvLoweringState& state,
ConvGemmPlan seedPlan = buildConvGemmPlan(
state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0, 1, forcedPackFactor);
Value weightMatrix = createWeightMatrix(state.w, seedPlan, rewriter, loc);
ConvGeometry geo = buildConvGeometry(state);
uint64_t chunkPositions = chooseStreamChunkPositions(geo, forcedPackFactor);
Value collectedRows = createChunkedConvRows(state,
Value collectedRows = createStreamedConvRows(state,
preparedInput,
weightMatrix,
biasMatrix,
wDenseAttr,
biasDenseAttr,
forcedPackFactor,
chunkPositions,
rewriter,
loc);
auto gemmOutType = cast<RankedTensorType>(collectedRows.getType());
@@ -2524,6 +2492,21 @@ static bool canConsumeNchwRowStripFragments(const ConvLoweringState& state, Stri
failureReason = "dilation_not_one";
return false;
}
ConvGeometry geometry = buildConvGeometry(state);
const bool pointwise = state.xHeight == 1 && state.xWidth == 1 && state.outHeight == 1 && state.outWidth == 1
&& state.wHeight == 1 && state.wWidth == 1 && state.padHeightBegin == 0
&& state.padHeightEnd == 0 && state.padWidthBegin == 0 && state.padWidthEnd == 0;
if (pointwise) {
if (!getHostConstDenseElementsAttr(state.w)) {
failureReason = "non_constant_weight";
return false;
}
if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) {
failureReason = "unsupported_bias";
return false;
}
return true;
}
if (state.wHeight != 3 || state.wWidth != 3) {
failureReason = "kernel_not_3x3";
return false;
@@ -2544,7 +2527,6 @@ static bool canConsumeNchwRowStripFragments(const ConvLoweringState& state, Stri
failureReason = "unsupported_bias";
return false;
}
ConvGeometry geometry = buildConvGeometry(state);
if (geometry.c > geometry.xbarSize) {
failureReason = "output_channels_exceed_crossbar";
return false;
@@ -2575,6 +2557,25 @@ static FailureOr<Value> createPaddedBiasRowConstant(const ConvLoweringState& sta
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), biasAttr, biasType);
}
static FailureOr<Value> createPaddedBiasTileConstant(const ConvLoweringState& state,
int64_t tileChannels,
PatternRewriter& rewriter) {
DenseElementsAttr denseAttr;
if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr))
return failure();
FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, state.outType);
if (failed(channelValues))
return failure();
const int64_t tileCount = ceilIntegerDivide(state.numChannelsOut, tileChannels);
auto tileType = RankedTensorType::get({tileCount, 1, tileChannels}, state.outType.getElementType());
SmallVector<Attribute> values(
tileType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(tileType.getElementType())));
for (int64_t channel = 0; channel < state.numChannelsOut; ++channel)
values[channel] = (*channelValues)[channel];
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
DenseElementsAttr::get(tileType, values), tileType);
}
static Value createHorizontallyPaddedRowStripFragment(Value fragment,
const ConvLoweringState& state,
PatternRewriter& rewriter,
@@ -2732,8 +2733,11 @@ static FailureOr<Value> createConvInputWindow(Value input,
? extractDenseConvWindowRow(input, sourceRowTable, state, outputHeight, kernelRow, rewriter, loc)
: extractProjectedRowStripWindowRow(
input, sourceRowTable, state, outputHeight, kernelRow, rewriter, loc);
Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, loc);
Value semanticRow = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, sourceRow, mask).getResult();
Value semanticRow = sourceRow;
if (state.padHeightBegin != 0 || state.padHeightEnd != 0) {
Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, loc);
semanticRow = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, sourceRow, mask).getResult();
}
Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, loc);
window = tensor::InsertSliceOp::create(rewriter,
loc,
@@ -2906,12 +2910,6 @@ static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
.getResult();
}
static bool rowStripOutputFitsOneCore(const ConvGeometry& geometry) {
const int64_t inputTileCount = ceilIntegerDivide(geometry.k, geometry.xbarSize);
const int64_t outputTileCount = ceilIntegerDivide(geometry.c, geometry.xbarSize);
return inputTileCount * outputTileCount <= static_cast<int64_t>(crossbarCountInCore.getValue());
}
static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) {
return ceilIntegerDivide(geometry.k, geometry.xbarSize)
<= static_cast<int64_t>(crossbarCountInCore.getValue());
@@ -2928,38 +2926,40 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
auto elementType = state.outType.getElementType();
auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
auto tilePixelType = RankedTensorType::get({1, xbarDim, 1, 1}, elementType);
auto tileFragmentType = RankedTensorType::get({1, xbarDim, 1, state.outWidth}, elementType);
auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.wType.getElementType());
SmallVector<Value> outputTiles;
outputTiles.reserve(outputTileCount);
for (int64_t outputTile = 0; outputTile < outputTileCount; ++outputTile) {
const int64_t channelOffset = outputTile * xbarDim;
const int64_t tileChannels = std::min(xbarDim, state.numChannelsOut - channelOffset);
auto tileRowType = RankedTensorType::get({1, tileChannels}, elementType);
auto tilePixelType = RankedTensorType::get({1, tileChannels, 1, 1}, elementType);
auto tileFragmentType = RankedTensorType::get({1, tileChannels, 1, state.outWidth}, elementType);
auto tileStorageType = spatial::getGraphBatchPhysicalResultType(state.outHeight, tileFragmentType);
SmallVector<OpFoldResult> weightOffsets {
rewriter.getIndexAttr(outputTile), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> weightSizes {
rewriter.getIndexAttr(1), rewriter.getIndexAttr(paddedK), rewriter.getIndexAttr(xbarDim)};
Value tileWeights = extractStaticSliceOrIdentity(
rewriter, loc, paddedWeights, tileWeightsType, weightOffsets, weightSizes, getUnitStrides(rewriter, 3));
auto tileBatch = createSpatComputeBatch(
rewriter,
loc,
TypeRange {tileStorageType},
state.outHeight,
ValueRange {tileWeights},
ValueRange {state.x},
[&](detail::SpatComputeBatchBodyArgs args) {
const int64_t laneCount = state.outHeight * outputTileCount;
auto tileStorageType = spatial::getGraphBatchPhysicalResultType(laneCount, tileFragmentType);
FailureOr<Value> paddedBias = failure();
if (state.hasBias)
paddedBias = createPaddedBiasTileConstant(state, xbarDim, rewriter);
if (state.hasBias && failed(paddedBias))
return failure();
auto tileBatch = createSpatComputeBatch(
rewriter, loc, TypeRange {tileStorageType}, laneCount, ValueRange {paddedWeights},
state.hasBias ? ValueRange {state.x, *paddedBias} : ValueRange {state.x},
[&](detail::SpatComputeBatchBodyArgs args) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
Value outputRow = affineFloorDivConst(rewriter, loc, args.lane, outputTileCount, anchorOp);
Value outputTile = affineModConst(rewriter, loc, args.lane, outputTileCount, anchorOp);
SmallVector<OpFoldResult> weightOffsets {
outputTile, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> weightSizes {
rewriter.getIndexAttr(1), rewriter.getIndexAttr(paddedK), rewriter.getIndexAttr(xbarDim)};
Value tileWeights = tensor::ExtractSliceOp::create(
rewriter, loc, tileWeightsType, args.weights.front(), weightOffsets, weightSizes, getUnitStrides(rewriter, 3));
FailureOr<Value> biasTile = failure();
if (state.hasBias)
biasTile = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[1], outputTile, paddedRowType);
if (state.hasBias && failed(biasTile))
return failure();
FailureOr<Value> inputWindow =
createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
createConvInputWindow(args.inputs.front(), state, outputRow, rewriter, loc);
if (failed(inputWindow))
return failure();
Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, tileFragmentType.getShape(), elementType);
@@ -2984,28 +2984,18 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
paddedPatchRow = createZeroPaddedTensor(
paddedPatchRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc);
FailureOr<Value> paddedOutputRow = createPaddedConvOutputTile(
paddedPatchRow, args.weights.front(), numKSlices, xbarDim, rewriter, widthLoc);
paddedPatchRow, tileWeights, numKSlices, xbarDim, rewriter, widthLoc);
if (failed(paddedOutputRow))
return failure();
Value outputRow = *paddedOutputRow;
if (tileChannels != xbarDim) {
SmallVector<OpFoldResult> rowOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> rowSizes {
rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels)};
outputRow = tensor::ExtractSliceOp::create(rewriter,
widthLoc,
tileRowType,
outputRow,
rowOffsets,
rowSizes,
getUnitStrides(rewriter, 2));
}
if (state.hasBias)
paddedOutputRow = spatial::SpatVAddOp::create(
rewriter, widthLoc, paddedRowType, *paddedOutputRow, *biasTile).getResult();
Value outputPixel = tensor::ExpandShapeOp::create(
rewriter, widthLoc, tilePixelType, outputRow, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
rewriter, widthLoc, tilePixelType, *paddedOutputRow, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
SmallVector<OpFoldResult> rowOffsets {
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(tileChannels),
rewriter.getIndexAttr(xbarDim),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1)};
Value nextFragment = tensor::InsertSliceOp::create(rewriter,
@@ -3024,58 +3014,9 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
rewriter, loc, widthLoop->results.front(), args.outputs.front(), args.lane);
return success();
});
if (failed(tileBatch))
return failure();
outputTiles.push_back(tileBatch->getResult(0));
}
auto fragmentType = getRowStripFragmentType(state.outType);
auto outputStorageType = getRowStripStorageType(state.outType);
auto assemblyBatch = createSpatComputeBatch(rewriter,
loc,
TypeRange {outputStorageType},
state.outHeight,
{},
ValueRange(outputTiles),
[&](detail::SpatComputeBatchBodyArgs args) {
Value fragment = tensor::EmptyOp::create(
rewriter, loc, fragmentType.getShape(), elementType);
for (int64_t outputTile = 0; outputTile < outputTileCount; ++outputTile) {
const int64_t channelOffset = outputTile * xbarDim;
const int64_t tileChannels =
std::min(xbarDim, state.numChannelsOut - channelOffset);
auto tileFragmentType = RankedTensorType::get(
{1, tileChannels, 1, state.outWidth}, elementType);
FailureOr<Value> tileFragment = extractGraphBatchPhysicalFragment(
rewriter, loc, args.inputs[outputTile], args.lane, tileFragmentType);
if (failed(tileFragment))
return failure();
SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0),
rewriter.getIndexAttr(channelOffset),
rewriter.getIndexAttr(0),
rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(tileChannels),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.outWidth)};
fragment = tensor::InsertSliceOp::create(rewriter,
loc,
*tileFragment,
fragment,
offsets,
sizes,
getUnitStrides(rewriter, 4));
}
insertRowStripFragment(
fragment, args.outputs.front(), state.outType, args.lane, rewriter, loc);
return success();
});
if (failed(assemblyBatch))
if (failed(tileBatch))
return failure();
Value output = assemblyBatch->getResult(0);
if (state.hasBias)
return applyRowStripBiasAdd(output, state.outType, state.b, rewriter, loc);
return output;
return tileBatch->getResult(0);
}
static FailureOr<Value>
@@ -3104,7 +3045,7 @@ createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRe
weightDenseAttr, state, paddedK, xbarDim, rewriter)
: standard::createPaddedOutputChannelTiledWeightConstant(
weightDenseAttr, state, paddedK, xbarDim, rewriter);
if (!rowStripOutputFitsOneCore(geometry))
if (state.numChannelsOut > xbarDim)
return createOutputChannelTiledRowStripConvOutput(
state, paddedWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
@@ -3279,11 +3220,106 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
return batchOp->getResult(0);
}
static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStripStorage,
const ConvLoweringState& state,
PatternRewriter& rewriter,
Location loc) {
FailureOr<RowStripPhysicalValue> input = describeRowStripPhysicalValue(rowStripStorage, state.xType);
if (failed(input)) return failure();
ConvGeometry geometry = buildConvGeometry(state);
const int64_t xbarDim = geometry.xbarSize;
const int64_t inputFragmentChannels = input->fragmentType.getDimSize(1);
if (inputFragmentChannels % xbarDim != 0 || state.numChannelsIn % xbarDim != 0)
return failure();
auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
if (!weightDenseAttr) return failure();
const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim);
const int64_t numKSlices = state.numChannelsIn / xbarDim;
auto elementType = state.outType.getElementType();
auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
auto inputRowType = RankedTensorType::get({1, inputFragmentChannels}, elementType);
auto weightTileType = RankedTensorType::get({state.numChannelsIn, xbarDim}, state.wType.getElementType());
auto weightSliceType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType());
auto outputFragmentType = RankedTensorType::get({1, xbarDim, 1, 1}, elementType);
auto outputStorageType = spatial::getGraphBatchPhysicalResultType(outputTileCount, outputFragmentType);
Value paddedWeights = standard::createPaddedOutputChannelTiledWeightConstant(
weightDenseAttr, state, state.numChannelsIn, xbarDim, rewriter);
FailureOr<Value> paddedBias = failure();
if (state.hasBias) paddedBias = createPaddedBiasTileConstant(state, xbarDim, rewriter);
if (state.hasBias && failed(paddedBias)) return failure();
auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {outputStorageType}, outputTileCount,
ValueRange {paddedWeights},
state.hasBias ? ValueRange {rowStripStorage, *paddedBias} : ValueRange {rowStripStorage},
[&](detail::SpatComputeBatchBodyArgs args) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
SmallVector<OpFoldResult> weightOffsets {args.lane, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> weightSizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.numChannelsIn), rewriter.getIndexAttr(xbarDim)};
Value weightTile = tensor::ExtractSliceOp::create(
rewriter, loc, weightTileType, args.weights.front(), weightOffsets, weightSizes, getUnitStrides(rewriter, 3));
auto createPiece = [&](Value kSlice, Location pieceLoc) -> FailureOr<Value> {
Value channelOffset = affineMulConst(rewriter, pieceLoc, kSlice, xbarDim, anchorOp);
Value sourceSlot = affineFloorDivConst(
rewriter, pieceLoc, channelOffset, inputFragmentChannels, anchorOp);
Value sourceOffset = affineModConst(
rewriter, pieceLoc, channelOffset, inputFragmentChannels, anchorOp);
FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
rewriter, pieceLoc, args.inputs.front(), sourceSlot, input->fragmentType);
if (failed(fragment)) return failure();
Value inputRow = tensor::CollapseShapeOp::create(rewriter, pieceLoc, inputRowType, *fragment,
SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
Value inputSlice = tensor::ExtractSliceOp::create(rewriter, pieceLoc, paddedRowType, inputRow,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), sourceOffset},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)},
getUnitStrides(rewriter, 2));
Value weightSlice = tensor::ExtractSliceOp::create(rewriter, pieceLoc, weightSliceType, weightTile,
SmallVector<OpFoldResult> {channelOffset, rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)},
getUnitStrides(rewriter, 2));
return spatial::SpatVMMOp::create(rewriter, pieceLoc, paddedRowType, weightSlice, inputSlice).getResult();
};
FailureOr<Value> result = createPiece(c0, loc);
if (failed(result)) return failure();
if (numKSlices > 1) {
auto reduction = buildNormalizedScfFor(rewriter, loc, c1, cNumKSlices, c1, ValueRange {*result},
[&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange iterArgs,
SmallVectorImpl<Value>& yielded) {
FailureOr<Value> piece = createPiece(kSlice, reduceLoc);
if (failed(piece)) return failure();
yielded.push_back(spatial::SpatVAddOp::create(
rewriter, reduceLoc, paddedRowType, iterArgs.front(), *piece).getResult());
return success();
});
if (failed(reduction)) return failure();
result = reduction->results.front();
}
if (state.hasBias) {
FailureOr<Value> bias = extractGraphBatchPhysicalFragment(
rewriter, loc, args.inputs[1], args.lane, paddedRowType);
if (failed(bias)) return failure();
result = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, *result, *bias).getResult();
}
Value fragment = tensor::ExpandShapeOp::create(rewriter, loc, outputFragmentType, *result,
SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
publishGraphBatchPhysicalFragment(rewriter, loc, fragment, args.outputs.front(), args.lane);
return success();
});
if (failed(batch)) return failure();
return batch->getResult(0);
}
static FailureOr<Value> createConvOutputFromRowStripInput(const ConvLoweringState& state,
[[maybe_unused]] const ConvLoweringDecision& decision,
Value rowStripInput,
PatternRewriter& rewriter,
Location loc) {
if (state.xHeight == 1 && state.xWidth == 1 && state.wHeight == 1 && state.wWidth == 1)
return createPointwiseOutputFromRowStripFragments(rowStripInput, state, rewriter, loc);
return createConvOutputFromNchwRowStripFragments(rowStripInput, state, rewriter, loc);
}
@@ -4187,22 +4223,9 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
}
if (failed(canLowerConvPlanToRowStrip(planOp)))
return planOp.emitOpError("selected row-strip layout is not supported for this Conv plan"), failure();
ConvLoweringState rowState = *state;
const bool applyBiasAfterStorage = rowState.hasBias;
Value originalBias = rowState.b;
if (applyBiasAfterStorage) {
rowState.b = Value();
rowState.hasBias = false;
}
FailureOr<Value> rowStripStorage = createRowStripConvOutputFromDenseInput(rowState, rewriter, planOp.getLoc());
FailureOr<Value> rowStripStorage = createRowStripConvOutputFromDenseInput(*state, rewriter, planOp.getLoc());
if (failed(rowStripStorage))
return planOp.emitOpError("failed to build row-strip fragment storage for the selected Conv plan"), failure();
if (applyBiasAfterStorage) {
rowStripStorage = applyRowStripBiasAdd(*rowStripStorage, state->outType, originalBias, rewriter, planOp.getLoc());
if (failed(rowStripStorage))
return planOp.emitOpError("failed to apply row-strip Conv bias per fragment"), failure();
}
return *rowStripStorage;
}
@@ -77,9 +77,16 @@ static LogicalResult eraseOldGraph(func::FuncOp funcOp,
rewriter.eraseOp(blueprint);
continue;
}
if (!op->use_empty())
return op->emitOpError(
"phase 2 cannot erase an old graph compute with live results");
if (!op->use_empty()) {
for (OpResult result : op->getResults()) {
if (!result.use_empty()) {
Operation *user = result.use_begin()->getOwner();
return op->emitOpError()
<< "phase 2 cannot erase old graph result "
<< result.getResultNumber() << " used by " << user->getName();
}
}
}
rewriter.eraseOp(op);
}
return success();
@@ -100,6 +107,25 @@ static LogicalResult eraseDeferredSourceSelectors(
return success();
}
static void eraseUnusedIdentityDeferredCommunications(
func::FuncOp funcOp, IRRewriter &rewriter) {
SmallVector<SpatDeferredCommunicationOp> unused;
funcOp.walk([&](SpatDeferredCommunicationOp deferred) {
if (!deferred.getOutput().use_empty() || !deferred.getBody().hasOneBlock())
return;
Block &body = deferred.getBody().front();
auto yield = dyn_cast<SpatYieldOp>(body.getTerminator());
auto argument = yield && yield.getOutputs().size() == 1
? dyn_cast<BlockArgument>(yield.getOutputs().front())
: BlockArgument();
if (argument && argument.getOwner() == &body
&& argument.getArgNumber() < deferred.getSources().size())
unused.push_back(deferred);
});
for (SpatDeferredCommunicationOp deferred : llvm::reverse(unused))
rewriter.eraseOp(deferred);
}
static LogicalResult verifyDominance(func::FuncOp funcOp) {
DominanceInfo dominance(funcOp);
WalkResult result = funcOp.walk([&](Operation *op) {
@@ -119,6 +145,9 @@ static LogicalResult verifyDominance(func::FuncOp funcOp) {
LogicalResult realizeDeferredCommunication(
func::FuncOp funcOp,
const ScheduledComputeMaterializationResult &materialization) {
IRRewriter rewriter(funcOp.getContext());
eraseUnusedIdentityDeferredCommunications(funcOp, rewriter);
auto transfers = buildDeferredTransferPlan(funcOp, materialization);
if (failed(transfers))
return funcOp.emitOpError(
@@ -134,7 +163,6 @@ LogicalResult realizeDeferredCommunication(
return funcOp.emitOpError(
"phase 2 failed to build sparse boundary programs");
IRRewriter rewriter(funcOp.getContext());
if (failed(retargetDeferredPublications(funcOp, *transfers))
|| failed(replaceFinalGraphPublications(funcOp, *transfers)))
return failure();
+270 -136
View File
@@ -1,174 +1,308 @@
# Validation Operations
# Operation Validation Suite
ONNX test models used by `validate.py` to verify the Raptor compiler + PIM simulator pipeline.
This directory contains the ONNX models used by `validation/validate.py` to
validate individual operations through compilation, PIM simulation, and
comparison with the ONNX-MLIR reference runtime.
Generated tests can be regenerated with:
## Naming
```
python3 validation/operations/gen_tests.py
Every model uses the same path convention:
```text
<category>/<case>/<category>_<case>.onnx
```
## Conv
The `<category>/<case>` pair is the operation ID printed by the validator.
Use lowercase `snake_case` for both components. Keep the generator function,
graph name, directory, and filename based on the same operation ID when adding
a test.
| Test | Directory | Input | Output | Kernel | Stride | Padding | Bias | Notes |
|------------------|-------------------------|-----------|-----------|--------|--------|------------|------|------------------------------------|
| Simple | `conv/simple` | [1,3,3,3] | [1,1,2,2] | 2x2 | 1 | none | no | Basic conv, hand-crafted |
| With constant | `conv/with_constant` | [1,3,3,3] | [1,1,3,3] | 2x2 | 1 | SAME_UPPER | yes | Hand-crafted, constant weight+bias |
| Batch 2 | `conv/batch_2` | [2,3,3,3] | [2,1,3,3] | 2x2 | 1 | SAME_UPPER | yes | Batched input |
| Kernel 3x3 | `conv/kernel_3x3` | [1,1,5,5] | [1,1,3,3] | 3x3 | 1 | none | no | Larger kernel |
| Stride 2 | `conv/stride_2` | [1,1,6,6] | [1,1,2,2] | 3x3 | 2 | none | no | Strided convolution |
| Multi channel | `conv/multi_channel` | [1,3,5,5] | [1,4,3,3] | 3x3 | 1 | none | no | 3 in channels, 4 out channels |
| Pointwise 1x1 | `conv/pointwise_1x1` | [1,8,4,4] | [1,4,4,4] | 1x1 | 1 | none | no | Channel mixing |
| SAME padding 3x3 | `conv/same_padding_3x3` | [1,1,5,5] | [1,1,5,5] | 3x3 | 1 | SAME_UPPER | no | Spatial dims preserved |
| Explicit padding | `conv/explicit_padding` | [1,1,4,4] | [1,1,4,4] | 3x3 | 1 | [1,1,1,1] | no | Symmetric explicit pads |
| With bias 3x3 | `conv/with_bias_3x3` | [1,3,5,5] | [1,2,3,3] | 3x3 | 1 | none | yes | Multi-channel with bias |
| Large spatial | `conv/large_spatial` | [1,1,8,8] | [1,1,6,6] | 3x3 | 1 | none | no | Larger spatial input |
| Grouped two groups | `conv/grouped_two_groups` | [1,4,4,4] | [1,4,4,4] | 1x1 | 1 | none | yes | group=2 channel partitioning |
| Depthwise grouped | `conv/depthwise_grouped` | [1,3,4,4] | [1,3,2,2] | 3x3 | 1 | none | no | group=3, one input channel per group |
| Dynamic | `conv/dynamic` | [1,1,4,4] | [1,1,2,2] | 3x3 | 1 | none | no | Runtime input and weight |
## Generate and validate
## Gemm
Regenerate all generated models from the repository root:
| Test | Directory | A (input) | B/W tensor | Output | transB | alpha | beta | Bias | Notes |
|---------------|-------------------------|-----------|------------|----------|--------|-------|------|-------|------------------------------|
| Simple | `gemm/simple` | [10,132] | [132,132] | [10,132] | no | 1 | 1 | no | Square weights |
| Non-square | `gemm/non_square` | [4,128] | [128,64] | [4,64] | no | 1 | 1 | no | K != N |
| With bias | `gemm/with_bias` | [4,128] | [128,128] | [4,128] | no | 1 | 1 | [128] | Bias vector |
| transB | `gemm/transB` | [4,128] | [64,128] | [4,64] | yes | 1 | 1 | no | Transposed weight |
| Alpha/beta | `gemm/alpha_beta` | [4,64] | [64,64] | [4,64] | no | 0.5 | 0.25 | [64] | Scaled matmul + bias |
| Small | `gemm/small` | [2,8] | [8,4] | [2,4] | no | 1 | 1 | no | Tiny matrices |
| Large | `gemm/large` | [8,256] | [256,128] | [8,128] | no | 1 | 1 | no | Larger matrices |
| transB + bias | `gemm/transB_with_bias` | [4,128] | [64,128] | [4,64] | yes | 1 | 1 | [64] | Combined |
| Dynamic | `gemm/dynamic` | [2,8] | [8,4] | [2,4] | no | 1 | 1 | no | Runtime matrix operands |
| Dynamic transB | `gemm/dynamic_transB` | [2,8] | [4,8] | [2,4] | yes | 1 | 1 | no | Runtime transpose handling |
| Dynamic bias | `gemm/dynamic_bias` | [2,8] | [8,4] | [2,4] | no | 1 | 1 | [4] | Runtime bias broadcast |
| Dynamic alpha | `gemm/dynamic_alpha` | [2,8] | [8,4] | [2,4] | no | 0.5 | 1 | no | Runtime alpha scaling |
| Dynamic beta | `gemm/dynamic_beta` | [2,8] | [8,4] | [2,4] | no | 1 | 2 | [4] | Runtime beta scaling |
| Dynamic bias + scale | `gemm/dynamic_bias_alpha_beta` | [2,8] | [8,4] | [2,4] | no | 0.5 | 2 | [4] | Runtime operands and bias |
```bash
.venv/bin/python validation/operations/gen_tests.py
```
## MatMul
Run the complete suite with deadlock detection:
| Test | Directory | A input | B tensor | Output | Notes |
|---------------------|----------------------------------|----------|----------|---------|-------------------------------------------------|
| Basic | `matmul/basic` | [2,3] | [3,4] | [2,4] | Direct 2D MatMul rewrite path |
| Left constant | `matmul/left_constant` | [2,3] | [3,4] | [2,4] | Constant LHS transpose rewrite path |
| Dynamic | `matmul/dynamic` | [2,3] | [3,4] | [2,4] | Runtime matrix operands |
| Batched 3D | `matmul/batched_3d` | [2,2,3] | [2,3,4] | [2,2,4] | Matching-batch direct batched lowering |
| Batched 3D dynamic | `matmul/batched_3d_dynamic` | [2,2,3] | [2,3,4] | [2,2,4] | Batched runtime operands |
| Batched left const | `matmul/batched_left_constant` | [2,2,3] | [2,3,4] | [2,2,4] | Batched constant-LHS transpose path |
| Batched RHS broadcast | `matmul/batched_rhs_broadcast` | [2,2,3] | [3,4] | [2,2,4] | Rank-2 RHS broadcast across batch |
| Batched LHS broadcast | `matmul/batched_lhs_broadcast` | [2,3] | [2,3,4] | [2,2,4] | Rank-2 LHS broadcast across batched RHS |
```bash
.venv/bin/python validation/validate.py \
--raptor-path build_release/Release/bin/onnx-mlir \
--onnx-include-dir onnx-mlir/include \
--operations-dir validation/operations \
--crossbar-count 64 \
--crossbar-size 128 \
--core-count 144 \
--raptor-extra-arg=--pim-detect-communication-deadlock \
--raptor-extra-arg=--pim-export-spatial-dataflow=none
```
## Gemv
Use `--compile-only` for compiler and deadlock checks, then `--run-only` to
reuse those artifacts for reference execution, simulation, and comparison.
The validator prints the complete operation results table before its summary.
| Test | Directory | Input | W (weight) | Output | Bias | Notes |
|---------------------|------------------------------------|----------|------------|---------|---------|----------------------------|
| Simple | `gemv/simple` | [1,132] | [132,132] | [1,132] | no | Single-sample matmul |
| Constant | `gemv/constant` | _(none)_ | [132,132] | [1,132] | no | All inputs constant |
| Homogeneous const | `gemv/with_homogeneous_constant` | [1,132] | [132,132] | [1,132] | [1,132] | Bias matches output shape |
| Heterogeneous const | `gemv/with_heterogeneous_constant` | [1,132] | [132,132] | [1,132] | [1,132] | Different constant pattern |
| Scalar const | `gemv/with_scalar_constant` | [1,132] | [132,132] | [1,132] | [1,1] | Scalar bias, broadcast |
## Complete inventory
## Pool
The suite contains 164 models. Tensor shapes, attributes, and constants are
defined in `gen_tests.py` and in the checked-in ONNX models.
| Test | Directory | Input | Output | Kernel | Stride | Padding | Notes |
|----------------------------|---------------------------------|-----------|-----------|------------------------|--------|------------|----------------------------------|
| Max basic | `pool/max_basic` | [1,1,4,4] | [1,1,3,3] | 2x2 | 1 | none | Basic max pooling |
| Max stride 2 multi-channel | `pool/max_stride2_multichannel` | [1,5,6,6] | [1,5,3,3] | 2x2 | 2 | none | Channel-preserving max pool |
| Max SAME_UPPER | `pool/max_same_upper` | [1,1,5,5] | [1,1,3,3] | 3x3 | 2 | SAME_UPPER | Deprecated auto_pad path |
| Avg basic | `pool/avg_basic` | [1,3,4,4] | [1,3,3,3] | 2x2 | 1 | none | Basic average pooling |
| Avg explicit padding | `pool/avg_explicit_padding` | [1,2,4,4] | [1,2,2,2] | 3x3 | 2 | [1,1,1,1] | `count_include_pad=0` |
| Avg include pad | `pool/avg_include_pad` | [1,2,4,4] | [1,2,2,2] | 3x3 | 2 | [1,1,1,1] | `count_include_pad=1` |
| Max after Conv | `pool/max_after_conv` | [1,3,6,6] | [1,4,2,2] | Conv 3x3 then Pool 2x2 | 2 | none | Regression for `pool(conv(...))` |
### Add (5)
## ReduceMean
| Case | Description |
|---|---|
| `after_gemm` | Gemm followed by Add with a broadcast bias vector. |
| `basic` | Elementwise Add on two inputs with identical shapes. |
| `broadcast_row` | Elementwise Add with row-vector broadcasting. |
| `channel_broadcast_1024` | NCHW per-channel broadcasting over 1024 channels. |
| `leading_dimension_broadcast` | Trailing-dimension broadcasting across leading dimensions. |
| Test | Directory | Input | Output | Axes | Keepdims | Notes |
|------------|--------------------------|-----------|-----------|-------|----------|-------------------------------------------------|
| Basic | `reduce_mean/basic` | [4,8] | [4,1] | [1] | 1 | Reduce feature dimension, preserving rank |
| Keepdims 0 | `reduce_mean/keepdims_0` | [4,8] | [4] | [1] | 0 | Reduce feature dimension, dropping reduced axis |
| 4D spatial | `reduce_mean/4d_spatial` | [1,3,4,4] | [1,3,1,1] | [2,3] | 1 | Reduce H and W on NCHW input |
| After Conv | `reduce_mean/after_conv` | [1,3,5,5] | [1,2,1,1] | [2,3] | 1 | Conv 3x3 + bias, then spatial ReduceMean |
### Concat (3)
## Relu
| Case | Description |
|---|---|
| `channel_axis` | Concatenates two runtime NCHW tensors along the channel axis. |
| `negative_axis` | Concatenates tensors using a negative axis. |
| `three_inputs_channel_axis` | Concatenates three runtime NCHW tensors along the channel axis. |
| Test | Directory | Input | Output | Notes |
|------------|-------------------|-----------|-----------|----------------------------|
| Basic | `relu/basic` | [4,8] | [4,8] | Standalone 2D Relu |
| 4D | `relu/4d` | [2,3,4,4] | [2,3,4,4] | Standalone NCHW Relu |
| After Conv | `relu/after_conv` | [1,3,5,5] | [1,2,3,3] | Conv 3x3 + bias, then Relu |
| After Gemm | `relu/after_gemm` | [4,64] | [4,32] | Gemm + bias, then Relu |
### Conv (31)
## Sigmoid
| Case | Description |
|---|---|
| `batch_2` | Batched Conv with SAME_UPPER padding and bias. |
| `batch_4_pointwise` | Pointwise Conv with batch size four. |
| `depthwise_1024_channels` | Depthwise pointwise Conv with 1024 groups. |
| `depthwise_grouped` | Depthwise-style grouped Conv with one input channel per group. |
| `dilated_3x3` | Conv with a dilated 3x3 kernel. |
| `dynamic` | Conv with runtime input and weight tensors. |
| `explicit_padding` | 3x3 Conv with symmetric explicit padding. |
| `grouped_many_groups` | Pointwise Conv with many groups and high channel counts. |
| `grouped_two_groups` | Two-group pointwise Conv with bias. |
| `huge_pointwise_1024` | Pointwise Conv with 1024 input and output channels. |
| `huge_pointwise_1024_dynamic` | The 1024-channel pointwise Conv with runtime weights. |
| `kernel_3x3` | Basic 3x3 Conv without padding. |
| `kernel_equals_input_spatial` | Conv whose kernel covers the full spatial input. |
| `large_input_channels_1x1` | Pointwise Conv with 1024 input channels and modest output width. |
| `large_output_channels_1x1` | Pointwise Conv with modest input width and 1024 output channels. |
| `large_spatial` | 3x3 Conv on a larger spatial input. |
| `multi_channel` | 3x3 Conv with multiple input and output channels. |
| `non_square_kernel_1x3` | Conv with a non-square 1x3 kernel. |
| `non_square_kernel_3x1` | Conv with a non-square 3x1 kernel. |
| `non_uniform_stride` | Conv with different height and width strides. |
| `pointwise_1x1` | Basic pointwise channel-mixing Conv. |
| `pointwise_tiled_chain` | Relu and chained pointwise Convs with a tiled intermediate. |
| `real_asymmetric_padding` | Conv with asymmetric explicit padding. |
| `relu_conv_store` | Conv, Relu, and a second Conv to validate an intermediate stored result. |
| `same_lower_3x3` | 3x3 Conv with SAME_LOWER padding. |
| `same_padding_3x3` | 3x3 Conv with SAME_UPPER padding. |
| `simple` | Hand-authored basic 2x2 Conv. |
| `stride_2` | 3x3 Conv with stride two. |
| `with_bias_3x3` | Multi-channel 3x3 Conv with bias. |
| `with_constant` | Hand-authored SAME_UPPER Conv with constant weight and bias. |
| `without_kernel_shape_attr` | Conv whose kernel shape is inferred from its weight tensor. |
| Test | Directory | Input | Output | Notes |
|------------|----------------------|-----------|-----------|---------------------------|
| Basic | `sigmoid/basic` | [4,8] | [4,8] | Standalone 2D Sigmoid |
| 4D | `sigmoid/4d` | [2,3,4,4] | [2,3,4,4] | Standalone NCHW Sigmoid |
| After Gemm | `sigmoid/after_gemm` | [4,64] | [4,32] | Gemm + bias, then Sigmoid |
### Div (6)
## Softmax
| Case | Description |
|---|---|
| `after_gemm` | Gemm followed by Div with a broadcast divisor vector. |
| `basic` | Elementwise Div by a same-shape constant tensor. |
| `channel_broadcast_1024` | Div with NCHW per-channel broadcasting over 1024 channels. |
| `leading_dimension_broadcast` | Div with trailing-dimension broadcasting. |
| `runtime_scalar_rhs` | Div of a runtime tensor by a scalar initializer. |
| `scalar_constant` | Div with scalar broadcasting on a 2D tensor. |
| Test | Directory | Input | Output | Axis | Notes |
|--------------|--------------------------|-------------|-------------|------|---------------------------------|
| Basic | `softmax/basic` | [3,5] | [3,5] | 1 | Row-wise softmax over features |
| 3D last axis | `softmax/3d_last_axis` | [2,3,4] | [2,3,4] | 2 | Last-dimension normalization |
| Channel axis | `softmax/channel_axis` | [1,3,2,2] | [1,3,2,2] | 1 | NCHW channel-wise softmax |
### Gather (5)
## Resize
| Case | Description |
|---|---|
| `3d_input_axis1` | Gathers along axis one of a 3D input. |
| `axis0_matrix_indices` | Gathers rows using a 2D indices tensor. |
| `axis1` | Gathers selected columns from a 2D tensor. |
| `negative_axis` | Gathers using a negative axis. |
| `negative_indices` | Gathers with negative indices along axis zero. |
| Test | Directory | Input | Output | Mode | Notes |
|---------------------|-------------------------|-----------|-----------|---------|-----------------------------------------|
| Nearest 2x | `resize/nearest_2x` | [1,1,2,3] | [1,1,4,6] | nearest | NCHW upsampling with scales [1,1,2,2] |
| Non-uniform scales | `resize/non_uniform` | [1,1,2,3] | [1,1,6,6] | nearest | Different height/width scaling factors |
| Explicit sizes | `resize/with_sizes` | [1,1,2,3] | [1,1,3,5] | nearest | Sizes input used instead of scales |
### Gemm (21)
## Split
| Case | Description |
|---|---|
| `alpha_beta` | Applies non-default alpha and beta scaling with bias. |
| `bias_rank2_broadcast` | Broadcasts a rank-2 bias across output rows. |
| `dynamic` | Uses both matrix operands at runtime. |
| `dynamic_alpha` | Uses runtime operands with non-default alpha scaling. |
| `dynamic_beta` | Uses runtime operands and bias with non-default beta scaling. |
| `dynamic_bias` | Uses runtime matrix operands and runtime bias. |
| `dynamic_bias_alpha_beta` | Combines runtime operands and bias with alpha and beta scaling. |
| `dynamic_transB` | Transposes a runtime right-hand matrix. |
| `huge_1024` | Uses 1024-wide inner and output dimensions. |
| `large` | Exercises larger rectangular matrices. |
| `large_k_small_n` | Uses a large reduction dimension and narrow output. |
| `non_square` | Uses different reduction and output widths. |
| `scalar_bias` | Broadcasts a scalar bias to the full output. |
| `simple` | Basic Gemm with square weights. |
| `small` | Tiny Gemm for fast focused validation. |
| `small_k_large_n` | Uses a modest reduction dimension and wide output. |
| `transA` | Transposes the left-hand matrix. |
| `transA_transB` | Transposes both matrix operands. |
| `transB` | Transposes the right-hand weight matrix. |
| `transB_with_bias` | Combines a transposed weight matrix with bias. |
| `with_bias` | Basic matrix product with vector bias. |
| Test | Directory | Input | Outputs | Axis | Notes |
|-----------------|---------------------------|-------|----------------------|------|-------------------------------------|
| Basic | `split/basic` | [2,6] | [2,2], [2,4] | 1 | Two-way split with explicit sizes |
| Equal three-way | `split/equal_three_way` | [2,6] | [2,2], [2,2], [2,2] | 1 | Optional split input omitted |
### Gemv (5)
## Gather
| Case | Description |
|---|---|
| `constant` | Vector-matrix product with all inputs constant. |
| `simple` | Basic single-row vector-matrix product. |
| `with_heterogeneous_constant` | Adds a non-uniform constant bias pattern. |
| `with_homogeneous_constant` | Adds a constant bias matching the output shape. |
| `with_scalar_constant` | Adds a scalar broadcast bias. |
| Test | Directory | Input | Indices | Output | Axis | Notes |
|----------------------|--------------------------------|-------|---------|----------|------|--------------------------------|
| Axis 1 | `gather/axis1` | [3,4] | [2] | [3,2] | 1 | Select two columns |
| Axis 0 matrix indices| `gather/axis0_matrix_indices` | [4,3] | [2,2] | [2,2,3] | 0 | Gather rows with 2D indices |
### MatMul (11)
## Concat
| Case | Description |
|---|---|
| `basic` | Direct 2D MatMul with constant right-hand matrix. |
| `batched_3d` | Batched 3D MatMul with matching batch dimensions. |
| `batched_3d_dynamic` | Batched 3D MatMul with both operands at runtime. |
| `batched_left_constant` | Batched 3D MatMul with constant left-hand matrix. |
| `batched_lhs_broadcast` | Broadcasts a 2D left-hand matrix across a batched right-hand tensor. |
| `batched_rhs_broadcast` | Broadcasts a 2D right-hand matrix across a batched left-hand tensor. |
| `dynamic` | Direct 2D MatMul with both operands at runtime. |
| `huge_1024` | Uses 1024-wide inner and output dimensions. |
| `left_constant` | Direct 2D MatMul with constant left-hand matrix. |
| `matrix_vector` | Matrix-vector multiplication producing a 1D output. |
| `vector_matrix` | Vector-matrix multiplication producing a 1D output. |
| Test | Directory | Input(s) | Output | Axis | Notes |
|--------------|-----------------------|---------------------------|-----------|------|-----------------------------|
| Channel axis | `concat/channel_axis` | A:[1,1,2,2], B:[1,2,2,2] | [1,3,2,2] | 1 | Runtime NCHW channel concat |
### Mul (5)
## Reshape
| Case | Description |
|---|---|
| `after_conv` | Conv followed by per-channel scaling. |
| `basic` | Elementwise Mul on two inputs with identical shapes. |
| `channel_broadcast_1024` | Mul with NCHW per-channel broadcasting over 1024 channels. |
| `leading_dimension_broadcast` | Mul with trailing-dimension broadcasting. |
| `scalar_constant` | Mul with scalar broadcasting. |
| Test | Directory | Input | Output | Notes |
|-----------|---------------------|-------|--------|----------------------------------------------|
| Same rank | `reshape/same_rank` | [2,3] | [3,2] | Runtime tensor with static shape initializer |
### Pool (15)
## Add
| Case | Description |
|---|---|
| `avg_basic` | AveragePool with a 2x2 kernel and unit stride. |
| `avg_ceil_mode` | AveragePool with ceil mode enabled. |
| `avg_explicit_padding` | Explicitly padded AveragePool excluding pad from the divisor. |
| `avg_include_pad` | Explicitly padded AveragePool including pad in the divisor. |
| `avg_large_channels` | AveragePool with a large channel count and small spatial extent. |
| `avg_non_uniform_stride` | AveragePool with different height and width strides. |
| `avg_real_asymmetric_padding` | AveragePool with asymmetric explicit padding. |
| `max_after_conv` | Conv followed by MaxPool. |
| `max_basic` | MaxPool with a 2x2 kernel and unit stride. |
| `max_ceil_mode` | MaxPool with ceil mode enabled. |
| `max_global_style_kernel_equals_input` | MaxPool whose kernel covers the full spatial input. |
| `max_non_square_kernel` | MaxPool with a non-square kernel. |
| `max_real_asymmetric_padding` | MaxPool with asymmetric explicit padding. |
| `max_same_upper` | MaxPool with SAME_UPPER padding. |
| `max_stride2_multichannel` | Multi-channel MaxPool with stride two. |
| Test | Directory | Input(s) | Output | Notes |
|---------------|---------------------|------------------|--------|---------------------------------------------|
| Basic | `add/basic` | A:[4,8], B:[4,8] | [4,8] | Elementwise add, same-shape inputs |
| Broadcast row | `add/broadcast_row` | A:[4,8], B:[8] | [4,8] | Row-vector broadcasting via initializer |
| After Gemm | `add/after_gemm` | A:[4,64], D:[32] | [4,32] | Gemm + bias, then Add with broadcast vector |
### ReduceMean (17)
## Mul
| Case | Description |
|---|---|
| `4d_spatial` | Reduces height and width of an NCHW tensor while preserving rank. |
| `4d_spatial_keepdims_0` | Reduces NCHW height and width while dropping those axes. |
| `after_conv` | Conv followed by a spatial ReduceMean. |
| `all_axes_keepdims_0` | Reduces all axes to a scalar. |
| `all_axes_keepdims_1` | Reduces all axes while preserving rank. |
| `basic` | Reduces a feature dimension while preserving rank. |
| `channel_axis_nchw` | Reduces the channel axis of an NCHW tensor. |
| `keepdims_0` | Reduces a feature dimension and drops that axis. |
| `large_dimension_1024` | Reduces a dimension of length 1024. |
| `legacy_axes_1_2_keepdims_1` | Opset-18 reduction over multiple positive axes. |
| `legacy_axis1_keepdims_0` | Opset-18 reduction over one axis while dropping it. |
| `legacy_axis1_keepdims_1` | Opset-18 reduction over one axis while preserving rank. |
| `legacy_empty_axes_noop` | Opset-18 empty-axes no-op followed by Relu. |
| `legacy_nchw_spatial` | Opset-18 spatial reduction on NCHW input. |
| `legacy_negative_axis` | Opset-18 reduction using a negative axis. |
| `legacy_reduce_all_keepdims_1` | Opset-18 all-axis reduction with the axes input omitted. |
| `negative_axis` | ReduceMean using a negative axis. |
| Test | Directory | Input(s) | Output | Notes |
|-----------------|-----------------------|--------------------------|-----------|-------------------------------------------|
| Basic | `mul/basic` | A:[4,8], B:[4,8] | [4,8] | Elementwise multiply, same-shape inputs |
| Scalar constant | `mul/scalar_constant` | X:[4,8], S:[1] | [4,8] | Scalar broadcasting via initializer |
| After Conv | `mul/after_conv` | X:[1,3,5,5], S:[1,2,1,1] | [1,2,3,3] | Conv 3x3 + bias, then per-channel scaling |
### Relu (4)
## Div
| Case | Description |
|---|---|
| `4d` | Standalone Relu on an NCHW tensor. |
| `after_conv` | Conv followed by Relu. |
| `after_gemm` | Gemm followed by Relu. |
| `basic` | Standalone Relu on a 2D tensor. |
| Test | Directory | Input(s) | Output | Notes |
|-----------------|-----------------------|------------------|--------|------------------------------------------------------|
| Basic | `div/basic` | X:[4,8], D:[4,8] | [4,8] | Elementwise divide by same-shape constant tensor |
| Scalar constant | `div/scalar_constant` | X:[4,8], S:[1] | [4,8] | Scalar broadcasting via initializer |
| After Gemm | `div/after_gemm` | A:[4,64], D:[32] | [4,32] | Gemm + bias, then Div with positive broadcast vector |
### Reshape (4)
| Case | Description |
|---|---|
| `4d_to_2d_flatten` | Flattens a 4D tensor into a 2D view. |
| `infer_dim_minus_one` | Uses `-1` to infer one output dimension. |
| `same_rank` | Changes shape without changing rank. |
| `zero_copies_input_dim` | Uses `0` to copy an input dimension. |
### Resize (6)
| Case | Description |
|---|---|
| `height_only` | Nearest-neighbor resize of only the height dimension. |
| `nearest_2x` | Nearest-neighbor upsampling by a factor of two. |
| `nearest_downsample` | Nearest-neighbor downsampling. |
| `non_uniform` | Nearest-neighbor resize with different spatial scales. |
| `width_only` | Nearest-neighbor resize of only the width dimension. |
| `with_sizes` | Resize using explicit output sizes instead of scales. |
### Sigmoid (3)
| Case | Description |
|---|---|
| `4d` | Standalone Sigmoid on an NCHW tensor. |
| `after_gemm` | Gemm followed by Sigmoid. |
| `basic` | Standalone Sigmoid on a 2D tensor. |
### Slice (8)
| Case | Description |
|---|---|
| `2d_basic` | Slices a 2D tensor with explicit axes and unit steps. |
| `after_conv` | Conv followed by a spatial crop. |
| `default_axes` | Omits axes and steps to use positional defaults. |
| `large_channel_1024` | Slices a channel range from a 1024-channel tensor. |
| `nchw_spatial_crop` | Crops the spatial axes of an NCHW tensor. |
| `negative_axis` | Slices using a negative axis. |
| `negative_indices` | Slices using negative indices. |
| `step2` | Slices using a positive step greater than one. |
### Softmax (5)
| Case | Description |
|---|---|
| `3d_last_axis` | Softmax over the last axis of a 3D tensor. |
| `basic` | Softmax over the last dimension of a 2D tensor. |
| `channel_axis` | Softmax over the channel axis of an NCHW tensor. |
| `large_dimension_1024` | Softmax over a last dimension of length 1024. |
| `negative_axis` | Softmax using a negative axis. |
### Split (4)
| Case | Description |
|---|---|
| `basic` | Splits a 2D tensor into two explicit output sizes. |
| `equal_three_way` | Splits a 2D tensor evenly into three outputs. |
| `negative_axis` | Splits using a negative axis. |
| `uneven_channel_axis_4d` | Splits an NCHW channel axis into uneven outputs. |
### Sub (6)
| Case | Description |
|---|---|
| `after_gemm` | Gemm followed by Sub with a broadcast constant vector. |
| `basic` | Elementwise Sub on two runtime inputs with identical shapes. |
| `broadcast_row` | Sub with a broadcast row-vector right-hand constant. |
| `channel_broadcast_1024` | Sub with NCHW per-channel broadcasting over 1024 channels. |
| `constant_lhs_broadcast` | Sub with a broadcast constant left-hand operand. |
| `leading_dimension_broadcast` | Sub with trailing-dimension broadcasting. |
+44 -16
View File
@@ -80,7 +80,7 @@ def conv_1x1():
kernel_shape=[1, 1], strides=[1, 1], pads=[0, 0, 0, 0])
graph = helper.make_graph([node], "conv_1x1", [X], [Y], initializer=[W])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "conv/pointwise_1x1", "conv_1x1.onnx")
save_model(model, "conv/pointwise_1x1", "conv_pointwise_1x1.onnx")
def conv_same_padding_3x3():
@@ -218,6 +218,33 @@ def conv_huge_pointwise_1024_dynamic():
save_model(model, "conv/huge_pointwise_1024_dynamic", "conv_huge_pointwise_1024_dynamic.onnx")
def conv_pointwise_tiled_chain():
"""Chained pointwise Convs with a tiled intermediate."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1024, 1, 1])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 256, 1, 1])
rng = np.random.default_rng(79)
W1 = numpy_helper.from_array(
rng.uniform(-1, 1, (1024, 1024, 1, 1)).astype(np.float32), name="W1")
B1 = numpy_helper.from_array(
rng.uniform(-1, 1, 1024).astype(np.float32), name="B1")
W2 = numpy_helper.from_array(
rng.uniform(-1, 1, (256, 1024, 1, 1)).astype(np.float32), name="W2")
B2 = numpy_helper.from_array(
rng.uniform(-1, 1, 256).astype(np.float32), name="B2")
nodes = [
helper.make_node("Relu", ["X"], ["X_relu"]),
helper.make_node("Conv", ["X_relu", "W1", "B1"], ["hidden"],
kernel_shape=[1, 1], strides=[1, 1], pads=[0, 0, 0, 0]),
helper.make_node("Relu", ["hidden"], ["hidden_relu"]),
helper.make_node("Conv", ["hidden_relu", "W2", "B2"], ["Y"],
kernel_shape=[1, 1], strides=[1, 1], pads=[0, 0, 0, 0]),
]
graph = helper.make_graph(
nodes, "conv_pointwise_tiled_chain", [X], [Y], initializer=[W1, B1, W2, B2])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "conv/pointwise_tiled_chain", "conv_pointwise_tiled_chain.onnx")
def conv_large_output_channels_1x1():
"""1x1 Conv with modest inputs and very large output channel count."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 64, 1, 1])
@@ -790,7 +817,7 @@ def maxpool_basic():
node = helper.make_node("MaxPool", ["X"], ["Y"], kernel_shape=[2, 2], strides=[1, 1], pads=[0, 0, 0, 0])
graph = helper.make_graph([node], "maxpool_basic", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/max_basic", "maxpool_basic.onnx")
save_model(model, "pool/max_basic", "pool_max_basic.onnx")
def maxpool_stride2_multichannel():
@@ -800,7 +827,7 @@ def maxpool_stride2_multichannel():
node = helper.make_node("MaxPool", ["X"], ["Y"], kernel_shape=[2, 2], strides=[2, 2], pads=[0, 0, 0, 0])
graph = helper.make_graph([node], "maxpool_stride2_multichannel", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/max_stride2_multichannel", "maxpool_stride2_multichannel.onnx")
save_model(model, "pool/max_stride2_multichannel", "pool_max_stride2_multichannel.onnx")
def maxpool_same_upper():
@@ -810,7 +837,7 @@ def maxpool_same_upper():
node = helper.make_node("MaxPool", ["X"], ["Y"], kernel_shape=[3, 3], strides=[2, 2], auto_pad="SAME_UPPER")
graph = helper.make_graph([node], "maxpool_same_upper", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/max_same_upper", "maxpool_same_upper.onnx")
save_model(model, "pool/max_same_upper", "pool_max_same_upper.onnx")
def avgpool_basic():
@@ -820,7 +847,7 @@ def avgpool_basic():
node = helper.make_node("AveragePool", ["X"], ["Y"], kernel_shape=[2, 2], strides=[1, 1], pads=[0, 0, 0, 0])
graph = helper.make_graph([node], "avgpool_basic", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/avg_basic", "avgpool_basic.onnx")
save_model(model, "pool/avg_basic", "pool_avg_basic.onnx")
def avgpool_explicit_padding():
@@ -831,7 +858,7 @@ def avgpool_explicit_padding():
kernel_shape=[3, 3], strides=[2, 2], pads=[1, 1, 1, 1], count_include_pad=0)
graph = helper.make_graph([node], "avgpool_explicit_padding", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/avg_explicit_padding", "avgpool_explicit_padding.onnx")
save_model(model, "pool/avg_explicit_padding", "pool_avg_explicit_padding.onnx")
def avgpool_include_pad():
@@ -842,7 +869,7 @@ def avgpool_include_pad():
kernel_shape=[3, 3], strides=[2, 2], pads=[1, 1, 1, 1], count_include_pad=1)
graph = helper.make_graph([node], "avgpool_include_pad", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/avg_include_pad", "avgpool_include_pad.onnx")
save_model(model, "pool/avg_include_pad", "pool_avg_include_pad.onnx")
def maxpool_after_conv():
@@ -855,7 +882,7 @@ def maxpool_after_conv():
pool = helper.make_node("MaxPool", ["C"], ["Y"], kernel_shape=[2, 2], strides=[2, 2], pads=[0, 0, 0, 0])
graph = helper.make_graph([conv, pool], "maxpool_after_conv", [X], [Y], initializer=[W])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/max_after_conv", "maxpool_after_conv.onnx")
save_model(model, "pool/max_after_conv", "pool_max_after_conv.onnx")
def maxpool_ceil_mode():
@@ -866,7 +893,7 @@ def maxpool_ceil_mode():
kernel_shape=[2, 2], strides=[2, 2], pads=[0, 0, 0, 0], ceil_mode=1)
graph = helper.make_graph([node], "maxpool_ceil_mode", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/max_ceil_mode", "maxpool_ceil_mode.onnx")
save_model(model, "pool/max_ceil_mode", "pool_max_ceil_mode.onnx")
def avgpool_ceil_mode():
@@ -877,7 +904,7 @@ def avgpool_ceil_mode():
kernel_shape=[2, 2], strides=[2, 2], pads=[0, 0, 0, 0], ceil_mode=1)
graph = helper.make_graph([node], "avgpool_ceil_mode", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/avg_ceil_mode", "avgpool_ceil_mode.onnx")
save_model(model, "pool/avg_ceil_mode", "pool_avg_ceil_mode.onnx")
def maxpool_real_asymmetric_padding():
@@ -888,7 +915,7 @@ def maxpool_real_asymmetric_padding():
kernel_shape=[3, 3], strides=[1, 2], pads=[0, 1, 2, 1])
graph = helper.make_graph([node], "maxpool_real_asymmetric_padding", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/max_real_asymmetric_padding", "maxpool_real_asymmetric_padding.onnx")
save_model(model, "pool/max_real_asymmetric_padding", "pool_max_real_asymmetric_padding.onnx")
def avgpool_real_asymmetric_padding():
@@ -899,7 +926,7 @@ def avgpool_real_asymmetric_padding():
kernel_shape=[3, 3], strides=[1, 2], pads=[0, 1, 2, 1], count_include_pad=0)
graph = helper.make_graph([node], "avgpool_real_asymmetric_padding", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/avg_real_asymmetric_padding", "avgpool_real_asymmetric_padding.onnx")
save_model(model, "pool/avg_real_asymmetric_padding", "pool_avg_real_asymmetric_padding.onnx")
def maxpool_non_square_kernel():
@@ -910,7 +937,7 @@ def maxpool_non_square_kernel():
kernel_shape=[2, 3], strides=[1, 2], pads=[0, 0, 0, 0])
graph = helper.make_graph([node], "maxpool_non_square_kernel", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/max_non_square_kernel", "maxpool_non_square_kernel.onnx")
save_model(model, "pool/max_non_square_kernel", "pool_max_non_square_kernel.onnx")
def avgpool_non_uniform_stride():
@@ -921,7 +948,7 @@ def avgpool_non_uniform_stride():
kernel_shape=[2, 3], strides=[1, 2], pads=[0, 0, 0, 0])
graph = helper.make_graph([node], "avgpool_non_uniform_stride", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/avg_non_uniform_stride", "avgpool_non_uniform_stride.onnx")
save_model(model, "pool/avg_non_uniform_stride", "pool_avg_non_uniform_stride.onnx")
def maxpool_global_style_kernel_equals_input():
@@ -931,7 +958,7 @@ def maxpool_global_style_kernel_equals_input():
node = helper.make_node("MaxPool", ["X"], ["Y"], kernel_shape=[4, 4], strides=[1, 1], pads=[0, 0, 0, 0])
graph = helper.make_graph([node], "maxpool_global_style_kernel_equals_input", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/max_global_style_kernel_equals_input", "maxpool_global_style_kernel_equals_input.onnx")
save_model(model, "pool/max_global_style_kernel_equals_input", "pool_max_global_style_kernel_equals_input.onnx")
def avgpool_large_channels():
@@ -941,7 +968,7 @@ def avgpool_large_channels():
node = helper.make_node("AveragePool", ["X"], ["Y"], kernel_shape=[2, 2], strides=[1, 1], pads=[0, 0, 0, 0])
graph = helper.make_graph([node], "avgpool_large_channels", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "pool/avg_large_channels", "avgpool_large_channels.onnx")
save_model(model, "pool/avg_large_channels", "pool_avg_large_channels.onnx")
# ---------------------------------------------------------------------------
@@ -2005,6 +2032,7 @@ if __name__ == "__main__":
conv_dynamic()
conv_huge_pointwise_1024()
conv_huge_pointwise_1024_dynamic()
conv_pointwise_tiled_chain()
conv_large_output_channels_1x1()
conv_large_input_channels_1x1()
conv_depthwise_1024_channels()
+12 -16
View File
@@ -174,22 +174,18 @@ def main():
# Summary
n_passed = sum(1 for passed in results.values() if passed)
n_total = len(results)
failing = [rel for rel, passed in results.items() if not passed]
if a.verbose or failing:
status_width = len("Result")
path_width = max(len("Operation"), *(len(rel) for rel in results))
separator = f"+-{'-' * path_width}-+-{'-' * status_width}-+"
print(separator)
print(f"| {'Operation'.ljust(path_width)} | {'Result'.ljust(status_width)} |")
print(separator)
for rel, passed in results.items():
if not a.verbose and passed:
continue
plain_status = "PASS" if passed else "FAIL"
status = Fore.GREEN + plain_status.ljust(status_width) + Style.RESET_ALL if passed else \
Fore.RED + plain_status.ljust(status_width) + Style.RESET_ALL
print(f"| {rel.ljust(path_width)} | {status} |")
print(separator)
status_width = len("Result")
path_width = max(len("Operation"), *(len(rel) for rel in results))
separator = f"+-{'-' * path_width}-+-{'-' * status_width}-+"
print(separator)
print(f"| {'Operation'.ljust(path_width)} | {'Result'.ljust(status_width)} |")
print(separator)
for rel, passed in results.items():
plain_status = "PASS" if passed else "FAIL"
status = Fore.GREEN + plain_status.ljust(status_width) + Style.RESET_ALL if passed else \
Fore.RED + plain_status.ljust(status_width) + Style.RESET_ALL
print(f"| {rel.ljust(path_width)} | {status} |")
print(separator)
print("\n" + Style.BRIGHT + Fore.CYAN + "Summary" + Style.RESET_ALL)
print(Style.BRIGHT + f"Passed: {n_passed}" + Style.RESET_ALL)
print(Style.BRIGHT + f"Failed: {n_total - n_passed}" + Style.RESET_ALL)