From 20daedacca803b81db6d8773b705345702bf0fc3 Mon Sep 17 00:00:00 2001 From: Benoit Jacob Date: Thu, 10 Jun 2021 13:57:44 -0700 Subject: [PATCH] 2d Arm Neon sdot op, and lowering to the intrinsic. This adds Sdot2d op, which is similar to the usual Neon intrinsic except that it takes 2d vector operands, reflecting the structure of the arithmetic that it's performing: 4 separate 4-dimensional dot products, whence the vector<4x4xi8> shape. This also adds a new pass, arm-neon-2d-to-intr, lowering this new 2d op to the 1d intrinsic. Reviewed By: nicolasvasilache Differential Revision: https://reviews.llvm.org/D102504 --- .../Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h | 30 +++++++++ mlir/include/mlir/Conversion/Passes.h | 1 + mlir/include/mlir/Conversion/Passes.td | 11 ++++ mlir/include/mlir/Dialect/ArmNeon/ArmNeon.td | 55 ++++++++++++++++ .../Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp | 75 ++++++++++++++++++++++ mlir/lib/Conversion/ArmNeon2dToIntr/CMakeLists.txt | 18 ++++++ mlir/lib/Conversion/CMakeLists.txt | 1 + mlir/lib/Conversion/PassDetail.h | 4 ++ mlir/test/Dialect/ArmNeon/invalid.mlir | 33 ++++++++++ mlir/test/Target/LLVMIR/arm-neon-2d.mlir | 17 +++++ 10 files changed, 245 insertions(+) create mode 100644 mlir/include/mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h create mode 100644 mlir/lib/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp create mode 100644 mlir/lib/Conversion/ArmNeon2dToIntr/CMakeLists.txt create mode 100644 mlir/test/Dialect/ArmNeon/invalid.mlir create mode 100644 mlir/test/Target/LLVMIR/arm-neon-2d.mlir diff --git a/mlir/include/mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h b/mlir/include/mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h new file mode 100644 index 0000000..a27e91f --- /dev/null +++ b/mlir/include/mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h @@ -0,0 +1,30 @@ +//===- ArmNeon2dToIntr.h - convert Arm Neon 2d ops to intrinsics ----------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_CONVERSION_ARMNEON2DTOINTR_ARMNEON2DTOINTR_H_ +#define MLIR_CONVERSION_ARMNEON2DTOINTR_ARMNEON2DTOINTR_H_ + +#include "mlir/Pass/Pass.h" + +namespace mlir { +class FuncOp; +template +class OperationPass; + +/// Populates patterns for the lowering of Arm NEON 2D ops to intrinsics. +/// See createConvertArmNeon2dToIntrPass. +void populateConvertArmNeon2dToIntrPatterns(RewritePatternSet &patterns); + +/// Creates a pass to lower Arm NEON 2D ops to intrinsics, i.e. +/// equivalent ops operating on flattened 1D vectors and mapping more +/// directly to the corresponding Arm NEON instruction. +std::unique_ptr> createConvertArmNeon2dToIntrPass(); + +} // namespace mlir + +#endif // MLIR_CONVERSION_ARMNEON2DTOINTR_ARMNEON2DTOINTR_H_ diff --git a/mlir/include/mlir/Conversion/Passes.h b/mlir/include/mlir/Conversion/Passes.h index d75e617..a78b728 100644 --- a/mlir/include/mlir/Conversion/Passes.h +++ b/mlir/include/mlir/Conversion/Passes.h @@ -10,6 +10,7 @@ #define MLIR_CONVERSION_PASSES_H #include "mlir/Conversion/AffineToStandard/AffineToStandard.h" +#include "mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h" #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h" #include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h" #include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h" diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td index 508a008..ba5e27a 100644 --- a/mlir/include/mlir/Conversion/Passes.td +++ b/mlir/include/mlir/Conversion/Passes.td @@ -607,4 +607,15 @@ def ConvertVectorToSPIRV : Pass<"convert-vector-to-spirv", "ModuleOp"> { let dependentDialects = ["spirv::SPIRVDialect"]; } +//===----------------------------------------------------------------------===// +// ArmNeon2dToIntr +//===----------------------------------------------------------------------===// + +def ConvertArmNeon2dToIntr : Pass<"arm-neon-2d-to-intr", "FuncOp"> { + let summary = "Convert Arm NEON structured ops to intrinsics"; + let constructor = "mlir::createConvertArmNeon2dToIntrPass()"; + let dependentDialects = ["arm_neon::ArmNeonDialect", "vector::VectorDialect"]; +} + + #endif // MLIR_CONVERSION_PASSES diff --git a/mlir/include/mlir/Dialect/ArmNeon/ArmNeon.td b/mlir/include/mlir/Dialect/ArmNeon/ArmNeon.td index a9d9f6f..b530c62 100644 --- a/mlir/include/mlir/Dialect/ArmNeon/ArmNeon.td +++ b/mlir/include/mlir/Dialect/ArmNeon/ArmNeon.td @@ -15,6 +15,7 @@ include "mlir/Dialect/LLVMIR/LLVMOpBase.td" include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/IR/OpBase.td" //===----------------------------------------------------------------------===// // ArmNeon dialect definition @@ -117,4 +118,58 @@ def SdotOp : ArmNeon_OverloadedOperandsWithOneResultIntrOp<"sdot",[1], [ "$a `,` $b `,` $c attr-dict `:` type($b) `,` type($c) `to` type($res)"; } +class ArmNeon_2dOp traits = []> + : Op; + +def Sdot2dOp : ArmNeon_2dOp<"sdot", [ + NoSideEffect, + AllTypesMatch<["b", "c"]>, + AllTypesMatch<["a", "res"]>, + PredOpTrait< + "operand `a` should be 1-dimensional", + CPred<"a().getType().cast().getShape().size() == 1"> + >, + PredOpTrait< + "operand `b` should be 2-dimensional", + CPred<"b().getType().cast().getShape().size() == 2"> + >, + PredOpTrait< + "operand `b` should have 4 columns", + CPred<"b().getType().cast().getShape()[1] == 4"> + >, + PredOpTrait< + "operand `b` should have as many rows as the size of operand `a`", + CPred<"b().getType().cast().getShape()[0] == a().getType().cast().getShape()[0]"> + >, + ] + > { + let summary = "sdot op"; + let description = [{ + The two input vectors `b` and `c` have a 2D shape, consisting of either 2 + or 4 rows, each row having length 4. This operation computes the pair-wise + dot-products of the rows of `b` and `c` and accumulates them with the + corresponding entry of `a`: + + ``` + res[i] := a[i] + dot_product(b[i, ...], c[i, ...]) + ``` + + }]; + // Supports either: + // (vector<2xi32>, vector<2x4xi8>, vector<2x4xi8>) -> vector<2xi32> + // (vector<4xi32>, vector<4x4xi8>, vector<4x4xi8>) -> vector<4xi32> + // TODO: how do we express 2D shape requirements here? + let arguments = (ins VectorOfLengthAndType<[4, 2], [I32]>:$a, + VectorOfLengthAndType<[16, 8], [I8]>:$b, + VectorOfLengthAndType<[16, 8], [I8]>:$c); + let results = (outs VectorOfLengthAndType<[4, 2], [I32]>:$res); + let assemblyFormat = + "$a `,` $b `,` $c attr-dict `:` type($b) `,` type($c) `to` type($res)"; + let extraClassDeclaration = [{ + static constexpr int kReductionSize = 4; + }]; +} + #endif // ARMNEON_OPS diff --git a/mlir/lib/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp b/mlir/lib/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp new file mode 100644 index 0000000..6d7e7c1 --- /dev/null +++ b/mlir/lib/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp @@ -0,0 +1,75 @@ +//===- ArmNeon2dToIntr.cpp - convert Arm Neon 2d ops to intrinsics --------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h" +#include "../PassDetail.h" +#include "mlir/Dialect/ArmNeon/ArmNeonDialect.h" +#include "mlir/Dialect/Vector/VectorOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassRegistry.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +using namespace mlir; +using namespace mlir::arm_neon; + +namespace { + +class Sdot2dLoweringPattern : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + /// Convert to 1-dimensional vector type to match the requirements of + /// arm.neon.intr.sdot + LogicalResult matchAndRewrite(Sdot2dOp op, + PatternRewriter &rewriter) const override { + Type elemType = op.b().getType().cast().getElementType(); + int length = op.b().getType().cast().getShape()[0] * + Sdot2dOp::kReductionSize; + VectorType flattenedVectorType = VectorType::get({length}, elemType); + Value b2d = op.b(); + Value c2d = op.c(); + Location loc = op.getLoc(); + Value b1d = + rewriter.create(loc, flattenedVectorType, b2d); + Value c1d = + rewriter.create(loc, flattenedVectorType, c2d); + Value newOp = + rewriter.create(loc, op.res().getType(), op.a(), b1d, c1d); + rewriter.replaceOp(op, {newOp}); + return success(); + } +}; + +class ConvertArmNeon2dToIntr + : public ConvertArmNeon2dToIntrBase { + void runOnOperation() override { + auto func = getOperation(); + auto *context = &getContext(); + + RewritePatternSet patterns(context); + populateConvertArmNeon2dToIntrPatterns(patterns); + + if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) + return signalPassFailure(); + } +}; + +} // namespace + +namespace mlir { + +void populateConvertArmNeon2dToIntrPatterns(RewritePatternSet &patterns) { + patterns.add(patterns.getContext()); +} + +std::unique_ptr> createConvertArmNeon2dToIntrPass() { + return std::make_unique(); +} + +} // namespace mlir diff --git a/mlir/lib/Conversion/ArmNeon2dToIntr/CMakeLists.txt b/mlir/lib/Conversion/ArmNeon2dToIntr/CMakeLists.txt new file mode 100644 index 0000000..5c729c8 --- /dev/null +++ b/mlir/lib/Conversion/ArmNeon2dToIntr/CMakeLists.txt @@ -0,0 +1,18 @@ +add_mlir_conversion_library(MLIRArmNeon2dToIntr + ArmNeon2dToIntr.cpp + + ADDITIONAL_HEADER_DIRS + ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ArmNeon2dToIntr + + DEPENDS + MLIRConversionPassIncGen + + LINK_COMPONENTS + Core + + LINK_LIBS PUBLIC + MLIRArmNeon + MLIRPass + MLIRTransforms + MLIRIR + ) diff --git a/mlir/lib/Conversion/CMakeLists.txt b/mlir/lib/Conversion/CMakeLists.txt index b89be2a..72cfb08 100644 --- a/mlir/lib/Conversion/CMakeLists.txt +++ b/mlir/lib/Conversion/CMakeLists.txt @@ -1,4 +1,5 @@ add_subdirectory(AffineToStandard) +add_subdirectory(ArmNeon2dToIntr) add_subdirectory(AsyncToLLVM) add_subdirectory(ComplexToLLVM) add_subdirectory(ComplexToStandard) diff --git a/mlir/lib/Conversion/PassDetail.h b/mlir/lib/Conversion/PassDetail.h index 5287c6c..993cac8 100644 --- a/mlir/lib/Conversion/PassDetail.h +++ b/mlir/lib/Conversion/PassDetail.h @@ -80,6 +80,10 @@ namespace vector { class VectorDialect; } // end namespace vector +namespace arm_neon { +class ArmNeonDialect; +} // end namespace arm_neon + #define GEN_PASS_CLASSES #include "mlir/Conversion/Passes.h.inc" diff --git a/mlir/test/Dialect/ArmNeon/invalid.mlir b/mlir/test/Dialect/ArmNeon/invalid.mlir new file mode 100644 index 0000000..460cdc0 --- /dev/null +++ b/mlir/test/Dialect/ArmNeon/invalid.mlir @@ -0,0 +1,33 @@ +// RUN: mlir-opt %s -split-input-file -verify-diagnostics + +// ----- + +func @a_is_2d(%a : vector<2x2xi32>, %b : vector<4x4xi8>) -> vector<2x2xi32> { + // expected-error@+1 {{operand `a` should be 1-dimensional}} + %0 = arm_neon.2d.sdot %a, %b, %b : vector<4x4xi8>, vector<4x4xi8> to vector<2x2xi32> + return %0 : vector<2x2xi32> +} + +// ----- + +func @b_is_3d(%a : vector<4xi32>, %b : vector<1x4x4xi8>) -> vector<4xi32> { + // expected-error@+1 {{operand `b` should be 2-dimensional}} + %0 = arm_neon.2d.sdot %a, %b, %b : vector<1x4x4xi8>, vector<1x4x4xi8> to vector<4xi32> + return %0 : vector<4xi32> +} + +// ----- + +func @b_has_2_columns(%a : vector<4xi32>, %b : vector<4x2xi8>) -> vector<4xi32> { + // expected-error@+1 {{operand `b` should have 4 columns}} + %0 = arm_neon.2d.sdot %a, %b, %b : vector<4x2xi8>, vector<4x2xi8> to vector<4xi32> + return %0 : vector<4xi32> +} + +// ----- + +func @b_has_2_rows_but_a_has_length_4(%a : vector<4xi32>, %b : vector<2x4xi8>) -> vector<4xi32> { + // expected-error@+1 {{operand `b` should have as many rows as the size of operand `a`}} + %0 = arm_neon.2d.sdot %a, %b, %b : vector<2x4xi8>, vector<2x4xi8> to vector<4xi32> + return %0 : vector<4xi32> +} diff --git a/mlir/test/Target/LLVMIR/arm-neon-2d.mlir b/mlir/test/Target/LLVMIR/arm-neon-2d.mlir new file mode 100644 index 0000000..b75afdc --- /dev/null +++ b/mlir/test/Target/LLVMIR/arm-neon-2d.mlir @@ -0,0 +1,17 @@ +// RUN: mlir-opt -arm-neon-2d-to-intr %s | FileCheck %s + +// CHECK-LABEL: arm_neon_sdot2d_4x4_i8i8 +func @arm_neon_sdot2d_4x4_i8i8(%a: vector<4xi32>, %b: vector<4x4xi8>, %c: vector<4x4xi8>) -> vector<4xi32> { + // CHECK: arm_neon.intr.sdot %{{.*}}, %{{.*}}, %{{.*}} : vector<16xi8>, vector<16xi8> to vector<4xi32> + // CHECK-NEXT: return %{{.*}} : vector<4xi32> + %0 = arm_neon.2d.sdot %a, %b, %c : vector<4x4xi8>, vector<4x4xi8> to vector<4xi32> + return %0 : vector<4xi32> +} + +// CHECK-LABEL: arm_neon_sdot2d_2x4_i8i8 +func @arm_neon_sdot2d_2x4_i8i8(%a: vector<2xi32>, %b: vector<2x4xi8>, %c: vector<2x4xi8>) -> vector<2xi32> { + // CHECK: arm_neon.intr.sdot %{{.*}}, %{{.*}}, %{{.*}} : vector<8xi8>, vector<8xi8> to vector<2xi32> + // CHECK-NEXT: return %{{.*}} : vector<2xi32> + %0 = arm_neon.2d.sdot %a, %b, %c : vector<2x4xi8>, vector<2x4xi8> to vector<2xi32> + return %0 : vector<2xi32> +} -- 2.7.4