nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
0aa3840a724fb84c557ca2de309201d2c2a8d038
Raptor/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp
T
ilgeco 0aa3840a72
Validate Operations / validate-operations (push) Has been cancelled
Details
Conv lowering report
2026-08-03 12:05:03 +02:00

3654 lines
184 KiB
C++
Raw Blame History

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/raw_ostream.h"
#include <algorithm>
#include <mutex>
#include <optional>
#include <string>
#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.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/MatrixProductLowering.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"
using namespace mlir;
namespace onnx_mlir {
namespace {
struct ConvToGemm : OpConversionPattern<ONNXConvOp> {
explicit ConvToGemm(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
: OpConversionPattern<ONNXConvOp>(ctx), target(target) {}
LogicalResult matchAndRewrite(ONNXConvOp convOp,
ONNXConvOpAdaptor convOpAdaptor,
ConversionPatternRewriter& rewriter) const override;
const spatial::SpatialTargetInfo& target;
};
struct PreparedConvInput {
Value value;
RankedTensorType type;
};
static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter);
static StringRef stringifyConvLoweringStrategy(spatial::ConvLoweringStrategy strategy) {
switch (strategy) {
case spatial::ConvLoweringStrategy::Auto: return "auto";
case spatial::ConvLoweringStrategy::Legacy: return "legacy";
case spatial::ConvLoweringStrategy::Depthwise: return "depthwise";
case spatial::ConvLoweringStrategy::PackedIm2Col: return "packed-im2col";
case spatial::ConvLoweringStrategy::StreamedPatch: return "streamed-patch";
case spatial::ConvLoweringStrategy::StreamedPacked: return "streamed-packed";
case spatial::ConvLoweringStrategy::OutputChannelTiled: return "output-channel-tiled";
case spatial::ConvLoweringStrategy::InputKTiled: return "input-k-tiled";
case spatial::ConvLoweringStrategy::Tiled2D: return "tiled-2d";
}
llvm_unreachable("unknown conv lowering strategy");
}
enum class ConvLoweringReportPhase { Planning, Realization };
struct ConvLoweringReportEntry {
size_t convId;
std::string phase;
std::string location;
std::string strategy;
std::string implementation;
};
struct ConvLoweringReportState {
std::mutex mutex;
llvm::SmallPtrSet<Operation*, 32> planned;
llvm::SmallPtrSet<Operation*, 32> realized;
llvm::DenseMap<Operation*, size_t> convIds;
llvm::SmallVector<ConvLoweringReportEntry, 256> entries;
size_t nextConvId = 1;
};
static StringRef stringifyConvLoweringReportPhase(ConvLoweringReportPhase phase) {
return phase == ConvLoweringReportPhase::Planning ? "planning" : "realization";
}
static std::string convReportLocation(Operation* op) {
std::string location;
llvm::raw_string_ostream stream(location);
op->getLoc().print(stream);
if (location.size() > 120)
location.replace(117, std::string::npos, "...");
return location;
}
static StringRef convLoweringImplementation(spatial::ConvLoweringStrategy strategy) {
switch (strategy) {
case spatial::ConvLoweringStrategy::Depthwise:
return "DW";
case spatial::ConvLoweringStrategy::Legacy:
case spatial::ConvLoweringStrategy::PackedIm2Col:
return "PIC";
case spatial::ConvLoweringStrategy::StreamedPatch:
case spatial::ConvLoweringStrategy::OutputChannelTiled:
case spatial::ConvLoweringStrategy::Tiled2D:
return "STR";
case spatial::ConvLoweringStrategy::InputKTiled:
return "IKT";
case spatial::ConvLoweringStrategy::StreamedPacked:
return "STP";
case spatial::ConvLoweringStrategy::Auto:
return "AUTO";
}
llvm_unreachable("unknown conv lowering implementation");
}
static StringRef convRowStripInputImplementation(const ConvLoweringState& state,
spatial::ConvLoweringStrategy strategy) {
if (strategy == spatial::ConvLoweringStrategy::Depthwise)
return "RSDW";
if (state.xHeight == 1 && state.xWidth == 1 && state.wHeight == 1 && state.wWidth == 1)
return "RSP";
return "RSM";
}
static constexpr size_t kConvReportIdWidth = 4;
static constexpr size_t kConvReportLocationWidth = 24;
static constexpr size_t kConvReportStrategyWidth = 20;
static constexpr size_t kConvReportCodeWidth = 8;
static std::string convReportCell(StringRef value, size_t width) {
std::string cell = value.str();
if (cell.size() > width) {
cell = width <= 3 ? std::string(width, '.') : cell.substr(0, width - 3) + "...";
}
cell.append(width - cell.size(), ' ');
return cell;
}
static void writeConvReportTableHeader(std::fstream& reportFile, StringRef fourthColumn) {
reportFile << "+------+--------------------------+----------------------+----------+\n";
reportFile << "| " << convReportCell("Conv", kConvReportIdWidth) << " | "
<< convReportCell("Location", kConvReportLocationWidth) << " | "
<< convReportCell("Strategy", kConvReportStrategyWidth) << " | "
<< convReportCell(fourthColumn, kConvReportCodeWidth) << " |\n";
reportFile << "+------+--------------------------+----------------------+----------+\n";
}
static void writeConvReportRow(std::fstream& reportFile,
const ConvLoweringReportEntry& entry) {
reportFile << "| " << convReportCell(std::to_string(entry.convId), kConvReportIdWidth) << " | "
<< convReportCell(entry.location, kConvReportLocationWidth) << " | "
<< convReportCell(entry.strategy, kConvReportStrategyWidth) << " | "
<< convReportCell(entry.implementation, kConvReportCodeWidth) << " |\n";
}
static void writeConvReportLegend(std::fstream& reportFile) {
reportFile << "Legend: Conv is shared by both sections; codes expand to:\n";
reportFile << " SEL selectConvLoweringPlan\n";
reportFile << " DW depthwise::rewriteConv\n";
reportFile << " PIC standard::rewritePackedIm2ColConv\n";
reportFile << " STR standard::rewriteStreamedConv(pack=1)\n";
reportFile << " IKT standard::rewriteInputKTiledConv\n";
reportFile << " STP standard::rewriteStreamedConv(pack=geo.pack)\n";
reportFile << " AUTO unresolved strategy\n";
reportFile << " RSD createRowStripConvOutputFromDenseInput -> createRowStripConvOutput\n";
reportFile << " RSDW createConvOutputFromRowStripInput -> createDepthwiseOutputFromRowStripFragments\n";
reportFile << " RSP createConvOutputFromRowStripInput -> createPointwiseOutputFromRowStripFragments\n";
reportFile << " RSM createConvOutputFromRowStripInput -> createConvOutputFromPixelMajorRowStripFragments\n\n";
}
static bool writeConvLoweringReport(const ConvLoweringReportEntry& entry,
ConvLoweringReportState& state) {
state.entries.push_back(entry);
std::fstream reportFile = openReportFile("conv_lowering_report");
if (!reportFile.is_open()) {
state.entries.pop_back();
return false;
}
reportFile << "# PIM Conv Lowering Report (bounded to 512 rows)\n\n";
reportFile << "## Plan selection\n";
writeConvReportTableHeader(reportFile, "Selector");
bool realizationSectionStarted = false;
for (const ConvLoweringReportEntry& reportEntry : state.entries) {
if (reportEntry.phase == "realization" && !realizationSectionStarted) {
reportFile << "\n## Realization\n";
writeConvReportTableHeader(reportFile, "Code");
realizationSectionStarted = true;
}
writeConvReportRow(reportFile, reportEntry);
}
reportFile << "\n";
writeConvReportLegend(reportFile);
if (!reportFile.good()) {
state.entries.pop_back();
return false;
}
return true;
}
static void recordConvLoweringReport(Operation* op,
ConvLoweringReportPhase phase,
spatial::ConvLoweringStrategy strategy,
StringRef implementation) {
if (!pimReportConvLowering)
return;
static ConvLoweringReportState state;
std::lock_guard<std::mutex> lock(state.mutex);
if (state.entries.size() >= 512)
return;
if (phase == ConvLoweringReportPhase::Planning) {
if (state.planned.contains(op))
return;
}
else {
if (state.realized.contains(op))
return;
}
size_t convId = state.convIds.lookup(op);
if (!convId) {
convId = state.nextConvId++;
state.convIds[op] = convId;
}
ConvLoweringReportEntry entry {convId,
stringifyConvLoweringReportPhase(phase).str(),
convReportLocation(op),
stringifyConvLoweringStrategy(strategy).str(),
implementation.str()};
if (!writeConvLoweringReport(entry, state))
return;
if (phase == ConvLoweringReportPhase::Planning)
state.planned.insert(op);
else
state.realized.insert(op);
}
static Value expandBiasIfNeeded(Value bias, PatternRewriter& rewriter, Location loc) {
auto biasType = cast<RankedTensorType>(bias.getType());
if (biasType.getRank() != 1)
return bias;
auto expandedBiasType = RankedTensorType::get({1, biasType.getDimSize(0)}, biasType.getElementType());
return tensor::ExpandShapeOp::create(rewriter,
loc,
expandedBiasType,
bias,
SmallVector<ReassociationIndices> {
{0, 1}
});
}
static int64_t findLargestDivisorAtMost(int64_t value, int64_t limit) {
assert(value > 0 && "expected positive value");
limit = std::min(value, limit);
for (int64_t candidate = limit; candidate >= 1; --candidate)
if (value % candidate == 0)
return candidate;
return 1;
}
static Value createZeroPaddedTensor(Value value,
RankedTensorType resultType,
ArrayRef<int64_t> lowPadValues,
ArrayRef<int64_t> highPadValues,
PatternRewriter& rewriter,
Location loc) {
auto valueType = cast<RankedTensorType>(value.getType());
if (valueType == resultType)
return value;
SmallVector<OpFoldResult> lowPads;
SmallVector<OpFoldResult> highPads;
lowPads.reserve(lowPadValues.size());
highPads.reserve(highPadValues.size());
for (auto lowPad : lowPadValues)
lowPads.push_back(rewriter.getIndexAttr(lowPad));
for (auto highPad : highPadValues)
highPads.push_back(rewriter.getIndexAttr(highPad));
auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
auto* padBlock = new Block();
for (int64_t dim = 0, rank = resultType.getRank(); dim < rank; ++dim)
padBlock->addArgument(rewriter.getIndexType(), loc);
padOp.getRegion().push_back(padBlock);
rewriter.setInsertionPointToStart(padBlock);
auto zero = getOrCreateConstant(
rewriter, padOp.getOperation(), rewriter.getZeroAttr(resultType.getElementType()), resultType.getElementType());
tensor::YieldOp::create(rewriter, loc, zero);
rewriter.setInsertionPointAfter(padOp);
return padOp.getResult();
}
static Value createConvInputPatch(Value input,
RankedTensorType patchType,
Value batchIndex,
Value channelOffset,
Value inputHeightOffset,
Value inputWidthOffset,
int64_t dilationHeight,
int64_t dilationWidth,
PatternRewriter& rewriter,
Location loc) {
const int64_t patchChannels = patchType.getDimSize(1);
const int64_t kernelHeight = patchType.getDimSize(2);
const int64_t kernelWidth = patchType.getDimSize(3);
if (dilationHeight == 1 && dilationWidth == 1) {
SmallVector<OpFoldResult> offsets {batchIndex, channelOffset, inputHeightOffset, inputWidthOffset};
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(patchChannels),
rewriter.getIndexAttr(kernelHeight),
rewriter.getIndexAttr(kernelWidth)};
return tensor::ExtractSliceOp::create(rewriter, loc, patchType, input, offsets, sizes, getUnitStrides(rewriter, 4));
}
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
auto elementType = patchType.getElementType();
auto pixelType = RankedTensorType::get({1, patchChannels, 1, 1}, elementType, patchType.getEncoding());
Value patch = tensor::EmptyOp::create(rewriter, loc, patchType.getShape(), elementType);
for (int64_t kernelH = 0; kernelH < kernelHeight; ++kernelH) {
Value sourceHeightOffset = affineAddConst(rewriter, loc, inputHeightOffset, kernelH * dilationHeight, anchorOp);
for (int64_t kernelW = 0; kernelW < kernelWidth; ++kernelW) {
Value sourceWidthOffset = affineAddConst(rewriter, loc, inputWidthOffset, kernelW * dilationWidth, anchorOp);
SmallVector<OpFoldResult> sourceOffsets {batchIndex, channelOffset, sourceHeightOffset, sourceWidthOffset};
SmallVector<OpFoldResult> sourceSizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(patchChannels),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1)};
Value sourcePixel = tensor::ExtractSliceOp::create(
rewriter, loc, pixelType, input, sourceOffsets, sourceSizes, getUnitStrides(rewriter, 4));
SmallVector<OpFoldResult> targetOffsets {
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(kernelH), rewriter.getIndexAttr(kernelW)};
patch = tensor::InsertSliceOp::create(
rewriter, loc, sourcePixel, patch, targetOffsets, sourceSizes, getUnitStrides(rewriter, 4));
}
}
return patch;
}
static Value createCollectedConvOutput(ValueRange gemmRows,
Type convType,
RankedTensorType gemmOutType,
RankedTensorType nhwcType,
RankedTensorType outType,
int64_t numPatches,
int64_t numChannelsOut,
int64_t packFactor,
PatternRewriter& rewriter,
Location loc);
static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp,
Value x,
Value w,
Value b,
const spatial::SpatialTargetInfo& target);
namespace depthwise {
struct Tiling {
int64_t outputMultiplier;
int64_t kernelElements;
int64_t channelsPerTile;
int64_t tileInputRows;
int64_t tileOutputChannels;
int64_t numChannelTiles;
int64_t spatialPatchesPerBatch;
int64_t totalPatches;
};
static std::optional<Tiling> computeTiling(int64_t batchSize,
int64_t numChannelsIn,
int64_t numChannelsOut,
int64_t wHeight,
int64_t wWidth,
int64_t outHeight,
int64_t outWidth,
int64_t xbarDim) {
const int64_t kernelElements = wHeight * wWidth;
const int64_t outputMultiplier = numChannelsOut / numChannelsIn;
if (kernelElements <= 0 || outputMultiplier <= 0 || kernelElements > xbarDim || outputMultiplier > xbarDim)
return std::nullopt;
const int64_t maxChannelsPerTile = std::min(xbarDim / kernelElements, xbarDim / outputMultiplier);
if (maxChannelsPerTile <= 0)
return std::nullopt;
const int64_t channelsPerTile = findLargestDivisorAtMost(numChannelsIn, maxChannelsPerTile);
const int64_t tileInputRows = channelsPerTile * kernelElements;
const int64_t tileOutputChannels = channelsPerTile * outputMultiplier;
if (tileInputRows > xbarDim || tileOutputChannels > xbarDim)
return std::nullopt;
return Tiling {
outputMultiplier,
kernelElements,
channelsPerTile,
tileInputRows,
tileOutputChannels,
numChannelsIn / channelsPerTile,
outHeight * outWidth,
batchSize * outHeight * outWidth,
};
}
static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
RankedTensorType wType,
const Tiling& tiling,
PatternRewriter& rewriter,
Location loc,
int64_t xbarDim,
int64_t paddedInputRows = -1) {
const int64_t paddedOutputChannels = xbarDim;
const int64_t packedInputRows = paddedInputRows > 0 ? paddedInputRows : tiling.tileInputRows;
auto packedWeightType = RankedTensorType::get(
{tiling.numChannelTiles, packedInputRows, paddedOutputChannels}, wType.getElementType());
SmallVector<Attribute> packedValues(packedWeightType.getNumElements(),
cast<Attribute>(rewriter.getZeroAttr(wType.getElementType())));
SmallVector<Attribute> sourceValues(wDenseAttr.getValues<Attribute>());
for (int64_t tileIndex = 0; tileIndex < tiling.numChannelTiles; ++tileIndex) {
const int64_t channelBase = tileIndex * tiling.channelsPerTile;
for (int64_t localChannel = 0; localChannel < tiling.channelsPerTile; ++localChannel) {
const int64_t globalChannel = channelBase + localChannel;
for (int64_t kernelIndex = 0; kernelIndex < tiling.kernelElements; ++kernelIndex) {
const int64_t kernelH = kernelIndex / wType.getDimSize(3);
const int64_t kernelW = kernelIndex % wType.getDimSize(3);
const int64_t targetRow = kernelIndex * tiling.channelsPerTile + localChannel;
for (int64_t multiplierIndex = 0; multiplierIndex < tiling.outputMultiplier; ++multiplierIndex) {
const int64_t globalOutChannel = globalChannel * tiling.outputMultiplier + multiplierIndex;
const int64_t sourceFlatIndex =
((globalOutChannel * wType.getDimSize(1) * wType.getDimSize(2)) + kernelH) * wType.getDimSize(3) + kernelW;
const int64_t targetCol = localChannel * tiling.outputMultiplier + multiplierIndex;
const int64_t targetFlatIndex =
((tileIndex * packedInputRows) + targetRow) * paddedOutputChannels + targetCol;
packedValues[targetFlatIndex] = sourceValues[sourceFlatIndex];
}
}
}
}
auto packedAttr = DenseElementsAttr::get(packedWeightType, packedValues);
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), packedAttr, packedWeightType);
}
static Value createPaddedInput(Value input,
RankedTensorType inputType,
int64_t padHeightBegin,
int64_t padHeightEnd,
int64_t padWidthBegin,
int64_t padWidthEnd,
PatternRewriter& rewriter,
Location loc) {
if (padHeightBegin == 0 && padHeightEnd == 0 && padWidthBegin == 0 && padWidthEnd == 0)
return input;
auto paddedInputType = RankedTensorType::get({inputType.getDimSize(0),
inputType.getDimSize(1),
inputType.getDimSize(2) + padHeightBegin + padHeightEnd,
inputType.getDimSize(3) + padWidthBegin + padWidthEnd},
inputType.getElementType());
auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
Value padded = createZeroPaddedTensor(computeInput,
paddedInputType,
{0, 0, padHeightBegin, padWidthBegin},
{0, 0, padHeightEnd, padWidthEnd},
rewriter,
loc);
spatial::SpatYieldOp::create(rewriter, loc, padded);
});
return computeOp.getResult(0);
}
static Value createInputTile(Value input,
Value patchIndex,
Value channelTileIndex,
RankedTensorType inputTileType,
const Tiling& tiling,
int64_t strideHeight,
int64_t strideWidth,
int64_t dilationHeight,
int64_t dilationWidth,
int64_t outWidth,
PatternRewriter& rewriter,
Location loc) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value batchIndex = affineFloorDivConst(rewriter, loc, patchIndex, tiling.spatialPatchesPerBatch, anchorOp);
Value batchPatchIndex = affineModConst(rewriter, loc, patchIndex, tiling.spatialPatchesPerBatch, anchorOp);
Value outHeightIndex = affineFloorDivConst(rewriter, loc, batchPatchIndex, outWidth, anchorOp);
Value outWidthIndex = affineModConst(rewriter, loc, batchPatchIndex, outWidth, anchorOp);
Value inputHeightOffset =
strideHeight == 1 ? outHeightIndex : affineMulConst(rewriter, loc, outHeightIndex, strideHeight, anchorOp);
Value inputWidthOffset =
strideWidth == 1 ? outWidthIndex : affineMulConst(rewriter, loc, outWidthIndex, strideWidth, anchorOp);
Value channelOffset = tiling.channelsPerTile == 1
? channelTileIndex
: affineMulConst(rewriter, loc, channelTileIndex, tiling.channelsPerTile, anchorOp);
Value tile4D;
if (dilationHeight == 1 && dilationWidth == 1) {
SmallVector<OpFoldResult> offsets {batchIndex, inputHeightOffset, inputWidthOffset, channelOffset};
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(inputTileType.getDimSize(1)),
rewriter.getIndexAttr(inputTileType.getDimSize(2)),
rewriter.getIndexAttr(tiling.channelsPerTile)};
tile4D = tensor::ExtractSliceOp::create(
rewriter, loc, inputTileType, input, offsets, sizes, getUnitStrides(rewriter, 4));
}
else {
auto pixelType = RankedTensorType::get(
{1, 1, 1, tiling.channelsPerTile}, inputTileType.getElementType(), inputTileType.getEncoding());
tile4D = tensor::EmptyOp::create(rewriter, loc, inputTileType.getShape(), inputTileType.getElementType());
for (int64_t kernelH = 0; kernelH < inputTileType.getDimSize(1); ++kernelH)
for (int64_t kernelW = 0; kernelW < inputTileType.getDimSize(2); ++kernelW) {
Value sourceHeight = affineAddConst(rewriter, loc, inputHeightOffset, kernelH * dilationHeight, anchorOp);
Value sourceWidth = affineAddConst(rewriter, loc, inputWidthOffset, kernelW * dilationWidth, anchorOp);
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(tiling.channelsPerTile)};
Value pixel = tensor::ExtractSliceOp::create(
rewriter,
loc,
pixelType,
input,
SmallVector<OpFoldResult> {batchIndex, sourceHeight, sourceWidth, channelOffset},
sizes,
getUnitStrides(rewriter, 4));
tile4D = tensor::InsertSliceOp::create(
rewriter,
loc,
pixel,
tile4D,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
rewriter.getIndexAttr(kernelH),
rewriter.getIndexAttr(kernelW),
rewriter.getIndexAttr(0)},
sizes,
getUnitStrides(rewriter, 4));
}
}
auto collapsedType = RankedTensorType::get({1, tiling.tileInputRows}, inputTileType.getElementType());
return tensor::CollapseShapeOp::create(rewriter,
loc,
collapsedType,
tile4D,
SmallVector<ReassociationIndices> {
{0},
{1, 2, 3}
});
}
static Value createWeightTile(Value packedWeights,
Value channelTileIndex,
RankedTensorType packedWeightType,
const Tiling& tiling,
PatternRewriter& rewriter,
Location loc) {
SmallVector<OpFoldResult> offsets {channelTileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
const int64_t paddedOutputChannels = packedWeightType.getDimSize(2);
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(tiling.tileInputRows),
rewriter.getIndexAttr(paddedOutputChannels)};
auto collapsedType =
RankedTensorType::get({tiling.tileInputRows, paddedOutputChannels}, packedWeightType.getElementType());
return extractMixedSliceOrIdentity(
rewriter, loc, packedWeights, collapsedType,
{offsets, sizes, getUnitStrides(rewriter, 3)});
}
static Value createBiasTile(
Value bias, Value channelTileIndex, const Tiling& tiling, PatternRewriter& rewriter, Location loc) {
auto biasType = cast<RankedTensorType>(bias.getType());
auto biasTileType = RankedTensorType::get({1, tiling.tileOutputChannels}, biasType.getElementType());
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value channelOffset = tiling.tileOutputChannels == 1
? channelTileIndex
: affineMulConst(rewriter, loc, channelTileIndex, tiling.tileOutputChannels, anchorOp);
SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), channelOffset};
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(tiling.tileOutputChannels)};
return tensor::ExtractSliceOp::create(rewriter, loc, biasTileType, bias, offsets, sizes, getUnitStrides(rewriter, 2));
}
static Value insertOutputTile(Value rowTile,
Value rowAccumulator,
Value channelTileIndex,
const Tiling& tiling,
PatternRewriter& rewriter,
Location loc) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value channelOffset = tiling.tileOutputChannels == 1
? channelTileIndex
: affineMulConst(rewriter, loc, channelTileIndex, tiling.tileOutputChannels, anchorOp);
SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), channelOffset};
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(tiling.tileOutputChannels)};
return tensor::InsertSliceOp::create(
rewriter, loc, rowTile, rowAccumulator, offsets, sizes, getUnitStrides(rewriter, 2));
}
static FailureOr<Value> reconstructDepthwiseGemmRows(Value pieces,
RankedTensorType piecesType,
RankedTensorType gemmOutType,
const Tiling& tiling,
PatternRewriter& rewriter,
Location loc) {
auto collectedOp = createSpatCompute<1>(rewriter, loc, TypeRange {gemmOutType}, {}, pieces, [&](Value piecesArg) {
auto rowType = RankedTensorType::get({1, gemmOutType.getDimSize(1)}, gemmOutType.getElementType());
Value outputInit = tensor::EmptyOp::create(rewriter, loc, gemmOutType.getShape(), gemmOutType.getElementType());
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, tiling.totalPatches);
Value cNumChannelTiles = getOrCreateIndexConstant(rewriter, anchorOp, tiling.numChannelTiles);
auto patchLoop = buildNormalizedScfFor(
rewriter,
loc,
c0,
cNumPatches,
c1,
ValueRange {outputInit},
[&](OpBuilder&,
Location nestedLoc,
Value patchIndex,
ValueRange patchIterArgs,
SmallVectorImpl<Value>& patchYielded) {
Value outputAcc = patchIterArgs.front();
Value rowInit = tensor::EmptyOp::create(rewriter, nestedLoc, rowType.getShape(), rowType.getElementType());
auto tileLoop = buildNormalizedScfFor(
rewriter,
nestedLoc,
c0,
cNumChannelTiles,
c1,
ValueRange {rowInit},
[&](OpBuilder&,
Location tileLoc,
Value channelTileIndex,
ValueRange tileIterArgs,
SmallVectorImpl<Value>& tileYielded) {
Value rowAcc = tileIterArgs.front();
MLIRContext* context = rewriter.getContext();
AffineExpr d0 = getAffineDimExpr(0, context);
AffineExpr d1 = getAffineDimExpr(1, context);
Value laneIndex = createOrFoldAffineApply(
rewriter, tileLoc, (d0 * tiling.totalPatches) + d1, ValueRange {channelTileIndex, patchIndex}, anchorOp);
auto rowTileType = RankedTensorType::get({1, tiling.tileOutputChannels}, piecesType.getElementType());
FailureOr<Value> rowTile = extractGraphBatchPhysicalFragment(rewriter, tileLoc, piecesArg, laneIndex, rowTileType);
if (failed(rowTile))
return failure();
Value rowNext = insertOutputTile(*rowTile, rowAcc, channelTileIndex, tiling, rewriter, tileLoc);
tileYielded.push_back(rowNext);
return success();
});
if (failed(tileLoop))
return failure();
SmallVector<OpFoldResult> rowOffsets {patchIndex, rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(gemmOutType.getDimSize(1))};
Value outputNext = tensor::InsertSliceOp::create(rewriter,
nestedLoc,
tileLoop->results.front(),
outputAcc,
rowOffsets,
rowSizes,
getUnitStrides(rewriter, 2))
.getResult();
patchYielded.push_back(outputNext);
return success();
});
if (failed(patchLoop))
return failure();
spatial::SpatYieldOp::create(rewriter, loc, patchLoop->results.front());
return success();
});
if (failed(collectedOp))
return failure();
return collectedOp->getResult(0);
}
static bool canUseStructuredRewrite(const ConvLoweringState& state) {
if (!getHostConstDenseElementsAttr(state.w))
return false;
auto tiling = computeTiling(state.batchSize,
state.numChannelsIn,
state.numChannelsOut,
state.wHeight,
state.wWidth,
state.outHeight,
state.outWidth,
state.targetInfo().matrixShape.rows);
if (!tiling)
return false;
if (!state.hasBias)
return true;
auto biasType = dyn_cast<RankedTensorType>(state.b.getType());
if (!biasType)
return false;
if (biasType.getRank() == 1)
return biasType.getDimSize(0) == state.numChannelsOut;
if (biasType.getRank() != 2)
return false;
return biasType.getDimSize(0) == 1 && biasType.getDimSize(1) == state.numChannelsOut;
}
static FailureOr<Value>
rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) {
auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
if (!wDenseAttr) {
convOp->emitOpError("requires constant-derived weights for structured depthwise Spatial lowering");
return failure();
}
auto tiling = computeTiling(state.xType.getDimSize(0),
state.xType.getDimSize(1),
state.outType.getDimSize(1),
state.wType.getDimSize(2),
state.wType.getDimSize(3),
state.outType.getDimSize(2),
state.outType.getDimSize(3),
state.targetInfo().matrixShape.rows);
if (!tiling) {
convOp->emitOpError("failed to derive a structured depthwise tiling that fits Spatial weighted VMM lowering");
return failure();
}
Value paddedInput = createPaddedInput(state.x,
state.xType,
state.padHeightBegin,
state.padHeightEnd,
state.padWidthBegin,
state.padWidthEnd,
rewriter,
loc);
auto paddedInputType = cast<RankedTensorType>(paddedInput.getType());
auto channelLastInputType = RankedTensorType::get({paddedInputType.getDimSize(0),
paddedInputType.getDimSize(2),
paddedInputType.getDimSize(3),
paddedInputType.getDimSize(1)},
paddedInputType.getElementType());
Value channelLastInput = createLinalgTranspose(
paddedInput, channelLastInputType, {0, 2, 3, 1}, rewriter, loc);
Value packedWeights = buildPackedWeights(
wDenseAttr, state.wType, *tiling, rewriter, loc, state.targetInfo().matrixShape.rows);
Value expandedBias;
SmallVector<Value> batchInputs {channelLastInput};
if (state.hasBias) {
expandedBias = expandBiasIfNeeded(state.b, rewriter, loc);
auto biasType = dyn_cast<RankedTensorType>(expandedBias.getType());
if (!biasType || biasType.getRank() != 2 || biasType.getDimSize(0) != 1
|| biasType.getDimSize(1) != state.outType.getDimSize(1)) {
convOp->emitOpError("requires bias sliceable as tensor<1xCout> for structured depthwise Spatial lowering");
return failure();
}
batchInputs.push_back(expandedBias);
}
auto gemmOutType =
RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType());
auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
auto paddedRowTileType = RankedTensorType::get(
{1, static_cast<int64_t>(state.targetInfo().matrixShape.rows)}, state.outType.getElementType());
auto piecesType = spatial::getGraphBatchPhysicalResultType(
tiling->totalPatches * tiling->numChannelTiles, rowTileType);
auto inputTileType =
RankedTensorType::get({1, state.wType.getDimSize(2), state.wType.getDimSize(3), tiling->channelsPerTile},
paddedInputType.getElementType());
SmallVector<Value> batchWeights;
if (tiling->numChannelTiles == 1) {
Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
batchWeights.push_back(createWeightTile(packedWeights,
c0,
cast<RankedTensorType>(packedWeights.getType()),
*tiling,
rewriter,
loc));
}
else {
batchWeights.push_back(packedWeights);
}
auto batchOp = createSpatComputeBatch(
rewriter,
loc,
TypeRange {piecesType},
tiling->totalPatches * tiling->numChannelTiles,
batchWeights,
batchInputs,
[&](detail::SpatComputeBatchBodyArgs args) {
auto pickInputByRank = [&](int64_t rank) -> Value {
for (Value input : args.inputs) {
auto inputType = dyn_cast<RankedTensorType>(input.getType());
if (inputType && inputType.getRank() == rank)
return input;
}
return Value();
};
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value patchIndex = tiling->numChannelTiles == 1
? args.lane
: affineModConst(rewriter, loc, args.lane, tiling->totalPatches, anchorOp);
Value channelTileIndex = tiling->numChannelTiles == 1
? getOrCreateIndexConstant(rewriter, anchorOp, 0)
: affineFloorDivConst(rewriter, loc, args.lane, tiling->totalPatches, anchorOp);
Value paddedInputArg = pickInputByRank(/*rank=*/4);
if (!paddedInputArg) {
convOp->emitOpError("structured depthwise batch body requires a rank-4 padded input block argument");
return failure();
}
Value inputTile = createInputTile(paddedInputArg,
patchIndex,
channelTileIndex,
inputTileType,
*tiling,
state.strideHeight,
state.strideWidth,
state.dilationHeight,
state.dilationWidth,
state.outType.getDimSize(3),
rewriter,
loc);
Value weightTile = tiling->numChannelTiles == 1
? args.weights.front()
: createWeightTile(args.weights.front(),
channelTileIndex,
cast<RankedTensorType>(args.weights.front().getType()),
*tiling,
rewriter,
loc);
Value paddedRowTile =
spatial::SpatVMMOp::create(rewriter, loc, paddedRowTileType, weightTile, inputTile).getResult();
Value rowTile = tensor::ExtractSliceOp::create(
rewriter,
loc,
rowTileType,
paddedRowTile,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(tiling->tileOutputChannels)},
getUnitStrides(rewriter, 2));
if (args.inputs.size() > 1) {
Value biasArg = pickInputByRank(/*rank=*/2);
if (!biasArg) {
convOp->emitOpError("structured depthwise batch body requires a rank-2 bias block argument when bias is present");
return failure();
}
Value biasTile = tiling->numChannelTiles == 1 ? biasArg : createBiasTile(biasArg, channelTileIndex, *tiling, rewriter, loc);
rowTile = spatial::SpatVAddOp::create(rewriter, loc, rowTileType, rowTile, biasTile).getResult();
}
publishGraphBatchPhysicalFragment(rewriter, loc, rowTile, args.outputs.front(), args.lane);
return success();
});
if (failed(batchOp))
return failure();
auto nhwcType = RankedTensorType::get(
{state.xType.getDimSize(0), state.outType.getDimSize(2), state.outType.getDimSize(3), state.outType.getDimSize(1)},
state.outType.getElementType());
auto reconstructedRows = reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc);
if (failed(reconstructedRows))
return failure();
Value collectedRows = *reconstructedRows;
return createCollectedConvOutput(ValueRange {collectedRows},
state.outType,
gemmOutType,
nhwcType,
state.outType,
tiling->totalPatches,
state.outType.getDimSize(1),
/*packFactor=*/1,
rewriter,
loc);
}
} // namespace depthwise
namespace standard {
struct ConvGemmPlan {
int64_t patchSize;
int64_t numPatchesPerBatch;
int64_t globalNumPatches;
int64_t chunkStart;
int64_t chunkNumPatches;
int64_t maxParallelPixels;
int64_t effectiveMaxParallelPixels;
int64_t packedNumRows;
RankedTensorType gemmInputRowsType;
RankedTensorType wFlatType;
RankedTensorType wTransType;
RankedTensorType gemmOutType;
RankedTensorType gemmOutputRowsType;
RankedTensorType nhwcType;
};
static ConvGemmPlan
buildConvGemmPlan(const ConvLoweringState& state,
bool canPackWeightsAsConstants,
bool canPackBiasAsConstants,
int64_t chunkStart,
int64_t chunkNumPatches,
std::optional<int64_t> forcedPackFactor = std::nullopt);
static PreparedConvInput prepareInputForIm2Col(const ConvLoweringState& state,
PatternRewriter& rewriter,
Location loc) {
if (state.padHeightBegin == 0 && state.padHeightEnd == 0 && state.padWidthBegin == 0 && state.padWidthEnd == 0)
return {state.x, state.xType};
auto paddedType = RankedTensorType::get({state.batchSize,
state.numChannelsIn,
state.xHeight + state.padHeightBegin + state.padHeightEnd,
state.xWidth + state.padWidthBegin + state.padWidthEnd},
state.xType.getElementType());
auto paddedInputOp =
createSpatCompute<1>(rewriter, loc, TypeRange {paddedType}, {}, state.x, [&](Value inputArg) {
Value paddedInput = createZeroPaddedTensor(inputArg,
paddedType,
{0, 0, state.padHeightBegin, state.padWidthBegin},
{0, 0, state.padHeightEnd, state.padWidthEnd},
rewriter,
loc);
spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
});
return {paddedInputOp.getResult(0), paddedType};
}
static Value unpackRowsFromParallelGemm(Value packedRows,
RankedTensorType packedRowsType,
int64_t unpackedRows,
int64_t rowWidth,
int64_t packFactor,
PatternRewriter& rewriter,
Location loc) {
if (packFactor == 1)
return packedRows;
const int64_t packedNumRows = packedRowsType.getDimSize(0);
const int64_t paddedNumRows = packedNumRows * packFactor;
auto expandedType = RankedTensorType::get(
{packedNumRows, packFactor, rowWidth}, packedRowsType.getElementType(), packedRowsType.getEncoding());
auto paddedType =
RankedTensorType::get({paddedNumRows, rowWidth}, packedRowsType.getElementType(), packedRowsType.getEncoding());
auto unpackedType =
RankedTensorType::get({unpackedRows, rowWidth}, packedRowsType.getElementType(), packedRowsType.getEncoding());
Value expanded = tensor::ExpandShapeOp::create(rewriter,
loc,
expandedType,
packedRows,
SmallVector<ReassociationIndices> {
{0},
{1, 2}
});
Value padded = tensor::CollapseShapeOp::create(rewriter,
loc,
paddedType,
expanded,
SmallVector<ReassociationIndices> {
{0, 1},
{2}
});
if (paddedNumRows == unpackedRows)
return padded;
SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(unpackedRows), rewriter.getIndexAttr(rowWidth)};
return tensor::ExtractSliceOp::create(rewriter, loc, unpackedType, padded, offsets, sizes, getUnitStrides(rewriter, 2));
}
static Value createWeightMatrix(
Value weights, const ConvGemmPlan& plan, bool transpose, PatternRewriter& rewriter, Location loc) {
auto buildWeightMatrix = [&](Value weight) -> Value {
Value flattened = tensor::CollapseShapeOp::create(rewriter,
loc,
plan.wFlatType,
weight,
SmallVector<ReassociationIndices> {
{0},
{1, 2, 3}
});
if (!transpose)
return flattened;
return createLinalgTranspose(flattened, plan.wTransType, {1, 0}, rewriter, loc);
};
if (isCompileTimeComputable(weights))
return buildWeightMatrix(weights);
RankedTensorType resultType = transpose ? plan.wTransType : plan.wFlatType;
auto computeOp =
createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, ValueRange {weights}, [&](Value weight) {
spatial::SpatYieldOp::create(rewriter, loc, buildWeightMatrix(weight));
});
return computeOp.getResult(0);
}
static Value createPaddedConvMatrix(Value matrix,
RankedTensorType sourceType,
RankedTensorType paddedType,
PatternRewriter& rewriter,
Location loc) {
if (sourceType == paddedType)
return matrix;
return createZeroPaddedTensor(matrix,
paddedType,
{0, 0},
{paddedType.getDimSize(0) - sourceType.getDimSize(0),
paddedType.getDimSize(1) - sourceType.getDimSize(1)},
rewriter,
loc);
}
static Value createPaddedConstantMatrix(DenseElementsAttr sourceAttr,
RankedTensorType sourceType,
RankedTensorType paddedType,
PatternRewriter& rewriter) {
SmallVector<Attribute> paddedValues(
paddedType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(paddedType.getElementType())));
SmallVector<Attribute> sourceValues(sourceAttr.getValues<Attribute>());
const int64_t sourceRows = sourceType.getDimSize(0);
const int64_t sourceCols = sourceType.getDimSize(1);
const int64_t paddedCols = paddedType.getDimSize(1);
for (int64_t row = 0; row < sourceRows; ++row)
for (int64_t col = 0; col < sourceCols; ++col)
paddedValues[row * paddedCols + col] = sourceValues[row * sourceCols + col];
auto paddedAttr = DenseElementsAttr::get(paddedType, paddedValues);
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), paddedAttr, paddedType);
}
static Value createPaddedInputKTiledWeightConstant(DenseElementsAttr sourceAttr,
const ConvLoweringState& state,
int64_t paddedK,
int64_t paddedC,
PatternRewriter& rewriter) {
auto paddedType = RankedTensorType::get({paddedK, paddedC}, state.wType.getElementType());
SmallVector<Attribute> sourceValues(sourceAttr.getValues<Attribute>());
SmallVector<Attribute> paddedValues(
paddedType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(paddedType.getElementType())));
for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel) {
for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH) {
for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) {
const int64_t sourceFlatIndex =
(((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
const int64_t patchIndex = ((inChannel * state.wHeight) + kernelH) * state.wWidth + kernelW;
paddedValues[patchIndex * paddedC + outChannel] = sourceValues[sourceFlatIndex];
}
}
}
}
auto paddedAttr = DenseElementsAttr::get(paddedType, paddedValues);
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), paddedAttr, paddedType);
}
static Value createPaddedPixelMajorWeightConstant(DenseElementsAttr sourceAttr,
const ConvLoweringState& state,
int64_t paddedK,
int64_t paddedC,
int64_t packFactor,
PatternRewriter& rewriter) {
auto paddedType = RankedTensorType::get({paddedK, paddedC}, state.wType.getElementType());
SmallVector<Attribute> sourceValues(sourceAttr.getValues<Attribute>());
SmallVector<Attribute> paddedValues(
paddedType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(paddedType.getElementType())));
const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
for (int64_t copy = 0; copy < packFactor; ++copy)
for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel)
for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH)
for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW)
for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
const int64_t sourceFlatIndex =
(((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
const int64_t patchIndex =
((kernelH * state.wWidth) + kernelW) * state.numChannelsIn + inChannel;
const int64_t packedRow = copy * patchSize + patchIndex;
const int64_t packedColumn = copy * state.numChannelsOut + outChannel;
paddedValues[packedRow * paddedC + packedColumn] = sourceValues[sourceFlatIndex];
}
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
DenseElementsAttr::get(paddedType, paddedValues), paddedType);
}
static Value createPaddedOutputChannelTiledWeightConstant(DenseElementsAttr sourceAttr,
const ConvLoweringState& state,
int64_t paddedK,
int64_t xbarDim,
PatternRewriter& rewriter) {
const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim);
auto paddedType =
RankedTensorType::get({outputTileCount, paddedK, xbarDim}, state.wType.getElementType());
SmallVector<Attribute> sourceValues(sourceAttr.getValues<Attribute>());
SmallVector<Attribute> paddedValues(
paddedType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(paddedType.getElementType())));
for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel) {
const int64_t outputTile = outChannel / xbarDim;
const int64_t tileChannel = outChannel % xbarDim;
for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH) {
for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) {
const int64_t sourceFlatIndex =
(((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
const int64_t patchIndex =
((kernelH * state.wWidth) + kernelW) * state.numChannelsIn + inChannel;
const int64_t destinationFlatIndex =
((outputTile * paddedK) + patchIndex) * xbarDim + tileChannel;
paddedValues[destinationFlatIndex] = sourceValues[sourceFlatIndex];
}
}
}
}
auto paddedAttr = DenseElementsAttr::get(paddedType, paddedValues);
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), paddedAttr, paddedType);
}
static FailureOr<Value> rewriteInputKTiledConv(const ConvLoweringState& state,
PatternRewriter& rewriter,
Location loc) {
PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc);
ConvGeometry geo = buildConvGeometry(state, state.targetInfo());
const int64_t xbarDim = geo.xbarSize;
const int64_t numKSlices = ceilIntegerDivide(geo.k, xbarDim);
const int64_t paddedK = numKSlices * xbarDim;
const uint64_t maxLanesPerBatch =
std::max<uint64_t>(1,
static_cast<uint64_t>(state.targetInfo().matrixUnitsPerProcessor)
/ static_cast<uint64_t>(std::max<int64_t>(1, numKSlices * 4)));
const uint64_t rowChunkWidth = std::max<uint64_t>(
1,
std::min<uint64_t>({chooseStreamChunkPositions(geo, /*packFactor=*/1, state.targetInfo()),
maxLanesPerBatch,
static_cast<uint64_t>(state.outWidth)}));
const auto elementType = state.outType.getElementType();
auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
if (!wDenseAttr)
return failure();
Value paddedWeight = createPaddedInputKTiledWeightConstant(wDenseAttr, state, paddedK, xbarDim, rewriter);
Value paddedBias;
RankedTensorType paddedBiasType;
if (state.hasBias) {
Value biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
auto biasMatrixType = cast<RankedTensorType>(biasMatrix.getType());
paddedBiasType = RankedTensorType::get({1, xbarDim}, elementType);
if (auto biasDenseAttr = getHostConstDenseElementsAttr(state.b))
paddedBias = createPaddedConstantMatrix(biasDenseAttr, biasMatrixType, paddedBiasType, rewriter);
else
paddedBias = materializeOrComputeUnary(
biasMatrix, paddedBiasType, rewriter, loc, [&](Value biasValue) {
return createPaddedConvMatrix(biasValue, biasMatrixType, paddedBiasType, rewriter, loc);
});
}
SmallVector<Value> chunkRows;
const int64_t totalPatches = state.batchSize * state.outHeight * state.outWidth;
chunkRows.reserve(
state.batchSize * state.outHeight * ceilIntegerDivide(state.outWidth, static_cast<int64_t>(rowChunkWidth)));
for (int64_t batchIndex = 0; batchIndex < state.batchSize; ++batchIndex) {
for (int64_t outHeightIndex = 0; outHeightIndex < state.outHeight; ++outHeightIndex) {
for (int64_t outWidthChunkStart = 0; outWidthChunkStart < state.outWidth;
outWidthChunkStart += static_cast<int64_t>(rowChunkWidth)) {
const int64_t chunkNumPatches =
std::min<int64_t>(static_cast<int64_t>(rowChunkWidth), state.outWidth - outWidthChunkStart);
auto chunkRowsType = RankedTensorType::get({chunkNumPatches, state.numChannelsOut}, elementType);
auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
auto paddedChunkRowType = RankedTensorType::get({1, paddedK}, elementType);
auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elementType);
auto collapsedPatchType = RankedTensorType::get({1, geo.k}, elementType);
auto weightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType());
auto rowType = RankedTensorType::get({1, state.numChannelsOut}, elementType);
SmallVector<Value> inputsStorage {preparedInput.value};
if (state.hasBias)
inputsStorage.push_back(paddedBias);
ValueRange inputs(inputsStorage);
auto chunkCompute = spatial::SpatCompute::create(rewriter, loc, TypeRange {chunkRowsType}, ValueRange {paddedWeight}, inputs);
auto* block = new Block();
block->addArgument(paddedWeight.getType(), loc);
for (Value input : inputs)
block->addArgument(input.getType(), loc);
chunkCompute.getBody().push_back(block);
rewriter.setInsertionPointToStart(block);
auto buildChunk = [&]() -> LogicalResult {
Value weightArg = block->getArgument(0);
Value inputArg = block->getArgument(1);
Value biasArg = state.hasBias ? block->getArgument(2) : Value();
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value cBatchIndex = getOrCreateIndexConstant(rewriter, anchorOp, batchIndex);
Value cZero = getOrCreateIndexConstant(rewriter, anchorOp, 0);
Value cKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
Value cOne = getOrCreateIndexConstant(rewriter, anchorOp, 1);
Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim);
Value cInputHeightOffset =
getOrCreateIndexConstant(rewriter, anchorOp, outHeightIndex * state.strideHeight);
Value chunkRowsValue = tensor::EmptyOp::create(rewriter, loc, chunkRowsType.getShape(), elementType);
auto widthLoop = buildNormalizedScfFor(
rewriter,
loc,
cZero,
getOrCreateIndexConstant(rewriter, anchorOp, chunkNumPatches),
cOne,
ValueRange {chunkRowsValue},
[&](OpBuilder&, Location nestedLoc, Value widthIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
Value laneWithChunkOffset = affineAddConst(rewriter, nestedLoc, widthIndex, outWidthChunkStart, anchorOp);
Value inputWidthOffset = createOrFoldAffineApply(rewriter,
nestedLoc,
getAffineDimExpr(0, rewriter.getContext()) * state.strideWidth,
ValueRange {laneWithChunkOffset},
anchorOp);
Value patch = createConvInputPatch(inputArg,
patchType,
cBatchIndex,
cZero,
cInputHeightOffset,
inputWidthOffset,
state.dilationHeight,
state.dilationWidth,
rewriter,
nestedLoc);
Value patchRow = tensor::CollapseShapeOp::create(rewriter,
nestedLoc,
collapsedPatchType,
patch,
SmallVector<ReassociationIndices> {
{0},
{1, 2, 3}
});
Value paddedPatchRow = createZeroPaddedTensor(
patchRow, paddedChunkRowType, {0, 0}, {0, paddedK - geo.k}, rewriter, nestedLoc);
auto zeroAttr = DenseElementsAttr::get(paddedRowType, rewriter.getZeroAttr(elementType));
Value zeroRow = getOrCreateConstant(rewriter, anchorOp, zeroAttr, paddedRowType);
auto kLoop = buildNormalizedScfFor(
rewriter,
nestedLoc,
cZero,
cKSlices,
cOne,
ValueRange {zeroRow},
[&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
Value acc = reduceIterArgs.front();
Value kOffset = arith::MulIOp::create(rewriter, reduceLoc, kSlice, cXbar);
SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
Value aTile = tensor::ExtractSliceOp::create(
rewriter, reduceLoc, paddedRowType, paddedPatchRow, aOffsets, aSizes, unitStrides);
SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
Value bTile = extractStaticSliceOrIdentity(
rewriter, reduceLoc, weightArg, weightTileType, bOffsets, bSizes, unitStrides);
Value piece = spatial::SpatVMMOp::create(rewriter, reduceLoc, paddedRowType, bTile, aTile).getResult();
reduceYielded.push_back(
spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, acc, piece).getResult());
return success();
});
if (failed(kLoop))
return failure();
Value reduced = kLoop->results.front();
if (state.hasBias)
reduced = spatial::SpatVAddOp::create(rewriter, nestedLoc, paddedRowType, reduced, biasArg).getResult();
Value row = reduced;
if (state.numChannelsOut != xbarDim) {
SmallVector<OpFoldResult> rowOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> rowSizes {
rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
row = tensor::ExtractSliceOp::create(
rewriter, nestedLoc, rowType, reduced, rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
}
SmallVector<OpFoldResult> outputOffsets {widthIndex, rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> outputSizes {
rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
Value updatedRows = tensor::InsertSliceOp::create(
rewriter, nestedLoc, row, iterArgs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
yielded.push_back(updatedRows);
return success();
});
if (failed(widthLoop))
return failure();
spatial::SpatYieldOp::create(rewriter, loc, widthLoop->results.front());
return success();
};
if (failed(buildChunk())) {
rewriter.setInsertionPointAfter(chunkCompute);
rewriter.eraseOp(chunkCompute);
return failure();
}
rewriter.setInsertionPointAfter(chunkCompute);
chunkRows.push_back(chunkCompute.getResult(0));
}
}
}
auto nhwcType = RankedTensorType::get({state.batchSize, state.outHeight, state.outWidth, state.numChannelsOut},
elementType);
return createCollectedConvOutput(
chunkRows, state.outType, cast<RankedTensorType>(chunkRows.front().getType()), nhwcType, state.outType, totalPatches,
state.numChannelsOut, /*packFactor=*/1, rewriter, loc);
}
static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
Value wTrans,
const ConvLoweringState& state,
const ConvGemmPlan& plan,
PatternRewriter& rewriter,
Location loc) {
if (plan.effectiveMaxParallelPixels == 1)
return wTrans;
auto packedWeightType = RankedTensorType::get(
{plan.effectiveMaxParallelPixels * plan.patchSize, plan.effectiveMaxParallelPixels * state.numChannelsOut},
state.wType.getElementType());
SmallVector<Attribute> sourceValues(wDenseAttr.getValues<Attribute>());
SmallVector<Attribute> packedValues(packedWeightType.getNumElements(),
cast<Attribute>(rewriter.getZeroAttr(state.wType.getElementType())));
for (int64_t copyId = 0; copyId < plan.effectiveMaxParallelPixels; ++copyId) {
for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel) {
for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH) {
for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) {
const int64_t sourceFlatIndex =
(((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
const int64_t patchIndex = ((inChannel * state.wHeight) + kernelH) * state.wWidth + kernelW;
const int64_t targetRow = copyId * plan.patchSize + patchIndex;
const int64_t targetCol = copyId * state.numChannelsOut + outChannel;
packedValues[targetRow * packedWeightType.getDimSize(1) + targetCol] = sourceValues[sourceFlatIndex];
}
}
}
}
}
auto packedAttr = DenseElementsAttr::get(packedWeightType, packedValues);
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), packedAttr, packedWeightType);
}
static Value buildPackedBias(Value gemmBias,
Value biasMatrix,
DenseElementsAttr biasDenseAttr,
const ConvLoweringState& state,
const ConvGemmPlan& plan,
PatternRewriter& rewriter,
Location loc) {
if (!state.hasBias)
return gemmBias;
if (plan.effectiveMaxParallelPixels == 1)
return biasMatrix;
SmallVector<Attribute> sourceValues(biasDenseAttr.getValues<Attribute>());
SmallVector<Attribute> packedValues;
packedValues.reserve(plan.effectiveMaxParallelPixels * state.numChannelsOut);
for (int64_t copyId = 0; copyId < plan.effectiveMaxParallelPixels; ++copyId)
packedValues.append(sourceValues.begin(), sourceValues.end());
auto packedBiasType =
RankedTensorType::get({1, plan.effectiveMaxParallelPixels * state.numChannelsOut}, state.outType.getElementType());
auto packedBiasAttr = DenseElementsAttr::get(packedBiasType, packedValues);
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), packedBiasAttr, packedBiasType);
}
static ConvGemmPlan
buildConvGemmPlan(const ConvLoweringState& state,
bool canPackWeightsAsConstants,
bool canPackBiasAsConstants,
int64_t chunkStart,
int64_t chunkNumPatches,
std::optional<int64_t> forcedPackFactor) {
ConvGemmPlan plan;
plan.patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
plan.numPatchesPerBatch = state.outHeight * state.outWidth;
plan.globalNumPatches = state.batchSize * plan.numPatchesPerBatch;
plan.chunkStart = chunkStart;
plan.chunkNumPatches = chunkNumPatches;
const int64_t wMaxDim = std::max(plan.patchSize, state.numChannelsOut);
plan.maxParallelPixels = forcedPackFactor
? *forcedPackFactor
: std::max<int64_t>(1, static_cast<int64_t>(state.targetInfo().matrixShape.rows) / wMaxDim);
plan.effectiveMaxParallelPixels =
(canPackWeightsAsConstants && canPackBiasAsConstants) ? plan.maxParallelPixels : 1;
plan.packedNumRows = ceilIntegerDivide(plan.chunkNumPatches, plan.effectiveMaxParallelPixels);
auto elemType = state.xType.getElementType();
auto outElemType = state.outType.getElementType();
plan.gemmInputRowsType =
RankedTensorType::get({plan.packedNumRows, plan.effectiveMaxParallelPixels * plan.patchSize}, elemType);
plan.wFlatType = RankedTensorType::get({state.numChannelsOut, plan.patchSize}, state.wType.getElementType());
plan.wTransType = RankedTensorType::get({plan.patchSize, state.numChannelsOut}, state.wType.getElementType());
plan.gemmOutType = RankedTensorType::get({plan.chunkNumPatches, state.numChannelsOut}, outElemType);
plan.gemmOutputRowsType =
RankedTensorType::get({plan.packedNumRows, plan.effectiveMaxParallelPixels * state.numChannelsOut}, outElemType);
plan.nhwcType =
RankedTensorType::get({state.batchSize, state.outHeight, state.outWidth, state.numChannelsOut}, outElemType);
return plan;
}
static Value createIm2colRows(const ConvLoweringState& state,
const PreparedConvInput& preparedInput,
const ConvGemmPlan& plan,
PatternRewriter& rewriter,
Location loc) {
if (plan.gemmInputRowsType.getDimSize(1)
> static_cast<int64_t>(state.targetInfo().matrixShape.rows)) {
assert(plan.effectiveMaxParallelPixels == 1 && "multi-crossbar im2col rows cannot pack pixels");
auto compute = createSpatCompute<1>(
rewriter, loc, TypeRange {plan.gemmInputRowsType}, {}, preparedInput.value, [&](Value input) {
auto elemType = preparedInput.type.getElementType();
Value empty = tensor::EmptyOp::create(rewriter, loc, plan.gemmInputRowsType.getShape(), elemType);
Operation *anchor = rewriter.getInsertionBlock()->getParentOp();
Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
Value upper = getOrCreateIndexConstant(rewriter, anchor, plan.chunkNumPatches);
auto patchType = RankedTensorType::get(
{1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
auto rowType = RankedTensorType::get({plan.patchSize}, elemType);
auto loop = buildNormalizedScfFor(
rewriter, loc, c0, upper, c1, ValueRange {empty},
[&](OpBuilder &, Location nestedLoc, Value patchIndex, ValueRange iterArgs,
SmallVectorImpl<Value> &yielded) {
Value batchIndex = affineAddFloorDivConst(
rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchor);
Value batchPatchIndex = affineAddModConst(
rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchor);
Value outHeight = affineFloorDivConst(
rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchor);
Value outWidth = affineModConst(
rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchor);
Value patch = createConvInputPatch(
input, patchType, batchIndex, c0,
affineMulConst(rewriter, nestedLoc, outHeight, state.strideHeight, anchor),
affineMulConst(rewriter, nestedLoc, outWidth, state.strideWidth, anchor),
state.dilationHeight, state.dilationWidth, rewriter, nestedLoc);
Value row = tensor::CollapseShapeOp::create(
rewriter, nestedLoc, rowType, patch,
SmallVector<ReassociationIndices> {{0, 1, 2, 3}});
Value next = tensor::InsertSliceOp::create(
rewriter, nestedLoc, row, iterArgs.front(),
SmallVector<OpFoldResult> {patchIndex, rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(plan.patchSize)},
getUnitStrides(rewriter, 2));
yielded.push_back(next);
return success();
});
if (failed(loop))
return failure();
spatial::SpatYieldOp::create(rewriter, loc, loop->results.front());
return success();
});
assert(succeeded(compute) && "Conv im2col compute construction must succeed");
return compute->getResult(0);
}
auto elemType = preparedInput.type.getElementType();
auto packedRowType = RankedTensorType::get(
{plan.effectiveMaxParallelPixels * plan.patchSize}, elemType, plan.gemmInputRowsType.getEncoding());
auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType);
bool hasPartialLane = plan.chunkNumPatches % plan.effectiveMaxParallelPixels != 0;
SmallVector<Value> im2colInputs {preparedInput.value};
auto im2colComputeOp = createSpatComputeBatch(
rewriter,
loc,
TypeRange {plan.gemmInputRowsType},
plan.packedNumRows,
{},
im2colInputs,
[&](detail::SpatComputeBatchBodyArgs args) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
Value cPack = getOrCreateIndexConstant(rewriter, anchorOp, plan.effectiveMaxParallelPixels);
Value laneStart = affineMulConst(rewriter, loc, args.lane, plan.effectiveMaxParallelPixels, anchorOp);
Value lanePatches = cPack;
if (hasPartialLane) {
Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);
Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart);
Value isPartial = arith::CmpIOp::create(
rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack);
lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack);
}
Value rowInit = tensor::EmptyOp::create(rewriter, loc, packedRowType.getShape(), elemType);
if (hasPartialLane) {
auto zeroAttr = cast<TypedAttr>(rewriter.getZeroAttr(elemType));
rowInit = linalg::MapOp::create(
rewriter, loc, ValueRange {}, rowInit,
[&](OpBuilder& builder, Location nestedLoc, ValueRange) {
Value zero = arith::ConstantOp::create(builder, nestedLoc, zeroAttr);
linalg::YieldOp::create(builder, nestedLoc, zero);
}).getResult().front();
}
auto rowLoop = buildNormalizedScfFor(
rewriter,
loc,
c0,
lanePatches,
c1,
ValueRange {rowInit},
[&](OpBuilder&, Location nestedLoc, Value copyIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
Value patchIndex = arith::AddIOp::create(rewriter, nestedLoc, laneStart, copyIndex);
Value batchIndex = state.batchSize == 1
? c0
: affineAddFloorDivConst(
rewriter, nestedLoc, patchIndex, plan.chunkStart,
plan.numPatchesPerBatch, anchorOp);
Value batchPatchIndex =
affineAddModConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
Value outWidthIndex = affineModConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
Value inputHeightOffset = affineMulConst(rewriter, nestedLoc, outHeightIndex, state.strideHeight, anchorOp);
Value inputWidthOffset = affineMulConst(rewriter, nestedLoc, outWidthIndex, state.strideWidth, anchorOp);
Value patch = createConvInputPatch(args.inputs.front(),
patchType,
batchIndex,
c0,
inputHeightOffset,
inputWidthOffset,
state.dilationHeight,
state.dilationWidth,
rewriter,
nestedLoc);
Value patchRow = tensor::CollapseShapeOp::create(rewriter,
nestedLoc,
patchRowType,
patch,
SmallVector<ReassociationIndices> {
{0, 1, 2, 3}
});
Value rowOffset = affineMulConst(rewriter, nestedLoc, copyIndex, plan.patchSize, anchorOp);
Value next = tensor::InsertSliceOp::create(rewriter,
nestedLoc,
patchRow,
iterArgs.front(),
SmallVector<OpFoldResult> {rowOffset},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(plan.patchSize)},
getUnitStrides(rewriter, 1));
yielded.push_back(next);
return success();
});
if (failed(rowLoop))
return failure();
Value row = rowLoop->results.front();
publishGraphBatchPhysicalFragment(rewriter, loc, row, args.outputs.front(), args.lane);
return success();
});
assert(succeeded(im2colComputeOp) && "Conv im2col compute construction must succeed");
return im2colComputeOp->getResult(0);
}
static Value maybeUnpackChunkRows(Value gemmRows,
const ConvGemmPlan& plan,
PatternRewriter& rewriter,
Location loc) {
if (plan.effectiveMaxParallelPixels == 1)
return gemmRows;
auto unpackedType = RankedTensorType::get(
{plan.chunkNumPatches, plan.gemmOutType.getDimSize(1)}, plan.gemmOutType.getElementType(), plan.gemmOutType.getEncoding());
auto unpackCompute = createSpatCompute<1>(rewriter, loc, TypeRange {unpackedType}, {}, gemmRows, [&](Value rowsArg) {
Value unpacked = unpackRowsFromParallelGemm(rowsArg,
cast<RankedTensorType>(rowsArg.getType()),
plan.chunkNumPatches,
plan.gemmOutType.getDimSize(1),
plan.effectiveMaxParallelPixels,
rewriter,
loc);
spatial::SpatYieldOp::create(rewriter, loc, unpacked);
});
return unpackCompute.getResult(0);
}
static FailureOr<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;
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);
FailureOr<Value> gemmRows = lowerGemmToSpatial(
state.diagnosticAnchor, inputRows, packedWeights, packedBias,
plan.gemmOutputRowsType, /*transA=*/false, /*transB=*/!wDenseAttr,
/*alpha=*/1.0f, /*beta=*/1.0f, state.targetInfo(), rewriter, loc);
if (failed(gemmRows))
return failure();
return maybeUnpackChunkRows(*gemmRows, plan, rewriter, loc);
}
static FailureOr<Value> rewritePackedIm2ColConv(const ConvLoweringState& state,
PatternRewriter& rewriter,
Location loc) {
auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc);
Value biasMatrix;
DenseElementsAttr biasDenseAttr;
if (state.hasBias) {
biasDenseAttr = getHostConstDenseElementsAttr(state.b);
biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
}
ConvGemmPlan plan =
buildConvGemmPlan(state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0,
state.batchSize * state.outHeight * state.outWidth);
// Static weights use the crossbar [patchSize, Cout] layout. Runtime weights
// stay in ONNX's contiguous [Cout, patchSize] layout and Gemm consumes them
// through transB without materializing a transpose.
Value weightMatrix = createWeightMatrix(state.w, plan, static_cast<bool>(wDenseAttr), rewriter, loc);
Value gemmInputRows = createIm2colRows(state, preparedInput, plan, rewriter, loc);
Value gemmB = buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc);
Value gemmBias = createZeroGemmBias(plan.gemmOutputRowsType, rewriter);
if (state.hasBias)
gemmBias = state.b;
Value gemmC = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc);
FailureOr<Value> gemmRows = lowerGemmToSpatial(
state.diagnosticAnchor, gemmInputRows, gemmB, gemmC,
plan.gemmOutputRowsType, /*transA=*/false, /*transB=*/!wDenseAttr,
/*alpha=*/1.0f, /*beta=*/1.0f, state.targetInfo(), rewriter, loc);
if (failed(gemmRows))
return failure();
return createCollectedConvOutput(ValueRange {*gemmRows},
state.outType,
plan.gemmOutType,
plan.nhwcType,
state.outType,
plan.chunkNumPatches,
state.numChannelsOut,
plan.effectiveMaxParallelPixels,
rewriter,
loc);
}
static FailureOr<Value> rewriteStreamedConv(const ConvLoweringState& state,
PatternRewriter& rewriter,
Location loc,
int64_t forcedPackFactor) {
auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc);
Value biasMatrix;
DenseElementsAttr biasDenseAttr;
if (state.hasBias) {
biasDenseAttr = getHostConstDenseElementsAttr(state.b);
biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
}
ConvGemmPlan seedPlan = buildConvGemmPlan(
state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0, 1, forcedPackFactor);
Value weightMatrix = createWeightMatrix(state.w, seedPlan, static_cast<bool>(wDenseAttr), rewriter, loc);
FailureOr<Value> collectedRows = createStreamedConvRows(state,
preparedInput,
weightMatrix,
biasMatrix,
wDenseAttr,
biasDenseAttr,
forcedPackFactor,
rewriter,
loc);
if (failed(collectedRows))
return failure();
auto gemmOutType = cast<RankedTensorType>(collectedRows->getType());
auto nhwcType = RankedTensorType::get({state.batchSize, state.outHeight, state.outWidth, state.numChannelsOut},
state.outType.getElementType());
return createCollectedConvOutput(
ValueRange {*collectedRows}, state.outType, gemmOutType, nhwcType, state.outType, gemmOutType.getDimSize(0),
state.numChannelsOut, /*packFactor=*/1, rewriter, loc);
}
} // namespace standard
static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter) {
auto zeroAttr = DenseElementsAttr::get(gemmResultType, rewriter.getZeroAttr(gemmResultType.getElementType()));
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, gemmResultType);
}
static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) {
return ceilIntegerDivide(geometry.k, geometry.xbarSize)
* ceilIntegerDivide(geometry.c, geometry.xbarSize)
<= geometry.matrixUnitsPerProcessor;
}
static bool rowStripOutputChannelTileFitsOneCore(const ConvGeometry& geometry) {
return ceilIntegerDivide(geometry.k, geometry.xbarSize)
<= geometry.matrixUnitsPerProcessor;
}
static int64_t chooseRowStripPixelPackFactor(const ConvLoweringState& state, int64_t xbarDim) {
const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
const int64_t baseWeightGroups = ceilIntegerDivide(patchSize, xbarDim)
* ceilIntegerDivide(state.numChannelsOut, xbarDim);
int64_t factor = std::min(state.outWidth, xbarDim / state.numChannelsOut);
while (factor > 1
&& (state.outWidth % factor != 0
|| ceilIntegerDivide(factor * patchSize, xbarDim)
* ceilIntegerDivide(factor * state.numChannelsOut, xbarDim)
> baseWeightGroups))
--factor;
return std::max<int64_t>(factor, 1);
}
static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
if (state.batchSize != 1) {
failureReason = "batch_not_one";
return false;
}
if (state.group != 1) {
failureReason = "grouped_conv";
return false;
}
if (!state.xType.hasStaticShape() || !state.wType.hasStaticShape() || !state.outType.hasStaticShape()) {
failureReason = "dynamic_shape";
return false;
}
if (!isa<FloatType>(state.xType.getElementType())) {
failureReason = "non_float_input";
return false;
}
if (state.dilationHeight != 1 || state.dilationWidth != 1) {
failureReason = "dilation_not_one";
return false;
}
if (!getHostConstDenseElementsAttr(state.w)) {
failureReason = "non_constant_weight";
return false;
}
if (!rowStripOutputChannelTileFitsOneCore(buildConvGeometry(state, state.targetInfo()))) {
failureReason = "output_channel_tile_does_not_fit_one_core";
return false;
}
if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) {
failureReason = "unsupported_bias";
return false;
}
return true;
}
static Value createZeroTensorConstant(RankedTensorType type, PatternRewriter& rewriter) {
auto zeroAttr = DenseElementsAttr::get(type, rewriter.getZeroAttr(type.getElementType()));
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, type);
}
static FailureOr<Value> createBiasRowConstant(const ConvLoweringState& state,
int64_t packFactor,
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();
SmallVector<Attribute> packedValues;
packedValues.reserve(packFactor * state.numChannelsOut);
for (int64_t copy = 0; copy < packFactor; ++copy)
packedValues.append(channelValues->begin(), channelValues->end());
auto biasType = RankedTensorType::get({1, packFactor * state.numChannelsOut}, state.outType.getElementType());
return getOrCreateConstant(rewriter,
rewriter.getInsertionBlock()->getParentOp(),
DenseElementsAttr::get(biasType, packedValues),
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,
Location loc) {
auto paddedType = RankedTensorType::get(
{1, 1, state.xWidth + state.padWidthBegin + state.padWidthEnd, state.numChannelsIn},
state.xType.getElementType(),
state.xType.getEncoding());
return createZeroPaddedTensor(fragment,
paddedType,
{0, 0, state.padWidthBegin, 0},
{0, 0, state.padWidthEnd, 0},
rewriter,
loc);
}
static Value createRowStripWindowSourceRowTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
auto tableType = RankedTensorType::get({state.outHeight * state.wHeight}, rewriter.getIndexType());
SmallVector<Attribute> values;
values.reserve(tableType.getNumElements());
for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) {
for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) {
int64_t sourceRow =
outputRow * state.strideHeight + kernelRow * state.dilationHeight - state.padHeightBegin;
sourceRow = std::clamp(sourceRow, int64_t {0}, state.xHeight - 1);
values.push_back(rewriter.getIndexAttr(sourceRow));
}
}
return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
}
static Value createRowStripWindowSourceSlotTable(const ConvLoweringState& state,
int64_t tilesPerRow,
PatternRewriter& rewriter) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
auto tableType =
RankedTensorType::get({state.outHeight * state.wHeight * tilesPerRow}, rewriter.getIndexType());
SmallVector<Attribute> values;
values.reserve(tableType.getNumElements());
for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow)
for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) {
int64_t sourceRow =
outputRow * state.strideHeight + kernelRow * state.dilationHeight - state.padHeightBegin;
sourceRow = std::clamp(sourceRow, int64_t {0}, state.xHeight - 1);
for (int64_t tile = 0; tile < tilesPerRow; ++tile)
values.push_back(rewriter.getIndexAttr(sourceRow * tilesPerRow + tile));
}
return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
}
static Value createRowStripWindowTableIndex(Value outputHeight,
Value kernelRow,
const ConvLoweringState& state,
PatternRewriter& rewriter,
Location loc) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
MLIRContext* ctx = rewriter.getContext();
AffineExpr outputRowExpr = getAffineDimExpr(0, ctx);
AffineExpr kernelRowExpr = getAffineDimExpr(1, ctx);
return createOrFoldAffineApply(
rewriter, loc, outputRowExpr * state.wHeight + kernelRowExpr, ValueRange {outputHeight, kernelRow}, anchorOp);
}
static FailureOr<Value> extractProjectedRowStripWindowRow(Value rowStripStorage,
Value sourceSlotTable,
const ConvLoweringState& state,
Value outputHeight,
Value kernelRow,
PatternRewriter& rewriter,
Location loc) {
FailureOr<RowStripPhysicalValue> physical = describeRowStripPhysicalValue(rowStripStorage, state.xType);
if (failed(physical))
return failure();
Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value tileTableIndex =
affineMulConst(rewriter, loc, tableIndex, physical->tilesPerRow, anchorOp);
if (physical->tilesPerRow == 1) {
Value sourceSlot =
tensor::ExtractOp::create(rewriter, loc, sourceSlotTable, ValueRange {tileTableIndex}).getResult();
return extractGraphBatchPhysicalFragment(
rewriter, loc, rowStripStorage, sourceSlot, physical->fragmentType);
}
auto fullFragmentType = getRowStripFragmentType(state.xType);
Value fullFragment = tensor::EmptyOp::create(
rewriter, loc, fullFragmentType.getShape(), fullFragmentType.getElementType());
const int64_t tileChannels = physical->fragmentType.getDimSize(3);
for (int64_t tile = 0; tile < physical->tilesPerRow; ++tile) {
Value slotTableIndex = affineAddConst(rewriter, loc, tileTableIndex, tile, anchorOp);
Value tileSlot =
tensor::ExtractOp::create(rewriter, loc, sourceSlotTable, ValueRange {slotTableIndex}).getResult();
FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
rewriter, loc, rowStripStorage, tileSlot, physical->fragmentType);
if (failed(fragment))
return failure();
const int64_t channelOffset = tile * tileChannels;
const int64_t validChannels = std::min(tileChannels, state.numChannelsIn - channelOffset);
auto validType = RankedTensorType::get(
{1, 1, state.xWidth, validChannels}, state.xType.getElementType(), state.xType.getEncoding());
Value validFragment = *fragment;
if (validChannels != tileChannels)
validFragment = tensor::ExtractSliceOp::create(
rewriter,
loc,
validType,
*fragment,
SmallVector<OpFoldResult>(4, rewriter.getIndexAttr(0)),
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.xWidth),
rewriter.getIndexAttr(validChannels)},
getUnitStrides(rewriter, 4));
fullFragment = tensor::InsertSliceOp::create(
rewriter,
loc,
validFragment,
fullFragment,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
rewriter.getIndexAttr(0),
rewriter.getIndexAttr(0),
rewriter.getIndexAttr(channelOffset)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.xWidth),
rewriter.getIndexAttr(validChannels)},
getUnitStrides(rewriter, 4));
}
return fullFragment;
}
static Value extractDenseConvWindowRow(Value denseInput,
Value sourceRowTable,
const ConvLoweringState& state,
Value outputHeight,
Value kernelRow,
PatternRewriter& rewriter,
Location loc) {
Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
Value sourceRow = tensor::ExtractOp::create(rewriter, loc, sourceRowTable, ValueRange {tableIndex}).getResult();
auto nchwType = RankedTensorType::get(
{1, state.numChannelsIn, 1, state.xWidth}, state.xType.getElementType(), state.xType.getEncoding());
auto fragmentType = getRowStripFragmentType(state.xType);
SmallVector<OpFoldResult> offsets {
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceRow, rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.numChannelsIn),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.xWidth)};
Value nchw = tensor::ExtractSliceOp::create(
rewriter, loc, nchwType, denseInput, offsets, sizes, getUnitStrides(rewriter, 4));
return createLinalgTranspose(nchw, fragmentType, {0, 2, 3, 1}, rewriter, loc);
}
static Value createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
auto elementType = cast<FloatType>(state.xType.getElementType());
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
auto tableType = RankedTensorType::get(
{2, 1, state.xWidth, state.numChannelsIn}, elementType, state.xType.getEncoding());
SmallVector<Attribute> values(tableType.getNumElements(), rewriter.getZeroAttr(elementType));
std::fill(values.begin() + tableType.getNumElements() / 2, values.end(), rewriter.getFloatAttr(elementType, 1.0));
return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
}
static Value createRowStripWindowMaskIndexTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
auto tableType = RankedTensorType::get({state.outHeight * state.wHeight}, rewriter.getIndexType());
SmallVector<Attribute> values;
values.reserve(tableType.getNumElements());
for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow)
for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) {
int64_t sourceRow =
outputRow * state.strideHeight + kernelRow * state.dilationHeight - state.padHeightBegin;
values.push_back(rewriter.getIndexAttr(sourceRow >= 0 && sourceRow < state.xHeight));
}
return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
}
static Value extractProjectedRowStripWindowMask(Value maskTable,
Value maskIndexTable,
const ConvLoweringState& state,
Value outputHeight,
Value kernelRow,
PatternRewriter& rewriter,
Location loc) {
Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
Value maskIndex = tensor::ExtractOp::create(rewriter, loc, maskIndexTable, ValueRange {tableIndex}).getResult();
auto fragmentType = getRowStripFragmentType(state.xType);
return tensor::ExtractSliceOp::create(
rewriter,
loc,
fragmentType,
maskTable,
SmallVector<OpFoldResult> {
maskIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.xWidth), rewriter.getIndexAttr(state.numChannelsIn)},
getUnitStrides(rewriter, 4));
}
static FailureOr<Value> createConvInputWindow(Value input,
const ConvLoweringState& state,
Value outputHeight,
PatternRewriter& rewriter,
Location loc) {
auto fragmentType = getRowStripFragmentType(state.xType);
auto inputType = dyn_cast<RankedTensorType>(input.getType());
const bool denseInput = inputType == state.xType;
FailureOr<RowStripPhysicalValue> physicalInput = describeRowStripPhysicalValue(input, state.xType);
if (!denseInput && failed(physicalInput))
return failure();
if (!denseInput && physicalInput->tilesPerRow == 1 && state.wHeight == 1 && state.wWidth == 1
&& state.strideHeight == 1 && state.strideWidth == 1 && state.padHeightBegin == 0
&& state.padHeightEnd == 0 && state.padWidthBegin == 0 && state.padWidthEnd == 0)
return extractGraphBatchPhysicalFragment(
rewriter, loc, input, outputHeight, physicalInput->fragmentType);
auto paddedWindowType = RankedTensorType::get(
{1, state.wHeight, state.xWidth + state.padWidthBegin + state.padWidthEnd, state.numChannelsIn},
state.xType.getElementType(),
state.xType.getEncoding());
Value sourceIndexTable =
denseInput ? createRowStripWindowSourceRowTable(state, rewriter)
: createRowStripWindowSourceSlotTable(state, physicalInput->tilesPerRow, rewriter);
Value maskTable = createRowStripWindowMaskTable(state, rewriter);
Value maskIndexTable = createRowStripWindowMaskIndexTable(state, rewriter);
Value initWindow = createZeroTensorConstant(paddedWindowType, rewriter);
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
Value cKernelRows = getOrCreateIndexConstant(rewriter, anchorOp, state.wHeight);
auto loop = buildNormalizedScfFor(
rewriter,
loc,
c0,
cKernelRows,
c1,
ValueRange {initWindow},
[&](OpBuilder&, Location rowLoc, Value kernelRow, ValueRange iterArgs,
SmallVectorImpl<Value>& yielded) -> LogicalResult {
FailureOr<Value> sourceRow =
denseInput
? FailureOr<Value>(
extractDenseConvWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, rowLoc))
: extractProjectedRowStripWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, rowLoc);
if (failed(sourceRow))
return failure();
Value semanticRow = *sourceRow;
if (state.padHeightBegin != 0 || state.padHeightEnd != 0) {
Value mask = extractProjectedRowStripWindowMask(
maskTable, maskIndexTable, state, outputHeight, kernelRow, rewriter, rowLoc);
semanticRow = spatial::SpatVMulOp::create(rewriter, rowLoc, fragmentType, semanticRow, mask).getResult();
}
Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, rowLoc);
yielded.push_back(tensor::InsertSliceOp::create(
rewriter,
rowLoc,
paddedRow,
iterArgs.front(),
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), kernelRow,
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.xWidth + state.padWidthBegin
+ state.padWidthEnd),
rewriter.getIndexAttr(state.numChannelsIn)},
getUnitStrides(rewriter, 4)));
return success();
});
return failed(loop) ? FailureOr<Value>(failure())
: FailureOr<Value>(loop->results.front());
}
static FailureOr<Value> createPixelMajorConvPatchRow(Value paddedWindow,
const ConvLoweringState& state,
Value outputWidth,
PatternRewriter& rewriter,
Location loc) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
auto patchType = RankedTensorType::get({1, state.wHeight, state.wWidth, state.numChannelsIn},
state.xType.getElementType(),
state.xType.getEncoding());
auto rowType = RankedTensorType::get({1, patchSize}, state.xType.getElementType(), state.xType.getEncoding());
Value inputWidthOffset = affineMulConst(rewriter, loc, outputWidth, state.strideWidth, anchorOp);
SmallVector<OpFoldResult> offsets {
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), inputWidthOffset, rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.wHeight),
rewriter.getIndexAttr(state.wWidth),
rewriter.getIndexAttr(state.numChannelsIn)};
Value patch;
if (state.dilationWidth == 1)
patch = tensor::ExtractSliceOp::create(
rewriter, loc, patchType, paddedWindow, offsets, sizes, getUnitStrides(rewriter, 4));
else {
auto columnType = RankedTensorType::get({1, state.wHeight, 1, state.numChannelsIn},
state.xType.getElementType(), state.xType.getEncoding());
patch = tensor::EmptyOp::create(rewriter, loc, patchType.getShape(), patchType.getElementType());
for (int64_t kernelColumn = 0; kernelColumn < state.wWidth; ++kernelColumn) {
Value sourceWidth =
affineAddConst(rewriter, loc, inputWidthOffset, kernelColumn * state.dilationWidth, anchorOp);
SmallVector<OpFoldResult> columnSizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.wHeight),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(state.numChannelsIn)};
Value column = tensor::ExtractSliceOp::create(
rewriter, loc, columnType, paddedWindow,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceWidth,
rewriter.getIndexAttr(0)},
columnSizes, getUnitStrides(rewriter, 4));
patch = tensor::InsertSliceOp::create(
rewriter, loc, column, patch,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
rewriter.getIndexAttr(kernelColumn), rewriter.getIndexAttr(0)},
columnSizes, getUnitStrides(rewriter, 4));
}
}
return tensor::CollapseShapeOp::create(
rewriter, loc, rowType, patch, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}})
.getResult();
}
static FailureOr<Value> createPackedPixelMajorConvPatchRow(Value paddedWindow,
const ConvLoweringState& state,
Value outputGroup,
int64_t packFactor,
PatternRewriter& rewriter,
Location loc) {
if (packFactor == 1)
return createPixelMajorConvPatchRow(paddedWindow, state, outputGroup, rewriter, loc);
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
auto packedType = RankedTensorType::get(
{1, packFactor * patchSize}, state.xType.getElementType(), state.xType.getEncoding());
Value packed = tensor::EmptyOp::create(rewriter, loc, packedType.getShape(), packedType.getElementType());
Value outputStart = affineMulConst(rewriter, loc, outputGroup, packFactor, anchorOp);
Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
Value cPackFactor = getOrCreateIndexConstant(rewriter, anchorOp, packFactor);
auto loop = buildNormalizedScfFor(
rewriter,
loc,
c0,
cPackFactor,
c1,
ValueRange {packed},
[&](OpBuilder&, Location copyLoc, Value copy, ValueRange iterArgs,
SmallVectorImpl<Value>& yielded) -> LogicalResult {
Value outputWidth = createOrFoldAffineApply(
rewriter, copyLoc, rewriter.getAffineDimExpr(0) + rewriter.getAffineDimExpr(1),
ValueRange {outputStart, copy}, anchorOp);
FailureOr<Value> patch =
createPixelMajorConvPatchRow(paddedWindow, state, outputWidth, rewriter, copyLoc);
if (failed(patch))
return failure();
Value packedOffset = affineMulConst(rewriter, copyLoc, copy, patchSize, anchorOp);
yielded.push_back(tensor::InsertSliceOp::create(
rewriter,
copyLoc,
*patch,
iterArgs.front(),
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), packedOffset},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)},
getUnitStrides(rewriter, 2)));
return success();
});
if (failed(loop))
return failure();
return loop->results.front();
}
static FailureOr<SmallVector<Value>> createConvInputTiles(Value paddedWindow,
const ConvLoweringState& state,
Value outputWidth,
int64_t packFactor,
Value& partialInputScratch,
int64_t patchSize,
int64_t numKSlices,
int64_t xbarDim,
PatternRewriter& rewriter,
Location loc) {
auto elementType = state.xType.getElementType();
auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
SmallVector<Value> inputTiles;
inputTiles.reserve(numKSlices);
if (packFactor == 1 && state.numChannelsIn % xbarDim == 0) {
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
auto inputTileType = RankedTensorType::get(
{1, 1, 1, xbarDim}, elementType, state.xType.getEncoding());
for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice) {
const int64_t linearOffset = kSlice * xbarDim;
const int64_t kernelPixel = linearOffset / state.numChannelsIn;
const int64_t kernelRow = kernelPixel / state.wWidth;
const int64_t kernelColumn = kernelPixel % state.wWidth;
const int64_t channelOffset = linearOffset % state.numChannelsIn;
Value inputWidthOffset =
affineMulConst(rewriter, loc, outputWidth, state.strideWidth, anchorOp);
inputWidthOffset = affineAddConst(
rewriter, loc, inputWidthOffset, kernelColumn * state.dilationWidth, anchorOp);
Value inputTile = tensor::ExtractSliceOp::create(
rewriter,
loc,
inputTileType,
paddedWindow,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
rewriter.getIndexAttr(kernelRow),
inputWidthOffset,
rewriter.getIndexAttr(channelOffset)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(xbarDim)},
getUnitStrides(rewriter, 4));
inputTiles.push_back(tensor::CollapseShapeOp::create(
rewriter, loc, paddedRowType, inputTile, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}})
.getResult());
}
return inputTiles;
}
FailureOr<Value> patchRow = createPackedPixelMajorConvPatchRow(
paddedWindow, state, outputWidth, packFactor, rewriter, loc);
if (failed(patchRow))
return failure();
for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice) {
const int64_t kOffset = kSlice * xbarDim;
const int64_t sliceSize = std::min(xbarDim, patchSize - kOffset);
Value inputTile;
if (sliceSize == xbarDim) {
inputTile = extractStaticSliceOrIdentity(
rewriter,
loc,
*patchRow,
paddedRowType,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)},
getUnitStrides(rewriter, 2));
}
else {
if (!partialInputScratch)
return failure();
auto partialType = RankedTensorType::get({1, sliceSize}, elementType);
Value partial = extractStaticSliceOrIdentity(
rewriter,
loc,
*patchRow,
partialType,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(kOffset)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)},
getUnitStrides(rewriter, 2));
partialInputScratch = tensor::InsertSliceOp::create(
rewriter,
loc,
partial,
partialInputScratch,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(sliceSize)},
getUnitStrides(rewriter, 2));
inputTile = partialInputScratch;
}
inputTiles.push_back(inputTile);
}
return inputTiles;
}
static FailureOr<Value> createConvOutputTile(ValueRange inputTiles,
Value tileWeights,
int64_t outputChannels,
int64_t xbarDim,
PatternRewriter& rewriter,
Location loc) {
auto elementType = cast<RankedTensorType>(inputTiles.front().getType()).getElementType();
auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
auto resultType = RankedTensorType::get({1, outputChannels}, elementType);
auto weightElementType = cast<RankedTensorType>(tileWeights.getType()).getElementType();
auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, weightElementType);
Value tileResult;
for (auto [kSlice, inputTile] : llvm::enumerate(inputTiles)) {
const int64_t kOffset = static_cast<int64_t>(kSlice) * xbarDim;
SmallVector<OpFoldResult> bOffsets {
rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
Value bTile = extractStaticSliceOrIdentity(
rewriter, loc, tileWeights, paddedWeightTileType, bOffsets, bSizes, getUnitStrides(rewriter, 2));
Value piece = spatial::SpatVMMOp::create(
rewriter, loc, paddedRowType, bTile, inputTile).getResult();
if (outputChannels != xbarDim)
piece = tensor::ExtractSliceOp::create(
rewriter,
loc,
resultType,
piece,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputChannels)},
getUnitStrides(rewriter, 2));
tileResult = tileResult
? spatial::SpatVAddOp::create(
rewriter, loc, resultType, tileResult, piece).getResult()
: piece;
}
return tileResult;
}
static FailureOr<Value> createConvOutputRow(ValueRange inputTiles,
int64_t paddedK,
int64_t outputChannels,
Value paddedWeights,
Value bias,
int64_t xbarDim,
PatternRewriter& rewriter,
Location loc) {
auto elementType = cast<RankedTensorType>(inputTiles.front().getType()).getElementType();
auto rowType = RankedTensorType::get({1, outputChannels}, elementType);
const int64_t outputTileCount = ceilIntegerDivide(outputChannels, xbarDim);
const int64_t paddedOutputChannels = outputTileCount * xbarDim;
auto paddedOutputType = RankedTensorType::get({1, paddedOutputChannels}, elementType);
auto weightSliceType = RankedTensorType::get(
{xbarDim, paddedOutputChannels},
cast<RankedTensorType>(paddedWeights.getType()).getElementType());
Value paddedOutput;
for (auto [kSlice, inputTile] : llvm::enumerate(inputTiles)) {
const int64_t kOffset = static_cast<int64_t>(kSlice) * xbarDim;
Value weightSlice = extractStaticSliceOrIdentity(
rewriter,
loc,
paddedWeights,
weightSliceType,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(kOffset), rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(xbarDim),
rewriter.getIndexAttr(paddedOutputChannels)},
getUnitStrides(rewriter, 2));
Value piece =
spatial::SpatVMMOp::create(rewriter, loc, paddedOutputType, weightSlice, inputTile).getResult();
paddedOutput = paddedOutput
? spatial::SpatVAddOp::create(
rewriter, loc, paddedOutputType, paddedOutput, piece).getResult()
: piece;
}
Value validRow = outputChannels == paddedOutputChannels
? paddedOutput
: tensor::ExtractSliceOp::create(
rewriter,
loc,
rowType,
paddedOutput,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(outputChannels)},
getUnitStrides(rewriter, 2))
.getResult();
if (bias)
validRow = spatial::SpatVAddOp::create(rewriter, loc, rowType, validRow, bias).getResult();
return validRow;
}
static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
Value input,
Value paddedWeights,
Value bias,
int64_t packFactor,
int64_t paddedK,
int64_t numKSlices,
int64_t xbarDim,
PatternRewriter& rewriter,
Location loc) {
const int64_t laneCount = state.outHeight;
const int64_t patchSize = packFactor * state.numChannelsIn * state.wHeight * state.wWidth;
const bool hasPartialInputTile = patchSize % xbarDim != 0;
auto elementType = state.outType.getElementType();
auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
auto outputPixelType = RankedTensorType::get({1, 1, packFactor, state.numChannelsOut}, elementType);
auto fragmentType = getRowStripFragmentType(state.outType);
auto storageType = getRowStripStorageType(state.outType);
auto batch = createSpatComputeBatch(
rewriter,
loc,
TypeRange {storageType},
laneCount,
ValueRange {paddedWeights},
bias ? ValueRange {input, bias} : ValueRange {input},
[&](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 / packFactor);
FailureOr<Value> inputWindow =
createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
if (failed(inputWindow))
return failure();
Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
SmallVector<Value> loopInit {fragmentInit};
if (hasPartialInputTile)
loopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
auto loop = buildNormalizedScfFor(
rewriter,
loc,
c0,
cOutWidth,
c1,
loopInit,
[&](OpBuilder&,
Location pixelLoc,
Value localColumn,
ValueRange iterArgs,
SmallVectorImpl<Value>& yielded) {
Value partialInputScratch = hasPartialInputTile ? iterArgs[1] : Value();
FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
state,
localColumn,
packFactor,
partialInputScratch,
patchSize,
numKSlices,
xbarDim,
rewriter,
pixelLoc);
if (failed(inputTiles))
return failure();
FailureOr<Value> output = createConvOutputRow(*inputTiles,
paddedK,
packFactor * state.numChannelsOut,
args.weights.front(),
bias ? args.inputs[1] : Value(),
xbarDim,
rewriter,
pixelLoc);
if (failed(output))
return failure();
Value outputPixel = tensor::ExpandShapeOp::create(
rewriter, pixelLoc, outputPixelType, *output, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
Value next = tensor::InsertSliceOp::create(
rewriter,
pixelLoc,
outputPixel,
iterArgs.front(),
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
rewriter.getIndexAttr(0),
affineMulConst(rewriter, pixelLoc, localColumn, packFactor, anchorOp),
rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(packFactor),
rewriter.getIndexAttr(state.numChannelsOut)},
getUnitStrides(rewriter, 4));
yielded.push_back(next);
if (hasPartialInputTile)
yielded.push_back(partialInputScratch);
return success();
});
if (failed(loop))
return failure();
publishGraphBatchPhysicalFragment(
rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
return success();
});
if (failed(batch))
return failure();
return batch->getResult(0);
}
static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLoweringState& state,
Value input,
Value paddedWeights,
int64_t paddedK,
int64_t numKSlices,
int64_t xbarDim,
PatternRewriter& rewriter,
Location loc) {
const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim);
const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
const bool hasPartialInputTile = patchSize % xbarDim != 0;
auto elementType = state.outType.getElementType();
auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
auto tilePixelType = RankedTensorType::get({1, 1, 1, xbarDim}, elementType);
auto tileFragmentType = RankedTensorType::get({1, 1, state.outWidth, xbarDim}, elementType);
auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.wType.getElementType());
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 {input, *paddedBias} : ValueRange {input},
[&](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, outputRow, rewriter, loc);
if (failed(inputWindow))
return failure();
Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, tileFragmentType.getShape(), elementType);
SmallVector<Value> widthLoopInit {fragmentInit};
if (hasPartialInputTile)
widthLoopInit.push_back(createZeroTensorConstant(partialInputScratchType, rewriter));
auto widthLoop = buildNormalizedScfFor(
rewriter,
loc,
c0,
cOutWidth,
c1,
widthLoopInit,
[&](OpBuilder&,
Location widthLoc,
Value widthIndex,
ValueRange widthIterArgs,
SmallVectorImpl<Value>& widthYielded) {
Value partialInputScratch = hasPartialInputTile ? widthIterArgs[1] : Value();
FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
state,
widthIndex,
/*packFactor=*/1,
partialInputScratch,
patchSize,
numKSlices,
xbarDim,
rewriter,
widthLoc);
if (failed(inputTiles))
return failure();
FailureOr<Value> paddedOutputRow =
createConvOutputTile(*inputTiles, tileWeights, xbarDim, xbarDim, rewriter, widthLoc);
if (failed(paddedOutputRow))
return failure();
if (state.hasBias)
paddedOutputRow =
spatial::SpatVAddOp::create(rewriter, widthLoc, paddedRowType, *paddedOutputRow, *biasTile).getResult();
Value outputPixel = tensor::ExpandShapeOp::create(
rewriter, widthLoc, tilePixelType, *paddedOutputRow, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
SmallVector<OpFoldResult> rowOffsets {
rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1),
rewriter.getIndexAttr(xbarDim)};
Value nextFragment = tensor::InsertSliceOp::create(rewriter,
widthLoc,
outputPixel,
widthIterArgs.front(),
rowOffsets,
rowSizes,
getUnitStrides(rewriter, 4));
widthYielded.push_back(nextFragment);
if (hasPartialInputTile)
widthYielded.push_back(partialInputScratch);
return success();
});
if (failed(widthLoop))
return failure();
publishGraphBatchPhysicalFragment(rewriter, loc, widthLoop->results.front(), args.outputs.front(), args.lane);
return success();
});
if (failed(tileBatch))
return failure();
return tileBatch->getResult(0);
}
static FailureOr<Value>
createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) {
ConvGeometry geometry = buildConvGeometry(state, state.targetInfo());
if (state.group != 1 || state.batchSize != 1 || !rowStripOutputChannelTileFitsOneCore(geometry))
return failure();
auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
if (!weightDenseAttr)
return failure();
if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType))
return failure();
const int64_t xbarDim = geometry.xbarSize;
const int64_t basePatchSize = state.numChannelsIn * state.wHeight * state.wWidth;
const int64_t baseNumKSlices = ceilIntegerDivide(basePatchSize, xbarDim);
const int64_t basePaddedK = baseNumKSlices * xbarDim;
if (!rowStripOutputTileFitsOneCore(geometry)) {
Value tiledWeights =
standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, basePaddedK, xbarDim, rewriter);
return createOutputChannelTiledRowStripConvOutput(
state, state.x, tiledWeights, basePaddedK, baseNumKSlices, xbarDim, rewriter, loc);
}
const int64_t packFactor = chooseRowStripPixelPackFactor(state, xbarDim);
const int64_t packedPatchSize = packFactor * basePatchSize;
const int64_t numKSlices = ceilIntegerDivide(packedPatchSize, xbarDim);
const int64_t paddedK = numKSlices * xbarDim;
const int64_t packedOutputChannels = packFactor * state.numChannelsOut;
const int64_t paddedOutputChannels = ceilIntegerDivide(packedOutputChannels, xbarDim) * xbarDim;
Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
weightDenseAttr, state, paddedK, paddedOutputChannels, packFactor, rewriter);
FailureOr<Value> bias = failure();
if (state.hasBias)
bias = createBiasRowConstant(state, packFactor, rewriter);
if (state.hasBias && failed(bias))
return failure();
return createRowStripConvOutput(
state, state.x, paddedWeights, state.hasBias ? *bias : Value(),
packFactor, paddedK, numKSlices, xbarDim, rewriter, loc);
}
static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage,
const ConvLoweringState& state,
PatternRewriter& rewriter,
Location loc) {
if (failed(describeRowStripPhysicalValue(rowStripStorage, state.xType)))
return failure();
StringRef failureReason;
if (!canConsumePixelMajorRowStripFragments(state, failureReason))
return failure();
ConvGeometry geometry = buildConvGeometry(state, state.targetInfo());
const int64_t xbarDim = geometry.xbarSize;
const int64_t basePatchSize = state.numChannelsIn * state.wHeight * state.wWidth;
const int64_t baseNumKSlices = ceilIntegerDivide(basePatchSize, xbarDim);
const int64_t basePaddedK = baseNumKSlices * xbarDim;
auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
if (!weightDenseAttr)
return failure();
if (!rowStripOutputTileFitsOneCore(geometry)) {
Value tiledWeights =
standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, basePaddedK, xbarDim, rewriter);
return createOutputChannelTiledRowStripConvOutput(
state, rowStripStorage, tiledWeights, basePaddedK, baseNumKSlices, xbarDim, rewriter, loc);
}
const int64_t packFactor = chooseRowStripPixelPackFactor(state, xbarDim);
const int64_t packedPatchSize = packFactor * basePatchSize;
const int64_t numKSlices = ceilIntegerDivide(packedPatchSize, xbarDim);
const int64_t paddedK = numKSlices * xbarDim;
const int64_t packedOutputChannels = packFactor * state.numChannelsOut;
const int64_t paddedOutputChannels = ceilIntegerDivide(packedOutputChannels, xbarDim) * xbarDim;
Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
weightDenseAttr, state, paddedK, paddedOutputChannels, packFactor, rewriter);
FailureOr<Value> bias = failure();
if (state.hasBias)
bias = createBiasRowConstant(state, packFactor, rewriter);
if (state.hasBias && failed(bias))
return failure();
return createRowStripConvOutput(
state, rowStripStorage, paddedWeights, state.hasBias ? *bias : Value(),
packFactor, paddedK, numKSlices, xbarDim, rewriter, loc);
}
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, state.targetInfo());
const int64_t xbarDim = geometry.xbarSize;
const int64_t inputFragmentChannels = input->fragmentType.getDimSize(3);
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, 1, 1, xbarDim}, 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 bool canConsumeDepthwiseRowStrip(const ConvLoweringState& state) {
if (state.batchSize != 1 || state.group != state.numChannelsIn
|| state.dilationHeight != 1 || state.dilationWidth != 1
|| !isa<FloatType>(state.xType.getElementType())
|| !getHostConstDenseElementsAttr(state.w)
|| (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)))
return false;
auto tiling = depthwise::computeTiling(state.batchSize,
state.numChannelsIn,
state.numChannelsOut,
state.wHeight,
state.wWidth,
state.outHeight,
state.outWidth,
state.targetInfo().matrixShape.rows);
return tiling && tiling->numChannelTiles
<= static_cast<int64_t>(state.targetInfo().matrixUnitsPerProcessor);
}
static Value insertDepthwiseInputSegment(Value inputWindow,
Value scratch,
Value tileIndex,
Value kernelRow,
Value sourceWidth,
Value scratchOffset,
int64_t inputChannel,
const depthwise::Tiling& tiling,
PatternRewriter& rewriter,
Location loc) {
auto inputWindowType = cast<RankedTensorType>(inputWindow.getType());
auto inputPixelType = RankedTensorType::get(
{1, 1, 1, tiling.channelsPerTile}, inputWindowType.getElementType(), inputWindowType.getEncoding());
Value inputPixel = tensor::ExtractSliceOp::create(
rewriter,
loc,
inputPixelType,
inputWindow,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), kernelRow, sourceWidth,
rewriter.getIndexAttr(inputChannel)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1), rewriter.getIndexAttr(tiling.channelsPerTile)},
getUnitStrides(rewriter, 4));
return tensor::InsertSliceOp::create(
rewriter,
loc,
inputPixel,
scratch,
SmallVector<OpFoldResult> {tileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), scratchOffset},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
rewriter.getIndexAttr(tiling.channelsPerTile)},
getUnitStrides(rewriter, 4));
}
static FailureOr<Value> assembleDepthwiseInputScratch(Value inputWindow,
Value scratch,
Value inputWidth,
const ConvLoweringState& state,
const depthwise::Tiling& tiling,
PatternRewriter& rewriter,
Location loc) {
Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
Value cKernelElements = getOrCreateIndexConstant(rewriter, anchor, tiling.kernelElements);
SmallVector<Value> tileIndices;
tileIndices.reserve(tiling.numChannelTiles);
for (int64_t tile = 0; tile < tiling.numChannelTiles; ++tile)
tileIndices.push_back(getOrCreateIndexConstant(rewriter, anchor, tile));
auto kernelLoop = buildNormalizedScfFor(
rewriter,
loc,
c0,
cKernelElements,
c1,
ValueRange {scratch},
[&](OpBuilder&, Location kernelLoc, Value kernelIndex, ValueRange iterArgs,
SmallVectorImpl<Value>& yielded) {
Value kernelRow = affineFloorDivConst(rewriter, kernelLoc, kernelIndex, state.wWidth, anchor);
Value kernelColumn = affineModConst(rewriter, kernelLoc, kernelIndex, state.wWidth, anchor);
Value sourceWidth = createOrFoldAffineApply(
rewriter,
kernelLoc,
getAffineDimExpr(0, rewriter.getContext()) + getAffineDimExpr(1, rewriter.getContext()),
ValueRange {inputWidth, kernelColumn},
anchor);
Value scratchOffset = affineMulConst(
rewriter, kernelLoc, kernelIndex, tiling.channelsPerTile, anchor);
Value nextScratch = iterArgs.front();
for (int64_t tile = 0; tile < tiling.numChannelTiles; ++tile)
nextScratch = insertDepthwiseInputSegment(inputWindow,
nextScratch,
tileIndices[tile],
kernelRow,
sourceWidth,
scratchOffset,
tile * tiling.channelsPerTile,
tiling,
rewriter,
kernelLoc);
yielded.push_back(nextScratch);
return success();
});
if (failed(kernelLoop))
return failure();
return kernelLoop->results.front();
}
static FailureOr<Value> createDepthwiseOutputFromRowStripFragments(Value rowStripStorage,
const ConvLoweringState& state,
PatternRewriter& rewriter,
Location loc) {
if (!canConsumeDepthwiseRowStrip(state)
|| failed(describeRowStripPhysicalValue(rowStripStorage, state.xType)))
return failure();
auto tiling = depthwise::computeTiling(state.batchSize,
state.numChannelsIn,
state.numChannelsOut,
state.wHeight,
state.wWidth,
state.outHeight,
state.outWidth,
state.targetInfo().matrixShape.rows);
auto weight = getHostConstDenseElementsAttr(state.w);
if (!tiling || !weight)
return failure();
Value packedWeights = depthwise::buildPackedWeights(
weight,
state.wType,
*tiling,
rewriter,
loc,
static_cast<int64_t>(state.targetInfo().matrixShape.rows),
static_cast<int64_t>(state.targetInfo().matrixShape.rows));
Value bias = state.hasBias ? expandBiasIfNeeded(state.b, rewriter, loc) : Value();
auto paddedOutputType = RankedTensorType::get(
{1, static_cast<int64_t>(state.targetInfo().matrixShape.rows)}, state.outType.getElementType());
auto outputTileType = RankedTensorType::get(
{1, tiling->tileOutputChannels}, state.outType.getElementType());
auto outputPixelType = RankedTensorType::get(
{1, 1, 1, tiling->tileOutputChannels}, state.outType.getElementType());
auto fragmentType = getRowStripFragmentType(state.outType);
auto storageType = getRowStripStorageType(state.outType);
auto batch = createSpatComputeBatch(
rewriter,
loc,
TypeRange {storageType},
state.outHeight,
ValueRange {packedWeights},
state.hasBias ? ValueRange {rowStripStorage, bias} : ValueRange {rowStripStorage},
[&](detail::SpatComputeBatchBodyArgs args) {
Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
FailureOr<Value> inputWindow =
createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
if (failed(inputWindow))
return failure();
Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
Value cOutWidth = getOrCreateIndexConstant(rewriter, anchor, state.outWidth);
const int64_t xbarDim = static_cast<int64_t>(state.targetInfo().matrixShape.rows);
auto paddedInputScratchType = RankedTensorType::get(
{tiling->numChannelTiles, 1, 1, xbarDim}, state.xType.getElementType(), state.xType.getEncoding());
auto tileScratchType = RankedTensorType::get(
{1, 1, 1, xbarDim}, state.xType.getElementType(), state.xType.getEncoding());
auto vmmInputType = RankedTensorType::get(
{1, xbarDim}, state.xType.getElementType(), state.xType.getEncoding());
Value zeroScratch = createZeroTensorConstant(paddedInputScratchType, rewriter);
Value fragment = tensor::EmptyOp::create(
rewriter, loc, fragmentType.getShape(), fragmentType.getElementType());
SmallVector<Value> weightTiles;
SmallVector<Value> biasTiles;
SmallVector<Value> tileIndices;
SmallVector<OpFoldResult> weightTileSizes {
rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim),
rewriter.getIndexAttr(xbarDim)};
for (int64_t tile = 0; tile < tiling->numChannelTiles; ++tile) {
Value tileIndex = getOrCreateIndexConstant(rewriter, anchor, tile);
tileIndices.push_back(tileIndex);
weightTiles.push_back(extractMixedSliceOrIdentity(
rewriter,
loc,
args.weights.front(),
RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType()),
{SmallVector<OpFoldResult> {tileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
weightTileSizes,
getUnitStrides(rewriter, 3)}));
if (state.hasBias)
biasTiles.push_back(depthwise::createBiasTile(args.inputs[1], tileIndex, *tiling, rewriter, loc));
}
auto widthLoop = buildNormalizedScfFor(
rewriter, loc, c0, cOutWidth, c1, ValueRange {fragment, zeroScratch},
[&](OpBuilder&, Location widthLoc, Value width, ValueRange iterArgs,
SmallVectorImpl<Value>& yielded) {
Value next = iterArgs.front();
Value scratch = iterArgs[1];
// Valid prefix entries are overwritten per tile; the padded tail stays zero.
Value inputWidth = affineMulConst(
rewriter, widthLoc, width, state.strideWidth, anchor);
FailureOr<Value> nextScratch = assembleDepthwiseInputScratch(
*inputWindow, scratch, inputWidth, state, *tiling, rewriter, widthLoc);
if (failed(nextScratch))
return failure();
scratch = *nextScratch;
for (int64_t tile = 0; tile < tiling->numChannelTiles; ++tile) {
Value tileScratch = tensor::ExtractSliceOp::create(
rewriter,
widthLoc,
tileScratchType,
scratch,
SmallVector<OpFoldResult> {tileIndices[tile], rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)},
getUnitStrides(rewriter, 4));
Value vmmInput = tensor::CollapseShapeOp::create(
rewriter,
widthLoc,
vmmInputType,
tileScratch,
SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
Value output = spatial::SpatVMMOp::create(
rewriter, widthLoc, paddedOutputType, weightTiles[tile], vmmInput);
Value validOutput = tensor::ExtractSliceOp::create(
rewriter,
widthLoc,
outputTileType,
output,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
rewriter.getIndexAttr(tiling->tileOutputChannels)},
getUnitStrides(rewriter, 2));
if (state.hasBias)
validOutput = spatial::SpatVAddOp::create(
rewriter, widthLoc, outputTileType, validOutput, biasTiles[tile]);
Value pixel = tensor::ExpandShapeOp::create(
rewriter, widthLoc, outputPixelType, validOutput,
SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
next = tensor::InsertSliceOp::create(
rewriter,
widthLoc,
pixel,
next,
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), width,
rewriter.getIndexAttr(tile * tiling->tileOutputChannels)},
SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
rewriter.getIndexAttr(1), rewriter.getIndexAttr(tiling->tileOutputChannels)},
getUnitStrides(rewriter, 4));
}
yielded.push_back(next);
yielded.push_back(scratch);
return success();
});
if (failed(widthLoop))
return failure();
publishGraphBatchPhysicalFragment(rewriter, loc, widthLoop->results.front(), args.outputs.front(), args.lane);
return success();
});
return failed(batch) ? FailureOr<Value>(failure()) : FailureOr<Value>(batch->getResult(0));
}
static FailureOr<Value> createConvOutputFromRowStripInput(const ConvLoweringState& state,
Value rowStripInput,
spatial::ConvLoweringStrategy strategy,
PatternRewriter& rewriter,
Location loc) {
if (strategy == spatial::ConvLoweringStrategy::Depthwise)
return createDepthwiseOutputFromRowStripFragments(rowStripInput, state, rewriter, loc);
if (state.xHeight == 1 && state.xWidth == 1 && state.wHeight == 1 && state.wWidth == 1)
return createPointwiseOutputFromRowStripFragments(rowStripInput, state, rewriter, loc);
return createConvOutputFromPixelMajorRowStripFragments(rowStripInput, state, rewriter, loc);
}
static Value createCollectedConvOutput(ValueRange gemmRows,
Type convType,
RankedTensorType gemmOutType,
RankedTensorType nhwcType,
RankedTensorType outType,
int64_t numPatches,
int64_t numChannelsOut,
int64_t packFactor,
PatternRewriter& rewriter,
Location loc) {
auto collectComputeOp = createSpatCompute(rewriter, loc, convType, {}, gemmRows, [&](ValueRange gemmRowArgs) {
Value gemmOut;
if (packFactor == 1) {
gemmOut = createSpatConcat(rewriter, loc, /*axis=*/0, gemmRowArgs);
}
else {
Value packedOutput = createSpatConcat(rewriter, loc, /*axis=*/0, gemmRowArgs);
gemmOut = standard::unpackRowsFromParallelGemm(
packedOutput, cast<RankedTensorType>(packedOutput.getType()), numPatches, numChannelsOut, packFactor, rewriter, loc);
}
// Restore output layout:
// [numPatches, numChannelsOut]
// -> [N, Hout, Wout, Cout]
// -> [N, Cout, Hout, Wout]
Value nhwcOut = tensor::ExpandShapeOp::create(rewriter,
loc,
nhwcType,
gemmOut,
SmallVector<ReassociationIndices> {
{0, 1, 2},
{3}
});
Value nchwOut = createLinalgTranspose(nhwcOut, outType, {0, 3, 1, 2}, rewriter, loc);
spatial::SpatYieldOp::create(rewriter, loc, nchwOut);
});
return collectComputeOp.getResult(0);
}
static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp,
Value x,
Value w,
Value b,
const spatial::SpatialTargetInfo& target) {
ConvLoweringState state;
state.diagnosticAnchor = convOp.getOperation();
state.x = x;
state.w = w;
state.b = b;
state.target = &target;
state.xType = cast<RankedTensorType>(state.x.getType());
state.wType = cast<RankedTensorType>(state.w.getType());
state.outType = cast<RankedTensorType>(convOp.getY().getType());
if (!state.xType.hasStaticShape()) {
pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv input");
return failure();
}
if (!state.wType.hasStaticShape()) {
pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv weight");
return failure();
}
if (!state.outType.hasStaticShape()) {
pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv result");
return failure();
}
if (state.xType.getRank() != 4) {
pim::emitUnsupportedRankDiagnostic(convOp, "conv input", state.xType.getRank(), {4});
return failure();
}
if (state.wType.getRank() != 4) {
pim::emitUnsupportedRankDiagnostic(convOp, "conv weight", state.wType.getRank(), {4});
return failure();
}
if (state.outType.getRank() != 4) {
pim::emitUnsupportedRankDiagnostic(convOp, "conv result", state.outType.getRank(), {4});
return failure();
}
state.group = convOp.getGroup();
if (state.group < 1) {
convOp.emitOpError("requires group >= 1 for Spatial lowering");
return failure();
}
state.batchSize = state.xType.getDimSize(0);
state.numChannelsIn = state.xType.getDimSize(1);
state.xHeight = state.xType.getDimSize(2);
state.xWidth = state.xType.getDimSize(3);
state.numChannelsOut = state.wType.getDimSize(0);
state.wHeight = state.wType.getDimSize(2);
state.wWidth = state.wType.getDimSize(3);
state.outHeight = state.outType.getDimSize(2);
state.outWidth = state.outType.getDimSize(3);
state.hasBias =
state.b && !isa<ONNXNoneOp>(state.b.getDefiningOp()) && !isZeroSplatHostConstant(state.b);
if (state.numChannelsIn % state.group != 0) {
convOp.emitOpError() << "requires input channels " << state.numChannelsIn << " to be divisible by group "
<< state.group << " for Spatial lowering";
return failure();
}
if (state.numChannelsOut % state.group != 0) {
convOp.emitOpError() << "requires output channels " << state.numChannelsOut << " to be divisible by group "
<< state.group << " for Spatial lowering";
return failure();
}
state.numChannelsInPerGroup = state.numChannelsIn / state.group;
state.numChannelsOutPerGroup = state.numChannelsOut / state.group;
if (state.wType.getDimSize(1) != state.numChannelsInPerGroup) {
convOp.emitOpError() << "requires grouped conv weight input channels " << state.wType.getDimSize(1)
<< " to match input channels per group " << state.numChannelsInPerGroup
<< " for Spatial lowering";
return failure();
}
if (state.wType.getDimSize(0) != state.numChannelsOut) {
convOp.emitOpError() << "requires weight output channels " << state.wType.getDimSize(0)
<< " to match result channels " << state.numChannelsOut << " for Spatial lowering";
return failure();
}
const auto stridesAttr = convOp.getStrides();
const auto dilationsAttr = convOp.getDilations();
const auto padsAttr = convOp.getPads();
if (stridesAttr && stridesAttr->size() != 2) {
convOp.emitOpError("requires exactly two stride values for Spatial lowering");
return failure();
}
if (dilationsAttr && dilationsAttr->size() != 2) {
convOp.emitOpError("requires exactly two dilation values for Spatial lowering");
return failure();
}
if (padsAttr && padsAttr->size() != 4) {
convOp.emitOpError("requires exactly four pad values for 2D Spatial lowering");
return failure();
}
state.strideHeight = getOptionalI64Attr(stridesAttr, 0, 1);
state.strideWidth = getOptionalI64Attr(stridesAttr, 1, 1);
state.dilationHeight = getOptionalI64Attr(dilationsAttr, 0, 1);
state.dilationWidth = getOptionalI64Attr(dilationsAttr, 1, 1);
state.padHeightBegin = 0;
state.padHeightEnd = 0;
state.padWidthBegin = 0;
state.padWidthEnd = 0;
if (padsAttr) {
state.padHeightBegin = getI64Attr(*padsAttr, 0);
state.padWidthBegin = getI64Attr(*padsAttr, 1);
state.padHeightEnd = getI64Attr(*padsAttr, 2);
state.padWidthEnd = getI64Attr(*padsAttr, 3);
classifyConvProblem(state);
return state;
}
const auto autoPad = convOp.getAutoPad();
if (autoPad == "SAME_UPPER" || autoPad == "SAME_LOWER") {
const int64_t effectiveKernelH = (state.wHeight - 1) * state.dilationHeight + 1;
const int64_t effectiveKernelW = (state.wWidth - 1) * state.dilationWidth + 1;
const int64_t totalPadH =
std::max(static_cast<int64_t>(0), (state.outHeight - 1) * state.strideHeight + effectiveKernelH - state.xHeight);
const int64_t totalPadW =
std::max(static_cast<int64_t>(0), (state.outWidth - 1) * state.strideWidth + effectiveKernelW - state.xWidth);
if (autoPad == "SAME_UPPER") {
state.padHeightBegin = totalPadH / 2;
state.padHeightEnd = totalPadH - state.padHeightBegin;
state.padWidthBegin = totalPadW / 2;
state.padWidthEnd = totalPadW - state.padWidthBegin;
}
else {
state.padHeightEnd = totalPadH / 2;
state.padHeightBegin = totalPadH - state.padHeightEnd;
state.padWidthEnd = totalPadW / 2;
state.padWidthBegin = totalPadW - state.padWidthEnd;
}
classifyConvProblem(state);
return state;
}
if (autoPad != "NOTSET" && autoPad != "VALID") {
convOp.emitOpError() << "unsupported auto_pad value `" << autoPad << "` for Spatial lowering";
return failure();
}
classifyConvProblem(state);
return state;
}
static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp,
ONNXConvOpAdaptor convOpAdaptor,
const spatial::SpatialTargetInfo& target) {
return analyzeConvLoweringState(
convOp, convOpAdaptor.getX(), convOpAdaptor.getW(), convOpAdaptor.getB(), target);
}
static FailureOr<ConvLoweringState> analyzeConvLoweringState(
spatial::SpatConv2DPlanOp planOp, const spatial::SpatialTargetInfo& target) {
ConvLoweringState state;
state.diagnosticAnchor = planOp.getOperation();
state.x = planOp.getInput();
state.w = planOp.getWeight();
state.b = planOp.getBias() ? planOp.getBias() : Value();
state.target = &target;
state.xType = dyn_cast<RankedTensorType>(state.x.getType());
state.wType = dyn_cast<RankedTensorType>(state.w.getType());
state.outType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
if (!state.xType || !state.wType || !state.outType)
return planOp.emitOpError("requires ranked tensor input, weight, and output"), failure();
if (!state.xType.hasStaticShape() || !state.wType.hasStaticShape() || !state.outType.hasStaticShape())
return planOp.emitOpError("requires static input, weight, and output shapes"), failure();
if (state.xType.getRank() != 4 || state.wType.getRank() != 4 || state.outType.getRank() != 4)
return planOp.emitOpError("requires rank-4 input, weight, and output tensors"), failure();
state.group = planOp.getGroup();
if (state.group < 1)
return planOp.emitOpError("requires group >= 1"), failure();
state.batchSize = state.xType.getDimSize(0);
state.numChannelsIn = state.xType.getDimSize(1);
state.xHeight = state.xType.getDimSize(2);
state.xWidth = state.xType.getDimSize(3);
state.numChannelsOut = state.wType.getDimSize(0);
state.wHeight = state.wType.getDimSize(2);
state.wWidth = state.wType.getDimSize(3);
state.outHeight = state.outType.getDimSize(2);
state.outWidth = state.outType.getDimSize(3);
state.hasBias = planOp.getBias() && !isZeroSplatHostConstant(planOp.getBias());
if (state.numChannelsIn % state.group != 0 || state.numChannelsOut % state.group != 0)
return planOp.emitOpError("requires input and output channels divisible by group"), failure();
state.numChannelsInPerGroup = state.numChannelsIn / state.group;
state.numChannelsOutPerGroup = state.numChannelsOut / state.group;
if (state.wType.getDimSize(1) != state.numChannelsInPerGroup)
return planOp.emitOpError("requires grouped conv weight channels to match input channels per group"), failure();
auto pads = planOp.getPads();
auto strides = planOp.getStrides();
auto dilations = planOp.getDilations();
if (pads.size() != 4 || strides.size() != 2 || dilations.size() != 2)
return planOp.emitOpError("requires 4 pads, 2 strides, and 2 dilations"), failure();
state.padHeightBegin = pads[0];
state.padWidthBegin = pads[1];
state.padHeightEnd = pads[2];
state.padWidthEnd = pads[3];
state.strideHeight = strides[0];
state.strideWidth = strides[1];
state.dilationHeight = dilations[0];
state.dilationWidth = dilations[1];
classifyConvProblem(state);
return state;
}
static FailureOr<spatial::ConvLoweringStrategy>
resolveRequestedConvLoweringStrategy(Operation* op, const spatial::SpatialTargetInfo& target) {
if (!target.useExperimentalConvImplementation)
return target.convLoweringStrategy;
if (target.convLoweringStrategy != spatial::ConvLoweringStrategy::Auto
&& target.convLoweringStrategy != spatial::ConvLoweringStrategy::PackedIm2Col) {
op->emitOpError() << "--use-experimental-conv-impl conflicts with --pim-conv-lowering="
<< stringifyConvLoweringStrategy(target.convLoweringStrategy);
return failure();
}
return spatial::ConvLoweringStrategy::PackedIm2Col;
}
static FailureOr<ConvPlan> selectConvLoweringPlan(
Operation* op, const ConvLoweringState& state, bool reportPlanning) {
FailureOr<spatial::ConvLoweringStrategy> requested =
resolveRequestedConvLoweringStrategy(op, state.targetInfo());
if (failed(requested))
return failure();
if (*requested == spatial::ConvLoweringStrategy::Auto) {
for (const ConvPlan& candidate : buildConvPlanCandidates(state, state.targetInfo())) {
if (candidate.strategy == spatial::ConvLoweringStrategy::Depthwise
&& !depthwise::canUseStructuredRewrite(state)) {
continue;
}
if (reportPlanning)
recordConvLoweringReport(op, ConvLoweringReportPhase::Planning, candidate.strategy, "SEL");
return candidate;
}
op->emitOpError("has no applicable Conv lowering candidate for the injected Spatial target");
return failure();
}
FailureOr<ConvPlan> candidate = makeConvPlan(state, *requested, state.targetInfo());
if (failed(candidate)) {
op->emitOpError() << "forced Conv lowering `"
<< stringifyConvLoweringStrategy(*requested)
<< "` is not applicable to this Conv problem";
return failure();
}
if (reportPlanning)
recordConvLoweringReport(op, ConvLoweringReportPhase::Planning, candidate->strategy, "SEL");
return *candidate;
}
static FailureOr<Value> lowerDenseSelectedConvPlan(Operation* op,
const ConvLoweringState& state,
spatial::ConvLoweringStrategy strategy,
PatternRewriter& rewriter,
Location loc);
static ConvLoweringState makeGroupedConvLoweringState(const ConvLoweringState& parent,
Value groupX,
Value groupW,
Value groupB,
RankedTensorType groupOutType);
static FailureOr<Value> buildConvValueForStrategy(Operation* op,
Location loc,
const ConvLoweringState& state,
spatial::ConvLoweringStrategy strategy,
PatternRewriter& rewriter);
static FailureOr<Value> buildGroupedConvValue(Operation* op,
Location loc,
const ConvLoweringState& state,
spatial::ConvLoweringStrategy strategy,
PatternRewriter& rewriter);
static FailureOr<Value> lowerGroupedSelectedConvPlan(Operation* op,
const ConvLoweringState& state,
spatial::ConvLoweringStrategy strategy,
PatternRewriter& rewriter,
Location loc) {
return buildGroupedConvValue(op, loc, state, strategy, rewriter);
}
static FailureOr<Value> lowerDenseSelectedConvPlan(Operation* op,
const ConvLoweringState& state,
spatial::ConvLoweringStrategy strategy,
PatternRewriter& rewriter,
Location loc) {
return buildConvValueForStrategy(op, loc, state, strategy, rewriter);
}
static FailureOr<Value> buildConvValueForStrategy(Operation* op,
Location loc,
const ConvLoweringState& state,
spatial::ConvLoweringStrategy strategy,
PatternRewriter& rewriter) {
const ConvGeometry geo = buildConvGeometry(state, state.targetInfo());
switch (strategy) {
case spatial::ConvLoweringStrategy::Depthwise: {
return depthwise::rewriteConv(op, state, rewriter, loc);
}
case spatial::ConvLoweringStrategy::Legacy:
case spatial::ConvLoweringStrategy::PackedIm2Col: {
return standard::rewritePackedIm2ColConv(state, rewriter, loc);
}
case spatial::ConvLoweringStrategy::StreamedPatch:
case spatial::ConvLoweringStrategy::OutputChannelTiled:
case spatial::ConvLoweringStrategy::Tiled2D: {
return standard::rewriteStreamedConv(state, rewriter, loc, /*forcedPackFactor=*/1);
}
case spatial::ConvLoweringStrategy::InputKTiled: {
return standard::rewriteInputKTiledConv(state, rewriter, loc);
}
case spatial::ConvLoweringStrategy::StreamedPacked: {
return standard::rewriteStreamedConv(state, rewriter, loc, geo.pack);
}
case spatial::ConvLoweringStrategy::Auto:
break;
}
op->emitOpError("unexpected auto strategy at Conv lowering dispatch");
return failure();
}
static ConvLoweringState makeGroupedConvLoweringState(const ConvLoweringState& parent,
Value groupX,
Value groupW,
Value groupB,
RankedTensorType groupOutType);
static ConvLoweringState makeGroupedConvLoweringState(
const ConvLoweringState& parent, Value groupX, Value groupW, Value groupB, RankedTensorType groupOutType) {
ConvLoweringState state = parent;
state.x = groupX;
state.w = groupW;
state.b = groupB;
state.xType = cast<RankedTensorType>(groupX.getType());
state.wType = cast<RankedTensorType>(groupW.getType());
state.outType = groupOutType;
state.batchSize = state.xType.getDimSize(0);
state.numChannelsIn = state.xType.getDimSize(1);
state.xHeight = state.xType.getDimSize(2);
state.xWidth = state.xType.getDimSize(3);
state.numChannelsOut = state.wType.getDimSize(0);
state.wHeight = state.wType.getDimSize(2);
state.wWidth = state.wType.getDimSize(3);
state.outHeight = state.outType.getDimSize(2);
state.outWidth = state.outType.getDimSize(3);
state.group = 1;
state.numChannelsInPerGroup = state.numChannelsIn;
state.numChannelsOutPerGroup = state.numChannelsOut;
state.hasBias = static_cast<bool>(groupB);
classifyConvProblem(state);
return state;
}
static FailureOr<Value> buildGroupedConvValue(Operation* op,
Location loc,
const ConvLoweringState& state,
spatial::ConvLoweringStrategy strategy,
PatternRewriter& rewriter) {
SmallVector<Value> xSlices = sliceTensor(state.x, /*axis=*/1, state.numChannelsInPerGroup, rewriter, loc);
SmallVector<Value> wSlices = sliceTensor(state.w, /*axis=*/0, state.numChannelsOutPerGroup, rewriter, loc);
SmallVector<Value> bSlices;
if (state.hasBias) {
auto biasType = cast<RankedTensorType>(state.b.getType());
int64_t biasAxis = -1;
if (biasType.getRank() == 1)
biasAxis = 0;
else if (biasType.getRank() == 2)
biasAxis = biasType.getDimSize(0) != 1 ? 0 : 1;
else {
op->emitOpError() << "requires rank-1 or rank-2 bias for grouped convolution Spatial lowering, but got rank "
<< biasType.getRank();
return failure();
}
bSlices = sliceTensor(state.b, biasAxis, state.numChannelsOutPerGroup, rewriter, loc);
}
if (xSlices.size() != static_cast<size_t>(state.group) || wSlices.size() != static_cast<size_t>(state.group)
|| (state.hasBias && bSlices.size() != static_cast<size_t>(state.group))) {
op->emitOpError("failed to partition grouped convolution operands for Spatial lowering");
return failure();
}
SmallVector<Value> groupResults;
groupResults.reserve(state.group);
auto groupOutType = RankedTensorType::get(
{state.batchSize, state.numChannelsOutPerGroup, state.outHeight, state.outWidth}, state.outType.getElementType());
for (int64_t groupId = 0; groupId < state.group; groupId++) {
Value groupX = xSlices[groupId];
Value groupW = wSlices[groupId];
Value groupB = state.hasBias ? bSlices[groupId] : Value();
ConvLoweringState groupState = makeGroupedConvLoweringState(state, groupX, groupW, groupB, groupOutType);
FailureOr<Value> groupResult = buildConvValueForStrategy(op, loc, groupState, strategy, rewriter);
if (failed(groupResult))
return failure();
groupResults.push_back(*groupResult);
}
if (llvm::all_of(groupResults, isCompileTimeComputable))
return createSpatConcat(rewriter, loc, /*axis=*/1, groupResults);
auto concatCompute = createSpatCompute(rewriter, loc, TypeRange {state.outType}, {}, groupResults, [&](ValueRange args) {
spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, /*axis=*/1, args));
});
return concatCompute.getResult(0);
}
} // namespace
LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
ONNXConvOpAdaptor convOpAdaptor,
ConversionPatternRewriter& rewriter) const {
FailureOr<ConvLoweringState> state = analyzeConvLoweringState(convOp, convOpAdaptor, target);
if (failed(state))
return failure();
SmallVector<int64_t> pads {
state->padHeightBegin, state->padWidthBegin, state->padHeightEnd, state->padWidthEnd};
SmallVector<int64_t> strides {state->strideHeight, state->strideWidth};
SmallVector<int64_t> dilations {state->dilationHeight, state->dilationWidth};
Value bias = state->hasBias ? convOpAdaptor.getB() : Value();
auto convPlan = spatial::SpatConv2DPlanOp::create(rewriter,
convOp.getLoc(),
convOp.getY().getType(),
convOpAdaptor.getX(),
convOpAdaptor.getW(),
bias,
rewriter.getDenseI64ArrayAttr(pads),
rewriter.getDenseI64ArrayAttr(strides),
rewriter.getDenseI64ArrayAttr(dilations),
rewriter.getI64IntegerAttr(state->group),
spatial::getNCHWLayout(rewriter.getContext()));
rewriter.replaceOp(convOp, convPlan.getResult());
return success();
}
void populateConvPatterns(RewritePatternSet& patterns,
MLIRContext* ctx,
const spatial::SpatialTargetInfo& target) {
patterns.insert<ConvToGemm>(ctx, target);
}
LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp,
const spatial::SpatialTargetInfo& target) {
FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp, target);
if (failed(state))
return failure();
if (state->group != 1 || state->batchSize != 1)
return failure();
if (state->outType.getRank() != 4 || !state->outType.hasStaticShape())
return failure();
if (!getHostConstDenseElementsAttr(state->w))
return failure();
if (state->hasBias && !isSupportedBiasAddValue(state->b, state->outType))
return failure();
ConvGeometry geometry = buildConvGeometry(*state, state->targetInfo());
if (!rowStripOutputChannelTileFitsOneCore(geometry))
return failure();
FailureOr<ConvPlan> plan =
selectConvLoweringPlan(planOp.getOperation(), *state, /*reportPlanning=*/true);
if (failed(plan))
return failure();
switch (plan->strategy) {
case spatial::ConvLoweringStrategy::Legacy:
case spatial::ConvLoweringStrategy::Depthwise:
case spatial::ConvLoweringStrategy::PackedIm2Col:
case spatial::ConvLoweringStrategy::StreamedPatch:
case spatial::ConvLoweringStrategy::OutputChannelTiled:
case spatial::ConvLoweringStrategy::Tiled2D:
case spatial::ConvLoweringStrategy::StreamedPacked:
return success();
case spatial::ConvLoweringStrategy::Auto:
case spatial::ConvLoweringStrategy::InputKTiled:
return failure();
}
llvm_unreachable("unknown conv lowering strategy");
}
LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp,
const spatial::SpatialTargetInfo& target) {
FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp, target);
if (failed(state))
return failure();
FailureOr<ConvPlan> plan =
selectConvLoweringPlan(planOp.getOperation(), *state, /*reportPlanning=*/true);
if (failed(plan))
return failure();
if (plan->strategy == spatial::ConvLoweringStrategy::Depthwise)
return canConsumeDepthwiseRowStrip(*state) ? success() : failure();
StringRef failureReason;
return canConsumePixelMajorRowStripFragments(*state, failureReason) ? success() : failure();
}
FailureOr<Value>
lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
std::optional<Value> rowStripInput,
bool emitRowStripLayout,
const spatial::SpatialTargetInfo& target,
PatternRewriter& rewriter) {
FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp, target);
if (failed(state))
return failure();
FailureOr<ConvPlan> plan =
selectConvLoweringPlan(planOp.getOperation(), *state, /*reportPlanning=*/false);
if (failed(plan))
return failure();
auto reportRealization = [&](StringRef implementation) {
recordConvLoweringReport(
planOp.getOperation(), ConvLoweringReportPhase::Realization, plan->strategy, implementation);
};
if (emitRowStripLayout) {
if (rowStripInput) {
if (failed(canConsumeAndProduceRowStrip(planOp, target)))
return planOp.emitOpError("selected row-strip input/output layout is not supported for this Conv plan"), failure();
reportRealization(convRowStripInputImplementation(*state, plan->strategy));
return createConvOutputFromRowStripInput(
*state, *rowStripInput, plan->strategy, rewriter, planOp.getLoc());
}
if (failed(canLowerConvPlanToRowStrip(planOp, target)))
return planOp.emitOpError("selected row-strip layout is not supported for this Conv plan"), failure();
reportRealization("RSD");
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();
return *rowStripStorage;
}
reportRealization(convLoweringImplementation(plan->strategy));
if (plan->strategy == spatial::ConvLoweringStrategy::Depthwise)
return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, plan->strategy, rewriter, planOp.getLoc());
if (state->group != 1)
return lowerGroupedSelectedConvPlan(planOp.getOperation(), *state, plan->strategy, rewriter, planOp.getLoc());
return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, plan->strategy, rewriter, planOp.getLoc());
}
} // namespace onnx_mlir
Reference in New Issue View Git Blame Copy Permalink