308 lines
13 KiB
TableGen
308 lines
13 KiB
TableGen
//===-- ArmNeonOps.td - ArmNeon dialect op definitions -----*- tablegen -*-===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file defines the basic operations for the ArmNeon dialect.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef ARMNEON_OPS
|
|
#define ARMNEON_OPS
|
|
|
|
include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
|
|
include "mlir/Interfaces/SideEffectInterfaces.td"
|
|
include "mlir/IR/OpBase.td"
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// ArmNeon dialect definition
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
def ArmNeon_Dialect : Dialect {
|
|
let name = "arm_neon";
|
|
let cppNamespace = "::mlir::arm_neon";
|
|
|
|
// Note: this does not need to depend on LLVMDialect as long as functions in
|
|
// this dialect (such as canonicalization) do not produce entities belonging
|
|
// to the LLVMDialect (ops or types).
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// ArmNeon type definition
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
class NeonVectorOfLength<int length, Type elementType> : ShapedContainerType<
|
|
[elementType], And<[IsVectorOfShape<[length]>, IsFixedVectorOfAnyRankTypePred]>,
|
|
"a vector with length " # length,
|
|
"::mlir::VectorType">;
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// ArmNeon op definitions
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
// ArmNeon dialect op that corresponds (and is convertible to) an LLVM IR
|
|
// intrinsic.
|
|
class ArmNeon_IntrOp<string mnemonic, list<int> overloadedResults,
|
|
list<int> overloadedOperands, int numResults,
|
|
list<Trait> traits = [], bit requiresAccessGroup = 0,
|
|
bit requiresAliasAnalysis = 0>
|
|
: LLVM_IntrOpBase</*dialect=*/ArmNeon_Dialect,
|
|
/*opName=*/"intr." # mnemonic,
|
|
/*enumName=*/"aarch64_neon_" # !subst(".", "_", mnemonic),
|
|
/*overloadedResults=*/overloadedResults,
|
|
/*overloadedOperands=*/overloadedOperands,
|
|
/*traits=*/traits,
|
|
/*numResults=*/numResults,
|
|
/*requiresAccessGroup=*/requiresAccessGroup,
|
|
/*requiresAliasAnalysis=*/requiresAliasAnalysis>;
|
|
|
|
// ArmNeon dialect op that corresponds to an LLVM IR intrinsic with one
|
|
// overloaded result.
|
|
class ArmNeon_OverloadedOneResultIntrOp<string mnemonic,
|
|
list<Trait> traits = []>
|
|
: ArmNeon_IntrOp<mnemonic, [0], [], 1, traits>;
|
|
|
|
// ArmNeon dialect op that corresponds to an LLVM IR intrinsic with one
|
|
// overloaded result and overloaded operands list.
|
|
class ArmNeon_OverloadedOperandsWithOneResultIntrOp<string mnemonic,
|
|
list<int> overloadedOperands,
|
|
list<Trait> traits = []>
|
|
: ArmNeon_IntrOp<mnemonic, [0], overloadedOperands, 1, traits>;
|
|
|
|
def SMullOp : ArmNeon_OverloadedOneResultIntrOp<"smull", [
|
|
Pure,
|
|
AllTypesMatch<["a", "b"]>,
|
|
TypesMatchWith<
|
|
"res has same vector shape and element bitwidth scaled by 2 as a",
|
|
"a", "res", "::llvm::cast<VectorType>($_self).scaleElementBitwidth(2)">
|
|
]> {
|
|
let summary = "smull roundscale op";
|
|
let description = [{
|
|
Signed Multiply Long (vector). This instruction multiplies corresponding
|
|
signed integer values in the lower or upper half of the vectors of the two
|
|
source SIMD&FP registers, places the results in a vector, and writes the
|
|
vector to the destination SIMD&FP register.
|
|
|
|
Source:
|
|
https://developer.arm.com/architectures/instruction-sets/simd-isas/neon/intrinsics
|
|
}];
|
|
|
|
// Supports either:
|
|
// (vector<8xi8>, vector<8xi8>) -> (vector<8xi16>)
|
|
// (vector<4xi16>, vector<4xi16>) -> (vector<4xi32>)
|
|
// (vector<2xi32>, vector<2xi32>) -> (vector<2xi64>)
|
|
let arguments = (ins VectorOfLengthAndType<[8, 4, 2], [I8, I16, I32]>:$a,
|
|
VectorOfLengthAndType<[8, 4, 2], [I8, I16, I32]>:$b);
|
|
let results = (outs VectorOfLengthAndType<[8, 4, 2], [I16, I32, I64]>:$res);
|
|
let assemblyFormat =
|
|
"$a `,` $b attr-dict `:` type($a) `to` type($res)";
|
|
}
|
|
|
|
def SdotOp : ArmNeon_OverloadedOperandsWithOneResultIntrOp<"sdot",[1], [
|
|
Pure,
|
|
AllTypesMatch<["b", "c"]>,
|
|
AllTypesMatch<["a", "res"]>,
|
|
TypesMatchWith<"res has the same number of elements as operand b",
|
|
"b", "res",
|
|
"VectorType::get({::llvm::cast<VectorType>($_self).getShape()[0] / 4},"
|
|
"IntegerType::get($_self.getContext(), 32))">]> {
|
|
let summary = "sdot op";
|
|
let description = [{
|
|
Signed integer addition of dot product (vector). This instruction performs
|
|
the following operation on signed integer vectors: res = dot(b, c) + a,
|
|
where vector operands are partitioned into groups of four elements.
|
|
|
|
Source:
|
|
https://developer.arm.com/architectures/instruction-sets/simd-isas/neon/intrinsics
|
|
}];
|
|
// Supports either:
|
|
// (vector<2xi32>, vector<8xi8>, vector<8xi8>) -> vector<2xi32>
|
|
// (vector<4xi32>, vector<16xi8>, vector<16xi8>) -> vector<4xi32>
|
|
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)";
|
|
}
|
|
|
|
def SmmlaOp : ArmNeon_OverloadedOperandsWithOneResultIntrOp<"smmla",[1], [
|
|
Pure,
|
|
AllTypesMatch<["src1", "src2"]>,
|
|
AllTypesMatch<["acc", "res"]>,
|
|
]> {
|
|
let summary = "Matrix-matrix multiply and accumulate op";
|
|
let description = [{
|
|
SMMLA: Signed integer matrix multiply-accumulate.
|
|
|
|
Signed 8-bit integer matrix multiply-accumulate. This instruction multiplies
|
|
the 2x8 matrix of signed 8-bit integer values in the first source vector by
|
|
the 8x2 matrix of signed 8-bit integer values in the second source vector.
|
|
The resulting 2x2 32-bit integer matrix product is destructively added to
|
|
the 32-bit integer matrix accumulator in the destination vector. This is
|
|
equivalent to performing an 8-way dot product per destination element.
|
|
|
|
Source:
|
|
https://developer.arm.com/architectures/instruction-sets/intrinsics/#f:@navigationhierarchiessimdisa=[Neon]&q=smmla
|
|
}];
|
|
// Supports (vector<16xi8>, vector<16xi8>) -> (vector<4xi32>)
|
|
let arguments = (ins
|
|
NeonVectorOfLength<4, I32>:$acc,
|
|
NeonVectorOfLength<16, I8>:$src1,
|
|
NeonVectorOfLength<16, I8>:$src2
|
|
);
|
|
let results = (outs NeonVectorOfLength<4, I32>:$res);
|
|
let assemblyFormat =
|
|
"$acc `,` $src1 `,` $src2 attr-dict `:` type($src1) `to` type($res)";
|
|
}
|
|
|
|
def UmmlaOp : ArmNeon_OverloadedOperandsWithOneResultIntrOp<"ummla",[1], [
|
|
Pure,
|
|
AllTypesMatch<["src1", "src2"]>,
|
|
AllTypesMatch<["acc", "res"]>,
|
|
]> {
|
|
let summary = "Unsinged matrix-matrix multiply and accumulate op";
|
|
let description = [{
|
|
UMMLA: Signed integer matrix multiply-accumulate.
|
|
|
|
Unsigned 8-bit integer matrix multiply-accumulate. This instruction
|
|
multiplies the 2x8 matrix of unsigned 8-bit integer values in the first
|
|
source vector by the 8x2 matrix of unsigned 8-bit integer values in the
|
|
second source vector. The resulting 2x2 32-bit integer matrix product is
|
|
destructively added to the 32-bit integer matrix accumulator in the
|
|
destination vector. This is equivalent to performing an 8-way dot product
|
|
per destination element.
|
|
|
|
Source:
|
|
https://developer.arm.com/architectures/instruction-sets/intrinsics/#f:@navigationhierarchiessimdisa=[Neon]&q=ummla
|
|
}];
|
|
// Supports (vector<16xi8>, vector<16xi8>) -> (vector<4xi32>)
|
|
let arguments = (ins
|
|
NeonVectorOfLength<4, I32>:$acc,
|
|
NeonVectorOfLength<16, I8>:$src1,
|
|
NeonVectorOfLength<16, I8>:$src2
|
|
);
|
|
let results = (outs NeonVectorOfLength<4, I32>:$res);
|
|
let assemblyFormat =
|
|
"$acc `,` $src1 `,` $src2 attr-dict `:` type($src1) `to` type($res)";
|
|
}
|
|
|
|
def UsmmlaOp : ArmNeon_OverloadedOperandsWithOneResultIntrOp<"usmmla",[1], [
|
|
Pure,
|
|
AllTypesMatch<["src1", "src2"]>,
|
|
AllTypesMatch<["acc", "res"]>,
|
|
]> {
|
|
let summary = "Unsignged and signed matrix-matrix multiply and accumulate op";
|
|
let description = [{
|
|
USMMLA: Signed integer matrix multiply-accumulate.
|
|
|
|
Unsigned and signed 8-bit integer matrix multiply-accumulate. This
|
|
instruction multiplies the 2x8 matrix of unsigned 8-bit integer values in
|
|
the first source vector by the 8x2 matrix of signed 8-bit integer values in
|
|
the second source vector. The resulting 2x2 32-bit integer matrix product is
|
|
destructively added to the 32-bit integer matrix accumulator in the
|
|
destination vector. This is equivalent to performing an 8-way dot product
|
|
per destination element.
|
|
|
|
|
|
Source:
|
|
https://developer.arm.com/architectures/instruction-sets/intrinsics/#f:@navigationhierarchiessimdisa=[Neon]&q=usmmla
|
|
}];
|
|
// Supports (vector<16xi8>, vector<16xi8>) -> (vector<4xi32>)
|
|
let arguments = (ins
|
|
NeonVectorOfLength<4, I32>:$acc,
|
|
NeonVectorOfLength<16, I8>:$src1,
|
|
NeonVectorOfLength<16, I8>:$src2
|
|
);
|
|
let results = (outs NeonVectorOfLength<4, I32>:$res);
|
|
let assemblyFormat =
|
|
"$acc `,` $src1 `,` $src2 attr-dict `:` type($src1) `to` type($res)";
|
|
}
|
|
|
|
def BfmmlaOp : ArmNeon_IntrOp<"bfmmla", [], [], 1, [
|
|
Pure,
|
|
AllTypesMatch<["src1", "src2"]>,
|
|
AllTypesMatch<["acc", "res"]>,
|
|
]> {
|
|
let summary = "BFloat16 matrix multiply-accumulate to single-precision";
|
|
let description = [{
|
|
BFMMLA: BFloat16 matrix multiply-accumulate to single-precision.
|
|
|
|
The operation multiplies the 2x4 BFloat16 matrix in the first source vector
|
|
with the 4x2 BFloat16 matrix in the second source vector, then accumulates
|
|
this intermediate result with the 2x2 Float32 matrix in the accumulator
|
|
vector, yielding the final 2x2 Float32 result.
|
|
|
|
Source:
|
|
https://developer.arm.com/architectures/instruction-sets/intrinsics/vbfmmlaq_f32
|
|
}];
|
|
// Supports (vector<8xbf16>, vector<8xbf16>) -> (vector<2xf32>)
|
|
let arguments = (ins
|
|
NeonVectorOfLength<4, F32>:$acc,
|
|
NeonVectorOfLength<8, BF16>:$src1,
|
|
NeonVectorOfLength<8, BF16>:$src2
|
|
);
|
|
let results = (outs NeonVectorOfLength<4, F32>:$res);
|
|
let assemblyFormat =
|
|
"$acc `,` $src1 `,` $src2 attr-dict `:` type($src1) `to` type($res)";
|
|
}
|
|
|
|
class ArmNeon_2dOp<string mnemonic, list<Trait> traits = []>
|
|
: Op</*dialect=*/ArmNeon_Dialect,
|
|
/*opName=*/"2d." # mnemonic,
|
|
/*traits=*/traits>;
|
|
|
|
def Sdot2dOp : ArmNeon_2dOp<"sdot", [
|
|
Pure,
|
|
AllTypesMatch<["b", "c"]>,
|
|
AllTypesMatch<["a", "res"]>,
|
|
PredOpTrait<
|
|
"operand `a` should be 1-dimensional",
|
|
CPred<"::llvm::cast<VectorType>(getA().getType()).getShape().size() == 1">
|
|
>,
|
|
PredOpTrait<
|
|
"operand `b` should be 2-dimensional",
|
|
CPred<"::llvm::cast<VectorType>(getB().getType()).getShape().size() == 2">
|
|
>,
|
|
PredOpTrait<
|
|
"operand `b` should have 4 columns",
|
|
CPred<"::llvm::cast<VectorType>(getB().getType()).getShape()[1] == 4">
|
|
>,
|
|
PredOpTrait<
|
|
"operand `b` should have as many rows as the size of operand `a`",
|
|
CPred<"::llvm::cast<VectorType>(getB().getType()).getShape()[0] == ::llvm::cast<VectorType>(getA().getType()).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
|