Commit 20daedac authored by Benoit Jacob's avatar Benoit Jacob Committed by Ahmed Taei
Browse files

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
parent 4f01122c
//===- 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 <typename T>
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<OperationPass<FuncOp>> createConvertArmNeon2dToIntrPass();
} // namespace mlir
#endif // MLIR_CONVERSION_ARMNEON2DTOINTR_ARMNEON2DTOINTR_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"
......
......@@ -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
......@@ -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<string mnemonic, list<OpTrait> traits = []>
: Op</*dialect=*/ArmNeon_Dialect,
/*opName=*/"2d." # mnemonic,
/*traits=*/traits>;
def Sdot2dOp : ArmNeon_2dOp<"sdot", [
NoSideEffect,
AllTypesMatch<["b", "c"]>,
AllTypesMatch<["a", "res"]>,
PredOpTrait<
"operand `a` should be 1-dimensional",
CPred<"a().getType().cast<VectorType>().getShape().size() == 1">
>,
PredOpTrait<
"operand `b` should be 2-dimensional",
CPred<"b().getType().cast<VectorType>().getShape().size() == 2">
>,
PredOpTrait<
"operand `b` should have 4 columns",
CPred<"b().getType().cast<VectorType>().getShape()[1] == 4">
>,
PredOpTrait<
"operand `b` should have as many rows as the size of operand `a`",
CPred<"b().getType().cast<VectorType>().getShape()[0] == a().getType().cast<VectorType>().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
//===- 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<Sdot2dOp> {
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<VectorType>().getElementType();
int length = op.b().getType().cast<VectorType>().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<vector::ShapeCastOp>(loc, flattenedVectorType, b2d);
Value c1d =
rewriter.create<vector::ShapeCastOp>(loc, flattenedVectorType, c2d);
Value newOp =
rewriter.create<SdotOp>(loc, op.res().getType(), op.a(), b1d, c1d);
rewriter.replaceOp(op, {newOp});
return success();
}
};
class ConvertArmNeon2dToIntr
: public ConvertArmNeon2dToIntrBase<ConvertArmNeon2dToIntr> {
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<Sdot2dLoweringPattern>(patterns.getContext());
}
std::unique_ptr<OperationPass<FuncOp>> createConvertArmNeon2dToIntrPass() {
return std::make_unique<ConvertArmNeon2dToIntr>();
}
} // namespace mlir
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
)
add_subdirectory(AffineToStandard)
add_subdirectory(ArmNeon2dToIntr)
add_subdirectory(AsyncToLLVM)
add_subdirectory(ComplexToLLVM)
add_subdirectory(ComplexToStandard)
......
......@@ -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"
......
// 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>
}
// 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>
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment