diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index 6b7d4402be..f7103c3b00 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -197,6 +197,7 @@ function(setup_mlas_source_for_windows) ${MLAS_SRC_DIR}/amd64/sgemma.asm ${MLAS_SRC_DIR}/amd64/cvtfp16a.asm ${MLAS_SRC_DIR}/amd64/SoftmaxKernelAvx.asm + ${MLAS_SRC_DIR}/amd64/SoftmaxKernelAvx512F.asm ${MLAS_SRC_DIR}/amd64/TransKernelFma3.asm ${MLAS_SRC_DIR}/amd64/TransKernelAvx512F.asm ${MLAS_SRC_DIR}/amd64/LogisticKernelFma3.asm @@ -536,6 +537,7 @@ else() ${MLAS_SRC_DIR}/x86_64/DgemmKernelAvx512F.S ${MLAS_SRC_DIR}/x86_64/SgemmKernelAvx512F.S ${MLAS_SRC_DIR}/x86_64/SconvKernelAvx512F.S + ${MLAS_SRC_DIR}/x86_64/SoftmaxKernelAvx512F.S ${MLAS_SRC_DIR}/x86_64/SpoolKernelAvx512F.S ${MLAS_SRC_DIR}/x86_64/TransKernelAvx512F.S ${MLAS_SRC_DIR}/intrinsics/avx512/quantize_avx512f.cpp diff --git a/onnxruntime/core/mlas/lib/amd64/SoftmaxKernelAvx512F.asm b/onnxruntime/core/mlas/lib/amd64/SoftmaxKernelAvx512F.asm new file mode 100644 index 0000000000..3e83bc852f --- /dev/null +++ b/onnxruntime/core/mlas/lib/amd64/SoftmaxKernelAvx512F.asm @@ -0,0 +1,103 @@ +;++ +; +;Copyright (c) Microsoft Corporation. All rights reserved. +; +;Licensed under the MIT License. +; +;Module Name: +; +; SoftmaxKernelAvx512F.asm +; +;Abstract: +; +; This module implements the kernels for the single precision softmax +; operation. +; +; This implementation uses AVX512F instructions. +; +;-- + + .xlist +INCLUDE mlasi.inc + .list + + EXTERN MlasMinimumF32Value:NEAR + +;++ +; +;Routine Description: +; +; This routine implements a vectorized kernel to find the maximum value of +; the supplied buffer. +; +;Arguments: +; +; Input (rcx) - Supplies the input buffer. +; +; N (rdx) - Supplies the number of elements to process. +; +;Return Value: +; +; Returns the maximum value of the supplied buffer. +; +;-- + + LEAF_ENTRY MlasReduceMaximumF32KernelAvx512F, _TEXT + + vbroadcastss zmm0,DWORD PTR [MlasMinimumF32Value] + test rdx,rdx + jz ExitKernel + cmp rdx,16 + jb ProcessRemainingCountBy1 + cmp rdx,64 + jb ProcessRemainingCountBy16 + vmovaps zmm1,zmm0 + vmovaps zmm2,zmm0 + vmovaps zmm3,zmm0 + +ProcessRemainingCountBy64: + vmaxps zmm0,zmm0,ZMMWORD PTR [rcx] + vmaxps zmm1,zmm1,ZMMWORD PTR [rcx+16*4] + sub rdx,64 + vmaxps zmm2,zmm2,ZMMWORD PTR [rcx+32*4] + vmaxps zmm3,zmm3,ZMMWORD PTR [rcx+48*4] + add rcx,64*4 ; advance input by 64 elements + cmp rdx,64 + jae ProcessRemainingCountBy64 + vmaxps zmm0,zmm0,zmm1 ; reduce to single vector + vmaxps zmm2,zmm2,zmm3 + vmaxps zmm0,zmm0,zmm2 + +ProcessRemainingCountBy16: + cmp rdx,16 + jb ProcessRemainingCountLessThan16 + vmaxps zmm0,zmm0,ZMMWORD PTR [rcx] + sub rdx,16 + add rcx,16*4 ; advance input by 16 elements + jmp ProcessRemainingCountBy16 + +ProcessRemainingCountLessThan16: + vextractf32x8 ymm1,zmm0,1 ; reduce to single scalar + vmaxps ymm0,ymm0,ymm1 + vextractf128 xmm1,ymm0,1 + vmaxps xmm0,xmm0,xmm1 + vshufps xmm1,xmm0,xmm0,0EEh + vmaxps xmm0,xmm0,xmm1 + vshufps xmm1,xmm0,xmm0,055h + vmaxss xmm0,xmm0,xmm1 + test rdx,rdx + jz ExitKernel + +ProcessRemainingCountBy1: + vmaxss xmm0,xmm0,DWORD PTR [rcx] + add rcx,4 ; advance input by 1 element + dec edx + jnz ProcessRemainingCountBy1 + +ExitKernel: + vzeroupper + ret + + LEAF_END MlasReduceMaximumF32KernelAvx512F, _TEXT + + END diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 624eb913d5..4b93dde1bc 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -846,6 +846,7 @@ extern "C" { MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL MlasReduceMinimumMaximumF32Kernel; #if defined(MLAS_TARGET_AMD64) MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32KernelAvx; + MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32KernelAvx512F; MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL MlasReduceMinimumMaximumF32KernelAvx; #endif diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index de092f7d1d..a53c5085b1 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -421,6 +421,7 @@ Return Value: this->PoolFloatKernel[MlasAveragePoolingIncludePad] = MlasPoolAverageIncludePadFloatKernelAvx512F; this->ComputeExpF32Kernel = MlasComputeExpF32KernelAvx512F; this->ComputeSumExpF32Kernel = MlasComputeSumExpF32KernelAvx512F; + this->ReduceMaximumF32Kernel = MlasReduceMaximumF32KernelAvx512F; this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8KernelAvx512F; this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8KernelAvx512F; this->NchwcBlockSize = 16; diff --git a/onnxruntime/core/mlas/lib/x86_64/SoftmaxKernelAvx512F.S b/onnxruntime/core/mlas/lib/x86_64/SoftmaxKernelAvx512F.S new file mode 100644 index 0000000000..db97286046 --- /dev/null +++ b/onnxruntime/core/mlas/lib/x86_64/SoftmaxKernelAvx512F.S @@ -0,0 +1,101 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + SoftmaxKernelAvx512F.s + +Abstract: + + This module implements the kernels for the single precision softmax + operation. + + This implementation uses AVX512F instructions. + +--*/ + +#include "asmmacro.h" + + .intel_syntax noprefix + + .text + +/*++ + +Routine Description: + + This routine implements a vectorized kernel to find the maximum value of + the supplied buffer. + +Arguments: + + Input (rdi) - Supplies the input buffer. + + N (rsi) - Supplies the number of elements to process. + +Return Value: + + Returns the maximum value of the supplied buffer. + +--*/ + + FUNCTION_ENTRY MlasReduceMaximumF32KernelAvx512F + + vbroadcastss zmm0,DWORD PTR C_UNDERSCORE(MlasMinimumF32Value)[rip] + test rsi,rsi + jz .LReduceMaximum.ExitKernel + cmp rsi,16 + jb .LReduceMaximum.ProcessRemainingCountBy1 + cmp rsi,64 + jb .LReduceMaximum.ProcessRemainingCountBy16 + vmovaps zmm1,zmm0 + vmovaps zmm2,zmm0 + vmovaps zmm3,zmm0 + +.LReduceMaximum.ProcessRemainingCountBy64: + vmaxps zmm0,zmm0,ZMMWORD PTR [rdi] + vmaxps zmm1,zmm1,ZMMWORD PTR [rdi+16*4] + sub rsi,64 + vmaxps zmm2,zmm2,ZMMWORD PTR [rdi+32*4] + vmaxps zmm3,zmm3,ZMMWORD PTR [rdi+48*4] + add rdi,64*4 # advance input by 64 elements + cmp rsi,64 + jae .LReduceMaximum.ProcessRemainingCountBy64 + vmaxps zmm0,zmm0,zmm1 # reduce to single vector + vmaxps zmm2,zmm2,zmm3 + vmaxps zmm0,zmm0,zmm2 + +.LReduceMaximum.ProcessRemainingCountBy16: + cmp rsi,16 + jb .LReduceMaximum.ProcessRemainingCountLessThan16 + vmaxps zmm0,zmm0,ZMMWORD PTR [rdi] + sub rsi,16 + add rdi,16*4 # advance input by 16 elements + jmp .LReduceMaximum.ProcessRemainingCountBy16 + +.LReduceMaximum.ProcessRemainingCountLessThan16: + vextractf32x8 ymm1,zmm0,1 # reduce to single scalar + vmaxps ymm0,ymm0,ymm1 + vextractf128 xmm1,ymm0,1 + vmaxps xmm0,xmm0,xmm1 + vshufps xmm1,xmm0,xmm0,0xEE + vmaxps xmm0,xmm0,xmm1 + vshufps xmm1,xmm0,xmm0,0x55 + vmaxss xmm0,xmm0,xmm1 + test rsi,rsi + jz .LReduceMaximum.ExitKernel + +.LReduceMaximum.ProcessRemainingCountBy1: + vmaxss xmm0,xmm0,DWORD PTR [rdi] + add rdi,4 # advance input by 1 element + dec esi + jnz .LReduceMaximum.ProcessRemainingCountBy1 + +.LReduceMaximum.ExitKernel: + vzeroupper + ret + + .end diff --git a/onnxruntime/test/mlas/bench/bench_computesoftmax.cpp b/onnxruntime/test/mlas/bench/bench_computesoftmax.cpp new file mode 100644 index 0000000000..f777a7cfc4 --- /dev/null +++ b/onnxruntime/test/mlas/bench/bench_computesoftmax.cpp @@ -0,0 +1,233 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/mlas/lib/mlasi.h" +#include "core/util/thread_utils.h" +#include "test/mlas/bench/bench_util.h" + +using onnxruntime::narrow; + +struct RestrictAlignedPtr { + float* ptr; // Aligned pointer within the underlying buffer + void* underlying_buffer; // Underlying buffer (including extra space for alignment) +}; + +// Return a RestrictAlignedPtr where the ptr is aligned to byte_aligned, but not to byte_aligned * 2 +RestrictAlignedPtr restrict_aligned_alloc(int D, int byte_aligned) { + if (byte_aligned <= 0 || (byte_aligned & (byte_aligned - 1)) != 0) { + throw std::invalid_argument("Alignment must be a power of 2"); + } + + const int byte_alignedx2 = byte_aligned << 1; + + void* buffer = malloc(D * sizeof(float) + byte_alignedx2 * 2); + if (buffer == nullptr) { + ORT_THROW_EX(std::bad_alloc); + } + + uintptr_t address = reinterpret_cast(buffer); + uintptr_t aligned_address = ((address + byte_alignedx2 - 1) & ~(byte_alignedx2 - 1)) + byte_aligned; + ORT_ENFORCE((aligned_address % byte_aligned == 0) && (aligned_address % byte_alignedx2 != 0)); + float* aligned_ptr = reinterpret_cast(aligned_address); + + return {aligned_ptr, buffer}; +} + +void COMPUTESOFTMAXINPLACE(benchmark::State& state) { + const auto byte_aligned = narrow(state.range(0)); + const auto N = narrow(state.range(1)); + const auto D = narrow(state.range(2)); + const auto threads = narrow(state.range(3)); + + if (N <= 0 || D <= 0 || threads <= 0) { + throw std::invalid_argument("N, D, and Threads must be greater than 0!"); + } + + OrtThreadPoolParams tpo; + tpo.thread_pool_size = threads; + tpo.auto_set_affinity = true; + + std::unique_ptr tp( + onnxruntime::concurrency::CreateThreadPool( + &onnxruntime::Env::Default(), tpo, onnxruntime::concurrency::ThreadPoolType::INTRA_OP)); + + auto data = RandomVectorUniform(static_cast(N * D), -1.0f, 1.0f); + RestrictAlignedPtr ptr = restrict_aligned_alloc(N * D, byte_aligned); + float* input = ptr.ptr; + float* output = input; + std::copy(data.begin(), data.end(), input); // Copy the data to the aligned memory + + // warming up run + MlasComputeSoftmax(input, output, N, D, false, tp.get()); + + for (auto _ : state) { + MlasComputeSoftmax(input, output, N, D, false, tp.get()); + } + + free(ptr.underlying_buffer); +} + +void REDUCEMAXIMUMF32KERNELAVX(benchmark::State& state) { + const auto byte_aligned = narrow(state.range(0)); + const auto D = narrow(state.range(1)); + + if (D <= 0) { + throw std::invalid_argument("D must be greater than 0!"); + } + + auto data = RandomVectorUniform(static_cast(D), -1.0f, 1.0f); + RestrictAlignedPtr ptr = restrict_aligned_alloc(D, byte_aligned); + float* input = ptr.ptr; + std::copy(data.begin(), data.end(), input); // Copy the data to the aligned memory + + // warming up run + float Maximum = MlasReduceMaximumF32KernelAvx(input, D); + + for (auto _ : state) { + Maximum = MlasReduceMaximumF32KernelAvx(input, D); + } + + free(ptr.underlying_buffer); + (void)Maximum; +} + +void REDUCEMAXIMUMF32KERNELAVX512F(benchmark::State& state) { + const auto byte_aligned = narrow(state.range(0)); + const auto D = narrow(state.range(1)); + + if (D <= 0) { + throw std::invalid_argument("D must be greater than 0!"); + } + + auto data = RandomVectorUniform(static_cast(D), -1.0f, 1.0f); + RestrictAlignedPtr ptr = restrict_aligned_alloc(D, byte_aligned); + float* input = ptr.ptr; + std::copy(data.begin(), data.end(), input); // Copy the data to the aligned memory + + // warming up run + float Maximum = MlasReduceMaximumF32KernelAvx512F(input, D); + + for (auto _ : state) { + Maximum = MlasReduceMaximumF32KernelAvx512F(input, D); + } + + free(ptr.underlying_buffer); + (void)Maximum; +} + +void COMPUTESUMEXPF32KERNELAVX512F(benchmark::State& state) { + const auto byte_aligned = narrow(state.range(0)); + const auto D = narrow(state.range(1)); + + if (D <= 0) { + throw std::invalid_argument("D must be greater than 0!"); + } + + auto data = RandomVectorUniform(static_cast(D), -1.0f, 1.0f); + RestrictAlignedPtr ptr = restrict_aligned_alloc(D, byte_aligned); + float* input = ptr.ptr; + float* output = input; + std::copy(data.begin(), data.end(), input); // Copy the data to the aligned memory + + float Maximum = MlasReduceMaximumF32KernelAvx(input, D); + float NegativeMaximum = -Maximum; + + // warming up run + float Accumulation = MlasComputeSumExpF32KernelAvx512F(input, output, D, &NegativeMaximum); + + for (auto _ : state) { + Accumulation = MlasComputeSumExpF32KernelAvx512F(input, output, D, &NegativeMaximum); + } + + free(ptr.underlying_buffer); + (void)Accumulation; +} + +void COMPUTESOFTMAXOUTPUTF32KERNELAVX(benchmark::State& state) { + const auto byte_aligned = narrow(state.range(0)); + const auto D = narrow(state.range(1)); + + if (D <= 0) { + throw std::invalid_argument("D must be greater than 0!"); + } + + auto data = RandomVectorUniform(static_cast(D), -1.0f, 1.0f); + RestrictAlignedPtr ptr = restrict_aligned_alloc(D, byte_aligned); + float* input = ptr.ptr; + float* output = input; + std::copy(data.begin(), data.end(), input); // Copy the data to the aligned memory + + float Maximum = MlasReduceMaximumF32KernelAvx(input, D); + float NegativeMaximum = -Maximum; + + float Accumulation = MlasComputeSumExpF32KernelAvx512F(input, output, D, &NegativeMaximum); + + float Parameters[] = {1.0f / Accumulation}; + + // warming up run + MlasComputeSoftmaxOutputF32KernelAvx(output, D, Parameters); + + for (auto _ : state) { + MlasComputeSoftmaxOutputF32KernelAvx(output, D, Parameters); + } + + free(ptr.underlying_buffer); +} + +static void ComputeSoftmaxInplaceArgs(benchmark::internal::Benchmark* b) { + b->ArgNames({"ByteAligned", "N", "D", "Threads"}); + for (int threads : {1, 8}) { + for (int byte_aligned : {64}) { // MLAS_DEFAULT_PREFERRED_BUFFER_ALIGNMENT is 64 + b->Args({byte_aligned, 16000, 4, threads}); + b->Args({byte_aligned, 16000, 500, threads}); + b->Args({byte_aligned, 48000, 3, threads}); + b->Args({byte_aligned, 48000, 2000, threads}); + b->Args({byte_aligned, 80000, 5, threads}); + b->Args({byte_aligned, 80000, 2000, threads}); + b->Args({byte_aligned, 112000, 7, threads}); + b->Args({byte_aligned, 112000, 2000, threads}); + b->Args({byte_aligned, 144000, 9, threads}); + b->Args({byte_aligned, 144000, 2000, threads}); + b->Args({byte_aligned, 176000, 11, threads}); + b->Args({byte_aligned, 176000, 2000, threads}); + b->Args({byte_aligned, 208000, 13, threads}); + b->Args({byte_aligned, 208000, 2000, threads}); + b->Args({byte_aligned, 240000, 15, threads}); + b->Args({byte_aligned, 240000, 2000, threads}); + } + } +} + +BENCHMARK(COMPUTESOFTMAXINPLACE)->Apply(ComputeSoftmaxInplaceArgs)->UseRealTime(); + +BENCHMARK(REDUCEMAXIMUMF32KERNELAVX) + ->ArgNames({"ByteAligned", "D"}) + ->ArgsProduct({ + {4, 8, 16, 32, 64, 128}, // ByteAligned + {3, 4, 5, 7, 9, 11, 13, 15, 16, 500, 2000}, // D + }) + ->UseRealTime(); + +BENCHMARK(REDUCEMAXIMUMF32KERNELAVX512F) + ->ArgNames({"ByteAligned", "D"}) + ->ArgsProduct({ + {4, 8, 16, 32, 64, 128}, // ByteAligned + {3, 4, 5, 7, 9, 11, 13, 15, 16, 500, 2000}, // D + }) + ->UseRealTime(); + +BENCHMARK(COMPUTESUMEXPF32KERNELAVX512F) + ->ArgNames({"ByteAligned", "D"}) + ->ArgsProduct({ + {4, 8, 16, 32, 64, 128}, // ByteAligned + {3, 4, 5, 7, 9, 11, 13, 15, 500, 2000}, // D + }) + ->UseRealTime(); + +BENCHMARK(COMPUTESOFTMAXOUTPUTF32KERNELAVX) + ->ArgNames({"ByteAligned", "D"}) + ->ArgsProduct({ + {4, 8, 16, 32, 64, 128}, // ByteAligned + {3, 4, 5, 7, 9, 11, 13, 15, 16, 500, 2000}, // D + }) + ->UseRealTime();