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-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #include "mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h" 10 #include "../PassDetail.h" 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/Pass/PassRegistry.h" 16 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" 17 18 using namespace mlir; 19 using namespace mlir::arm_neon; 20 21 namespace { 22 23 class Sdot2dLoweringPattern : public OpRewritePattern<Sdot2dOp> { 24 public: 25 using OpRewritePattern::OpRewritePattern; 26 27 /// Convert to 1-dimensional vector type to match the requirements of 28 /// arm.neon.intr.sdot 29 LogicalResult matchAndRewrite(Sdot2dOp op, 30 PatternRewriter &rewriter) const override { 31 Type elemType = op.b().getType().cast<VectorType>().getElementType(); 32 int length = op.b().getType().cast<VectorType>().getShape()[0] * 33 Sdot2dOp::kReductionSize; 34 VectorType flattenedVectorType = VectorType::get({length}, elemType); 35 Value b2d = op.b(); 36 Value c2d = op.c(); 37 Location loc = op.getLoc(); 38 Value b1d = 39 rewriter.create<vector::ShapeCastOp>(loc, flattenedVectorType, b2d); 40 Value c1d = 41 rewriter.create<vector::ShapeCastOp>(loc, flattenedVectorType, c2d); 42 Value newOp = 43 rewriter.create<SdotOp>(loc, op.res().getType(), op.a(), b1d, c1d); 44 rewriter.replaceOp(op, {newOp}); 45 return success(); 46 } 47 }; 48 49 class ConvertArmNeon2dToIntr 50 : public ConvertArmNeon2dToIntrBase<ConvertArmNeon2dToIntr> { 51 void runOnOperation() override { 52 auto func = getOperation(); 53 auto *context = &getContext(); 54 55 RewritePatternSet patterns(context); 56 populateConvertArmNeon2dToIntrPatterns(patterns); 57 58 if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) 59 return signalPassFailure(); 60 } 61 }; 62 63 } // namespace 64 65 namespace mlir { 66 67 void populateConvertArmNeon2dToIntrPatterns(RewritePatternSet &patterns) { 68 patterns.add<Sdot2dLoweringPattern>(patterns.getContext()); 69 } 70 71 std::unique_ptr<OperationPass<FuncOp>> createConvertArmNeon2dToIntrPass() { 72 return std::make_unique<ConvertArmNeon2dToIntr>(); 73 } 74 75 } // namespace mlir 76