222 lines
9.2 KiB
C++
222 lines
9.2 KiB
C++
//=- Transforms.h - X86 Dialect Transformation Entrypoints --------*- C++ -*-=//
|
|
//
|
|
// 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_DIALECT_X86_TRANSFORMS_H
|
|
#define MLIR_DIALECT_X86_TRANSFORMS_H
|
|
|
|
#include "mlir/IR/Value.h"
|
|
|
|
namespace mlir {
|
|
|
|
class ImplicitLocOpBuilder;
|
|
class LLVMConversionTarget;
|
|
class LLVMTypeConverter;
|
|
class RewritePatternSet;
|
|
|
|
namespace x86 {
|
|
|
|
/// Helper class to factor out the creation and extraction of masks from nibs.
|
|
struct MaskHelper {
|
|
/// b0 captures the lowest bit, b7 captures the highest bit.
|
|
/// Meant to be used with instructions such as mm256BlendPs.
|
|
template <uint8_t b0, uint8_t b1, uint8_t b2, uint8_t b3, uint8_t b4,
|
|
uint8_t b5, uint8_t b6, uint8_t b7>
|
|
static uint8_t blend() {
|
|
static_assert(b0 <= 1 && b1 <= 1 && b2 <= 1 && b3 <= 1, "overflow");
|
|
static_assert(b4 <= 1 && b5 <= 1 && b6 <= 1 && b7 <= 1, "overflow");
|
|
return static_cast<uint8_t>((b7 << 7) | (b6 << 6) | (b5 << 5) | (b4 << 4) |
|
|
(b3 << 3) | (b2 << 2) | (b1 << 1) | b0);
|
|
}
|
|
/// b0 captures the lowest bit, b7 captures the highest bit.
|
|
/// Meant to be used with instructions such as mm256BlendPs.
|
|
static void extractBlend(uint8_t mask, uint8_t &b0, uint8_t &b1, uint8_t &b2,
|
|
uint8_t &b3, uint8_t &b4, uint8_t &b5, uint8_t &b6,
|
|
uint8_t &b7) {
|
|
b7 = mask & (1 << 7);
|
|
b6 = mask & (1 << 6);
|
|
b5 = mask & (1 << 5);
|
|
b4 = mask & (1 << 4);
|
|
b3 = mask & (1 << 3);
|
|
b2 = mask & (1 << 2);
|
|
b1 = mask & (1 << 1);
|
|
b0 = mask & 1;
|
|
}
|
|
/// b01 captures the lower 2 bits, b67 captures the higher 2 bits.
|
|
/// Meant to be used with instructions such as mm256ShufflePs.
|
|
template <unsigned b67, unsigned b45, unsigned b23, unsigned b01>
|
|
static uint8_t shuffle() {
|
|
static_assert(b01 <= 0x03, "overflow");
|
|
static_assert(b23 <= 0x03, "overflow");
|
|
static_assert(b45 <= 0x03, "overflow");
|
|
static_assert(b67 <= 0x03, "overflow");
|
|
return static_cast<uint8_t>((b67 << 6) | (b45 << 4) | (b23 << 2) | b01);
|
|
}
|
|
/// b01 captures the lower 2 bits, b67 captures the higher 2 bits.
|
|
static void extractShuffle(uint8_t mask, uint8_t &b01, uint8_t &b23,
|
|
uint8_t &b45, uint8_t &b67) {
|
|
b67 = (mask & (0x03 << 6)) >> 6;
|
|
b45 = (mask & (0x03 << 4)) >> 4;
|
|
b23 = (mask & (0x03 << 2)) >> 2;
|
|
b01 = mask & 0x03;
|
|
}
|
|
/// b03 captures the lower 4 bits, b47 captures the higher 4 bits.
|
|
/// Meant to be used with instructions such as mm256Permute2f128Ps.
|
|
template <unsigned b47, unsigned b03>
|
|
static uint8_t permute() {
|
|
static_assert(b03 <= 0x0f, "overflow");
|
|
static_assert(b47 <= 0x0f, "overflow");
|
|
return static_cast<uint8_t>((b47 << 4) + b03);
|
|
}
|
|
/// b03 captures the lower 4 bits, b47 captures the higher 4 bits.
|
|
static void extractPermute(uint8_t mask, uint8_t &b03, uint8_t &b47) {
|
|
b47 = (mask & (0x0f << 4)) >> 4;
|
|
b03 = mask & 0x0f;
|
|
}
|
|
};
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
// A set of patterns for specialized lowering of vector contraction
|
|
// operation to vector fused multiply and add (FMA) operation.
|
|
void populateVectorContractToFMAPatterns(RewritePatternSet &patterns);
|
|
|
|
// A set of patterns for lowering 32-bit packed vector contraction operations
|
|
// to their corresponding packed-type dot-product operations, ultimately
|
|
// targeting the relevant x86 LLVM intrinsics (e.g., BF16 and Int8).
|
|
void populateVectorContractToPackedTypeDotProductPatterns(
|
|
RewritePatternSet &patterns);
|
|
|
|
// A set of patterns for lowering 32-bit packed BF16 vector contraction
|
|
// operations to vector fused multiply-add (FMA) operations, following
|
|
// the emulation-based approach using BF16 packed operations.
|
|
void populateVectorContractBF16ToFMAPatterns(RewritePatternSet &patterns);
|
|
|
|
// Performs forward scheduling of vector producer ops to minimize their live
|
|
// range by placing them at their earliest legal use site.
|
|
void populateSinkVectorProducerOpsPatterns(RewritePatternSet &patterns);
|
|
|
|
// Shuffles FMAs with x86 operations as operands such that FMAs are
|
|
// grouped with respect to odd/even packed index.
|
|
void populateShuffleVectorFMAOpsPatterns(RewritePatternSet &patterns);
|
|
|
|
// A set of patterns for lowering 32-bit packed vector contraction operations
|
|
// to their corresponding packed-type tiled dot-product operations, using
|
|
// AMX ultimately targeting the relevant x86 LLVM intrinsics (e.g., BF16 and
|
|
// Int8).
|
|
void populateVectorContractToAMXDotProductPatterns(RewritePatternSet &patterns);
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
/// Helpers extracted from:
|
|
/// - clang/lib/Headers/avxintrin.h
|
|
/// - clang/test/CodeGen/X86/avx-builtins.c
|
|
/// - clang/test/CodeGen/X86/avx2-builtins.c
|
|
/// - clang/test/CodeGen/X86/avx-shuffle-builtins.c
|
|
/// as well as the Intel Intrinsics Guide
|
|
/// (https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html)
|
|
/// make it easier to just implement known good lowerings.
|
|
/// All intrinsics correspond 1-1 to the Intel definition.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace avx2 {
|
|
|
|
namespace inline_asm {
|
|
//===----------------------------------------------------------------------===//
|
|
/// Methods in the inline_asm namespace emit calls to LLVM::InlineAsmOp.
|
|
//===----------------------------------------------------------------------===//
|
|
/// If bit i of `mask` is zero, take f32@i from v1 else take it from v2.
|
|
Value mm256BlendPsAsm(ImplicitLocOpBuilder &b, Value v1, Value v2,
|
|
uint8_t mask);
|
|
|
|
} // namespace inline_asm
|
|
|
|
namespace intrin {
|
|
//===----------------------------------------------------------------------===//
|
|
/// Methods in the intrin namespace emulate clang's impl. of X86 intrinsics.
|
|
//===----------------------------------------------------------------------===//
|
|
/// Lower to vector.shuffle v1, v2, [0, 8, 1, 9, 4, 12, 5, 13].
|
|
Value mm256UnpackLoPs(ImplicitLocOpBuilder &b, Value v1, Value v2);
|
|
|
|
/// Lower to vector.shuffle v1, v2, [0, 8, 1, 9, 4, 12, 5, 13].
|
|
Value mm256UnpackHiPs(ImplicitLocOpBuilder &b, Value v1, Value v2);
|
|
|
|
/// a a b b a a b b
|
|
/// Take an 8 bit mask, 2 bit for each position of a[0, 3) **and** b[0, 4):
|
|
/// 0:127 | 128:255
|
|
/// b01 b23 C8 D8 | b01+4 b23+4 C8+4 D8+4
|
|
Value mm256ShufflePs(ImplicitLocOpBuilder &b, Value v1, Value v2, uint8_t mask);
|
|
|
|
// imm[0:1] out of imm[0:3] is:
|
|
// 0 1 2 3
|
|
// a[0:127] or a[128:255] or b[0:127] or b[128:255] |
|
|
// a[0:127] or a[128:255] or b[0:127] or b[128:255]
|
|
// 0 1 2 3
|
|
// imm[0:1] out of imm[4:7].
|
|
Value mm256Permute2f128Ps(ImplicitLocOpBuilder &b, Value v1, Value v2,
|
|
uint8_t mask);
|
|
|
|
/// If bit i of `mask` is zero, take f32@i from v1 else take it from v2.
|
|
Value mm256BlendPs(ImplicitLocOpBuilder &b, Value v1, Value v2, uint8_t mask);
|
|
} // namespace intrin
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
/// Generic lowerings may either use intrin or inline_asm depending on needs.
|
|
//===----------------------------------------------------------------------===//
|
|
/// 4x8xf32-specific AVX2 transpose lowering.
|
|
void transpose4x8xf32(ImplicitLocOpBuilder &ib, MutableArrayRef<Value> vs);
|
|
|
|
/// 8x8xf32-specific AVX2 transpose lowering.
|
|
void transpose8x8xf32(ImplicitLocOpBuilder &ib, MutableArrayRef<Value> vs);
|
|
|
|
/// Structure to control the behavior of specialized AVX2 transpose lowering.
|
|
struct TransposeLoweringOptions {
|
|
bool lower4x8xf32_ = false;
|
|
TransposeLoweringOptions &lower4x8xf32(bool lower = true) {
|
|
lower4x8xf32_ = lower;
|
|
return *this;
|
|
}
|
|
bool lower8x8xf32_ = false;
|
|
TransposeLoweringOptions &lower8x8xf32(bool lower = true) {
|
|
lower8x8xf32_ = lower;
|
|
return *this;
|
|
}
|
|
};
|
|
|
|
/// Options for controlling specialized AVX2 lowerings.
|
|
struct LoweringOptions {
|
|
/// Configure specialized vector lowerings.
|
|
TransposeLoweringOptions transposeOptions;
|
|
LoweringOptions &setTransposeOptions(TransposeLoweringOptions options) {
|
|
transposeOptions = options;
|
|
return *this;
|
|
}
|
|
};
|
|
|
|
/// Insert specialized transpose lowering patterns.
|
|
void populateSpecializedTransposeLoweringPatterns(
|
|
RewritePatternSet &patterns, LoweringOptions options = LoweringOptions(),
|
|
int benefit = 10);
|
|
|
|
} // namespace avx2
|
|
} // namespace x86
|
|
|
|
/// Collect a set of patterns to lower X86 ops to ops that map to LLVM
|
|
/// intrinsics.
|
|
void populateX86LegalizeForLLVMExportPatterns(LLVMTypeConverter &converter,
|
|
RewritePatternSet &patterns);
|
|
|
|
/// Configure the target to support lowering X86 ops to ops that map to
|
|
/// LLVM intrinsics.
|
|
void configureX86LegalizeForExportTarget(LLVMConversionTarget &target);
|
|
|
|
/// Register LLVM conversion interface for X86 dialect.
|
|
void registerConvertX86ToLLVMInterface(DialectRegistry ®istry);
|
|
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_DIALECT_X86_TRANSFORMS_H
|