71 lines · cpp
1//===- ArmNeon2dToIntr.cpp - convert Arm Neon 2d ops to intrinsics --------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h"10 11#include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"12#include "mlir/Dialect/Vector/IR/VectorOps.h"13#include "mlir/IR/PatternMatch.h"14#include "mlir/Pass/Pass.h"15#include "mlir/Transforms/GreedyPatternRewriteDriver.h"16 17namespace mlir {18#define GEN_PASS_DEF_CONVERTARMNEON2DTOINTRPASS19#include "mlir/Conversion/Passes.h.inc"20} // namespace mlir21 22using namespace mlir;23using namespace mlir::arm_neon;24 25namespace {26 27class Sdot2dLoweringPattern : public OpRewritePattern<Sdot2dOp> {28public:29 using OpRewritePattern::OpRewritePattern;30 31 /// Convert to 1-dimensional vector type to match the requirements of32 /// arm.neon.intr.sdot33 LogicalResult matchAndRewrite(Sdot2dOp op,34 PatternRewriter &rewriter) const override {35 Type elemType = cast<VectorType>(op.getB().getType()).getElementType();36 int length = cast<VectorType>(op.getB().getType()).getShape()[0] *37 Sdot2dOp::kReductionSize;38 VectorType flattenedVectorType = VectorType::get({length}, elemType);39 Value b2d = op.getB();40 Value c2d = op.getC();41 Location loc = op.getLoc();42 Value b1d =43 vector::ShapeCastOp::create(rewriter, loc, flattenedVectorType, b2d);44 Value c1d =45 vector::ShapeCastOp::create(rewriter, loc, flattenedVectorType, c2d);46 Value newOp = SdotOp::create(rewriter, loc, op.getRes().getType(),47 op.getA(), b1d, c1d);48 rewriter.replaceOp(op, {newOp});49 return success();50 }51};52 53class ConvertArmNeon2dToIntr54 : public impl::ConvertArmNeon2dToIntrPassBase<ConvertArmNeon2dToIntr> {55 void runOnOperation() override {56 auto *context = &getContext();57 58 RewritePatternSet patterns(context);59 populateConvertArmNeon2dToIntrPatterns(patterns);60 61 if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))62 return signalPassFailure();63 }64};65 66} // namespace67 68void mlir::populateConvertArmNeon2dToIntrPatterns(RewritePatternSet &patterns) {69 patterns.add<Sdot2dLoweringPattern>(patterns.getContext());70}71