Skip to content

Commit eea0458

Browse files
jorickertjorickert
committed
Add lowering of gridsample to linalg
1 parent ca7817b commit eea0458

4 files changed

Lines changed: 439 additions & 2 deletions

File tree

‎lib/Conversion/XTenNNToLinalg.cpp‎

Lines changed: 265 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -565,6 +565,268 @@ class ResizeToLinalg : public OpConversionPattern<ResizeOp> {
565565
}
566566
};
567567

568+
// Lower grid_sample to an output-indexed linalg.generic. Each output element
569+
// extracts the normalized x/y grid coordinate, converts it to input-space pixel
570+
// coordinates, samples the input with the requested interpolation rule, and
571+
// applies zeros or border padding while materializing the sample.
572+
class GridSampleToLinalg : public OpConversionPattern<GridSampleOp> {
573+
public:
574+
using OpConversionPattern<GridSampleOp>::OpConversionPattern;
575+
576+
struct GridSampleParams {
577+
Value input;
578+
Value grid;
579+
RankedTensorType inputTy;
580+
RankedTensorType gridTy;
581+
RankedTensorType resultTy;
582+
uint64_t alignCorners;
583+
uint64_t mode;
584+
uint64_t paddingMode;
585+
};
586+
587+
static Value gridCoordinateToInput(const GridSampleParams &params,
588+
OpBuilder &b, Location loc,
589+
Value gridValue, int64_t dim,
590+
Type calcType) {
591+
gridValue = convertValue(b, loc, gridValue, calcType);
592+
const double inputSize = params.inputTy.getDimSize(dim);
593+
const Value one = createFloatConstant(b, loc, calcType, 1.0);
594+
const Value half = createFloatConstant(b, loc, calcType, 0.5);
595+
const Value shifted = b.create<arith::AddFOp>(loc, gridValue, one);
596+
597+
if (params.alignCorners) {
598+
const Value sizeMinusOne =
599+
createFloatConstant(b, loc, calcType, inputSize - 1.0);
600+
const Value scaled = b.create<arith::MulFOp>(loc, shifted, sizeMinusOne);
601+
return b.create<arith::MulFOp>(loc, scaled, half);
602+
}
603+
604+
const Value size = createFloatConstant(b, loc, calcType, inputSize);
605+
const Value scaled = b.create<arith::MulFOp>(loc, shifted, size);
606+
const Value shiftedBack = b.create<arith::SubFOp>(loc, scaled, one);
607+
return b.create<arith::MulFOp>(loc, shiftedBack, half);
608+
}
609+
610+
static Value isIndexInBounds(OpBuilder &b, Location loc, Value index,
611+
int64_t size) {
612+
const Value zero = createIndexConstant(b, loc, 0);
613+
const Value upper = createIndexConstant(b, loc, size - 1);
614+
const Value aboveLower =
615+
b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sge, index, zero);
616+
const Value belowUpper =
617+
b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sle, index, upper);
618+
return b.create<arith::AndIOp>(loc, aboveLower, belowUpper);
619+
}
620+
621+
static Value sampleInput(const GridSampleParams &params, OpBuilder &b,
622+
Location loc, Value n, Value c, Value h, Value w,
623+
Type resultElementTy) {
624+
Value sampleH = h;
625+
Value sampleW = w;
626+
Value inBounds;
627+
if (params.paddingMode == 0) {
628+
Value hInBounds =
629+
isIndexInBounds(b, loc, h, params.inputTy.getDimSize(2));
630+
Value wInBounds =
631+
isIndexInBounds(b, loc, w, params.inputTy.getDimSize(3));
632+
inBounds = b.create<arith::AndIOp>(loc, hInBounds, wInBounds);
633+
sampleH = clampIndex(b, loc, h, params.inputTy.getDimSize(2) - 1);
634+
sampleW = clampIndex(b, loc, w, params.inputTy.getDimSize(3) - 1);
635+
} else {
636+
sampleH = clampIndex(b, loc, h, params.inputTy.getDimSize(2) - 1);
637+
sampleW = clampIndex(b, loc, w, params.inputTy.getDimSize(3) - 1);
638+
}
639+
640+
Value sample =
641+
b.create<tensor::ExtractOp>(loc, params.input,
642+
ValueRange{n, c, sampleH, sampleW});
643+
if (params.paddingMode == 0) {
644+
const Value zero = createFloatConstant(b, loc, resultElementTy, 0.0);
645+
sample = b.create<arith::SelectOp>(loc, inBounds, sample, zero);
646+
}
647+
return sample;
648+
}
649+
650+
static Value getNearestIndex(OpBuilder &b, Location loc, Value coord,
651+
Type calcType) {
652+
const Value floorCoord = b.create<math::FloorOp>(loc, coord);
653+
const Value floorIndex = castFloatToIndex(b, loc, floorCoord);
654+
const Value one = createIndexConstant(b, loc, 1);
655+
const Value upperIndex = b.create<arith::AddIOp>(loc, floorIndex, one);
656+
657+
const Value fraction = b.create<arith::SubFOp>(loc, coord, floorCoord);
658+
const Value half = createFloatConstant(b, loc, calcType, 0.5);
659+
const Value takeUpper =
660+
b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OGT, fraction, half);
661+
const Value isTie =
662+
b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OEQ, fraction, half);
663+
664+
const Value floorInt =
665+
b.create<arith::IndexCastOp>(loc, b.getI64Type(), floorIndex);
666+
const Value two = b.create<arith::ConstantIntOp>(loc, 2, 64);
667+
const Value remainder = b.create<arith::RemSIOp>(loc, floorInt, two);
668+
const Value zero = b.create<arith::ConstantIntOp>(loc, 0, 64);
669+
const Value floorIsOdd =
670+
b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, remainder, zero);
671+
const Value tieTakesUpper =
672+
b.create<arith::AndIOp>(loc, isTie, floorIsOdd);
673+
const Value selectUpper =
674+
b.create<arith::OrIOp>(loc, takeUpper, tieTakesUpper);
675+
return b.create<arith::SelectOp>(loc, selectUpper, upperIndex, floorIndex);
676+
}
677+
678+
static Value buildNearest(const GridSampleParams &params, OpBuilder &b,
679+
Location loc, Type calcType,
680+
Type resultElementTy) {
681+
const Value n = b.create<linalg::IndexOp>(loc, 0);
682+
const Value c = b.create<linalg::IndexOp>(loc, 1);
683+
const Value outH = b.create<linalg::IndexOp>(loc, 2);
684+
const Value outW = b.create<linalg::IndexOp>(loc, 3);
685+
const Value zero = createIndexConstant(b, loc, 0);
686+
const Value one = createIndexConstant(b, loc, 1);
687+
const Value gridX =
688+
b.create<tensor::ExtractOp>(loc, params.grid,
689+
ValueRange{n, outH, outW, zero});
690+
const Value gridY =
691+
b.create<tensor::ExtractOp>(loc, params.grid,
692+
ValueRange{n, outH, outW, one});
693+
const Value inputX =
694+
gridCoordinateToInput(params, b, loc, gridX, 3, calcType);
695+
const Value inputY =
696+
gridCoordinateToInput(params, b, loc, gridY, 2, calcType);
697+
const Value x = getNearestIndex(b, loc, inputX, calcType);
698+
const Value y = getNearestIndex(b, loc, inputY, calcType);
699+
return sampleInput(params, b, loc, n, c, y, x, resultElementTy);
700+
}
701+
702+
static Value buildBilinear(const GridSampleParams &params, OpBuilder &b,
703+
Location loc, Type calcType,
704+
Type resultElementTy) {
705+
const Value n = b.create<linalg::IndexOp>(loc, 0);
706+
const Value c = b.create<linalg::IndexOp>(loc, 1);
707+
const Value outH = b.create<linalg::IndexOp>(loc, 2);
708+
const Value outW = b.create<linalg::IndexOp>(loc, 3);
709+
const Value zeroIndex = createIndexConstant(b, loc, 0);
710+
const Value oneIndex = createIndexConstant(b, loc, 1);
711+
const Value gridX =
712+
b.create<tensor::ExtractOp>(loc, params.grid,
713+
ValueRange{n, outH, outW, zeroIndex});
714+
const Value gridY =
715+
b.create<tensor::ExtractOp>(loc, params.grid,
716+
ValueRange{n, outH, outW, oneIndex});
717+
const Value inputX =
718+
gridCoordinateToInput(params, b, loc, gridX, 3, calcType);
719+
const Value inputY =
720+
gridCoordinateToInput(params, b, loc, gridY, 2, calcType);
721+
722+
const Value x0Float = b.create<math::FloorOp>(loc, inputX);
723+
const Value y0Float = b.create<math::FloorOp>(loc, inputY);
724+
const Value x0 = castFloatToIndex(b, loc, x0Float);
725+
const Value y0 = castFloatToIndex(b, loc, y0Float);
726+
const Value x1 = b.create<arith::AddIOp>(loc, x0, oneIndex);
727+
const Value y1 = b.create<arith::AddIOp>(loc, y0, oneIndex);
728+
729+
const Value xLerp = b.create<arith::SubFOp>(loc, inputX, x0Float);
730+
const Value yLerp = b.create<arith::SubFOp>(loc, inputY, y0Float);
731+
const Value oneFloat = createFloatConstant(b, loc, calcType, 1.0);
732+
const Value x0Weight = b.create<arith::SubFOp>(loc, oneFloat, xLerp);
733+
const Value y0Weight = b.create<arith::SubFOp>(loc, oneFloat, yLerp);
734+
735+
auto weightedSample = [&](Value h, Value w, Value hWeight,
736+
Value wWeight) -> Value {
737+
Value sample = sampleInput(params, b, loc, n, c, h, w, resultElementTy);
738+
sample = convertValue(b, loc, sample, calcType);
739+
const Value weight = b.create<arith::MulFOp>(loc, hWeight, wWeight);
740+
return b.create<arith::MulFOp>(loc, sample, weight);
741+
};
742+
743+
Value acc = weightedSample(y0, x0, y0Weight, x0Weight);
744+
acc = b.create<arith::AddFOp>(
745+
loc, acc, weightedSample(y0, x1, y0Weight, xLerp));
746+
acc = b.create<arith::AddFOp>(
747+
loc, acc, weightedSample(y1, x0, yLerp, x0Weight));
748+
acc = b.create<arith::AddFOp>(
749+
loc, acc, weightedSample(y1, x1, yLerp, xLerp));
750+
return convertValue(b, loc, acc, resultElementTy);
751+
}
752+
753+
LogicalResult
754+
matchAndRewrite(GridSampleOp op, OpAdaptor adaptor,
755+
ConversionPatternRewriter &rewriter) const override {
756+
const Location loc = op->getLoc();
757+
const auto inputTy = cast<RankedTensorType>(adaptor.getX().getType());
758+
const auto gridTy = cast<RankedTensorType>(adaptor.getGrid().getType());
759+
const auto resultTy = cast<RankedTensorType>(op->getResult(0).getType());
760+
const Type resultElementTy = resultTy.getElementType();
761+
762+
if (!inputTy.hasStaticShape() || !gridTy.hasStaticShape() ||
763+
!resultTy.hasStaticShape())
764+
return rewriter.notifyMatchFailure(
765+
op, "grid_sample lowering requires static shapes");
766+
if (inputTy.getRank() != 4 || gridTy.getRank() != 4 ||
767+
resultTy.getRank() != 4)
768+
return rewriter.notifyMatchFailure(
769+
op, "grid_sample lowering supports rank-4 tensors");
770+
if (!isa<FloatType>(inputTy.getElementType()) ||
771+
!isa<FloatType>(gridTy.getElementType()) ||
772+
!isa<FloatType>(resultElementTy))
773+
return rewriter.notifyMatchFailure(
774+
op, "grid_sample lowering only supports floating point tensors");
775+
// align_corners encoding: 0=false, 1=true.
776+
if (adaptor.getAlignCorners() > 1)
777+
return rewriter.notifyMatchFailure(
778+
op, "grid_sample align_corners must be 0 or 1");
779+
780+
// mode encoding: 0=bilinear, 1=nearest, 2=cubic.
781+
if (adaptor.getMode() > 2)
782+
return rewriter.notifyMatchFailure(op, "invalid grid_sample mode");
783+
if (adaptor.getMode() == 2)
784+
return rewriter.notifyMatchFailure(
785+
op, "grid_sample cubic mode is not supported");
786+
787+
// padding_mode encoding: 0=zeros, 1=border, 2=reflection.
788+
if (adaptor.getPaddingMode() > 2)
789+
return rewriter.notifyMatchFailure(
790+
op, "invalid grid_sample padding mode");
791+
if (adaptor.getPaddingMode() == 2)
792+
return rewriter.notifyMatchFailure(
793+
op, "grid_sample reflection padding is not supported");
794+
795+
Type calcType = resultElementTy;
796+
if (auto floatTy = dyn_cast<FloatType>(calcType)) {
797+
if (floatTy.getWidth() < 32)
798+
calcType = rewriter.getF32Type();
799+
}
800+
801+
const GridSampleParams params{adaptor.getX(),
802+
adaptor.getGrid(),
803+
inputTy,
804+
gridTy,
805+
resultTy,
806+
adaptor.getAlignCorners(),
807+
adaptor.getMode(),
808+
adaptor.getPaddingMode()};
809+
const Value output = getEmptyTensor(rewriter, loc, resultTy, {});
810+
const AffineMap outputMap =
811+
rewriter.getMultiDimIdentityMap(resultTy.getRank());
812+
auto generic = rewriter.create<linalg::GenericOp>(
813+
loc, TypeRange{resultTy}, ValueRange{}, output,
814+
SmallVector<AffineMap>{outputMap},
815+
mlir::tosa::getNParallelLoopsAttrs(resultTy.getRank()),
816+
[&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange) {
817+
Value result = params.mode == 0
818+
? buildBilinear(params, nestedBuilder, nestedLoc,
819+
calcType, resultElementTy)
820+
: buildNearest(params, nestedBuilder, nestedLoc,
821+
calcType, resultElementTy);
822+
nestedBuilder.create<linalg::YieldOp>(nestedLoc, result);
823+
});
824+
825+
rewriter.replaceOp(op, generic.getResults());
826+
return success();
827+
}
828+
};
829+
568830
struct ConvertXtenNNtoLinalg
569831
: public xilinx::xten::impl::ConvertXTenNNToLinalgBase<
570832
ConvertXtenNNtoLinalg> {
@@ -582,14 +844,15 @@ struct ConvertXtenNNtoLinalg
582844
auto funcOp = getOperation();
583845

584846
ConversionTarget target(*context);
585-
target.addIllegalOp<EluOp, ReduceMeanOp, ResizeOp, SignOp>();
847+
target.addIllegalOp<EluOp, GridSampleOp, ReduceMeanOp, ResizeOp, SignOp>();
586848
target.addLegalDialect<linalg::LinalgDialect, scf::SCFDialect,
587849
complex::ComplexDialect, math::MathDialect,
588850
shape::ShapeDialect, tensor::TensorDialect,
589851
arith::ArithDialect>();
590852

591853
RewritePatternSet patterns(context);
592-
patterns.add<EluToLinalg, ReduceMeanToLinalg, ResizeToLinalg>(context);
854+
patterns.add<EluToLinalg, GridSampleToLinalg, ReduceMeanToLinalg,
855+
ResizeToLinalg>(context);
593856

594857
if (failed(applyPartialConversion(funcOp, target, std::move(patterns))))
595858
signalPassFailure();

0 commit comments

Comments
 (0)