diff --git a/csharp/src/Microsoft.ML.OnnxRuntime/InferenceSession.cs b/csharp/src/Microsoft.ML.OnnxRuntime/InferenceSession.cs index af59d13df8114..bfc2f7c839d85 100644 --- a/csharp/src/Microsoft.ML.OnnxRuntime/InferenceSession.cs +++ b/csharp/src/Microsoft.ML.OnnxRuntime/InferenceSession.cs @@ -47,6 +47,21 @@ public InferenceSession(string modelPath) Init(modelPath, _builtInSessionOptions); } + /// + /// Constructs an InferenceSession from a model file and it will use + /// the provided pre-packed weights container to store and share pre-packed buffers + /// of shared initializers across sessions if any. + /// + /// Model path + /// Instance of PrepackedWeightsContainer. + /// Lifetime of 'prepackedWeightsContainer' must be + /// managed by the user and it must outlive any sessions reliant on it + public InferenceSession(string modelPath, PrePackedWeightsContainer prepackedWeightsContainer) + { + _builtInSessionOptions = new SessionOptions(); // need to be disposed + Init(modelPath, _builtInSessionOptions, prepackedWeightsContainer); + } + /// /// Constructs an InferenceSession from a model file, with some additional session options @@ -58,6 +73,23 @@ public InferenceSession(string modelPath, SessionOptions options) Init(modelPath, options); } + + /// + /// Constructs an InferenceSession from a model file, with some additional session options + /// and it will use the provided pre-packed weights container to store and share pre-packed buffers + /// of shared initializers across sessions if any. + /// + /// Model path + /// Session options + /// Instance of PrepackedWeightsContainer. + /// Lifetime of 'prepackedWeightsContainer' must be + /// managed by the user and it must outlive any sessions reliant on it + public InferenceSession(string modelPath, SessionOptions options, + PrePackedWeightsContainer prepackedWeightsContainer) + { + Init(modelPath, options, prepackedWeightsContainer); + } + /// /// Constructs an InferenceSession from a model data in byte array /// @@ -68,6 +100,21 @@ public InferenceSession(byte[] model) Init(model, _builtInSessionOptions); } + /// + /// Constructs an InferenceSession from a model data (in byte array) and it will use + /// the provided pre-packed weights container to store and share pre-packed buffers + /// of shared initializers across sessions if any. + /// + /// Model as byte array + /// Instance of PrepackedWeightsContainer. + /// Lifetime of 'prepackedWeightsContainer' must be + /// managed by the user and it must outlive any sessions reliant on it + public InferenceSession(byte[] model, PrePackedWeightsContainer prepackedWeightsContainer) + { + _builtInSessionOptions = new SessionOptions(); // need to be disposed + Init(model, _builtInSessionOptions, prepackedWeightsContainer); + } + /// /// Constructs an InferenceSession from a model data in byte array, with some additional session options /// @@ -78,6 +125,21 @@ public InferenceSession(byte[] model, SessionOptions options) Init(model, options); } + /// + /// Constructs an InferenceSession from a model data (in byte array) with some additional + /// session options and it will use the provided pre-packed weights container to store + /// and share pre-packed buffers of shared initializers across sessions if any. + /// + /// Model as byte array + /// Session Options + /// Instance of PrepackedWeightsContainer. + /// Lifetime of 'prepackedWeightsContainer' must be + /// managed by the user and it must outlive any sessions reliant on it + public InferenceSession(byte[] model, SessionOptions options, + PrePackedWeightsContainer prepackedWeightsContainer) + { + Init(model, options, prepackedWeightsContainer); + } /// /// Meta data regarding the input nodes, keyed by input names /// @@ -310,7 +372,7 @@ public void Run( IReadOnlyCollection outputs, RunOptions options) { - using(var cleanupList = new DisposableList()) + using (var cleanupList = new DisposableList()) { var inputNamesArray = ConvertNamesToUtf8(inputs, i => i.Name, cleanupList); var inputValuesArray = GetOrtValuesHandles(inputs, cleanupList); @@ -367,7 +429,7 @@ public void Run( throw new ArgumentException($"Length of {nameof(outputNames)} ({outputNames.Count}) must match that of {nameof(outputValues)} ({outputValues.Count})."); } - using(var cleanupList = new DisposableList()) + using (var cleanupList = new DisposableList()) { // prepare inputs var inputNamesArray = ConvertNamesToUtf8(inputs, i => i.Name, cleanupList); @@ -428,7 +490,7 @@ public void Run( throw new ArgumentException($"Length of {nameof(inputNames)} ({inputNames.Count}) must match that of {nameof(inputValues)} ({inputValues.Count})."); } - using(var cleanupList = new DisposableList()) + using (var cleanupList = new DisposableList()) { // prepare inputs var inputNamesArray = ConvertNamesToUtf8(inputNames, n => n, cleanupList); @@ -515,7 +577,8 @@ public IDisposableReadOnlyCollection RunWithBindingAnd var ortValue = ortValues.ElementAt(i); result.Add(DisposableNamedOnnxValue.CreateFromOrtValue(outputNames[i], ortValue)); } - } catch(Exception e) + } + catch (Exception e) { result.Dispose(); throw e; @@ -535,7 +598,7 @@ public string EndProfiling() NativeApiStatus.VerifySuccess(NativeMethods.OrtSessionEndProfiling(_nativeHandle, allocator.Pointer, out nameHandle)); - using(var allocation = new OrtMemoryAllocation(allocator, nameHandle, 0)) + using (var allocation = new OrtMemoryAllocation(allocator, nameHandle, 0)) { return NativeOnnxValueHelper.StringFromNativeUtf8(nameHandle); } @@ -609,8 +672,8 @@ private IntPtr[] GetOrtValuesHandles(IReadOnlyCollection v } - private DisposableList RunImpl(RunOptions options, IntPtr[] inputNames, IntPtr[] inputValues, IntPtr[] outputNames, - DisposableList cleanupList) + private DisposableList RunImpl(RunOptions options, IntPtr[] inputNames, IntPtr[] inputValues, IntPtr[] outputNames, + DisposableList cleanupList) { var ortValues = new DisposableList(outputNames.Length); cleanupList.Add(ortValues); @@ -680,31 +743,57 @@ public ModelMetadata ModelMetadata /// public ulong ProfilingStartTimeNs { - get - { - return _profilingStartTimeNs; - } - } + get + { + return _profilingStartTimeNs; + } + } #endregion #region private methods - private void Init(string modelPath, SessionOptions options) + private void Init(string modelPath, SessionOptions options, + PrePackedWeightsContainer prepackedWeightsContainer = null) { var envHandle = OrtEnv.Handle; var session = IntPtr.Zero; - NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSession(envHandle, NativeMethods.GetPlatformSerializedString(modelPath), options.Handle, out session)); + + if (prepackedWeightsContainer == null) + { + NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSession(envHandle, NativeMethods.GetPlatformSerializedString(modelPath), + options.Handle, out session)); + } + + else + { + NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSessionWithPrepackedWeightsContainer( + envHandle, NativeMethods.GetPlatformSerializedString(modelPath), + options.Handle, prepackedWeightsContainer.Pointer, out session)); + } InitWithSessionHandle(session, options); } - private void Init(byte[] modelData, SessionOptions options) + private void Init(byte[] modelData, SessionOptions options, + PrePackedWeightsContainer prepackedWeightsContainer = null) { var envHandle = OrtEnv.Handle; var session = IntPtr.Zero; - NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSessionFromArray(envHandle, modelData, (UIntPtr)modelData.Length, options.Handle, out session)); + if (prepackedWeightsContainer == null) + { + + NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSessionFromArray(envHandle, modelData, (UIntPtr)modelData.Length, options.Handle, out session)); + } + + else + { + NativeApiStatus.VerifySuccess(NativeMethods.OrtCreateSessionFromArrayWithPrepackedWeightsContainer( + envHandle, modelData, (UIntPtr)modelData.Length, options.Handle, prepackedWeightsContainer.Pointer, + out session)); + + } InitWithSessionHandle(session, options); } @@ -757,8 +846,8 @@ private void InitWithSessionHandle(IntPtr session, SessionOptions options) // set profiling's start time UIntPtr startTime = UIntPtr.Zero; NativeApiStatus.VerifySuccess(NativeMethods.OrtSessionGetProfilingStartTimeNs(_nativeHandle, - out startTime)); - _profilingStartTimeNs = (ulong) startTime; + out startTime)); + _profilingStartTimeNs = (ulong)startTime; } catch (OnnxRuntimeException e) { @@ -821,7 +910,7 @@ private string GetOverridableInitializerName(ulong index) (UIntPtr)index, allocator.Pointer, out nameHandle)); - using(var ortAllocation = new OrtMemoryAllocation(allocator, nameHandle, 0)) + using (var ortAllocation = new OrtMemoryAllocation(allocator, nameHandle, 0)) { str = NativeOnnxValueHelper.StringFromNativeUtf8(nameHandle); } @@ -963,7 +1052,7 @@ public void Dispose() /// true if invoked from Dispose() method protected virtual void Dispose(bool disposing) { - if(_disposed) + if (_disposed) { return; } @@ -1137,7 +1226,7 @@ internal ModelMetadata(InferenceSession session) } // Process each key via the stored key handles - foreach(var allocation in ortAllocationKeys) + foreach (var allocation in ortAllocationKeys) { IntPtr keyHandle = allocation.Pointer; IntPtr valueHandle = IntPtr.Zero; @@ -1160,9 +1249,9 @@ internal ModelMetadata(InferenceSession session) { // Free ModelMetadata handle - NativeMethods.OrtReleaseModelMetadata(modelMetadataHandle); + NativeMethods.OrtReleaseModelMetadata(modelMetadataHandle); - } + } } diff --git a/csharp/src/Microsoft.ML.OnnxRuntime/NativeMethods.cs b/csharp/src/Microsoft.ML.OnnxRuntime/NativeMethods.cs index 0df4c77404898..c8d6864e9a478 100644 --- a/csharp/src/Microsoft.ML.OnnxRuntime/NativeMethods.cs +++ b/csharp/src/Microsoft.ML.OnnxRuntime/NativeMethods.cs @@ -190,6 +190,17 @@ public struct OrtApi public IntPtr CreateArenaCfg; public IntPtr ReleaseArenaCfg; public IntPtr ModelMetadataGetGraphDescription; + public IntPtr SessionOptionsAppendExecutionProvider_TensorRT; + public IntPtr SetCurrentGpuDeviceId; + public IntPtr GetCurrentGpuDeviceId; + public IntPtr KernelInfoGetAttributeArray_float; + public IntPtr KernelInfoGetAttributeArray_int64; + public IntPtr CreateArenaCfgV2; + public IntPtr AddRunConfigEntry; + public IntPtr CreatePrepackedWeightsContainer; + public IntPtr ReleasePrepackedWeightsContainer; + public IntPtr CreateSessionWithPrepackedWeightsContainer; + public IntPtr CreateSessionFromArrayWithPrepackedWeightsContainer; } internal static class NativeMethods @@ -218,7 +229,11 @@ static NativeMethods() OrtReleaseStatus = (DOrtReleaseStatus)Marshal.GetDelegateForFunctionPointer(api_.ReleaseStatus, typeof(DOrtReleaseStatus)); OrtCreateSession = (DOrtCreateSession)Marshal.GetDelegateForFunctionPointer(api_.CreateSession, typeof(DOrtCreateSession)); + OrtCreateSessionWithPrepackedWeightsContainer = + (DOrtCreateSessionWithPrepackedWeightsContainer)Marshal.GetDelegateForFunctionPointer(api_.CreateSessionWithPrepackedWeightsContainer, typeof(DOrtCreateSessionWithPrepackedWeightsContainer)); OrtCreateSessionFromArray = (DOrtCreateSessionFromArray)Marshal.GetDelegateForFunctionPointer(api_.CreateSessionFromArray, typeof(DOrtCreateSessionFromArray)); + OrtCreateSessionFromArrayWithPrepackedWeightsContainer = + (DOrtCreateSessionFromArrayWithPrepackedWeightsContainer)Marshal.GetDelegateForFunctionPointer(api_.CreateSessionFromArrayWithPrepackedWeightsContainer, typeof(DOrtCreateSessionFromArrayWithPrepackedWeightsContainer)); OrtRun = (DOrtRun)Marshal.GetDelegateForFunctionPointer(api_.Run, typeof(DOrtRun)); OrtRunWithBinding = (DOrtRunWithBinding)Marshal.GetDelegateForFunctionPointer(api_.RunWithBinding, typeof(DOrtRunWithBinding)); OrtSessionGetInputCount = (DOrtSessionGetInputCount)Marshal.GetDelegateForFunctionPointer(api_.SessionGetInputCount, typeof(DOrtSessionGetInputCount)); @@ -335,6 +350,10 @@ static NativeMethods() OrtGetAvailableProviders = (DOrtGetAvailableProviders)Marshal.GetDelegateForFunctionPointer(api_.GetAvailableProviders, typeof(DOrtGetAvailableProviders)); OrtReleaseAvailableProviders = (DOrtReleaseAvailableProviders)Marshal.GetDelegateForFunctionPointer(api_.ReleaseAvailableProviders, typeof(DOrtReleaseAvailableProviders)); + + OrtCreatePrepackedWeightsContainer = (DOrtCreatePrepackedWeightsContainer)Marshal.GetDelegateForFunctionPointer(api_.CreatePrepackedWeightsContainer, typeof(DOrtCreatePrepackedWeightsContainer)); + OrtReleasePrepackedWeightsContainer = (DOrtReleasePrepackedWeightsContainer)Marshal.GetDelegateForFunctionPointer(api_.ReleasePrepackedWeightsContainer, typeof(DOrtReleasePrepackedWeightsContainer)); + } [DllImport(nativeLib, CharSet = charSet)] @@ -381,14 +400,48 @@ static NativeMethods() out IntPtr /**/ session); public static DOrtCreateSession OrtCreateSession; + /// + /// Creates an instance of OrtSession with provided parameters + /// + /// Native OrtEnv instance + /// UTF-8 bytes corresponding to model string path + /// Native SessionOptions instance + /// Native OrtPrepackedWeightsContainer instance + /// (Output) Created native OrtSession instance + public delegate IntPtr /* OrtStatus* */DOrtCreateSessionWithPrepackedWeightsContainer( + IntPtr /* (OrtEnv*) */ environment, + byte[] modelPath, + IntPtr /* (OrtSessionOptions*) */sessionOptions, + IntPtr /* (OrtPrepackedWeightsContainer*) */prepackedWeightsContainer, + out IntPtr /* (OrtSession**) */ session); + public static DOrtCreateSessionWithPrepackedWeightsContainer OrtCreateSessionWithPrepackedWeightsContainer; + public delegate IntPtr /* OrtStatus* */DOrtCreateSessionFromArray( IntPtr /* (OrtEnv*) */ environment, byte[] modelData, UIntPtr modelSize, - IntPtr /* (OrtSessionOptions*) */sessionOptions, + IntPtr /* (OrtSessionOptions*) */ sessionOptions, out IntPtr /**/ session); public static DOrtCreateSessionFromArray OrtCreateSessionFromArray; + /// + /// Creates an instance of OrtSession with provided parameters + /// + /// Native OrtEnv instance + /// Byte array correspoonding to the model + /// Size of the model in bytes + /// Native SessionOptions instance + /// Native OrtPrepackedWeightsContainer instance + /// (Output) Created native OrtSession instance + public delegate IntPtr /* OrtStatus* */DOrtCreateSessionFromArrayWithPrepackedWeightsContainer( + IntPtr /* (OrtEnv*) */ environment, + byte[] /* (void*) */ modelData, + UIntPtr /* (size_t) */ modelSize, + IntPtr /* (OrtSessionOptions*) */ sessionOptions, + IntPtr /* (OrtPrepackedWeightsContainer*) */prepackedWeightsContainer, + out IntPtr /* (OrtSession**) */ session); + public static DOrtCreateSessionFromArrayWithPrepackedWeightsContainer OrtCreateSessionFromArrayWithPrepackedWeightsContainer; + public delegate IntPtr /*(ONNStatus*)*/ DOrtRun( IntPtr /*(OrtSession*)*/ session, IntPtr /*(OrtSessionRunOptions*)*/ runOptions, // can be null to use the default options @@ -632,7 +685,6 @@ IntPtr[] outputValues /* An array of output value pointers. Array must be alloca IntPtr /*(const char*)*/ name, IntPtr /*(OrtValue*)*/ ortValue); public static DOrtAddInitializer OrtAddInitializer; - #endregion #region RunOptions API @@ -1161,6 +1213,20 @@ IntPtr[] outputValues /* An array of output value pointers. Array must be alloca public delegate IntPtr /* (OrtStatus*) */ DOrtReleaseAvailableProviders(IntPtr /* (char**) */ providers, int /* (int) */ numProviders); public static DOrtReleaseAvailableProviders OrtReleaseAvailableProviders; + + /// + /// Create an instance of PrepackedWeightsContainer + /// + /// (output) Created native OrtPrepackedWeightsContainer instance + public delegate IntPtr /*(OrtStatus*)*/ DOrtCreatePrepackedWeightsContainer(out IntPtr /*(OrtPrepackedWeightsContainer**)*/ prepackedWeightsContainer); + public static DOrtCreatePrepackedWeightsContainer OrtCreatePrepackedWeightsContainer; + + /// + /// Destroy an instance of PrepackedWeightsContainer + /// + /// Native OrtPrepackedWeightsContainer instance to be destroyed + public delegate void DOrtReleasePrepackedWeightsContainer(IntPtr /*(OrtPrepackedWeightsContainer*)*/ prepackedWeightsContainer); + public static DOrtReleasePrepackedWeightsContainer OrtReleasePrepackedWeightsContainer; #endregion public static byte[] GetPlatformSerializedString(string str) diff --git a/csharp/src/Microsoft.ML.OnnxRuntime/PrepackedWeightsContainer.cs b/csharp/src/Microsoft.ML.OnnxRuntime/PrepackedWeightsContainer.cs new file mode 100644 index 0000000000000..bfeebe411b6b5 --- /dev/null +++ b/csharp/src/Microsoft.ML.OnnxRuntime/PrepackedWeightsContainer.cs @@ -0,0 +1,53 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +using System; +using System.Runtime.InteropServices; + +namespace Microsoft.ML.OnnxRuntime +{ + + /// + /// This class holds pre-packed weights of shared initializers to be shared across sessions using these initializers + /// and thereby provide memory savings by sharing the same pre-packed versions of these shared initializers + /// + public class PrePackedWeightsContainer : SafeHandle + { + + /// + /// Constructs an empty PrePackedWeightsContainer + /// + public PrePackedWeightsContainer() + : base(IntPtr.Zero, true) + { + NativeApiStatus.VerifySuccess(NativeMethods.OrtCreatePrepackedWeightsContainer(out handle)); + } + + /// + /// Internal accessor to call native methods + /// + internal IntPtr Pointer { get { return handle; } } + + /// + /// Overrides SafeHandle.IsInvalid + /// + /// returns true if handle is equal to Zero + public override bool IsInvalid { get { return handle == IntPtr.Zero; } } + + #region SafeHandle + /// + /// Overrides SafeHandle.ReleaseHandle() to deallocate + /// a chunk of memory using the specified allocator. + /// + /// always returns true + protected override bool ReleaseHandle() + { + NativeMethods.OrtReleasePrepackedWeightsContainer(handle); + handle = IntPtr.Zero; + return true; + } + + #endregion + } + +} diff --git a/csharp/test/Microsoft.ML.OnnxRuntime.Tests/InferenceTest.cs b/csharp/test/Microsoft.ML.OnnxRuntime.Tests/InferenceTest.cs index edf2a299dd2a1..17546cbd21f12 100644 --- a/csharp/test/Microsoft.ML.OnnxRuntime.Tests/InferenceTest.cs +++ b/csharp/test/Microsoft.ML.OnnxRuntime.Tests/InferenceTest.cs @@ -342,7 +342,7 @@ private void CanRunInferenceOnAModel(GraphOptimizationLevel graphOptimizationLev // Run inference with outputs pinned from buffers using (var pinnedInputs = new DisposableListTest()) - using(var pinnedOutputs = new DisposableListTest()) + using (var pinnedOutputs = new DisposableListTest()) { var memInfo = OrtMemoryInfo.DefaultInstance; // CPU @@ -367,7 +367,7 @@ private void CanRunInferenceOnAModel(GraphOptimizationLevel graphOptimizationLev longShape = Array.ConvertAll(outputMeta[outputName].Dimensions, d => d); byteSize = longShape.Aggregate(1L, (a, b) => a * b) * sizeof(float); float[] outputBuffer = new float[expectedOutput.Length]; - pinnedOutputs.Add(FixedBufferOnnxValue.CreateFromMemory(memInfo, outputBuffer, + pinnedOutputs.Add(FixedBufferOnnxValue.CreateFromMemory(memInfo, outputBuffer, TensorElementType.Float, longShape, byteSize)); session.Run(inputNames, pinnedInputs, outputNames, pinnedOutputs); @@ -1068,7 +1068,7 @@ private void TestPreTrainedModels(string opset, string modelName) catch (Exception ex) { var msg = $"Opset {opset}, Model {modelName}: ModelFile = {onnxModelFileName} error = {ex.Message}"; - if(ex.Message.Contains("ONNX Runtime only *guarantees* support for models stamped with official released onnx opset versions")) + if (ex.Message.Contains("ONNX Runtime only *guarantees* support for models stamped with official released onnx opset versions")) { // If the exception is thrown because the opset version of the test model is // not supported by ONNXRuntime yet, then ignore the test and proceed. @@ -1138,7 +1138,7 @@ private void UnloadLibrary(IntPtr libraryHandle) { if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows)) { - if(!FreeLibrary(libraryHandle)) + if (!FreeLibrary(libraryHandle)) { throw new Exception("Could not unload the provided shared library using its handle"); } @@ -2225,12 +2225,12 @@ private void TestIOBinding() [Fact] private void TestWeightSharingBetweenSessions() { - string modelPath = Path.Combine(Directory.GetCurrentDirectory(), "mul_1.onnx"); + string modelPath = Path.Combine(Directory.GetCurrentDirectory(), "matmul_1.onnx"); // create initializer to share var ortCpuMemInfo = OrtMemoryInfo.DefaultInstance; - var dims = new long[] { 3, 2 }; - var dataBuffer = new float[] { 1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F }; + var dims = new long[] { 2, 1 }; + var dataBuffer = new float[] { 2.0F, 1.0F }; var dataHandle = GCHandle.Alloc(dataBuffer, GCHandleType.Pinned); try @@ -2244,47 +2244,61 @@ private void TestWeightSharingBetweenSessions() } } var dataBufferNumBytes = (uint)dataBuffer.Length * sizeof(float); - var sharedInitializer = OrtValue.CreateTensorValueWithData(ortCpuMemInfo, Tensors.TensorElementType.Float, - dims, dataHandle.AddrOfPinnedObject(), dataBufferNumBytes); - - SessionOptions options = new SessionOptions(); - options.AddInitializer("W", sharedInitializer); - - float[] expectedOutput = { 1.0F, 4.0F, 9.0F, 16.0F, 25.0F, 36.0F }; - int[] expectedDimensions = { 3, 2 }; - using (var session = new InferenceSession(modelPath, options)) - using (var session2 = new InferenceSession(modelPath, options)) + using (var sharedInitializer = OrtValue.CreateTensorValueWithData(ortCpuMemInfo, Tensors.TensorElementType.Float, + dims, dataHandle.AddrOfPinnedObject(), dataBufferNumBytes)) { - var inputMeta = session.InputMetadata; - var container = new List(); - foreach (var name in inputMeta.Keys) + using (var prepackedWeightsContainer = new PrePackedWeightsContainer()) { - Assert.Equal(typeof(float), inputMeta[name].ElementType); - Assert.True(inputMeta[name].IsTensor); - var tensor = new DenseTensor(dataBuffer, inputMeta[name].Dimensions); - container.Add(NamedOnnxValue.CreateFromTensor(name, tensor)); - } + using (var options = new SessionOptions()) + { + // We want to share initializers between the two sessions + options.AddInitializer("W", sharedInitializer); - ReadOnlySpan expectedOutputDimensions = new int[] { 1, 1000, 1, 1 }; - string[] expectedOutputNames = new string[] { "Y" }; + float[] expectedOutput = { 4.0F, 10.0F, 16.0F }; + int[] expectedDimensions = { 3, 1 }; - // Run inference with named inputs and outputs created with in Run() - using (var results = session.Run(container)) // results is an IReadOnlyList container - { - foreach (var r in results) - { - validateRunResultData(r.AsTensor(), expectedOutput, expectedDimensions); - } - } + // We want the pre-packed weights of the shared initializer to be shared between sessions (memory savings) + // and hence we pass in the 'prepackedWeightsContainer' at session creation time + byte[] modelData = File.ReadAllBytes(modelPath); - // Run inference with named inputs and outputs created with in Run() - using (var results2 = session2.Run(container)) // results is an IReadOnlyList container - { - foreach (var r in results2) - { - validateRunResultData(r.AsTensor(), expectedOutput, expectedDimensions); + // Test both InferenceSession ctors that take PrePackedWeightsContainer instances + using (var session = new InferenceSession(modelPath, options, prepackedWeightsContainer)) + using (var session2 = new InferenceSession(modelData, options, prepackedWeightsContainer)) + { + var inputMeta = session.InputMetadata; + var container = new List(); + + foreach (var name in inputMeta.Keys) + { + Assert.Equal(typeof(float), inputMeta[name].ElementType); + Assert.True(inputMeta[name].IsTensor); + var tensor = new DenseTensor(new float[] { 1, 2, 3, 4, 5, 6 }, inputMeta[name].Dimensions); + container.Add(NamedOnnxValue.CreateFromTensor(name, tensor)); + } + + ReadOnlySpan expectedOutputDimensions = new int[] { 1, 1000, 1, 1 }; + string[] expectedOutputNames = new string[] { "Y" }; + + // Run inference with named inputs and outputs created with in Run() + using (var results = session.Run(container)) // results is an IReadOnlyList container + { + foreach (var r in results) + { + validateRunResultData(r.AsTensor(), expectedOutput, expectedDimensions); + } + } + + // Run inference with named inputs and outputs created with in Run() + using (var results2 = session2.Run(container)) // results is an IReadOnlyList container + { + foreach (var r in results2) + { + validateRunResultData(r.AsTensor(), expectedOutput, expectedDimensions); + } + } + } } } } @@ -2617,7 +2631,7 @@ static NamedOnnxValue CreateNamedOnnxValueFromRawData(string name, byte[] raw { T[] typedArr = new T[rawData.Length / elemWidth]; var typeOf = typeof(T); - if(typeOf == typeof(Float16) || typeOf == typeof(BFloat16)) + if (typeOf == typeof(Float16) || typeOf == typeof(BFloat16)) { using (var memSrcHandle = new Memory(rawData).Pin()) using (var memDstHandle = new Memory(typedArr).Pin()) diff --git a/csharp/testdata/matmul_1.onnx b/csharp/testdata/matmul_1.onnx new file mode 100644 index 0000000000000..257aa3029d76d Binary files /dev/null and b/csharp/testdata/matmul_1.onnx differ diff --git a/include/onnxruntime/core/framework/buffer_deleter.h b/include/onnxruntime/core/framework/buffer_deleter.h new file mode 100644 index 0000000000000..d8322af77c80c --- /dev/null +++ b/include/onnxruntime/core/framework/buffer_deleter.h @@ -0,0 +1,36 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/framework/allocator.h" + +namespace onnxruntime { + +// TODO: Do we need this class or is IAllocator::MakeUniquePtr sufficient/better +class BufferDeleter { + public: + BufferDeleter() : alloc_(nullptr) {} + BufferDeleter(AllocatorPtr alloc) + : alloc_(alloc) {} + + void operator()(void* p) const { + if (alloc_) + alloc_->Free(p); + } + + private: + // TODO: we may need consider the lifetime of alloc carefully + // The alloc_ here is the allocator that used to allocate the buffer + // And need go with the unique_ptr together. If it is using our internal + // allocator, it is ok as our allocators are global managed. But if it + // is provide by user, user need to be very careful about it. + // A weak_ptr may be a choice to reduce the impact, but that require to + // change our current allocator mgr to use shared_ptr. Will revisit it + // later. + AllocatorPtr alloc_; +}; + +using BufferUniquePtr = std::unique_ptr; +using BufferNakedPtr = void*; +} // namespace onnxruntime \ No newline at end of file diff --git a/include/onnxruntime/core/framework/op_kernel.h b/include/onnxruntime/core/framework/op_kernel.h index fdce63cad0edf..ee0464e035472 100644 --- a/include/onnxruntime/core/framework/op_kernel.h +++ b/include/onnxruntime/core/framework/op_kernel.h @@ -5,6 +5,10 @@ #include "boost/mp11.hpp" +// It is safe to include the below header even if SHARED_PROVIDER macro is enabled +// as it doesn't include any pb headers. +#include "core/framework/prepacked_weights_container.h" + #ifndef SHARED_PROVIDER #include #include "core/common/exceptions.h" @@ -48,28 +52,68 @@ class OpKernel { // Override this function to PrePack initialized constant tensor to the format as needed. // For example, MatMul kernel can pack the input B if it is constant like code below. - // Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override { + // Status PrePack(const Tensor& tensor, int input_idx, /*out*/ bool& is_packed, + // /*out*/ PrePackedWeights* prepacked_weight_for_caching, + // AllocatorPtr alloc) override { // is_packed = false; // if (input_idx == 1) { - // this.Pack(tensor, this.buffer_); // is_packed = true; + // this.Pack(tensor, this.buffer_, alloc); + // if (prepacked_weight_for_caching) { + // // LOGIC TO CACHE `this.buffer_` SINCE THE KERNEL DOESN"T OWN THE PACKED WEIGHT + // } // } // return Status::OK(); // } // Please refer to MatMulIntegerToFloatBase for a complete example // @param tensor: The initialized constant tensor // @param input_idx: The input index of the tensor in this kernel + // @param alloc: The kernel's PrePack() method MUST use this allocator for allocating the pre-packed + // weights' buffers. The alloc that the PrePack() method will receive will be either + // the allocator tied to the session if the kernel owns the pre-packed buffer or an + // allocator shared between sessions if the pre-packed buffer is to be shared across sessions + // (i.e.) the kernel does not own the buffer. // @param is_packed: Set it to true if the kernel packed the tensor or to false // The kernel is responsible for keeping the packed data and related metadata if is_packed is true, // and the original initialized constant tensor will be released and not accessible anymore in // the Compute function. - virtual Status PrePack(const Tensor& /*tensor*/, int /*input_idx*/, bool& is_packed) { + // @param prepacked_weights: A PrePackedWeights instance will be provided to the kernel IF the pre-packed weights + // are meant to be stored in a shared container. + + virtual Status + PrePack(const Tensor& /*tensor*/, int /*input_idx*/, AllocatorPtr /*alloc*/, + /*out*/ bool& is_packed, /*out*/ PrePackedWeights* /*prepacked_weights*/) { is_packed = false; return Status::OK(); } + // Override this function to use provided pre-packed weight. + // Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + // int input_idx, + // /*out*/ bool& used_shared_buffers) { + // used_shared_buffers = true; + // this.buffer_ = std::move(prepacked_buffers[0]); + // return Status::OK(); + // } + // Please refer to MatMulIntegerToFloatBase for a complete example + // @param prepacked_buffers: The pre-packed buffers to be used by this kernel for the provided input index + // (Sometimes a single constant initializer may have multiple pre-packed buffers associated + // with it and it upto the kernel developer to store it in any order of their choice in PrePack() + // and must use the same order for retrieval in UseSharedPrePackedBuffers(). + // @param input_idx: The input index of the tensor in this kernel + // @param used_shared_buffers: Boolean flag set by the kernel implementation indicating + // that the provided weight has been used by the kernel. + virtual Status UseSharedPrePackedBuffers(std::vector& /*prepacked_buffers*/, + int /*input_idx*/, + /*out*/ bool& used_shared_buffers) { + used_shared_buffers = false; + return Status::OK(); + } + const OrtMemoryInfo& Allocator(int id, OrtMemType mem_type) const; - const OpKernelInfo& Info() const { return *op_kernel_info_; } + const OpKernelInfo& Info() const { + return *op_kernel_info_; + } private: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(OpKernel); diff --git a/include/onnxruntime/core/framework/tensor.h b/include/onnxruntime/core/framework/tensor.h index 92a952ca383e3..1df19f40cea0c 100644 --- a/include/onnxruntime/core/framework/tensor.h +++ b/include/onnxruntime/core/framework/tensor.h @@ -12,37 +12,13 @@ #include "core/common/common.h" #include "core/framework/allocator.h" #include "core/framework/tensor_shape.h" +#include "core/framework/buffer_deleter.h" #include "onnxruntime_config.h" #include "core/framework/data_types.h" #include "core/framework/data_types_internal.h" namespace onnxruntime { -// TODO: Do we need this class or is IAllocator::MakeUniquePtr sufficient/better -class BufferDeleter { - public: - BufferDeleter() : alloc_(nullptr) {} - BufferDeleter(AllocatorPtr alloc) - : alloc_(alloc) {} - - void operator()(void* p) const { - if (alloc_) - alloc_->Free(p); - } - - private: - // TODO: we may need consider the lifetime of alloc carefully - // The alloc_ here is the allocator that used to allocate the buffer - // And need go with the unique_ptr together. If it is using our internal - // allocator, it is ok as our allocators are global managed. But if it - // is provide by user, user need to be very careful about it. - // A weak_ptr may be a choice to reduce the impact, but that require to - // change our current allocator mgr to use shared_ptr. Will revisit it - // later. - AllocatorPtr alloc_; -}; -using BufferUniquePtr = std::unique_ptr; -using BufferNakedPtr = void*; //TODO:ensure dtype_!=nullptr #ifdef __GNUC__ #pragma GCC diagnostic push diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 78110d84ad89c..927ca7016113d 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -170,6 +170,7 @@ ORT_RUNTIME_CLASS(ModelMetadata); ORT_RUNTIME_CLASS(ThreadPoolParams); ORT_RUNTIME_CLASS(ThreadingOptions); ORT_RUNTIME_CLASS(ArenaCfg); +ORT_RUNTIME_CLASS(PrepackedWeightsContainer); #ifdef _WIN32 typedef _Return_type_success_(return == 0) OrtStatus* OrtStatusPtr; @@ -268,10 +269,13 @@ typedef enum OrtCudnnConvAlgoSearch { typedef struct OrtCUDAProviderOptions { int device_id; // cuda device with id=0 as default device. OrtCudnnConvAlgoSearch cudnn_conv_algo_search; // cudnn conv algo search option - size_t gpu_mem_limit; // default cuda memory limitation to maximum finite value of size_t. - // (will be overridden by "max_mem" value used while creating `arena_cfg` if `arena_cfg` is provided) - int arena_extend_strategy; // default area extend strategy to KNextPowerOfTwo. - // (will be overridden by "arena_extend_strategy" value used while creating `arena_cfg` if `arena_cfg` is provided) + + size_t gpu_mem_limit; // default cuda memory limitation to maximum finite value of size_t. + // (will be overridden by "max_mem" value used while creating `arena_cfg` if `arena_cfg` is provided) + + int arena_extend_strategy; // default area extend strategy to KNextPowerOfTwo. + // (will be overridden by "arena_extend_strategy" value used while creating `arena_cfg` if `arena_cfg` is provided) + int do_copy_in_default_stream; int has_user_compute_stream; void* user_compute_stream; @@ -1316,6 +1320,59 @@ struct OrtApi { */ ORT_API2_STATUS(AddRunConfigEntry, _Inout_ OrtRunOptions* options, _In_z_ const char* config_key, _In_z_ const char* config_value); + + /* + * Creates an OrtPrepackedWeightsContainer instance. + * This container will hold pre-packed buffers of shared initializers for sharing between sessions + * (i.e.) if there are shared initializers that can be shared between sessions, the pre-packed buffers + * of these (if any) may possibly be shared to provide memory footprint savings. Pass this container + * to sessions that you would like to share pre-packed buffers of shared initializers at session + * creation time. + * \out - created OrtPrepackedWeightsContainer instance + */ + ORT_API2_STATUS(CreatePrepackedWeightsContainer, _Outptr_ OrtPrepackedWeightsContainer** out); + + /* + * Release OrtPrepackedWeightsContainer instance + * Note: The OrtPrepackedWeightsContainer instance must not be released until the sessions using it are released + */ + ORT_CLASS_RELEASE(PrepackedWeightsContainer); + + /** + * Same functionality offered by CreateSession() API except that a container that contains + pre-packed weights' buffers is written into/read from by the created session. + This is useful when used in conjunction with the AddInitializer() API which injects + shared initializer info into sessions. Wherever possible, the pre-packed versions of these + shared initializers are cached in this container so that multiple sessions can just re-use + these instead of duplicating these in memory. + * \env - OrtEnv instance instance + * \model_path - model path + * \options - OrtSessionOptions instance + * \prepacked_weights_container - OrtPrepackedWeightsContainer instance + * \out - created session instance + */ + ORT_API2_STATUS(CreateSessionWithPrepackedWeightsContainer, _In_ const OrtEnv* env, _In_ const ORTCHAR_T* model_path, + _In_ const OrtSessionOptions* options, _Inout_ OrtPrepackedWeightsContainer* prepacked_weights_container, + _Outptr_ OrtSession** out); + + /** + * Same functionality offered by CreateSessionFromArray() API except that a container that contains + pre-packed weights' buffers is written into/read from by the created session. + This is useful when used in conjunction with the AddInitializer() API which injects + shared initializer info into sessions. Wherever possible, the pre-packed versions of these + shared initializers are cached in this container so that multiple sessions can just re-use + these instead of duplicating these in memory. + * \env - OrtEnv instance instance + * \model_data - model byte array + * \model_data_length - the size of the model byte array + * \options - OrtSessionOptions instance + * \prepacked_weights_container - OrtPrepackedWeightsContainer instance + * \out - created session instance + */ + ORT_API2_STATUS(CreateSessionFromArrayWithPrepackedWeightsContainer, _In_ const OrtEnv* env, + _In_ const void* model_data, size_t model_data_length, + _In_ const OrtSessionOptions* options, _Inout_ OrtPrepackedWeightsContainer* prepacked_weights_container, + _Outptr_ OrtSession** out); }; /* diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index 3cba4ca776ce4..ffd4a9d6c1681 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -349,6 +349,7 @@ struct ModelMetadata : Base { struct Session : Base { explicit Session(std::nullptr_t) {} Session(Env& env, const ORTCHAR_T* model_path, const SessionOptions& options); + Session(Env& env, const ORTCHAR_T* model_path, const SessionOptions& options, OrtPrepackedWeightsContainer* prepacked_weights_container); Session(Env& env, const void* model_data, size_t model_data_length, const SessionOptions& options); // Run that will allocate the output values diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h index a51070665da1a..d08b890e80bc1 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h @@ -514,6 +514,11 @@ inline Session::Session(Env& env, const ORTCHAR_T* model_path, const SessionOpti ThrowOnError(GetApi().CreateSession(env, model_path, options, &p_)); } +inline Session::Session(Env& env, const ORTCHAR_T* model_path, const SessionOptions& options, + OrtPrepackedWeightsContainer* prepacked_weights_container) { + ThrowOnError(GetApi().CreateSessionWithPrepackedWeightsContainer(env, model_path, options, prepacked_weights_container, &p_)); +} + inline Session::Session(Env& env, const void* model_data, size_t model_data_length, const SessionOptions& options) { ThrowOnError(GetApi().CreateSessionFromArray(env, model_data, model_data_length, options, &p_)); } diff --git a/onnxruntime/contrib_ops/cpu/bert/attention.cc b/onnxruntime/contrib_ops/cpu/bert/attention.cc index 827692ede88c6..d992eb6143c68 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/attention.cc @@ -21,7 +21,14 @@ class Attention : public OpKernel, public AttentionCPUBase { explicit Attention(const OpKernelInfo& info); Status Compute(OpKernelContext* context) const override; - Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; + + Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) override; + + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) override; private: BufferUniquePtr packed_weights_; @@ -205,7 +212,9 @@ Attention::Attention(const OpKernelInfo& info) : OpKernel(info), AttentionCPU } template -Status Attention::PrePack(const Tensor& weights, int input_idx, bool& is_packed) { +Status Attention::PrePack(const Tensor& weights, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) { is_packed = false; if (1 != input_idx) { @@ -236,8 +245,14 @@ Status Attention::PrePack(const Tensor& weights, int input_idx, bool& is_pack } const size_t loop_len = static_cast(3) * num_heads_; - auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); + size_t packed_weights_data_size = packed_weights_size_ * loop_len; // The same size would be computed by AllocArray() below auto* packed_weights_data = static_cast(alloc->AllocArray(packed_weights_size_, loop_len)); + + // Initialize memory to 0 as there could be some padding associated with pre-packed + // buffer memory and we don not want it uninitialized and generate different hashes + // if and when we try to cache this pre-packed buffer for sharing between sessions. + memset(packed_weights_data, 0, packed_weights_data_size); + packed_weights_ = BufferUniquePtr(packed_weights_data, BufferDeleter(alloc)); for (size_t i = 0; i < loop_len; i++) { @@ -246,10 +261,30 @@ Status Attention::PrePack(const Tensor& weights, int input_idx, bool& is_pack weights_data += head_size; } + bool share_prepacked_weights = (prepacked_weights != nullptr); + if (share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(packed_weights_)); + prepacked_weights->buffer_sizes_.push_back(packed_weights_data_size); + } + is_packed = true; return Status::OK(); } +template +Status Attention::UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) { + if (1 != input_idx) { + return Status::OK(); + } + + used_shared_buffers = true; + packed_weights_ = std::move(prepacked_buffers[0]); + + return Status::OK(); +} + template Status Attention::Compute(OpKernelContext* context) const { const Tensor* input = context->Input(0); diff --git a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc index b83ff424a5147..7b8189c32e367 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc @@ -22,7 +22,14 @@ class QAttention : public OpKernel, public AttentionCPUBase { QAttention(const OpKernelInfo& info); Status Compute(OpKernelContext* context) const override; - Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; + + Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + bool& /*out*/ is_packed, + /*out*/ PrePackedWeights* prepacked_weights) override; + + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) override; private: BufferUniquePtr packed_weights_; @@ -49,9 +56,9 @@ template QAttention::QAttention(const OpKernelInfo& info) : OpKernel(info), AttentionCPUBase(info) {} template -Status QAttention::PrePack(const Tensor& weights, int input_idx, bool& is_packed) { - is_packed = false; - +Status QAttention::PrePack(const Tensor& weights, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) { if (1 != input_idx) { return Status::OK(); } @@ -81,8 +88,14 @@ Status QAttention::PrePack(const Tensor& weights, int input_idx, bool& is_pac } const size_t loop_len = 3 * num_heads_; - auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); - auto* packed_weights_data = static_cast(alloc->Alloc(packed_weights_size_ * loop_len)); + size_t packed_weights_data_size = packed_weights_size_ * loop_len; + auto* packed_weights_data = static_cast(alloc->Alloc(packed_weights_data_size)); + + // Initialize memory to 0 as there could be some padding associated with pre-packed + // buffer memory and we don not want it uninitialized and generate different hashes + // if and when we try to cache this pre-packed buffer for sharing between sessions. + memset(packed_weights_data, 0, packed_weights_data_size); + packed_weights_ = BufferUniquePtr(packed_weights_data, BufferDeleter(alloc)); for (size_t i = 0; i < loop_len; i++) { @@ -91,10 +104,30 @@ Status QAttention::PrePack(const Tensor& weights, int input_idx, bool& is_pac weights_data += head_size; } + bool share_prepacked_weights = (prepacked_weights != nullptr); + if (share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(packed_weights_)); + prepacked_weights->buffer_sizes_.push_back(packed_weights_data_size); + } + is_packed = true; return Status::OK(); } +template +Status QAttention::UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) { + if (1 != input_idx) { + return Status::OK(); + } + + used_shared_buffers = true; + packed_weights_ = std::move(prepacked_buffers[0]); + + return Status::OK(); +} + template Status QAttention::Compute(OpKernelContext* context) const { // Input and output shapes: diff --git a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc index 4511a66f53da9..5637461dd723d 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc @@ -10,14 +10,22 @@ using namespace rnn::detail; class DynamicQuantizeLSTM : public OpKernel, public LSTMBase { public: DynamicQuantizeLSTM(const OpKernelInfo& info) : OpKernel(info), LSTMBase(info) {} - Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; + + Status PrePack(const Tensor& tensor, int input_idx, + AllocatorPtr alloc, /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) override; + + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) override; Status Compute(OpKernelContext* context) const override; ~DynamicQuantizeLSTM() override = default; private: - Status TryPackWeights(const Tensor& weights, PackedWeights& packed_weights, bool& is_packed, bool& is_weight_signed); + Status TryPackWeights(const Tensor& weights, PackedWeights& packed_weights, bool& is_packed, + bool& is_weight_signed, AllocatorPtr& alloc); template Status ComputeImpl(OpKernelContext& context) const; @@ -28,7 +36,8 @@ class DynamicQuantizeLSTM : public OpKernel, public LSTMBase { bool is_R_signed_; }; -Status DynamicQuantizeLSTM::TryPackWeights(const Tensor& weights, PackedWeights& packed_weights, bool& is_packed, bool& is_weight_signed) { +Status DynamicQuantizeLSTM::TryPackWeights(const Tensor& weights, PackedWeights& packed_weights, + bool& is_packed, bool& is_weight_signed, AllocatorPtr& alloc) { const auto& shape = weights.Shape(); if (shape.NumDimensions() != 3) { return Status::OK(); @@ -49,9 +58,16 @@ Status DynamicQuantizeLSTM::TryPackWeights(const Tensor& weights, PackedWeights& return Status::OK(); } - auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); - auto* packed_weights_data = alloc->Alloc(SafeInt(packed_weights_size) * num_directions_); + size_t packed_weights_data_size = SafeInt(packed_weights_size) * num_directions_; + auto* packed_weights_data = alloc->Alloc(packed_weights_data_size); + + // Initialize memory to 0 as there could be some padding associated with pre-packed + // buffer memory and we don not want it uninitialized and generate different hashes + // if and when we try to cache this pre-packed buffer for sharing between sessions. + memset(packed_weights_data, 0, packed_weights_data_size); + packed_weights.buffer_ = BufferUniquePtr(packed_weights_data, BufferDeleter(alloc)); + packed_weights.buffer_size_ = packed_weights_data_size; packed_weights.weights_size_ = packed_weights_size; packed_weights.shape_ = shape; @@ -66,13 +82,48 @@ Status DynamicQuantizeLSTM::TryPackWeights(const Tensor& weights, PackedWeights& return Status::OK(); } -Status DynamicQuantizeLSTM::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) { +static void UseSharedPrePackedBuffersImpl(std::vector& prepacked_buffers, + rnn::detail::PackedWeights& packed_tensor) { + packed_tensor.buffer_ = std::move(prepacked_buffers[0]); +} + +Status DynamicQuantizeLSTM::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) { is_packed = false; if (input_idx == 1) { - return TryPackWeights(tensor, packed_W_, is_packed, is_W_signed_); + ORT_RETURN_IF_ERROR(TryPackWeights(tensor, packed_W_, is_packed, is_W_signed_, alloc)); + + bool share_prepacked_weights = (prepacked_weights != nullptr); + if (is_packed && share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(packed_W_.buffer_)); + prepacked_weights->buffer_sizes_.push_back(packed_W_.buffer_size_); + } + } else if (input_idx == 2) { + ORT_RETURN_IF_ERROR(TryPackWeights(tensor, packed_R_, is_packed, is_R_signed_, alloc)); + + bool share_prepacked_weights = (prepacked_weights != nullptr); + if (is_packed && share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(packed_R_.buffer_)); + prepacked_weights->buffer_sizes_.push_back(packed_R_.buffer_size_); + } + } + + return Status::OK(); +} + +Status DynamicQuantizeLSTM::UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) { + used_shared_buffers = false; + + if (input_idx == 1) { + used_shared_buffers = true; + UseSharedPrePackedBuffersImpl(prepacked_buffers, packed_W_); } else if (input_idx == 2) { - return TryPackWeights(tensor, packed_R_, is_packed, is_R_signed_); + used_shared_buffers = true; + UseSharedPrePackedBuffersImpl(prepacked_buffers, packed_R_); } return Status::OK(); diff --git a/onnxruntime/core/framework/prepacked_weights.cc b/onnxruntime/core/framework/prepacked_weights.cc new file mode 100644 index 0000000000000..563c3a8bd5402 --- /dev/null +++ b/onnxruntime/core/framework/prepacked_weights.cc @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "prepacked_weights.h" +#include "core/framework/murmurhash3.h" + +namespace onnxruntime { + +uint64_t PrePackedWeights::GetHash() const { + // Adaptation of the hashing logic of the KernelDef class + + uint32_t hash[4] = {0, 0, 0, 0}; + + auto hash_int8_t_buffer = [&hash](void* data, int len) { MurmurHash3::x86_128(data, len, hash[0], &hash); }; + + ORT_ENFORCE(buffers_.size() == buffer_sizes_.size()); + for (size_t iter = 0; iter < buffers_.size(); ++iter) { + // some pre-packed buffers may be null if they were just "place-holders" occupying an index + // in the "buffers_" vector + if (buffers_[iter].get() != nullptr) { + hash_int8_t_buffer(buffers_[iter].get(), static_cast(buffer_sizes_[iter])); + } + } + + uint64_t returned_hash = hash[0] & 0xfffffff8; // save low 3 bits for hash version info in case we need it in the future + returned_hash |= uint64_t(hash[1]) << 32; + + return returned_hash; +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/framework/prepacked_weights.h b/onnxruntime/core/framework/prepacked_weights.h new file mode 100644 index 0000000000000..7167c53fbcac3 --- /dev/null +++ b/onnxruntime/core/framework/prepacked_weights.h @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include "core/framework/buffer_deleter.h" +#include "core/framework/tensor_shape.h" + +namespace onnxruntime { + +struct PrePackedWeights final { + // Some weights may be associated with multiple pre-packed buffers (e.g.) QLinearConv. + // Hence we hold them in container. It is upto the developer implementing each PrePack() + // method to define what gets stored in which position of the container. + + std::vector> buffers_; // cache pre-packed buffers associated with the kernel + std::vector buffer_sizes_; // cache sizes of pre-packed buffers (in bytes) + + // Produces a hash of the buffers stored in the given instance of this class + uint64_t GetHash() const; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/framework/prepacked_weights_container.cc b/onnxruntime/core/framework/prepacked_weights_container.cc new file mode 100644 index 0000000000000..cc6b45e00f148 --- /dev/null +++ b/onnxruntime/core/framework/prepacked_weights_container.cc @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/framework/prepacked_weights_container.h" +#include "core/framework/allocatormgr.h" + +namespace onnxruntime { + +AllocatorPtr PrepackedWeightsContainer::GetOrCreateAllocator(const std::string& device_name) { + auto iter = allocators_.find(device_name); + + if (iter != allocators_.end()) + return iter->second; + + // Support only CPU based allocators for now. + // as pre-packing is only supported by CPU kernels for now. + if (device_name == CPU) { + // TODO: Investigate benefits of using an arena based allocator + // For now, we go with a non-arena based allocator + AllocatorCreationInfo device_info{[](int) { return std::make_unique(); }, + 0, false}; + auto allocator = CreateAllocator(device_info); + + allocators_[device_name] = allocator; + + return allocator; + + } else { + ORT_THROW("Unsupported device allocator in the context of pre-packed weights caching: ", device_name); + } +} + +const PrePackedWeights& PrepackedWeightsContainer::GetWeight(const std::string& key) const { + // .at() will throw if th key doesn't exist + return prepacked_weights_map_.at(key); +} + +bool PrepackedWeightsContainer::WriteWeight(const std::string& key, PrePackedWeights&& packed_weight) { + auto ret = prepacked_weights_map_.insert({key, std::move(packed_weight)}); + return ret.second; +} + +bool PrepackedWeightsContainer::HasWeight(const std::string& key) const { + return prepacked_weights_map_.find(key) != + prepacked_weights_map_.end(); +} + +size_t PrepackedWeightsContainer::GetNumberOfElements() const { + return prepacked_weights_map_.size(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/framework/prepacked_weights_container.h b/onnxruntime/core/framework/prepacked_weights_container.h new file mode 100644 index 0000000000000..7fe317b6c4317 --- /dev/null +++ b/onnxruntime/core/framework/prepacked_weights_container.h @@ -0,0 +1,69 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include "core/framework/buffer_deleter.h" + +#include "core/framework/allocator.h" +#include "core/platform/ort_mutex.h" +#include "prepacked_weights.h" + +namespace onnxruntime { + +class PrepackedWeightsContainer final { + public: + PrepackedWeightsContainer() { + } + + ~PrepackedWeightsContainer() = default; + + // Returns an allocator keyed by device name. + // If an allocator doesn't exist for that specific device, an allocator + // is created and stored in a member to be returned on subsequent calls. + // Currently, the only supported device is "Cpu". + AllocatorPtr GetOrCreateAllocator(const std::string& device_name); + + // Returns the PrePackedWeights instance pertaining to the provided key. + // The key is : op_type + "+" + hash_of_prepacked_buffers_in_the_PrepackedWeights_instance. + // Throws an exception if the key doesn't exist + const PrePackedWeights& GetWeight(const std::string& key) const; + + // Writes the PrePackedWeights instance pertaining to the provided key. + // The key is : op_type + "+" + hash_of_prepacked_buffers_in_the_PrepackedWeights_instance. + // Returns a boolean indicating if the insertion took place. + bool WriteWeight(const std::string& key, PrePackedWeights&& packed_weight); + + // Returns a boolean indicating if there is a PrePackedWeights instance + // pertaining to the provided key. + // The key is : op_type + "+" + hash_of_prepacked_buffers_in_the_PrepackedWeights_instance. + bool HasWeight(const std::string& key) const; + + // Returns the number of elements in the container + size_t GetNumberOfElements() const; + + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(PrepackedWeightsContainer); + + // Resource to be acquired by the method that is going to invoke calls to the kernels' + // PrePack() methods and does the read/write into the pre-packed weights' container. + // We only want to invoke PrePack() on a kernel that doesn't have a cached version + // of its pre-packed weight. + OrtMutex mutex_; + + // Define allocators ahead of the container containing tensors because the allocators + // needs to destructed after the container containing the pre-packed cached tensors + // because the Tensor buffers will be de-allocated using these allocators + std::unordered_map allocators_; + + // This is an unordered map that holds a mapping between a composite key + // to PrePackedWeights instances. + // The key is : op_type + "+" + hash_of_prepacked_buffers_in_the_PrepackedWeights_instance. + std::unordered_map prepacked_weights_map_; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc index 1683e53076c79..000d604961159 100644 --- a/onnxruntime/core/framework/session_state.cc +++ b/onnxruntime/core/framework/session_state.cc @@ -5,6 +5,7 @@ #include +#include "core/platform/ort_mutex.h" #include "core/common/logging/logging.h" #include "core/common/safeint.h" #include "core/flatbuffers/schema/ort.fbs.h" @@ -258,45 +259,167 @@ void SessionState::CleanInitializedTensorsFromGraph() { graph_.CleanAllInitializedTensors(); } -Status SessionState::PrepackConstantInitializedTensors(std::unordered_map& constant_initializers_use_count) { - for (auto& node : GetGraphViewer().Nodes()) { - auto kernel = GetMutableKernel(node.Index()); - int input_idx = 0; - for (auto& input_def : node.InputDefs()) { - if (input_def->Exists()) { - const std::string& input_name = input_def->Name(); - SessionState* st = this; - // subgraph can use the value from outer scope, - // so it needs to check if current node uses constant initialized tensor from current and outer graphs - do { - int ort_value_idx; - if (st->GetOrtValueNameIdxMap().GetIdx(input_name, ort_value_idx).IsOK()) { - std::unordered_map& constant_initialized_tensors = st->constant_initialized_tensors_; - if (constant_initialized_tensors.count(ort_value_idx)) { - bool is_packed = false; - const Tensor& const_initialized_tensor = constant_initialized_tensors[ort_value_idx].Get(); - ORT_RETURN_IF_ERROR(kernel->PrePack(const_initialized_tensor, input_idx, is_packed)); - if (is_packed && constant_initializers_use_count.count(input_name) && --constant_initializers_use_count[input_name] == 0) { - // release the constant initialized tensor - st->initialized_tensors_.erase(ort_value_idx); - constant_initialized_tensors.erase(ort_value_idx); +static Status KernelUseSharedPrePackedBuffers(OpKernel& kernel, int input_idx, + const PrePackedWeights& prepacked_weights, + const std::string& node_name) { + std::vector shared_prepacked_buffers; + shared_prepacked_buffers.reserve(4); // Unlikely to see more than 4 prepacked buffers per initializer + + for (const auto& prepacked_buffer : prepacked_weights.buffers_) { + // BufferDeleter is nullptr because the kernel should not delete the shared buffer - it can only use it + shared_prepacked_buffers.emplace_back(prepacked_buffer.get(), BufferDeleter(nullptr)); + } + + bool used_shared_buffers = false; + ORT_RETURN_IF_ERROR(kernel.UseSharedPrePackedBuffers(shared_prepacked_buffers, input_idx, used_shared_buffers)); + + // BUG CHECK: Ensure that the kernel used the provided shared buffers + // Mostly a debug check to ensure that the kernel has an overridden implementation of the + // base UseSharedPrePackedBuffers() which is basically a no-op. + if (!used_shared_buffers) + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "The kernel corresponding to the node ", node_name, + " doesn't have an implementation that can consume provided pre-packed weights"); + + return Status::OK(); +} + +static std::string GenerateKeyForPrepackedWeightsMap(const std::string& op_type, + const PrePackedWeights& pre_packed_weights) { + std::ostringstream ss_1; + ss_1 << op_type; + ss_1 << "+"; + ss_1 << std::to_string(pre_packed_weights.GetHash()); + + return ss_1.str(); +} + +Status SessionState::PrepackConstantInitializedTensors(std::unordered_map& constant_initializers_use_count, + const std::unordered_map& initializers_to_share_map) { + auto prepacked_constant_weights = [this, &constant_initializers_use_count, &initializers_to_share_map]( + bool should_cache_prepacked_weights_for_shared_initializers) -> Status { + for (auto& node : GetGraphViewer().Nodes()) { + auto kernel = GetMutableKernel(node.Index()); + int input_idx = 0; + for (auto& input_def : node.InputDefs()) { + if (input_def->Exists()) { + const std::string& input_name = input_def->Name(); + SessionState* st = this; + // subgraph can use the value from outer scope, + // so it needs to check if current node uses constant initialized tensor from current and outer graphs + do { + int ort_value_idx; + if (st->GetOrtValueNameIdxMap().GetIdx(input_name, ort_value_idx).IsOK()) { + std::unordered_map& constant_initialized_tensors = st->constant_initialized_tensors_; + + if (constant_initialized_tensors.count(ort_value_idx)) { + bool is_packed = false; + const Tensor& const_initialized_tensor = constant_initialized_tensors[ort_value_idx].Get(); + + auto iter = initializers_to_share_map.find(input_name); + bool is_shared_initializer = (iter != initializers_to_share_map.end()); + + // Caching pre-packed weights is limited to shared initializers associated with the CPU EP for now + if (is_shared_initializer && should_cache_prepacked_weights_for_shared_initializers && + node.GetExecutionProviderType() == kCpuExecutionProvider) { // caching of pre-packed weights' turned ON + + AllocatorPtr allocator_for_caching = prepacked_weights_container_->GetOrCreateAllocator(CPU); + ORT_ENFORCE(allocator_for_caching.get() != nullptr); + + PrePackedWeights weights_to_be_filled_in; + // The reason we invoke PrePack() before looking into the container for any pre-packed weight + // cached by another instance of the same op_type (for the same constant initializer) is because + // to truly know if we can use a cached pre-packed weight, we would have to compare the cached pre-packed + // weight with the pre-packed weight generated by this instance of the same op_type because other static + // properties of the node like node attributes could play a role in the pre-packed weights' contents. + ORT_RETURN_IF_ERROR(kernel->PrePack(const_initialized_tensor, input_idx, allocator_for_caching, + is_packed, + &weights_to_be_filled_in)); + + if (is_packed) { + // BUG CHECK: Ensure that the kernel has filled in the pre-packed weight to be cached if the weight was pre-packed + ORT_ENFORCE(weights_to_be_filled_in.buffers_.size() > 0, "The kernel corresponding to the node ", node.Name(), + " doesn't have an implementation that can cache computed pre-packed weights"); + + const auto& op_type = node.OpType(); + + // Sanity check + // TODO: Check if some version of the ONNX IR allows op_type to be empty + ORT_ENFORCE(!op_type.empty(), "The op type of a node cannot be empty"); + + // The key for the pre-packed weights container lookup is the op_type + hash of the prepacked-weight + // that we just got by invoking PrePack() on this kernel. + + const std::string& prepacked_weights_container_key = GenerateKeyForPrepackedWeightsMap(op_type, + weights_to_be_filled_in); + + bool container_contains_packed_weight = prepacked_weights_container_->HasWeight(prepacked_weights_container_key); + + if (container_contains_packed_weight) { + LOGS(logger_, INFO) << "Using cached version of pre-packed weight for constant initializer: " << input_name + << " used in the node: " << node.Name() << " which is of op type: " << node.OpType(); + + ORT_RETURN_IF_ERROR(KernelUseSharedPrePackedBuffers(*kernel, input_idx, + prepacked_weights_container_->GetWeight(prepacked_weights_container_key), + node.Name())); + + ++used_shared_pre_packed_weights_counter_; + } else { // container doesn't contain the pre-packed weight - so write into it for sharing across kernel instances + + if (!prepacked_weights_container_->WriteWeight(prepacked_weights_container_key, std::move(weights_to_be_filled_in))) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unable to write the provided PrePackedWeights instance into the container"); + } + + ORT_RETURN_IF_ERROR(KernelUseSharedPrePackedBuffers(*kernel, input_idx, + prepacked_weights_container_->GetWeight(prepacked_weights_container_key), + node.Name())); + } + } + + } else { // caching of pre-packed weights' turned OFF + AllocatorPtr session_cpu_alloc = kernel->Info().GetAllocator(0, OrtMemType::OrtMemTypeDefault); + ORT_RETURN_IF_ERROR(kernel->PrePack(const_initialized_tensor, input_idx, + session_cpu_alloc, // use allocator tied to this session + is_packed, + nullptr // no caching required + )); + } + if (is_packed) { + ++number_of_prepacks_counter_; + + if (constant_initializers_use_count.count(input_name) && --constant_initializers_use_count[input_name] == 0) { + // release the constant initialized tensor + st->initialized_tensors_.erase(ort_value_idx); + constant_initialized_tensors.erase(ort_value_idx); + } + } + } + // stop searching in 2 cases: + // 1. value is not from OuterScope + // 2. value is from OuterScope and the current OuterScope has the value + if (st != this || !st->graph_.IsOuterScopeValue(input_name)) { + break; } } - // stop searching in 2 cases: - // 1. value is not from OuterScope - // 2. value is from OuterScope and the current OuterScope has the value - if (st != this || !st->graph_.IsOuterScopeValue(input_name)) { - break; - } - } - st = st->Parent(); - } while (st); + st = st->Parent(); + } while (st); + } + input_idx++; } - input_idx++; } - } - return Status::OK(); + return Status::OK(); + }; + + bool should_cache_prepacked_weights_for_shared_initializers = (prepacked_weights_container_ != nullptr); + + if (should_cache_prepacked_weights_for_shared_initializers) { + // serialize calls to the method that looks up the container, calls UseCachedPrePackedWeight/PrePack + // and writes pre-packed weights to the container + std::lock_guard l(prepacked_weights_container_->mutex_); + return prepacked_constant_weights(true); + } else { + return prepacked_constant_weights(false); + } } static int64_t CalculateMemoryPatternsKey(const std::vector>& shapes) { @@ -1065,7 +1188,8 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string& constant_initializers_use_count); + Status PrepackConstantInitializedTensors(std::unordered_map& constant_initializers_use_count, + const std::unordered_map& initializers_to_share_map); SessionState* GetMutableSubgraphSessionState(onnxruntime::NodeIndex index, const std::string& attribute_name); @@ -450,6 +462,12 @@ class SessionState { std::unique_ptr node_index_info_; std::multimap> cached_feeds_fetches_managers_; + // Container to store pre-packed weights to share between sessions. + // The life-cycle of the cache itself is maintained by the user and the user will ensure + // the cache is valid until any session reliant on it is still in scope. + // prepacked_weights_container_ can be nullptr if no caching is required for prepacked weights + PrepackedWeightsContainer* const prepacked_weights_container_{}; + #if !defined(ORT_MINIMAL_BUILD) std::map, std::unordered_set> to_be_executed_nodes_; #endif @@ -466,6 +484,14 @@ class SessionState { graph_id_ = p->next_graph_id_++; } #endif + + // Counter for number of times pre-packing of weights was performed across kernels + // part the model + size_t number_of_prepacks_counter_ = 0; + + // Counter for number of times a shared version of the pre-packed weight corresponding to + // a constant initialized weight was used by the session state + size_t used_shared_pre_packed_weights_counter_ = 0; }; } // namespace onnxruntime diff --git a/onnxruntime/core/mlas/lib/qgemm.cpp b/onnxruntime/core/mlas/lib/qgemm.cpp index 6015dc0f8680a..b4714b8130dc2 100644 --- a/onnxruntime/core/mlas/lib/qgemm.cpp +++ b/onnxruntime/core/mlas/lib/qgemm.cpp @@ -21,28 +21,20 @@ Module Name: // Quantized integer matrix/matrix dispatch structure. // -typedef -void -(MLAS_GEMM_U8X8_OPERATION)( - const MLAS_GEMM_U8X8_SHAPE_PARAMS* Shape, - const MLAS_GEMM_U8X8_DATA_PARAMS* Data, - const size_t RangeStartM, - const size_t RangeCountM, - const size_t RangeStartN, - const size_t RangeCountN - ); - -typedef -void -(MLAS_GEMM_U8X8_COPY_PACKB_ROUTINE)( - uint8_t* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer, - bool BIsSigned - ); +typedef void(MLAS_GEMM_U8X8_OPERATION)(const MLAS_GEMM_U8X8_SHAPE_PARAMS* Shape, + const MLAS_GEMM_U8X8_DATA_PARAMS* Data, + const size_t RangeStartM, + const size_t RangeCountM, + const size_t RangeStartN, + const size_t RangeCountN); + +typedef void(MLAS_GEMM_U8X8_COPY_PACKB_ROUTINE)(uint8_t* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer, + bool BIsSigned); struct MLAS_GEMM_U8X8_DISPATCH { MLAS_GEMM_U8X8_OPERATION* Operation; @@ -53,9 +45,7 @@ struct MLAS_GEMM_U8X8_DISPATCH { }; const MLAS_GEMM_U8X8_DISPATCH* -MlasGemmU8X8GetDispatch( - bool BIsSigned - ) +MlasGemmU8X8GetDispatch(bool BIsSigned) { const MLAS_GEMM_U8X8_DISPATCH* GemmU8X8Dispatch; @@ -102,12 +92,7 @@ struct MLAS_GEMM_U8X8_STRIDES { }; void -MlasGemmU8X8ScaleSumBuffer( - int32_t* Output, - const int32_t* Input, - size_t N, - int32_t Scale - ) +MlasGemmU8X8ScaleSumBuffer(int32_t* Output, const int32_t* Input, size_t N, int32_t Scale) { for (size_t n = 0; n < N; n++) { Output[n] = Input[n] * Scale; @@ -116,27 +101,20 @@ MlasGemmU8X8ScaleSumBuffer( MLAS_FORCEINLINE void -MlasGemmU8X8ScaleSumBuffer( - int32_t* SumBuffer, - size_t N, - int32_t Scale - ) +MlasGemmU8X8ScaleSumBuffer(int32_t* SumBuffer, size_t N, int32_t Scale) { return MlasGemmU8X8ScaleSumBuffer(SumBuffer, SumBuffer, N, Scale); } -template -MLAS_FORCEINLINE -bool -MlasGemmU8X8TryGemvKernel( - const uint8_t* A, - const uint8_t* B, - size_t ldb, - int32_t* C, - size_t CountK, - size_t CountN, - bool BIsSigned - ) +template +MLAS_FORCEINLINE bool +MlasGemmU8X8TryGemvKernel(const uint8_t* A, + const uint8_t* B, + size_t ldb, + int32_t* C, + size_t CountK, + size_t CountN, + bool BIsSigned) { MLAS_UNREFERENCED_PARAMETER(A); MLAS_UNREFERENCED_PARAMETER(B); @@ -149,32 +127,25 @@ MlasGemmU8X8TryGemvKernel( return false; } -template +template int32_t -MlasGemmU8X8FixupZeroPointB( - int32_t ZeroPointB, - bool BIsSigned - ) +MlasGemmU8X8FixupZeroPointB(int32_t ZeroPointB, bool BIsSigned) { MLAS_UNREFERENCED_PARAMETER(BIsSigned); return ZeroPointB; } -template -MLAS_FORCEINLINE -void -MlasGemmU8X8FixupZeroPointB( - const uint8_t* PackedZeroPointB, - int32_t* ZeroPointBBuffer, - size_t N, - bool BIsSigned - ) +template +MLAS_FORCEINLINE void +MlasGemmU8X8FixupZeroPointB(const uint8_t* PackedZeroPointB, + int32_t* ZeroPointBBuffer, + size_t N, + bool BIsSigned) { int32_t ZeroPointB; for (size_t n = 0; n < N; n++) { - ZeroPointB = typename KernelType::OffsetBType(PackedZeroPointB[n]); ZeroPointB = MlasGemmU8X8FixupZeroPointB(ZeroPointB, BIsSigned); @@ -186,62 +157,55 @@ MlasGemmU8X8FixupZeroPointB( // against tools that check for uninitialized data usage. // - size_t AlignedN = (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1); + size_t AlignedN = + (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1); for (size_t n = N; n < AlignedN; n++) { ZeroPointBBuffer[n] = 0; } } -template +template void -MlasGemmU8X8CopyPackA( - typename KernelType::PackedAType* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumBuffer - ); - -template +MlasGemmU8X8CopyPackA(typename KernelType::PackedAType* D, + const uint8_t* A, + size_t lda, + size_t CountM, + size_t CountK, + int32_t* RowSumBuffer); + +template void -MlasGemmU8X8CopyPackB( - typename KernelType::PackedBType* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer, - bool BIsSigned - ); - -template +MlasGemmU8X8CopyPackB(typename KernelType::PackedBType* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer, + bool BIsSigned); + +template size_t -MlasGemmU8X8Kernel( - const typename KernelType::PackedAType* A, - const typename KernelType::PackedBType* B, - int32_t* C, - size_t PackedCountK, - size_t CountM, - size_t CountN, - size_t ldc, - const int32_t* RowSumBuffer, - const int32_t* ColumnSumBuffer, - const int32_t* ZeroPointB, - bool ZeroMode - ); - -template +MlasGemmU8X8Kernel(const typename KernelType::PackedAType* A, + const typename KernelType::PackedBType* B, + int32_t* C, + size_t PackedCountK, + size_t CountM, + size_t CountN, + size_t ldc, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, + const int32_t* ZeroPointB, + bool ZeroMode); + +template void -MlasGemmU8X8Operation( - const MLAS_GEMM_U8X8_SHAPE_PARAMS* Shape, - const MLAS_GEMM_U8X8_DATA_PARAMS* Data, - const size_t RangeStartM, - const size_t RangeCountM, - const size_t RangeStartN, - const size_t RangeCountN - ) +MlasGemmU8X8Operation(const MLAS_GEMM_U8X8_SHAPE_PARAMS* Shape, + const MLAS_GEMM_U8X8_DATA_PARAMS* Data, + const size_t RangeStartM, + const size_t RangeCountM, + const size_t RangeStartN, + const size_t RangeCountN) /*++ Routine Description: @@ -287,8 +251,8 @@ Return Value: const uint8_t* A = Data->A + RangeStartM * lda; const uint8_t* B = (const uint8_t*)Data->B + RangeStartN; int32_t* C = Data->C + RangeStartM * ldc + RangeStartN; - const uint8_t* PackedZeroPointB = Data->PerColumnZeroPoints ? - Data->ZeroPointB + RangeStartN : nullptr; + const uint8_t* PackedZeroPointB = + Data->PerColumnZeroPoints ? Data->ZeroPointB + RangeStartN : nullptr; int32_t ZeroPointA = Data->ZeroPointA; int32_t ZeroPointB = typename KernelType::OffsetBType(*Data->ZeroPointB); @@ -297,9 +261,8 @@ Return Value: // Try to use a GEMV kernel if supported by this kernel type. // - if ((RangeCountM == 1) && - (ZeroPointA == 0) && (PackedZeroPointB == nullptr) && (ZeroPointB == 0) && - (Data->OutputProcessor == nullptr)) { + if ((RangeCountM == 1) && (ZeroPointA == 0) && (PackedZeroPointB == nullptr) && + (ZeroPointB == 0) && (Data->OutputProcessor == nullptr)) { if (MlasGemmU8X8TryGemvKernel(A, B, ldb, C, K, RangeCountN, Shape->BIsSigned)) { return; } @@ -320,7 +283,6 @@ Return Value: size_t CountK; for (size_t k = 0; k < K; k += CountK) { - CountK = std::min(K - k, Strides.K); const size_t PackedCountK = (CountK + KernelType::PackedK - 1) / KernelType::PackedK; @@ -332,7 +294,6 @@ Return Value: size_t CountN; for (size_t n = 0; n < RangeCountN; n += CountN) { - CountN = std::min(RangeCountN - n, Strides.N); // @@ -341,25 +302,16 @@ Return Value: // if (PackedZeroPointB != nullptr) { - MlasGemmU8X8FixupZeroPointB( - PackedZeroPointB + n, - ZeroPointBBuffer, - CountN, - Shape->BIsSigned); + MlasGemmU8X8FixupZeroPointB(PackedZeroPointB + n, ZeroPointBBuffer, + CountN, Shape->BIsSigned); } // // Copy a panel of matrix B to a local packed buffer. // - MlasGemmU8X8CopyPackB( - PanelB, - B + n, - ldb, - CountN, - CountK, - ColumnSumBuffer, - Shape->BIsSigned); + MlasGemmU8X8CopyPackB(PanelB, B + n, ldb, CountN, CountK, ColumnSumBuffer, + Shape->BIsSigned); MlasGemmU8X8ScaleSumBuffer(ColumnSumBuffer, CountN, -ZeroPointA); @@ -371,20 +323,14 @@ Return Value: size_t CountM; for (size_t m = 0; m < RangeCountM; m += CountM) { - CountM = std::min(RangeCountM - m, Strides.M); // // Copy a panel of matrix A to a local packed buffer. // - MlasGemmU8X8CopyPackA( - PanelA, - A + m * lda, - lda, - CountM, - CountK, - RowSumBuffer); + MlasGemmU8X8CopyPackA(PanelA, A + m * lda, lda, CountM, CountK, + RowSumBuffer); // // Apply the global depth value constant without the ZeroPointB scaling from: @@ -421,28 +367,15 @@ Return Value: bool PostProcess = (k + CountK == K); while (RowsRemaining > 0) { - size_t RowsHandled = MlasGemmU8X8Kernel( - pa, - PanelB, - c, - PackedCountK, - RowsRemaining, - CountN, - ldc, - RowSums, - ColumnSumBuffer, - (PackedZeroPointB != nullptr) ? ZeroPointBBuffer : nullptr, + pa, PanelB, c, PackedCountK, RowsRemaining, CountN, ldc, RowSums, + ColumnSumBuffer, (PackedZeroPointB != nullptr) ? ZeroPointBBuffer : nullptr, ZeroMode); if (PostProcess && Data->OutputProcessor != nullptr) { Data->OutputProcessor->Process( - Data->C, - RangeStartM + m + CountM - RowsRemaining, - RangeStartN + n, - RowsHandled, - CountN, - Data->ldc); + Data->C, RangeStartM + m + CountM - RowsRemaining, RangeStartN + n, + RowsHandled, CountN, Data->ldc); } c += ldc * RowsHandled; @@ -458,16 +391,14 @@ Return Value: } } -template +template void -MlasGemmU8X8PackedOperation( - const MLAS_GEMM_U8X8_SHAPE_PARAMS* Shape, - const MLAS_GEMM_U8X8_DATA_PARAMS* Data, - const size_t RangeStartM, - const size_t RangeCountM, - const size_t RangeStartN, - const size_t RangeCountN - ) +MlasGemmU8X8PackedOperation(const MLAS_GEMM_U8X8_SHAPE_PARAMS* Shape, + const MLAS_GEMM_U8X8_DATA_PARAMS* Data, + const size_t RangeStartM, + const size_t RangeCountM, + const size_t RangeStartN, + const size_t RangeCountN) /*++ Routine Description: @@ -511,8 +442,8 @@ Return Value: const uint8_t* A = Data->A + RangeStartM * lda; const uint8_t* PackedB = (const uint8_t*)Data->B; int32_t* C = Data->C + RangeStartM * ldc + RangeStartN; - const uint8_t* PackedZeroPointB = Data->PerColumnZeroPoints ? - Data->ZeroPointB + RangeStartN : nullptr; + const uint8_t* PackedZeroPointB = + Data->PerColumnZeroPoints ? Data->ZeroPointB + RangeStartN : nullptr; int32_t ZeroPointA = Data->ZeroPointA; int32_t ZeroPointB = typename KernelType::OffsetBType(*Data->ZeroPointB); @@ -542,7 +473,6 @@ Return Value: size_t CountK; for (size_t k = 0; k < K; k += CountK) { - CountK = std::min(K - k, Strides.K); const size_t PackedCountK = (CountK + KernelType::PackedK - 1) / KernelType::PackedK; @@ -558,12 +488,11 @@ Return Value: size_t CountN; for (size_t n = 0; n < RangeCountN; n += CountN) { - CountN = std::min(RangeCountN - n, Strides.N); if (k == 0) { - MlasGemmU8X8ScaleSumBuffer(ColumnSumBuffer, PackedColumnSumBuffer + n, - CountN, -ZeroPointA); + MlasGemmU8X8ScaleSumBuffer(ColumnSumBuffer, PackedColumnSumBuffer + n, CountN, + -ZeroPointA); } // @@ -572,37 +501,27 @@ Return Value: // if (PackedZeroPointB != nullptr) { - MlasGemmU8X8FixupZeroPointB( - PackedZeroPointB + n, - ZeroPointBBuffer, - CountN, - Shape->BIsSigned); + MlasGemmU8X8FixupZeroPointB(PackedZeroPointB + n, ZeroPointBBuffer, + CountN, Shape->BIsSigned); } // // Step through each slice of matrix A along the M dimension. // - const uint8_t* b = PackedB + (RangeStartN + n) * - KernelType::PackedK * PackedCountK; + const uint8_t* b = PackedB + (RangeStartN + n) * KernelType::PackedK * PackedCountK; int32_t* c = C + n; size_t CountM; for (size_t m = 0; m < RangeCountM; m += CountM) { - CountM = std::min(RangeCountM - m, Strides.M); // // Copy a panel of matrix A to a local packed buffer. // - MlasGemmU8X8CopyPackA( - PanelA, - A + m * lda, - lda, - CountM, - CountK, - RowSumBuffer); + MlasGemmU8X8CopyPackA(PanelA, A + m * lda, lda, CountM, CountK, + RowSumBuffer); // // Apply the global depth value constant without the ZeroPointB scaling from: @@ -639,28 +558,15 @@ Return Value: bool PostProcess = (k + CountK == K); while (RowsRemaining > 0) { - size_t RowsHandled = MlasGemmU8X8Kernel( - pa, - b, - c, - PackedCountK, - RowsRemaining, - CountN, - ldc, - RowSums, - ColumnSumBuffer, - (PackedZeroPointB != nullptr) ? ZeroPointBBuffer : nullptr, + pa, b, c, PackedCountK, RowsRemaining, CountN, ldc, RowSums, + ColumnSumBuffer, (PackedZeroPointB != nullptr) ? ZeroPointBBuffer : nullptr, ZeroMode); if (PostProcess && Data->OutputProcessor != nullptr) { Data->OutputProcessor->Process( - Data->C, - RangeStartM + m + CountM - RowsRemaining, - RangeStartN + n, - RowsHandled, - CountN, - Data->ldc); + Data->C, RangeStartM + m + CountM - RowsRemaining, RangeStartN + n, + RowsHandled, CountN, Data->ldc); } c += ldc * RowsHandled; @@ -678,8 +584,7 @@ Return Value: #if defined(MLAS_SSE2_INTRINSICS) -struct MLAS_GEMM_U8X8_KERNEL_SSE -{ +struct MLAS_GEMM_U8X8_KERNEL_SSE { typedef int16_t PackedAType; typedef int16_t PackedBType; typedef int8_t OffsetBType; @@ -691,13 +596,9 @@ struct MLAS_GEMM_U8X8_KERNEL_SSE constexpr size_t MLAS_GEMM_U8X8_KERNEL_SSE::PackedK; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8X8_KERNEL_SSE::Strides; -template<> -MLAS_FORCEINLINE -int32_t -MlasGemmU8X8FixupZeroPointB( - int32_t ZeroPointB, - bool BIsSigned - ) +template <> +MLAS_FORCEINLINE int32_t +MlasGemmU8X8FixupZeroPointB(int32_t ZeroPointB, bool BIsSigned) { if (!BIsSigned) { ZeroPointB = MLAS_GEMM_U8X8_KERNEL_SSE::OffsetBType(ZeroPointB ^ 0x80); @@ -706,27 +607,24 @@ MlasGemmU8X8FixupZeroPointB( return ZeroPointB; } -template<> +template <> void -MlasGemmU8X8CopyPackA( - MLAS_GEMM_U8X8_KERNEL_SSE::PackedAType* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumBuffer - ) +MlasGemmU8X8CopyPackA(MLAS_GEMM_U8X8_KERNEL_SSE::PackedAType* D, + const uint8_t* A, + size_t lda, + size_t CountM, + size_t CountK, + int32_t* RowSumBuffer) { const __m128i ZeroVector = _mm_setzero_si128(); const __m128i OnesWordBroadcast = _mm_set1_epi16(1); - uint8_t PaddedMatrixAData[8] = { 0 }; + uint8_t PaddedMatrixAData[8] = {0}; // // Process a single row of matrix A in a loop. // while (CountM > 0) { - const uint8_t* a = A; size_t k = CountK; __m128i ReductionVector = ZeroVector; @@ -745,7 +643,6 @@ MlasGemmU8X8CopyPackA( // while (k >= 8) { - __m128i Bytes = _mm_loadl_epi64((const __m128i*)&a[0]); __m128i Words = _mm_unpacklo_epi8(Bytes, ZeroVector); @@ -759,7 +656,6 @@ MlasGemmU8X8CopyPackA( } if (k > 0) { - // // Copy the remaining bytes to the zero padded stack buffer. // @@ -795,10 +691,10 @@ MlasGemmU8X8CopyPackA( // ReductionVector = _mm_madd_epi16(ReductionVector, OnesWordBroadcast); - ReductionVector = _mm_add_epi32(ReductionVector, - _mm_shuffle_epi32(ReductionVector, _MM_SHUFFLE(3, 2, 3, 2))); - ReductionVector = _mm_add_epi32(ReductionVector, - _mm_shuffle_epi32(ReductionVector, _MM_SHUFFLE(0, 1, 0, 1))); + ReductionVector = _mm_add_epi32( + ReductionVector, _mm_shuffle_epi32(ReductionVector, _MM_SHUFFLE(3, 2, 3, 2))); + ReductionVector = _mm_add_epi32( + ReductionVector, _mm_shuffle_epi32(ReductionVector, _MM_SHUFFLE(0, 1, 0, 1))); *RowSumBuffer++ = _mm_cvtsi128_si32(ReductionVector); @@ -809,20 +705,20 @@ MlasGemmU8X8CopyPackA( MLAS_FORCEINLINE void -MlasGemmU8X8CopyPackBProcessSse( - MLAS_GEMM_U8X8_KERNEL_SSE::PackedBType* D, - __m128i BytesRow0, - __m128i BytesRow1, - __m128i BitFlipVector, - __m128i ColumnSums[2] - ) +MlasGemmU8X8CopyPackBProcessSse(MLAS_GEMM_U8X8_KERNEL_SSE::PackedBType* D, + __m128i BytesRow0, + __m128i BytesRow1, + __m128i BitFlipVector, + __m128i ColumnSums[2]) { __m128i BytesInterleaved = _mm_unpacklo_epi8(BytesRow0, BytesRow1); BytesInterleaved = _mm_xor_si128(BytesInterleaved, BitFlipVector); - __m128i WordsInterleaved0 = _mm_srai_epi16(_mm_unpacklo_epi8(BytesInterleaved, BytesInterleaved), 8); - __m128i WordsInterleaved1 = _mm_srai_epi16(_mm_unpackhi_epi8(BytesInterleaved, BytesInterleaved), 8); + __m128i WordsInterleaved0 = + _mm_srai_epi16(_mm_unpacklo_epi8(BytesInterleaved, BytesInterleaved), 8); + __m128i WordsInterleaved1 = + _mm_srai_epi16(_mm_unpackhi_epi8(BytesInterleaved, BytesInterleaved), 8); ColumnSums[0] = _mm_add_epi16(ColumnSums[0], WordsInterleaved0); ColumnSums[1] = _mm_add_epi16(ColumnSums[1], WordsInterleaved1); @@ -831,17 +727,15 @@ MlasGemmU8X8CopyPackBProcessSse( _mm_storeu_si128((__m128i*)&D[8], WordsInterleaved1); } -template<> +template <> void -MlasGemmU8X8CopyPackB( - MLAS_GEMM_U8X8_KERNEL_SSE::PackedBType* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer, - bool BIsSigned - ) +MlasGemmU8X8CopyPackB(MLAS_GEMM_U8X8_KERNEL_SSE::PackedBType* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer, + bool BIsSigned) { const __m128i OnesWordBroadcast = _mm_set1_epi16(1); const __m128i BitFlipVector = _mm_set1_epi32(BIsSigned ? 0 : 0x80808080); @@ -851,7 +745,6 @@ MlasGemmU8X8CopyPackB( // while (CountN >= 8) { - const uint8_t* b = B; size_t k = CountK; __m128i ColumnSums[2]; @@ -868,7 +761,6 @@ MlasGemmU8X8CopyPackB( // while (k >= MLAS_GEMM_U8X8_KERNEL_SSE::PackedK) { - __m128i BytesRow0 = _mm_loadl_epi64((const __m128i*)&b[0]); __m128i BytesRow1 = _mm_loadl_epi64((const __m128i*)&b[ldb]); @@ -880,7 +772,6 @@ MlasGemmU8X8CopyPackB( } if (k > 0) { - __m128i BytesRow0 = _mm_loadl_epi64((const __m128i*)&b[0]); MlasGemmU8X8CopyPackBProcessSse(D, BytesRow0, BitFlipVector, BitFlipVector, ColumnSums); @@ -904,7 +795,6 @@ MlasGemmU8X8CopyPackB( // if (CountN > 0) { - const uint8_t* b = B; size_t k = CountK; __m128i ColumnSums[2]; @@ -921,7 +811,6 @@ MlasGemmU8X8CopyPackB( // while (k >= MLAS_GEMM_U8X8_KERNEL_SSE::PackedK) { - const uint8_t* bcopy = b; uint8_t* padded = PaddedMatrixBData; uint8_t* padded_end = padded + CountN; @@ -944,7 +833,6 @@ MlasGemmU8X8CopyPackB( } if (k > 0) { - const uint8_t* bcopy = b; uint8_t* padded = PaddedMatrixBData; uint8_t* padded_end = padded + CountN; @@ -970,11 +858,7 @@ MlasGemmU8X8CopyPackB( MLAS_FORCEINLINE void -MlasGemmU8X8MultiplyAccumulateRowSse( - __m128i ABroadcast, - const int16_t* B, - __m128i Accumulators[2] - ) +MlasGemmU8X8MultiplyAccumulateRowSse(__m128i ABroadcast, const int16_t* B, __m128i Accumulators[2]) { __m128i BElements0 = _mm_load_si128((__m128i*)&B[0]); __m128i BElements1 = _mm_load_si128((__m128i*)&B[8]); @@ -983,27 +867,24 @@ MlasGemmU8X8MultiplyAccumulateRowSse( Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_madd_epi16(BElements1, ABroadcast)); } -template<> +template <> size_t -MlasGemmU8X8Kernel( - const MLAS_GEMM_U8X8_KERNEL_SSE::PackedAType* A, - const MLAS_GEMM_U8X8_KERNEL_SSE::PackedBType* B, - int32_t* C, - size_t PackedCountK, - size_t CountM, - size_t CountN, - size_t ldc, - const int32_t* RowSumBuffer, - const int32_t* ColumnSumBuffer, - const int32_t* ZeroPointB, - bool ZeroMode - ) +MlasGemmU8X8Kernel(const MLAS_GEMM_U8X8_KERNEL_SSE::PackedAType* A, + const MLAS_GEMM_U8X8_KERNEL_SSE::PackedBType* B, + int32_t* C, + size_t PackedCountK, + size_t CountM, + size_t CountN, + size_t ldc, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, + const int32_t* ZeroPointB, + bool ZeroMode) { MLAS_UNREFERENCED_PARAMETER(CountM); MLAS_UNREFERENCED_PARAMETER(ldc); while (CountN > 0) { - __m128i Accumulators[2]; // @@ -1013,7 +894,6 @@ MlasGemmU8X8Kernel( int32_t RowSumValue = RowSumBuffer[0]; if (ZeroPointB != nullptr) { - int32_t ScaledRowSumBuffer[8]; for (size_t i = 0; i < 8; i++) { @@ -1026,13 +906,14 @@ MlasGemmU8X8Kernel( Accumulators[1] = _mm_loadu_si128((__m128i*)&ScaledRowSumBuffer[4]); } else { - Accumulators[0] = _mm_set1_epi32(RowSumValue); Accumulators[1] = Accumulators[0]; } - Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[0])); - Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[4])); + Accumulators[0] = + _mm_add_epi32(Accumulators[0], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[0])); + Accumulators[1] = + _mm_add_epi32(Accumulators[1], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[4])); ColumnSumBuffer += 8; // @@ -1045,7 +926,6 @@ MlasGemmU8X8Kernel( size_t k = PackedCountK; while (k >= 4) { - __m128i AElements = _mm_loadu_si128((__m128i*)a); __m128i ABroadcast; @@ -1067,7 +947,6 @@ MlasGemmU8X8Kernel( } while (k > 0) { - __m128i ABroadcast = _mm_set1_epi32(*((int32_t*)a)); MlasGemmU8X8MultiplyAccumulateRowSse(ABroadcast, &B[0], Accumulators); @@ -1082,7 +961,6 @@ MlasGemmU8X8Kernel( // if (CountN >= 8) { - if (!ZeroMode) { Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&C[0])); Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((__m128i*)&C[4])); @@ -1095,15 +973,14 @@ MlasGemmU8X8Kernel( CountN -= 8; } else { - // // Output the remaining partial output block. // if ((CountN & 4) != 0) { - if (!ZeroMode) { - Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&C[0])); + Accumulators[0] = + _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&C[0])); } _mm_storeu_si128((__m128i*)&C[0], Accumulators[0]); @@ -1113,9 +990,9 @@ MlasGemmU8X8Kernel( } if ((CountN & 2) != 0) { - if (!ZeroMode) { - Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadl_epi64((__m128i*)&C[0])); + Accumulators[0] = + _mm_add_epi32(Accumulators[0], _mm_loadl_epi64((__m128i*)&C[0])); } _mm_storel_epi64((__m128i*)&C[0], Accumulators[0]); @@ -1125,7 +1002,6 @@ MlasGemmU8X8Kernel( } if ((CountN & 1) != 0) { - int32_t AccumulatorValue = _mm_cvtsi128_si32(Accumulators[0]); if (!ZeroMode) { @@ -1156,8 +1032,7 @@ const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8X8DispatchSse = { // for this code is Windows only, so restrict this kernel to that environment. #if defined(MLAS_SSE2_INTRINSICS) && defined(_MSC_VER) -struct MLAS_GEMM_U8S8_KERNEL_SSE41 -{ +struct MLAS_GEMM_U8S8_KERNEL_SSE41 { typedef uint8_t PackedAType; typedef uint8_t PackedBType; typedef int8_t OffsetBType; @@ -1171,16 +1046,14 @@ constexpr size_t MLAS_GEMM_U8S8_KERNEL_SSE41::PackedK; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8S8_KERNEL_SSE41::Strides; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8S8_KERNEL_SSE41::PackedStrides; -template<> +template <> void -MlasGemmU8X8CopyPackA( - MLAS_GEMM_U8S8_KERNEL_SSE41::PackedAType* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumBuffer - ) +MlasGemmU8X8CopyPackA(MLAS_GEMM_U8S8_KERNEL_SSE41::PackedAType* D, + const uint8_t* A, + size_t lda, + size_t CountM, + size_t CountK, + int32_t* RowSumBuffer) { const __m128i ZeroVector = _mm_setzero_si128(); const __m128i OnesWordBroadcast = _mm_set1_epi16(1); @@ -1190,7 +1063,6 @@ MlasGemmU8X8CopyPackA( // while (CountM > 0) { - const uint8_t* a = A; size_t k = CountK; __m128i ReductionVector = ZeroVector; @@ -1204,11 +1076,11 @@ MlasGemmU8X8CopyPackA( // while (k >= 8) { - __m128i Bytes = _mm_loadl_epi64((const __m128i*)&a[0]); __m128i Words = _mm_unpacklo_epi8(Bytes, ZeroVector); - ReductionVector = _mm_add_epi32(ReductionVector, _mm_madd_epi16(Words, OnesWordBroadcast)); + ReductionVector = + _mm_add_epi32(ReductionVector, _mm_madd_epi16(Words, OnesWordBroadcast)); _mm_storel_epi64((__m128i*)&D[0], Bytes); @@ -1218,7 +1090,6 @@ MlasGemmU8X8CopyPackA( } if (k > 0) { - // // Copy the remaining bytes to the zero padded stack buffer. // @@ -1231,7 +1102,8 @@ MlasGemmU8X8CopyPackA( D += (k + 3) & ~3; __m128i Words = _mm_unpacklo_epi8(Bytes, ZeroVector); - ReductionVector = _mm_add_epi32(ReductionVector, _mm_madd_epi16(Words, OnesWordBroadcast)); + ReductionVector = + _mm_add_epi32(ReductionVector, _mm_madd_epi16(Words, OnesWordBroadcast)); } // @@ -1250,13 +1122,11 @@ MlasGemmU8X8CopyPackA( MLAS_FORCEINLINE void -MlasGemmU8X8CopyPackBProcessSse41( - MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* D, - __m128i BytesRows[4], - __m128i OnesByteBroadcast, - __m128i OnesWordBroadcast, - __m128i ColumnSums[2] - ) +MlasGemmU8X8CopyPackBProcessSse41(MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* D, + __m128i BytesRows[4], + __m128i OnesByteBroadcast, + __m128i OnesWordBroadcast, + __m128i ColumnSums[2]) { __m128i PairsInterleaved0 = _mm_unpacklo_epi8(BytesRows[0], BytesRows[1]); __m128i PairsInterleaved1 = _mm_unpacklo_epi8(BytesRows[2], BytesRows[3]); @@ -1277,17 +1147,15 @@ MlasGemmU8X8CopyPackBProcessSse41( _mm_storeu_si128((__m128i*)&D[16], QuadsInterleaved1); } -template<> +template <> void -MlasGemmU8X8CopyPackB( - MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer, - bool BIsSigned - ) +MlasGemmU8X8CopyPackB(MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer, + bool BIsSigned) { const __m128i OnesByteBroadcast = _mm_set1_epi8(1); const __m128i OnesWordBroadcast = _mm_set1_epi16(1); @@ -1300,7 +1168,6 @@ MlasGemmU8X8CopyPackB( // while (CountN >= 8) { - const uint8_t* b = B; size_t k = CountK; __m128i ColumnSums[2]; @@ -1313,13 +1180,13 @@ MlasGemmU8X8CopyPackB( // while (k >= MLAS_GEMM_U8S8_KERNEL_SSE41::PackedK) { - BytesRows[0] = _mm_loadl_epi64((const __m128i*)&b[ldb * 0]); BytesRows[1] = _mm_loadl_epi64((const __m128i*)&b[ldb * 1]); BytesRows[2] = _mm_loadl_epi64((const __m128i*)&b[ldb * 2]); BytesRows[3] = _mm_loadl_epi64((const __m128i*)&b[ldb * 3]); - MlasGemmU8X8CopyPackBProcessSse41(D, BytesRows, OnesByteBroadcast, OnesWordBroadcast, ColumnSums); + MlasGemmU8X8CopyPackBProcessSse41(D, BytesRows, OnesByteBroadcast, OnesWordBroadcast, + ColumnSums); b += ldb * 4; D += 32; @@ -1327,7 +1194,6 @@ MlasGemmU8X8CopyPackB( } if (k > 0) { - BytesRows[0] = _mm_loadl_epi64((const __m128i*)&b[ldb * 0]); BytesRows[1] = _mm_setzero_si128(); BytesRows[2] = _mm_setzero_si128(); @@ -1341,7 +1207,8 @@ MlasGemmU8X8CopyPackB( BytesRows[2] = _mm_loadl_epi64((const __m128i*)&b[ldb * 2]); } - MlasGemmU8X8CopyPackBProcessSse41(D, BytesRows, OnesByteBroadcast, OnesWordBroadcast, ColumnSums); + MlasGemmU8X8CopyPackBProcessSse41(D, BytesRows, OnesByteBroadcast, OnesWordBroadcast, + ColumnSums); D += 32; } @@ -1359,7 +1226,6 @@ MlasGemmU8X8CopyPackB( // if (CountN > 0) { - const __m128i ZeroVector = _mm_setzero_si128(); __m128i ColumnSums[2]; @@ -1369,7 +1235,6 @@ MlasGemmU8X8CopyPackB( ColumnSums[1] = _mm_setzero_si128(); while (CountK > 0) { - size_t k = std::min(CountK, MLAS_GEMM_U8S8_KERNEL_SSE41::PackedK); CountK -= k; @@ -1379,7 +1244,6 @@ MlasGemmU8X8CopyPackB( uint8_t* padded = PaddedMatrixBData; do { - std::copy_n(B, CountN, padded); padded += 8; @@ -1393,7 +1257,8 @@ MlasGemmU8X8CopyPackB( BytesRows[2] = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[16]); BytesRows[3] = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[24]); - MlasGemmU8X8CopyPackBProcessSse41(D, BytesRows, OnesByteBroadcast, OnesWordBroadcast, ColumnSums); + MlasGemmU8X8CopyPackBProcessSse41(D, BytesRows, OnesByteBroadcast, OnesWordBroadcast, + ColumnSums); D += 32; } @@ -1405,12 +1270,10 @@ MlasGemmU8X8CopyPackB( MLAS_FORCEINLINE void -MlasGemmU8X8MultiplyAccumulateRowSse41( - __m128i ABroadcast, - const MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* B, - __m128i OnesWordBroadcast, - __m128i Accumulators[2] - ) +MlasGemmU8X8MultiplyAccumulateRowSse41(__m128i ABroadcast, + const MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* B, + __m128i OnesWordBroadcast, + __m128i Accumulators[2]) { __m128i BElements0 = _mm_load_si128((__m128i*)&B[0]); __m128i BElements1 = _mm_load_si128((__m128i*)&B[16]); @@ -1418,25 +1281,25 @@ MlasGemmU8X8MultiplyAccumulateRowSse41( __m128i Intermediate0 = _mm_maddubs_epi16(ABroadcast, BElements0); __m128i Intermediate1 = _mm_maddubs_epi16(ABroadcast, BElements1); - Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_madd_epi16(Intermediate0, OnesWordBroadcast)); - Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_madd_epi16(Intermediate1, OnesWordBroadcast)); + Accumulators[0] = + _mm_add_epi32(Accumulators[0], _mm_madd_epi16(Intermediate0, OnesWordBroadcast)); + Accumulators[1] = + _mm_add_epi32(Accumulators[1], _mm_madd_epi16(Intermediate1, OnesWordBroadcast)); } -template<> +template <> size_t -MlasGemmU8X8Kernel( - const MLAS_GEMM_U8S8_KERNEL_SSE41::PackedAType* A, - const MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* B, - int32_t* C, - size_t PackedCountK, - size_t CountM, - size_t CountN, - size_t ldc, - const int32_t* RowSumBuffer, - const int32_t* ColumnSumBuffer, - const int32_t* ZeroPointB, - bool ZeroMode - ) +MlasGemmU8X8Kernel(const MLAS_GEMM_U8S8_KERNEL_SSE41::PackedAType* A, + const MLAS_GEMM_U8S8_KERNEL_SSE41::PackedBType* B, + int32_t* C, + size_t PackedCountK, + size_t CountM, + size_t CountN, + size_t ldc, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, + const int32_t* ZeroPointB, + bool ZeroMode) { const __m128i OnesWordBroadcast = _mm_set1_epi16(1); @@ -1444,7 +1307,6 @@ MlasGemmU8X8Kernel( MLAS_UNREFERENCED_PARAMETER(ldc); while (CountN > 0) { - __m128i Accumulators[2]; // @@ -1455,13 +1317,17 @@ MlasGemmU8X8Kernel( Accumulators[1] = Accumulators[0]; if (ZeroPointB != nullptr) { - Accumulators[0] = _mm_mullo_epi32(Accumulators[0], _mm_loadu_si128((const __m128i*)&ZeroPointB[0])); - Accumulators[1] = _mm_mullo_epi32(Accumulators[1], _mm_loadu_si128((const __m128i*)&ZeroPointB[4])); + Accumulators[0] = + _mm_mullo_epi32(Accumulators[0], _mm_loadu_si128((const __m128i*)&ZeroPointB[0])); + Accumulators[1] = + _mm_mullo_epi32(Accumulators[1], _mm_loadu_si128((const __m128i*)&ZeroPointB[4])); ZeroPointB += 8; } - Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[0])); - Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[4])); + Accumulators[0] = + _mm_add_epi32(Accumulators[0], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[0])); + Accumulators[1] = + _mm_add_epi32(Accumulators[1], _mm_loadu_si128((const __m128i*)&ColumnSumBuffer[4])); ColumnSumBuffer += 8; // @@ -1474,21 +1340,24 @@ MlasGemmU8X8Kernel( size_t k = PackedCountK; while (k >= 4) { - __m128i AElements = _mm_loadu_si128((__m128i*)a); __m128i ABroadcast; ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(0, 0, 0, 0)); - MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[0], OnesWordBroadcast, Accumulators); + MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[0], OnesWordBroadcast, + Accumulators); ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(1, 1, 1, 1)); - MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[32], OnesWordBroadcast, Accumulators); + MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[32], OnesWordBroadcast, + Accumulators); ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(2, 2, 2, 2)); - MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[64], OnesWordBroadcast, Accumulators); + MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[64], OnesWordBroadcast, + Accumulators); ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(3, 3, 3, 3)); - MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[96], OnesWordBroadcast, Accumulators); + MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[96], OnesWordBroadcast, + Accumulators); a += 4 * 4; B += 4 * 32; @@ -1496,9 +1365,9 @@ MlasGemmU8X8Kernel( } while (k > 0) { - __m128i ABroadcast = _mm_set1_epi32(*((int32_t*)a)); - MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[0], OnesWordBroadcast, Accumulators); + MlasGemmU8X8MultiplyAccumulateRowSse41(ABroadcast, &B[0], OnesWordBroadcast, + Accumulators); a += 4; B += 32; @@ -1511,7 +1380,6 @@ MlasGemmU8X8Kernel( // if (CountN >= 8) { - if (!ZeroMode) { Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&C[0])); Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((__m128i*)&C[4])); @@ -1524,15 +1392,14 @@ MlasGemmU8X8Kernel( CountN -= 8; } else { - // // Output the remaining partial output block. // if ((CountN & 4) != 0) { - if (!ZeroMode) { - Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&C[0])); + Accumulators[0] = + _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&C[0])); } _mm_storeu_si128((__m128i*)&C[0], Accumulators[0]); @@ -1542,9 +1409,9 @@ MlasGemmU8X8Kernel( } if ((CountN & 2) != 0) { - if (!ZeroMode) { - Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadl_epi64((__m128i*)&C[0])); + Accumulators[0] = + _mm_add_epi32(Accumulators[0], _mm_loadl_epi64((__m128i*)&C[0])); } _mm_storel_epi64((__m128i*)&C[0], Accumulators[0]); @@ -1554,7 +1421,6 @@ MlasGemmU8X8Kernel( } if ((CountN & 1) != 0) { - int32_t AccumulatorValue = _mm_cvtsi128_si32(Accumulators[0]); if (!ZeroMode) { @@ -1587,8 +1453,9 @@ const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8S8DispatchSse41 = { // Stores a vector to transpose a 4x4 byte vector using vpshufb. // -MLAS_INTERNAL_DATA MLAS_DECLSPEC_ALIGN(const uint8_t MlasTranspose4x4BytesAvx[16], 16) = - { 0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15 }; +MLAS_INTERNAL_DATA +MLAS_DECLSPEC_ALIGN(const uint8_t MlasTranspose4x4BytesAvx[16], 16) = {0, 4, 8, 12, 1, 5, 9, 13, + 2, 6, 10, 14, 3, 7, 11, 15}; // // Define the prototypes of the AVX2/AVX512 routines written in assembly. @@ -1596,54 +1463,33 @@ MLAS_INTERNAL_DATA MLAS_DECLSPEC_ALIGN(const uint8_t MlasTranspose4x4BytesAvx[16 extern "C" { - void - MLASCALL - MlasGemmU8S8CopyPackAAvx2( - uint8_t* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumBuffer - ); - - void - MLASCALL - MlasGemmU8S8CopyPackBAvx2( - uint8_t* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer, - bool BIsSigned - ); - - void - MLASCALL - MlasGemmU8U8CopyPackAAvx2( - int16_t* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumBuffer - ); - - void - MLASCALL - MlasGemmU8U8CopyPackBAvx2( - uint8_t* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer - ); +void MLASCALL +MlasGemmU8S8CopyPackAAvx2( + uint8_t* D, const uint8_t* A, size_t lda, size_t CountM, size_t CountK, int32_t* RowSumBuffer); + +void MLASCALL +MlasGemmU8S8CopyPackBAvx2(uint8_t* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer, + bool BIsSigned); + +void MLASCALL +MlasGemmU8U8CopyPackAAvx2( + int16_t* D, const uint8_t* A, size_t lda, size_t CountM, size_t CountK, int32_t* RowSumBuffer); + +void MLASCALL +MlasGemmU8U8CopyPackBAvx2(uint8_t* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer); } -struct MLAS_GEMM_U8S8_KERNEL_AVX2 -{ +struct MLAS_GEMM_U8S8_KERNEL_AVX2 { typedef uint8_t PackedAType; typedef uint8_t PackedBType; typedef int8_t OffsetBType; @@ -1657,18 +1503,15 @@ constexpr size_t MLAS_GEMM_U8S8_KERNEL_AVX2::PackedK; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8S8_KERNEL_AVX2::Strides; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8S8_KERNEL_AVX2::PackedStrides; -template<> -MLAS_FORCEINLINE -bool -MlasGemmU8X8TryGemvKernel( - const uint8_t* A, - const uint8_t* B, - size_t ldb, - int32_t* C, - size_t CountK, - size_t CountN, - bool BIsSigned - ) +template <> +MLAS_FORCEINLINE bool +MlasGemmU8X8TryGemvKernel(const uint8_t* A, + const uint8_t* B, + size_t ldb, + int32_t* C, + size_t CountK, + size_t CountN, + bool BIsSigned) { if (BIsSigned) { MlasPlatform.GemvU8S8Kernel(A, B, C, CountK, CountN, ldb); @@ -1678,13 +1521,9 @@ MlasGemmU8X8TryGemvKernel( return false; } -template<> -MLAS_FORCEINLINE -int32_t -MlasGemmU8X8FixupZeroPointB( - int32_t ZeroPointB, - bool BIsSigned - ) +template <> +MLAS_FORCEINLINE int32_t +MlasGemmU8X8FixupZeroPointB(int32_t ZeroPointB, bool BIsSigned) { if (!BIsSigned) { ZeroPointB = MLAS_GEMM_U8S8_KERNEL_AVX2::OffsetBType(ZeroPointB ^ 0x80); @@ -1693,56 +1532,47 @@ MlasGemmU8X8FixupZeroPointB( return ZeroPointB; } -template<> -MLAS_FORCEINLINE -void -MlasGemmU8X8CopyPackA( - MLAS_GEMM_U8S8_KERNEL_AVX2::PackedAType* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumBuffer - ) +template <> +MLAS_FORCEINLINE void +MlasGemmU8X8CopyPackA(MLAS_GEMM_U8S8_KERNEL_AVX2::PackedAType* D, + const uint8_t* A, + size_t lda, + size_t CountM, + size_t CountK, + int32_t* RowSumBuffer) { MlasGemmU8S8CopyPackAAvx2(D, A, lda, CountM, CountK, RowSumBuffer); } -template<> -MLAS_FORCEINLINE -void -MlasGemmU8X8CopyPackB( - MLAS_GEMM_U8S8_KERNEL_AVX2::PackedBType* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer, - bool BIsSigned - ) +template <> +MLAS_FORCEINLINE void +MlasGemmU8X8CopyPackB(MLAS_GEMM_U8S8_KERNEL_AVX2::PackedBType* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer, + bool BIsSigned) { MlasGemmU8S8CopyPackBAvx2(D, B, ldb, CountN, CountK, ColumnSumBuffer, BIsSigned); } -template<> -MLAS_FORCEINLINE -size_t -MlasGemmU8X8Kernel( - const MLAS_GEMM_U8S8_KERNEL_AVX2::PackedAType* A, - const MLAS_GEMM_U8S8_KERNEL_AVX2::PackedBType* B, - int32_t* C, - size_t PackedCountK, - size_t CountM, - size_t CountN, - size_t ldc, - const int32_t* RowSumBuffer, - const int32_t* ColumnSumBuffer, - const int32_t* ZeroPointB, - bool ZeroMode - ) +template <> +MLAS_FORCEINLINE size_t +MlasGemmU8X8Kernel(const MLAS_GEMM_U8S8_KERNEL_AVX2::PackedAType* A, + const MLAS_GEMM_U8S8_KERNEL_AVX2::PackedBType* B, + int32_t* C, + size_t PackedCountK, + size_t CountM, + size_t CountN, + size_t ldc, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, + const int32_t* ZeroPointB, + bool ZeroMode) { - return MlasPlatform.GemmU8S8Kernel(A, B, C, PackedCountK, CountM, CountN, ldc, - RowSumBuffer, ColumnSumBuffer, ZeroPointB, ZeroMode); + return MlasPlatform.GemmU8S8Kernel(A, B, C, PackedCountK, CountM, CountN, ldc, RowSumBuffer, + ColumnSumBuffer, ZeroPointB, ZeroMode); } const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8S8DispatchAvx2 = { @@ -1753,8 +1583,7 @@ const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8S8DispatchAvx2 = { MLAS_GEMM_U8S8_KERNEL_AVX2::PackedStrides.K, }; -struct MLAS_GEMM_U8U8_KERNEL_AVX2 -{ +struct MLAS_GEMM_U8U8_KERNEL_AVX2 { typedef int16_t PackedAType; typedef uint8_t PackedBType; typedef uint8_t OffsetBType; @@ -1768,58 +1597,49 @@ constexpr size_t MLAS_GEMM_U8U8_KERNEL_AVX2::PackedK; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8U8_KERNEL_AVX2::Strides; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8U8_KERNEL_AVX2::PackedStrides; -template<> -MLAS_FORCEINLINE -void -MlasGemmU8X8CopyPackA( - MLAS_GEMM_U8U8_KERNEL_AVX2::PackedAType* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumBuffer - ) +template <> +MLAS_FORCEINLINE void +MlasGemmU8X8CopyPackA(MLAS_GEMM_U8U8_KERNEL_AVX2::PackedAType* D, + const uint8_t* A, + size_t lda, + size_t CountM, + size_t CountK, + int32_t* RowSumBuffer) { MlasGemmU8U8CopyPackAAvx2(D, A, lda, CountM, CountK, RowSumBuffer); } -template<> -MLAS_FORCEINLINE -void -MlasGemmU8X8CopyPackB( - MLAS_GEMM_U8U8_KERNEL_AVX2::PackedBType* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer, - bool BIsSigned - ) +template <> +MLAS_FORCEINLINE void +MlasGemmU8X8CopyPackB(MLAS_GEMM_U8U8_KERNEL_AVX2::PackedBType* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer, + bool BIsSigned) { MLAS_UNREFERENCED_PARAMETER(BIsSigned); MlasGemmU8U8CopyPackBAvx2(D, B, ldb, CountN, CountK, ColumnSumBuffer); } -template<> -MLAS_FORCEINLINE -size_t -MlasGemmU8X8Kernel( - const MLAS_GEMM_U8U8_KERNEL_AVX2::PackedAType* A, - const MLAS_GEMM_U8U8_KERNEL_AVX2::PackedBType* B, - int32_t* C, - size_t PackedCountK, - size_t CountM, - size_t CountN, - size_t ldc, - const int32_t* RowSumBuffer, - const int32_t* ColumnSumBuffer, - const int32_t* ZeroPointB, - bool ZeroMode - ) +template <> +MLAS_FORCEINLINE size_t +MlasGemmU8X8Kernel(const MLAS_GEMM_U8U8_KERNEL_AVX2::PackedAType* A, + const MLAS_GEMM_U8U8_KERNEL_AVX2::PackedBType* B, + int32_t* C, + size_t PackedCountK, + size_t CountM, + size_t CountN, + size_t ldc, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, + const int32_t* ZeroPointB, + bool ZeroMode) { - return MlasPlatform.GemmU8U8Kernel(A, B, C, PackedCountK, CountM, CountN, ldc, - RowSumBuffer, ColumnSumBuffer, ZeroPointB, ZeroMode); + return MlasPlatform.GemmU8U8Kernel(A, B, C, PackedCountK, CountM, CountN, ldc, RowSumBuffer, + ColumnSumBuffer, ZeroPointB, ZeroMode); } const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8U8DispatchAvx2 = { @@ -1842,25 +1662,21 @@ const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8U8DispatchAvx2 = { extern "C" { - size_t - MLASCALL - MlasGemmU8X8KernelNeon( - const uint8_t* A, - const uint8_t* B, - int32_t* C, - size_t PackedCountK, - size_t CountM, - size_t CountN, - size_t ldc, - const int32_t* RowSumVector, - const int32_t* ColumnSumVector, - const int32_t* ZeroPointB, - bool ZeroMode - ); +size_t MLASCALL +MlasGemmU8X8KernelNeon(const uint8_t* A, + const uint8_t* B, + int32_t* C, + size_t PackedCountK, + size_t CountM, + size_t CountN, + size_t ldc, + const int32_t* RowSumVector, + const int32_t* ColumnSumVector, + const int32_t* ZeroPointB, + bool ZeroMode); } -struct MLAS_GEMM_U8X8_KERNEL_NEON -{ +struct MLAS_GEMM_U8X8_KERNEL_NEON { typedef uint8_t PackedAType; typedef uint8_t PackedBType; typedef uint8_t OffsetBType; @@ -1874,13 +1690,9 @@ constexpr size_t MLAS_GEMM_U8X8_KERNEL_NEON::PackedK; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8X8_KERNEL_NEON::Strides; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8X8_KERNEL_NEON::PackedStrides; -template<> -MLAS_FORCEINLINE -int32_t -MlasGemmU8X8FixupZeroPointB( - int32_t ZeroPointB, - bool BIsSigned - ) +template <> +MLAS_FORCEINLINE int32_t +MlasGemmU8X8FixupZeroPointB(int32_t ZeroPointB, bool BIsSigned) { if (BIsSigned) { ZeroPointB = MLAS_GEMM_U8X8_KERNEL_NEON::OffsetBType(ZeroPointB ^ 0x80); @@ -1889,16 +1701,14 @@ MlasGemmU8X8FixupZeroPointB( return ZeroPointB; } -template<> +template <> void -MlasGemmU8X8CopyPackA( - MLAS_GEMM_U8X8_KERNEL_NEON::PackedAType* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumBuffer - ) +MlasGemmU8X8CopyPackA(MLAS_GEMM_U8X8_KERNEL_NEON::PackedAType* D, + const uint8_t* A, + size_t lda, + size_t CountM, + size_t CountK, + int32_t* RowSumBuffer) { uint8_t PaddedMatrixAData[16]; @@ -1918,7 +1728,6 @@ MlasGemmU8X8CopyPackA( // while (CountM >= 4) { - const uint8_t* a0 = A; const uint8_t* a1 = a0 + lda; const uint8_t* a2 = a1 + lda; @@ -1928,7 +1737,6 @@ MlasGemmU8X8CopyPackA( uint32x4_t RowSums = vmovq_n_u32(0); while (k >= 16) { - uint32x4_t v0 = vld1q_u32(reinterpret_cast(a0)); a0 += 16; uint32x4_t v1 = vld1q_u32(reinterpret_cast(a1)); @@ -1981,7 +1789,6 @@ MlasGemmU8X8CopyPackA( } while (k >= 4) { - uint32_t v0 = *reinterpret_cast(a0); a0 += 4; uint32_t v1 = *reinterpret_cast(a1); @@ -2003,7 +1810,6 @@ MlasGemmU8X8CopyPackA( } if (k > 0) { - // // Copy the remaining bytes to the zero padded stack buffer. // @@ -2013,7 +1819,6 @@ MlasGemmU8X8CopyPackA( vst1q_u8(PaddedMatrixAData, vmovq_n_u8(0)); while (k > 0) { - d[0] = *a0++; d[4] = *a1++; d[8] = *a2++; @@ -2054,7 +1859,6 @@ MlasGemmU8X8CopyPackA( // if ((CountM & 2) != 0) { - const uint8_t* a0 = A; const uint8_t* a1 = a0 + lda; @@ -2062,7 +1866,6 @@ MlasGemmU8X8CopyPackA( uint32x2_t RowSums = vmov_n_u32(0); while (k >= 4) { - uint32_t v0 = *reinterpret_cast(a0); a0 += 4; uint32_t v1 = *reinterpret_cast(a1); @@ -2078,7 +1881,6 @@ MlasGemmU8X8CopyPackA( } if (k > 0) { - // // Copy the remaining bytes to the zero padded stack buffer. // @@ -2088,7 +1890,6 @@ MlasGemmU8X8CopyPackA( vst1q_u8(PaddedMatrixAData, vmovq_n_u8(0)); while (k > 0) { - d[0] = *a0++; d[4] = *a1++; @@ -2126,13 +1927,11 @@ MlasGemmU8X8CopyPackA( // if ((CountM & 1) != 0) { - const uint8_t* a = A; size_t k = CountK; uint32x4_t RowSums = vmovq_n_u32(0); while (k >= 16) { - uint8x16_t v = vld1q_u8(a); a += 16; @@ -2145,7 +1944,6 @@ MlasGemmU8X8CopyPackA( } if (k > 0) { - // // Copy the remaining bytes to the zero padded stack buffer. // @@ -2182,12 +1980,10 @@ MlasGemmU8X8CopyPackA( MLAS_FORCEINLINE void -MlasGemmU8X8CopyPackBProcessNeon( - uint8_t* D, - const uint8_t* B, - uint8x8_t BitFlipVector, - uint32x4_t ColumnSums[2] - ) +MlasGemmU8X8CopyPackBProcessNeon(uint8_t* D, + const uint8_t* B, + uint8x8_t BitFlipVector, + uint32x4_t ColumnSums[2]) { uint8x8_t BytesRow = veor_u8(vld1_u8(B), BitFlipVector); vst1_u8(D, BytesRow); @@ -2201,22 +1997,20 @@ MlasGemmU8X8CopyPackBProcessNeon( #endif } -template<> +template <> void -MlasGemmU8X8CopyPackB( - MLAS_GEMM_U8X8_KERNEL_NEON::PackedBType* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer, - bool BIsSigned - ) +MlasGemmU8X8CopyPackB(MLAS_GEMM_U8X8_KERNEL_NEON::PackedBType* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer, + bool BIsSigned) { const uint8x8_t BitFlipVector = vdup_n_u8(BIsSigned ? 0x80 : 0); const uint8x8_t ZeroVector = vmov_n_u8(0); - const size_t AlignedCountK = - (CountK + MLAS_GEMM_U8X8_KERNEL_NEON::PackedK - 1) & ~(MLAS_GEMM_U8X8_KERNEL_NEON::PackedK - 1); + const size_t AlignedCountK = (CountK + MLAS_GEMM_U8X8_KERNEL_NEON::PackedK - 1) & + ~(MLAS_GEMM_U8X8_KERNEL_NEON::PackedK - 1); // // Process 8 columns of matrix B in a loop. @@ -2232,7 +2026,6 @@ MlasGemmU8X8CopyPackB( // while (CountN >= 8) { - const uint8_t* b = B; uint32x4_t ColumnSums[2]; @@ -2240,7 +2033,6 @@ MlasGemmU8X8CopyPackB( ColumnSums[1] = vmovq_n_u32(0); for (size_t k = CountK; k > 0; k--) { - MlasGemmU8X8CopyPackBProcessNeon(D, b, BitFlipVector, ColumnSums); b += ldb; @@ -2265,7 +2057,6 @@ MlasGemmU8X8CopyPackB( // if (CountN > 0) { - const uint8_t* b = B; uint8_t PaddedMatrixBData[8]; uint32x4_t ColumnSums[2]; @@ -2276,7 +2067,6 @@ MlasGemmU8X8CopyPackB( ColumnSums[1] = vmovq_n_u32(0); for (size_t k = CountK; k > 0; k--) { - for (size_t n = 0; n < CountN; n++) { PaddedMatrixBData[n] = b[n]; } @@ -2297,25 +2087,22 @@ MlasGemmU8X8CopyPackB( } } -template<> -MLAS_FORCEINLINE -size_t -MlasGemmU8X8Kernel( - const MLAS_GEMM_U8X8_KERNEL_NEON::PackedAType* A, - const MLAS_GEMM_U8X8_KERNEL_NEON::PackedBType* B, - int32_t* C, - size_t PackedCountK, - size_t CountM, - size_t CountN, - size_t ldc, - const int32_t* RowSumBuffer, - const int32_t* ColumnSumBuffer, - const int32_t* ZeroPointB, - bool ZeroMode - ) +template <> +MLAS_FORCEINLINE size_t +MlasGemmU8X8Kernel(const MLAS_GEMM_U8X8_KERNEL_NEON::PackedAType* A, + const MLAS_GEMM_U8X8_KERNEL_NEON::PackedBType* B, + int32_t* C, + size_t PackedCountK, + size_t CountM, + size_t CountN, + size_t ldc, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, + const int32_t* ZeroPointB, + bool ZeroMode) { - return MlasGemmU8X8KernelNeon(A, B, C, PackedCountK, CountM, CountN, ldc, - RowSumBuffer, ColumnSumBuffer, ZeroPointB, ZeroMode); + return MlasGemmU8X8KernelNeon(A, B, C, PackedCountK, CountM, CountN, ldc, RowSumBuffer, + ColumnSumBuffer, ZeroPointB, ZeroMode); } const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8X8DispatchNeon = { @@ -2336,25 +2123,21 @@ const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8X8DispatchNeon = { extern "C" { - size_t - MLASCALL - MlasGemmU8X8KernelUdot( - const uint8_t* A, - const uint8_t* B, - int32_t* C, - size_t PackedCountK, - size_t CountM, - size_t CountN, - size_t ldc, - const int32_t* RowSumVector, - const int32_t* ColumnSumVector, - const int32_t* ZeroPointB, - bool ZeroMode - ); +size_t MLASCALL +MlasGemmU8X8KernelUdot(const uint8_t* A, + const uint8_t* B, + int32_t* C, + size_t PackedCountK, + size_t CountM, + size_t CountN, + size_t ldc, + const int32_t* RowSumVector, + const int32_t* ColumnSumVector, + const int32_t* ZeroPointB, + bool ZeroMode); } -struct MLAS_GEMM_U8X8_KERNEL_UDOT -{ +struct MLAS_GEMM_U8X8_KERNEL_UDOT { typedef uint8_t PackedAType; typedef uint8_t PackedBType; typedef uint8_t OffsetBType; @@ -2368,13 +2151,9 @@ constexpr size_t MLAS_GEMM_U8X8_KERNEL_UDOT::PackedK; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8X8_KERNEL_UDOT::Strides; constexpr MLAS_GEMM_U8X8_STRIDES MLAS_GEMM_U8X8_KERNEL_UDOT::PackedStrides; -template<> -MLAS_FORCEINLINE -int32_t -MlasGemmU8X8FixupZeroPointB( - int32_t ZeroPointB, - bool BIsSigned - ) +template <> +MLAS_FORCEINLINE int32_t +MlasGemmU8X8FixupZeroPointB(int32_t ZeroPointB, bool BIsSigned) { if (BIsSigned) { ZeroPointB = MLAS_GEMM_U8X8_KERNEL_UDOT::OffsetBType(ZeroPointB ^ 0x80); @@ -2383,16 +2162,14 @@ MlasGemmU8X8FixupZeroPointB( return ZeroPointB; } -template<> +template <> void -MlasGemmU8X8CopyPackA( - MLAS_GEMM_U8X8_KERNEL_NEON::PackedAType* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumBuffer - ) +MlasGemmU8X8CopyPackA(MLAS_GEMM_U8X8_KERNEL_NEON::PackedAType* D, + const uint8_t* A, + size_t lda, + size_t CountM, + size_t CountK, + int32_t* RowSumBuffer) { uint8_t PaddedMatrixAData[16]; @@ -2412,7 +2189,6 @@ MlasGemmU8X8CopyPackA( // while (CountM >= 4) { - const uint8_t* a0 = A; const uint8_t* a1 = a0 + lda; const uint8_t* a2 = a1 + lda; @@ -2422,7 +2198,6 @@ MlasGemmU8X8CopyPackA( uint32x4_t RowSums = vmovq_n_u32(0); while (k >= 16) { - uint32x4_t v0 = vld1q_u32(reinterpret_cast(a0)); a0 += 16; uint32x4_t v1 = vld1q_u32(reinterpret_cast(a1)); @@ -2457,7 +2232,6 @@ MlasGemmU8X8CopyPackA( } while (k >= 4) { - uint32_t v0 = *reinterpret_cast(a0); a0 += 4; uint32_t v1 = *reinterpret_cast(a1); @@ -2479,7 +2253,6 @@ MlasGemmU8X8CopyPackA( } if (k > 0) { - // // Copy the remaining bytes to the zero padded stack buffer. // @@ -2489,7 +2262,6 @@ MlasGemmU8X8CopyPackA( vst1q_u8(PaddedMatrixAData, vmovq_n_u8(0)); while (k > 0) { - d[0] = *a0++; d[4] = *a1++; d[8] = *a2++; @@ -2508,7 +2280,6 @@ MlasGemmU8X8CopyPackA( } if (((CountK - 1) & 7) < 4) { - vst1q_u8(D, vmovq_n_u8(0)); D += 16; @@ -2537,7 +2308,6 @@ MlasGemmU8X8CopyPackA( // if (CountM >= 2) { - const uint8_t* a0 = A; const uint8_t* a1 = a0 + lda; @@ -2545,7 +2315,6 @@ MlasGemmU8X8CopyPackA( uint32x2_t RowSums = vmov_n_u32(0); while (k >= 4) { - uint32_t v0 = *reinterpret_cast(a0); a0 += 4; uint32_t v1 = *reinterpret_cast(a1); @@ -2561,7 +2330,6 @@ MlasGemmU8X8CopyPackA( } if (k > 0) { - // // Copy the remaining bytes to the zero padded stack buffer. // @@ -2571,7 +2339,6 @@ MlasGemmU8X8CopyPackA( vst1_u8(PaddedMatrixAData, vmov_n_u8(0)); while (k > 0) { - d[0] = *a0++; d[4] = *a1++; @@ -2588,7 +2355,6 @@ MlasGemmU8X8CopyPackA( } if (((CountK - 1) & 7) < 4) { - vst1_u8(D, vmov_n_u8(0)); D += 8; @@ -2616,13 +2382,11 @@ MlasGemmU8X8CopyPackA( // if (CountM > 0) { - const uint8_t* a = A; size_t k = CountK; uint32x4_t RowSums = vmovq_n_u32(0); while (k >= 16) { - uint8x16_t v = vld1q_u8(a); a += 16; @@ -2635,7 +2399,6 @@ MlasGemmU8X8CopyPackA( } if (k > 0) { - // // Copy the remaining bytes to the zero padded stack buffer. // @@ -2668,12 +2431,10 @@ MlasGemmU8X8CopyPackA( MLAS_FORCEINLINE void -MlasGemmU8X8CopyPackBProcessUdot( - MLAS_GEMM_U8X8_KERNEL_UDOT::PackedBType* D, - uint8x8_t BytesRow[4], - uint8x16_t BitFlipVector, - uint32x4_t ColumnSums[2] - ) +MlasGemmU8X8CopyPackBProcessUdot(MLAS_GEMM_U8X8_KERNEL_UDOT::PackedBType* D, + uint8x8_t BytesRow[4], + uint8x16_t BitFlipVector, + uint32x4_t ColumnSums[2]) { uint8x16_t v02 = veorq_u8(vcombine_u8(BytesRow[0], BytesRow[2]), BitFlipVector); uint8x16_t v13 = veorq_u8(vcombine_u8(BytesRow[1], BytesRow[3]), BitFlipVector); @@ -2688,17 +2449,15 @@ MlasGemmU8X8CopyPackBProcessUdot( ColumnSums[1] = vpadalq_u16(ColumnSums[1], vpaddlq_u8(vreinterpretq_u8_u16(zd.val[1]))); } -template<> +template <> void -MlasGemmU8X8CopyPackB( - MLAS_GEMM_U8X8_KERNEL_UDOT::PackedBType* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer, - bool BIsSigned - ) +MlasGemmU8X8CopyPackB(MLAS_GEMM_U8X8_KERNEL_UDOT::PackedBType* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer, + bool BIsSigned) { const uint8x16_t ZeroVector = vmovq_n_u8(0); const uint8x16_t BitFlipVector = vdupq_n_u8(BIsSigned ? 0x80 : 0); @@ -2726,7 +2485,6 @@ MlasGemmU8X8CopyPackB( // while (CountN >= 8) { - const uint8_t* b = B; size_t k = CountK; uint32x4_t ColumnSums[2]; @@ -2739,7 +2497,6 @@ MlasGemmU8X8CopyPackB( // while (k >= 4) { - BytesRow[0] = vld1_u8(&b[ldb * 0]); BytesRow[1] = vld1_u8(&b[ldb * 1]); BytesRow[2] = vld1_u8(&b[ldb * 2]); @@ -2753,7 +2510,6 @@ MlasGemmU8X8CopyPackB( } if (k > 0) { - BytesRow[0] = vld1_u8(&b[ldb * 0]); BytesRow[1] = (k >= 2) ? vld1_u8(&b[ldb * 1]) : vget_low_u8(BitFlipVector); BytesRow[2] = (k > 2) ? vld1_u8(&b[ldb * 2]) : vget_low_u8(BitFlipVector); @@ -2770,7 +2526,6 @@ MlasGemmU8X8CopyPackB( // if (((CountK - 1) & 7) < 4) { - vst1q_u8(&D[0], ZeroVector); vst1q_u8(&D[16], ZeroVector); @@ -2790,7 +2545,6 @@ MlasGemmU8X8CopyPackB( // if (CountN > 0) { - const uint8_t* b = B; size_t k = CountK; uint8_t PaddedMatrixBData[32]; @@ -2808,19 +2562,16 @@ MlasGemmU8X8CopyPackB( // while (k > 0) { - const uint8_t* bcopy0 = &b[ldb * 0]; const uint8_t* bcopy1 = &b[ldb * 1]; const uint8_t* bcopy2 = &b[ldb * 2]; const uint8_t* bcopy3 = &b[ldb * 3]; if (k >= 4) { - b += ldb * 4; k -= 4; } else { - vst1q_u8(&PaddedMatrixBData[0], BitFlipVector); vst1q_u8(&PaddedMatrixBData[16], BitFlipVector); @@ -2857,7 +2608,6 @@ MlasGemmU8X8CopyPackB( // if (((CountK - 1) & 7) < 4) { - vst1q_u8(&D[0], ZeroVector); vst1q_u8(&D[16], ZeroVector); @@ -2869,25 +2619,22 @@ MlasGemmU8X8CopyPackB( } } -template<> -MLAS_FORCEINLINE -size_t -MlasGemmU8X8Kernel( - const MLAS_GEMM_U8X8_KERNEL_UDOT::PackedAType* A, - const MLAS_GEMM_U8X8_KERNEL_UDOT::PackedBType* B, - int32_t* C, - size_t PackedCountK, - size_t CountM, - size_t CountN, - size_t ldc, - const int32_t* RowSumBuffer, - const int32_t* ColumnSumBuffer, - const int32_t* ZeroPointB, - bool ZeroMode - ) +template <> +MLAS_FORCEINLINE size_t +MlasGemmU8X8Kernel(const MLAS_GEMM_U8X8_KERNEL_UDOT::PackedAType* A, + const MLAS_GEMM_U8X8_KERNEL_UDOT::PackedBType* B, + int32_t* C, + size_t PackedCountK, + size_t CountM, + size_t CountN, + size_t ldc, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, + const int32_t* ZeroPointB, + bool ZeroMode) { - return MlasGemmU8X8KernelUdot(A, B, C, PackedCountK, CountM, CountN, ldc, - RowSumBuffer, ColumnSumBuffer, ZeroPointB, ZeroMode); + return MlasGemmU8X8KernelUdot(A, B, C, PackedCountK, CountM, CountN, ldc, RowSumBuffer, + ColumnSumBuffer, ZeroPointB, ZeroMode); } const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8X8DispatchUdot = { @@ -2900,8 +2647,7 @@ const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8X8DispatchUdot = { #endif -struct MLAS_GEMM_U8X8_KERNEL_DEFAULT -{ +struct MLAS_GEMM_U8X8_KERNEL_DEFAULT { typedef uint8_t PackedAType; typedef uint8_t PackedBType; typedef uint8_t OffsetBType; @@ -2911,13 +2657,9 @@ struct MLAS_GEMM_U8X8_KERNEL_DEFAULT static constexpr MLAS_GEMM_U8X8_STRIDES PackedStrides{16, 128, 128}; }; -template<> -MLAS_FORCEINLINE -int32_t -MlasGemmU8X8FixupZeroPointB( - int32_t ZeroPointB, - bool BIsSigned - ) +template <> +MLAS_FORCEINLINE int32_t +MlasGemmU8X8FixupZeroPointB(int32_t ZeroPointB, bool BIsSigned) { if (BIsSigned) { ZeroPointB = MLAS_GEMM_U8X8_KERNEL_DEFAULT::OffsetBType(ZeroPointB ^ 0x80); @@ -2926,30 +2668,26 @@ MlasGemmU8X8FixupZeroPointB( return ZeroPointB; } -template<> +template <> void -MlasGemmU8X8CopyPackA( - MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedAType* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumBuffer - ) +MlasGemmU8X8CopyPackA(MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedAType* D, + const uint8_t* A, + size_t lda, + size_t CountM, + size_t CountK, + int32_t* RowSumBuffer) { - const size_t AlignedCountK = - (CountK + MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedK - 1) & ~(MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedK - 1); + const size_t AlignedCountK = (CountK + MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedK - 1) & + ~(MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedK - 1); // // Process a single row of matrix A in a loop. // while (CountM-- > 0) { - int32_t RowSum = 0; for (size_t k = 0; k < CountK; k++) { - uint8_t a0 = A[k]; D[k] = a0; @@ -2967,20 +2705,18 @@ MlasGemmU8X8CopyPackA( } } -template<> +template <> void -MlasGemmU8X8CopyPackB( - MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedBType* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumBuffer, - bool BIsSigned - ) +MlasGemmU8X8CopyPackB(MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedBType* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumBuffer, + bool BIsSigned) { - const size_t AlignedCountK = - (CountK + MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedK - 1) & ~(MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedK - 1); + const size_t AlignedCountK = (CountK + MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedK - 1) & + ~(MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedK - 1); const uint8_t BitFlipValue = (BIsSigned ? 0x80 : 0); // @@ -2988,7 +2724,6 @@ MlasGemmU8X8CopyPackB( // while (CountN-- > 0) { - const uint8_t* b = B; int32_t ColumnSum = 0; @@ -2997,7 +2732,6 @@ MlasGemmU8X8CopyPackB( // for (size_t k = 0; k < CountK; k++) { - uint8_t b0 = b[0] ^ BitFlipValue; D[k] = b0; @@ -3017,7 +2751,7 @@ MlasGemmU8X8CopyPackB( } } -template<> +template <> size_t MlasGemmU8X8Kernel( const MLAS_GEMM_U8X8_KERNEL_DEFAULT::PackedAType* A, @@ -3030,8 +2764,7 @@ MlasGemmU8X8Kernel( const int32_t* RowSumBuffer, const int32_t* ColumnSumBuffer, const int32_t* ZeroPointB, - bool ZeroMode - ) + bool ZeroMode) { MLAS_UNREFERENCED_PARAMETER(CountM); MLAS_UNREFERENCED_PARAMETER(ldc); @@ -3041,7 +2774,6 @@ MlasGemmU8X8Kernel( // while (CountN-- > 0) { - int32_t Accumulator = *RowSumBuffer; if (ZeroPointB != nullptr) { @@ -3053,7 +2785,6 @@ MlasGemmU8X8Kernel( const auto* a = A; for (size_t k = 0; k < PackedCountK; k++) { - Accumulator += a[0] * B[0]; Accumulator += a[1] * B[1]; Accumulator += a[2] * B[2]; @@ -3082,14 +2813,11 @@ const MLAS_GEMM_U8X8_DISPATCH MlasGemmU8X8DispatchDefault = { 0, }; - void -MlasGemmU8X8Threaded( - const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock, - const MLAS_GEMM_U8X8_SHAPE_PARAMS* Shape, - const MLAS_GEMM_U8X8_DATA_PARAMS* Data, - ptrdiff_t ThreadId - ) +MlasGemmU8X8Threaded(const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock, + const MLAS_GEMM_U8X8_SHAPE_PARAMS* Shape, + const MLAS_GEMM_U8X8_DATA_PARAMS* Data, + ptrdiff_t ThreadId) /*++ Routine Description: @@ -3136,11 +2864,10 @@ Return Value: const size_t N = Shape->N; - const size_t BlockedN = (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) / - MLAS_QGEMM_STRIDEN_THREAD_ALIGN; + const size_t BlockedN = + (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) / MLAS_QGEMM_STRIDEN_THREAD_ALIGN; - MlasPartitionWork(ThreadIdN, WorkBlock->ThreadCountN, BlockedN, - &RangeStartN, &RangeCountN); + MlasPartitionWork(ThreadIdN, WorkBlock->ThreadCountN, BlockedN, &RangeStartN, &RangeCountN); RangeStartN *= MLAS_QGEMM_STRIDEN_THREAD_ALIGN; RangeCountN *= MLAS_QGEMM_STRIDEN_THREAD_ALIGN; @@ -3163,13 +2890,10 @@ Return Value: GemmU8X8Operation(Shape, Data, RangeStartM, RangeCountM, RangeStartN, RangeCountN); } - -void -MLASCALL -MlasGemm( - const MLAS_GEMM_U8X8_SHAPE_PARAMS &Shape, - const MLAS_GEMM_U8X8_DATA_PARAMS &DataParams, - MLAS_THREADPOOL *ThreadPool) +void MLASCALL +MlasGemm(const MLAS_GEMM_U8X8_SHAPE_PARAMS& Shape, + const MLAS_GEMM_U8X8_DATA_PARAMS& DataParams, + MLAS_THREADPOOL* ThreadPool) /*++ Routine Description: @@ -3195,13 +2919,11 @@ Return Value: MlasGemmBatch(Shape, &DataParams, 1, ThreadPool); } -void -MLASCALL -MlasGemmBatch( - const MLAS_GEMM_U8X8_SHAPE_PARAMS& Shape, - const MLAS_GEMM_U8X8_DATA_PARAMS* DataParams, - const size_t BatchN, - MLAS_THREADPOOL* ThreadPool) +void MLASCALL +MlasGemmBatch(const MLAS_GEMM_U8X8_SHAPE_PARAMS& Shape, + const MLAS_GEMM_U8X8_DATA_PARAMS* DataParams, + const size_t BatchN, + MLAS_THREADPOOL* ThreadPool) { const size_t M = Shape.M; const size_t N = Shape.N; @@ -3243,9 +2965,8 @@ MlasGemmBatch( MLAS_GEMM_U8X8_WORK_BLOCK WorkBlock; if (N > M) { - - const size_t BlockedN = (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) / - MLAS_QGEMM_STRIDEN_THREAD_ALIGN; + const size_t BlockedN = + (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) / MLAS_QGEMM_STRIDEN_THREAD_ALIGN; if (size_t(ThreadsPerGemm) > BlockedN) { ThreadsPerGemm = ptrdiff_t(BlockedN); @@ -3255,7 +2976,6 @@ MlasGemmBatch( WorkBlock.ThreadCountN = ThreadsPerGemm; } else { - if (size_t(ThreadsPerGemm) > M) { ThreadsPerGemm = ptrdiff_t(M); } @@ -3272,14 +2992,8 @@ MlasGemmBatch( }); } - -size_t -MLASCALL -MlasGemmPackBSize( - size_t N, - size_t K, - bool BIsSigned - ) +size_t MLASCALL +MlasGemmPackBSize(size_t N, size_t K, bool BIsSigned) /*++ Routine Description: @@ -3327,22 +3041,14 @@ Return Value: const size_t BytesRequired = (AlignedN * sizeof(int32_t)) + (AlignedN * AlignedK * sizeof(uint8_t)); const size_t BufferAlignment = MlasGetPreferredBufferAlignment(); - const size_t AlignedBytesRequired = (BytesRequired + BufferAlignment - 1) & - ~(BufferAlignment - 1); + const size_t AlignedBytesRequired = + (BytesRequired + BufferAlignment - 1) & ~(BufferAlignment - 1); return AlignedBytesRequired; } -void -MLASCALL -MlasGemmPackB( - size_t N, - size_t K, - const uint8_t* B, - size_t ldb, - bool BIsSigned, - void* PackedB - ) +void MLASCALL +MlasGemmPackB(size_t N, size_t K, const uint8_t* B, size_t ldb, bool BIsSigned, void* PackedB) /*++ Routine Description: @@ -3399,7 +3105,6 @@ Return Value: size_t CountK; for (size_t k = 0; k < K; k += CountK) { - CountK = std::min(K - k, PackedStrideK); // @@ -3411,13 +3116,13 @@ Return Value: size_t CountN; for (size_t n = 0; n < N; n += CountN) { - constexpr size_t BatchedN = 128; MLAS_DECLSPEC_ALIGN(int32_t ColumnSumBuffer[BatchedN], 64); CountN = std::min(N - n, BatchedN); - GemmU8X8Dispatch->CopyPackBRoutine(pb, B + n, ldb, CountN, CountK, ColumnSumBuffer, BIsSigned); + GemmU8X8Dispatch->CopyPackBRoutine(pb, B + n, ldb, CountN, CountK, ColumnSumBuffer, + BIsSigned); // // Accumulate this batch of the column sum buffer into the packed diff --git a/onnxruntime/core/providers/cpu/math/gemm.cc b/onnxruntime/core/providers/cpu/math/gemm.cc index 15a7034d1c768..c23398ee51c3b 100644 --- a/onnxruntime/core/providers/cpu/math/gemm.cc +++ b/onnxruntime/core/providers/cpu/math/gemm.cc @@ -70,10 +70,11 @@ ONNX_CPU_OPERATOR_TYPED_KERNEL( KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), Gemm); -bool GemmPackBFp32(const OpKernelInfo& info, +bool GemmPackBFp32(AllocatorPtr& alloc, const Tensor& tensor_b, bool trans_b, BufferUniquePtr& packed_b, + size_t& packed_b_size, TensorShape& b_shape) { // Only handle the common case of a 2D weight matrix. Additional matrices // could be handled by stacking the packed buffers. @@ -85,13 +86,18 @@ bool GemmPackBFp32(const OpKernelInfo& info, const size_t K = trans_b ? static_cast(b_shape[1]) : static_cast(b_shape[0]); const size_t N = trans_b ? static_cast(b_shape[0]) : static_cast(b_shape[1]); - const size_t packed_b_size = MlasGemmPackBSize(N, K); + packed_b_size = MlasGemmPackBSize(N, K); if (packed_b_size == 0) { return false; } - auto alloc = info.GetAllocator(0, OrtMemTypeDefault); auto* packed_b_data = alloc->Alloc(packed_b_size); + + // Initialize memory to 0 as there could be some padding associated with pre-packed + // buffer memory and we don not want it uninitialized and generate different hashes + // if and when we try to cache this pre-packed buffer for sharing between sessions. + memset(packed_b_data, 0, packed_b_size); + packed_b = BufferUniquePtr(packed_b_data, BufferDeleter(alloc)); MlasGemmPackB(trans_b ? CblasTrans : CblasNoTrans, N, @@ -164,18 +170,49 @@ template void Gemm::ComputeGemm(CBLAS_TRANSPOSE trans_a, CBLAS_TRANSPOSE concurrency::ThreadPool* thread_pool); template -Status Gemm::PrePack(const Tensor& /* tensor */, int /* input_idx */, bool& is_packed) { +Status Gemm::PrePack(const Tensor& /* tensor */, int /* input_idx */, AllocatorPtr /*alloc_for_caching*/, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* /*prepacked_weight_for_caching*/) { is_packed = false; return Status::OK(); } template <> -Status Gemm::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) { +Status Gemm::PrePack(const Tensor& tensor, int input_idx, + AllocatorPtr alloc, /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) { is_packed = false; // only pack Matrix B if (input_idx == 1) { - is_packed = GemmPackBFp32(Info(), tensor, trans_B_ != CblasNoTrans, packed_b_, b_shape_); + size_t packed_b_size; + is_packed = GemmPackBFp32(alloc, tensor, trans_B_ != CblasNoTrans, packed_b_, packed_b_size, b_shape_); + bool share_prepacked_weights = (prepacked_weights != nullptr); + if (is_packed && share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(packed_b_)); + prepacked_weights->buffer_sizes_.push_back(packed_b_size); + } + } + return Status::OK(); +} + +template +Status Gemm::UseSharedPrePackedBuffers(std::vector& /*prepacked_buffers*/, + int /*input_idx*/, + /*out*/ bool& used_shared_buffers) { + used_shared_buffers = false; + return Status::OK(); +} + +template <> +Status Gemm::UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) { + used_shared_buffers = false; + + if (input_idx == 1) { + used_shared_buffers = true; + packed_b_ = std::move(prepacked_buffers[0]); } return Status::OK(); } diff --git a/onnxruntime/core/providers/cpu/math/gemm.h b/onnxruntime/core/providers/cpu/math/gemm.h index cc91d756f0228..d0cf142cd0231 100644 --- a/onnxruntime/core/providers/cpu/math/gemm.h +++ b/onnxruntime/core/providers/cpu/math/gemm.h @@ -27,7 +27,13 @@ class Gemm : public OpKernel { Status Compute(OpKernelContext* context) const override; - Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; + Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) override; + + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) override; static void ComputeGemm(CBLAS_TRANSPOSE trans_a, CBLAS_TRANSPOSE trans_b, int64_t M, int64_t N, int64_t K, diff --git a/onnxruntime/core/providers/cpu/math/gemm_matmul_common.h b/onnxruntime/core/providers/cpu/math/gemm_matmul_common.h index ffd229e825629..57d0e3d0d116d 100644 --- a/onnxruntime/core/providers/cpu/math/gemm_matmul_common.h +++ b/onnxruntime/core/providers/cpu/math/gemm_matmul_common.h @@ -7,10 +7,11 @@ namespace onnxruntime { -bool GemmPackBFp32(const OpKernelInfo& info, +bool GemmPackBFp32(AllocatorPtr& alloc, const Tensor& tensor_b, bool trans_b, BufferUniquePtr& packed_b, + size_t& packed_b_size, TensorShape& b_shape); }; // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/matmul.cc b/onnxruntime/core/providers/cpu/math/matmul.cc index 6165d922fa847..702c63bb8f89f 100644 --- a/onnxruntime/core/providers/cpu/math/matmul.cc +++ b/onnxruntime/core/providers/cpu/math/matmul.cc @@ -126,16 +126,37 @@ Status MatMul::Compute(OpKernelContext* ctx) const { return Status::OK(); } -Status MatMul::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) { +Status MatMul::PrePack(const Tensor& tensor, int input_idx, /*out*/ AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) { is_packed = false; // only pack Matrix B if (input_idx == 1) { - is_packed = GemmPackBFp32(Info(), tensor, trans_b_attr_, packed_b_, b_shape_); + size_t packed_b_size; + is_packed = GemmPackBFp32(alloc, tensor, trans_b_attr_, packed_b_, packed_b_size, b_shape_); + bool share_prepacked_weights = (prepacked_weights != nullptr); + if (is_packed && share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(packed_b_)); + prepacked_weights->buffer_sizes_.push_back(packed_b_size); + } } return Status::OK(); } +Status MatMul::UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) { + used_shared_buffers = false; + + if (input_idx == 1) { + used_shared_buffers = true; + packed_b_ = std::move(prepacked_buffers[0]); + } + + return Status::OK(); +} + Status MatMul::Compute(OpKernelContext* ctx) const { concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool(); @@ -179,7 +200,7 @@ Status MatMul::Compute(OpKernelContext* ctx) const { data[i].beta = 0.0f; } MlasGemmBatch(trans_a ? CblasTrans : CblasNoTrans, trans_b ? CblasTrans : CblasNoTrans, - M, N, K, data.data(), max_len, thread_pool); + M, N, K, data.data(), max_len, thread_pool); return Status::OK(); } diff --git a/onnxruntime/core/providers/cpu/math/matmul.h b/onnxruntime/core/providers/cpu/math/matmul.h index 2bf12c666ca00..4f312855f9adf 100644 --- a/onnxruntime/core/providers/cpu/math/matmul.h +++ b/onnxruntime/core/providers/cpu/math/matmul.h @@ -24,7 +24,13 @@ class MatMul final : public OpKernel { info.GetAttrOrDefault("alpha", &alpha_attr_, 1.0); } - Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; + Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) override; + + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) override; Status Compute(OpKernelContext* context) const override; diff --git a/onnxruntime/core/providers/cpu/math/matmul_integer_base.h b/onnxruntime/core/providers/cpu/math/matmul_integer_base.h index 9f6fa16307511..920f7cac71193 100644 --- a/onnxruntime/core/providers/cpu/math/matmul_integer_base.h +++ b/onnxruntime/core/providers/cpu/math/matmul_integer_base.h @@ -11,7 +11,9 @@ class MatMulIntegerBase : public OpKernel { public: MatMulIntegerBase(const OpKernelInfo& info) : OpKernel(info) {} - Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override { + Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) override { is_packed = false; // only pack Matrix B @@ -23,26 +25,52 @@ class MatMulIntegerBase : public OpKernel { return Status::OK(); } + b_is_signed_ = tensor.IsDataType(); + const size_t K = static_cast(b_shape_[0]); const size_t N = static_cast(b_shape_[1]); const auto* b_data = static_cast(tensor.DataRaw()); - b_is_signed_ = tensor.IsDataType(); const size_t packed_b_size = MlasGemmPackBSize(N, K, b_is_signed_); if (packed_b_size == 0) { return Status::OK(); } - auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); auto* packed_b_data = alloc->Alloc(packed_b_size); + + // Initialize memory to 0 as there could be some padding associated with pre-packed + // buffer memory and we don not want it uninitialized and generate different hashes + // if and when we try to cache this pre-packed buffer for sharing between sessions. + memset(packed_b_data, 0, packed_b_size); + packed_b_ = BufferUniquePtr(packed_b_data, BufferDeleter(alloc)); MlasGemmPackB(N, K, b_data, N, b_is_signed_, packed_b_data); + + bool share_prepacked_weights = (prepacked_weights != nullptr); + if (share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(packed_b_)); + prepacked_weights->buffer_sizes_.push_back(packed_b_size); + } + is_packed = true; } return Status::OK(); } + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) override { + used_shared_buffers = false; + + if (input_idx == GetBIdx()) { + used_shared_buffers = true; + packed_b_ = std::move(prepacked_buffers[0]); + } + + return Status::OK(); + } + protected: /** * @return input index of Matrix B, the weight tensor diff --git a/onnxruntime/core/providers/cpu/nn/conv_transpose.cc b/onnxruntime/core/providers/cpu/nn/conv_transpose.cc index d28a63268b275..d7e39c16b6d58 100644 --- a/onnxruntime/core/providers/cpu/nn/conv_transpose.cc +++ b/onnxruntime/core/providers/cpu/nn/conv_transpose.cc @@ -37,13 +37,18 @@ ONNX_CPU_OPERATOR_KERNEL( ConvTranspose); template -Status ConvTranspose::PrePack(const Tensor& /* tensor */, int /* input_idx */, bool& is_packed) { +Status ConvTranspose::PrePack(const Tensor& /*tensor*/, int /*input_idx*/, AllocatorPtr /*alloc*/, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* /*prepacked_weights*/ +) { is_packed = false; return Status::OK(); } template <> -Status ConvTranspose::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) { +Status ConvTranspose::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) { is_packed = false; // only pack filter tensor @@ -56,12 +61,18 @@ Status ConvTranspose::PrePack(const Tensor& tensor, int input_idx, bool& const size_t K = static_cast(filter_shape_[0]) / conv_transpose_attrs_.group; const size_t N = filter_shape_.SizeFromDimension(1); auto packed_elements_per_group = N * K; - if (packed_elements_per_group == 0 || N == 1 || K == 1) { // No need for single row or single col case + if (packed_elements_per_group == 0 || N == 1 || K == 1) { // No need for single row or single col case return Status::OK(); } - auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); - auto* packed_filter_data = alloc->Alloc(packed_elements_per_group * sizeof(float) * conv_transpose_attrs_.group); + size_t packed_filter_data_size = packed_elements_per_group * sizeof(float) * conv_transpose_attrs_.group; + auto* packed_filter_data = alloc->Alloc(packed_filter_data_size); + + // Initialize memory to 0 as there could be some padding associated with pre-packed + // buffer memory and we don not want it uninitialized and generate different hashes + // if and when we try to cache this pre-packed buffer for sharing between sessions. + memset(packed_filter_data, 0, packed_filter_data_size); + transposed_filter_ = BufferUniquePtr(packed_filter_data, BufferDeleter(alloc)); for (int64_t group_id = 0; group_id < conv_transpose_attrs_.group; ++group_id) { @@ -70,11 +81,39 @@ Status ConvTranspose::PrePack(const Tensor& tensor, int input_idx, bool& K, N); } + bool share_prepacked_weights = (prepacked_weights != nullptr); + if (share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(transposed_filter_)); + prepacked_weights->buffer_sizes_.push_back(packed_filter_data_size); + } + is_packed = true; } return Status::OK(); } +template +Status ConvTranspose::UseSharedPrePackedBuffers(std::vector& /*prepacked_buffers*/, + int /*input_idx*/, + /*out*/ bool& used_shared_buffers) { + used_shared_buffers = false; + return Status::OK(); +} + +template <> +Status ConvTranspose::UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) { + used_shared_buffers = false; + + if (input_idx == 1) { + used_shared_buffers = true; + transposed_filter_ = std::move(prepacked_buffers[0]); + } + + return Status::OK(); +} + template Status ConvTranspose::Compute(OpKernelContext* context) const { return ConvTranspose::DoConvTranspose(context, false); diff --git a/onnxruntime/core/providers/cpu/nn/conv_transpose.h b/onnxruntime/core/providers/cpu/nn/conv_transpose.h index 6de3499e627b2..49c174969b38b 100644 --- a/onnxruntime/core/providers/cpu/nn/conv_transpose.h +++ b/onnxruntime/core/providers/cpu/nn/conv_transpose.h @@ -27,7 +27,13 @@ class ConvTranspose : public OpKernel { public: ConvTranspose(const OpKernelInfo& info) : OpKernel(info), conv_transpose_attrs_(info) {} - Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; + Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) override; + + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) override; Status Compute(OpKernelContext* context) const override; diff --git a/onnxruntime/core/providers/cpu/nn/qlinearconv.cc b/onnxruntime/core/providers/cpu/nn/qlinearconv.cc index 52e734cedde6d..f5fe4bf4fe8c4 100644 --- a/onnxruntime/core/providers/cpu/nn/qlinearconv.cc +++ b/onnxruntime/core/providers/cpu/nn/qlinearconv.cc @@ -23,7 +23,14 @@ class QLinearConv : public OpKernel { } Status Compute(OpKernelContext* context) const override; - Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; + + Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) override; + + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) override; private: static void ReorderFilter(const uint8_t* input, @@ -83,7 +90,9 @@ ONNX_OPERATOR_KERNEL_EX( #endif -Status QLinearConv::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) { +Status QLinearConv::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) { is_packed = false; // Support packing the weight matrix. @@ -91,6 +100,8 @@ Status QLinearConv::PrePack(const Tensor& tensor, int input_idx, bool& is_packed return Status::OK(); } + is_W_signed_ = tensor.IsDataType(); + const auto& shape = tensor.Shape().GetDims(); size_t rank = shape.size(); if (rank <= 2) { @@ -110,20 +121,26 @@ Status QLinearConv::PrePack(const Tensor& tensor, int input_idx, bool& is_packed const auto* Wdata = static_cast(tensor.DataRaw()); W_shape_ = shape; - is_W_signed_ = tensor.IsDataType(); - - auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); const size_t group_count = static_cast(conv_attrs_.group); const size_t group_output_channels = output_channels / group_count; const size_t kernel_dim = group_input_channels * kernel_size; + bool share_prepacked_weights = (prepacked_weights != nullptr); + // Don't pack the filter buffer if the MlasConvDepthwise path is used. if (group_input_channels != 1 && group_output_channels != 1) { packed_W_size_ = MlasGemmPackBSize(group_output_channels, kernel_dim, is_W_signed_); if (packed_W_size_ != 0) { - auto* packed_W = static_cast(alloc->Alloc(SafeInt(group_count) * packed_W_size_)); + size_t packed_W_data_size = SafeInt(group_count) * packed_W_size_; + auto* packed_W = static_cast(alloc->Alloc(packed_W_data_size)); + + // Initialize memory to 0 as there could be some padding associated with pre-packed + // buffer memory and we don not want it uninitialized and generate different hashes + // if and when we try to cache this pre-packed buffer for sharing between sessions. + memset(packed_W, 0, packed_W_data_size); + packed_W_buffer_ = BufferUniquePtr(packed_W, BufferDeleter(alloc)); // Allocate a temporary buffer to hold the reordered oihw->hwio filter for @@ -143,22 +160,69 @@ Status QLinearConv::PrePack(const Tensor& tensor, int input_idx, bool& is_packed Wdata += W_offset; } + if (share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(packed_W_buffer_)); + prepacked_weights->buffer_sizes_.push_back(packed_W_data_size); + } + is_W_packed_ = true; is_packed = true; return Status::OK(); + } else { + if (share_prepacked_weights) { + prepacked_weights->buffers_.push_back(nullptr); // packed_W_buffer_ is nullptr + prepacked_weights->buffer_sizes_.push_back(0); + } + } + } else { + if (share_prepacked_weights) { + prepacked_weights->buffers_.push_back(nullptr); // packed_W_buffer_ is nullptr + prepacked_weights->buffer_sizes_.push_back(0); } } - auto* reordered_W = static_cast(alloc->Alloc(SafeInt(sizeof(uint8_t)) * output_channels * group_input_channels * kernel_size)); + size_t reordered_w_data_size = SafeInt(sizeof(uint8_t)) * output_channels * group_input_channels * kernel_size; + auto* reordered_W = static_cast(alloc->Alloc(reordered_w_data_size)); + + // Initialize memory to 0 as there could be some padding associated with pre-packed + // buffer memory and we don not want it uninitialized and generate different hashes + // if and when we try to cache this pre-packed buffer for sharing between sessions. + memset(reordered_W, 0, reordered_w_data_size); + reordered_W_buffer_ = BufferUniquePtr(reordered_W, BufferDeleter(alloc)); ReorderFilter(Wdata, reordered_W, output_channels, group_input_channels, kernel_size); + if (share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(reordered_W_buffer_)); + prepacked_weights->buffer_sizes_.push_back(reordered_w_data_size); + } + is_W_packed_ = true; is_packed = true; return Status::OK(); } +Status QLinearConv::UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) { + if (input_idx != 3) { + return Status::OK(); + } + + used_shared_buffers = true; + + if (prepacked_buffers.size() == 1) { // This means that only packed_W_ exists + packed_W_buffer_ = std::move(prepacked_buffers[0]); + } else if (prepacked_buffers.size() == 2) { // This means that only reordered_W_ exists + // Enforce that the first "placeholder" buffer is nullptr + ORT_ENFORCE(prepacked_buffers[0].get() == nullptr); + reordered_W_buffer_ = std::move(prepacked_buffers[1]); + } + + return Status::OK(); +} + Status QLinearConv::Compute(OpKernelContext* context) const { const Tensor* X = context->Input(0); const Tensor* W = is_W_packed_ ? nullptr : context->Input(3); diff --git a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc index cb1185853637b..a711a59bd1ed8 100644 --- a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc +++ b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc @@ -174,7 +174,7 @@ using namespace rnn::detail; // LSTM details -Status DeepCpuLstmOp::TryPackWeights(const Tensor& weights, PackedWeights& packed_weights, bool& is_packed) { +Status DeepCpuLstmOp::TryPackWeights(const Tensor& weights, PackedWeights& packed_weights, bool& is_packed, AllocatorPtr& alloc) { const auto& shape = weights.Shape(); if (shape.NumDimensions() != 3) { return Status::OK(); @@ -194,9 +194,16 @@ Status DeepCpuLstmOp::TryPackWeights(const Tensor& weights, PackedWeights& packe return Status::OK(); } - auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); - auto* packed_weights_data = alloc->Alloc(SafeInt(packed_weights_size) * num_directions_); + size_t packed_weights_data_size = SafeInt(packed_weights_size) * num_directions_; + auto* packed_weights_data = alloc->Alloc(packed_weights_data_size); + + // Initialize memory to 0 as there could be some padding associated with pre-packed + // buffer memory and we don not want it uninitialized and generate different hashes + // if and when we try to cache this pre-packed buffer for sharing between sessions. + memset(packed_weights_data, 0, packed_weights_data_size); + packed_weights.buffer_ = BufferUniquePtr(packed_weights_data, BufferDeleter(alloc)); + packed_weights.buffer_size_ = packed_weights_data_size; packed_weights.weights_size_ = packed_weights_size; packed_weights.shape_ = shape; @@ -211,20 +218,55 @@ Status DeepCpuLstmOp::TryPackWeights(const Tensor& weights, PackedWeights& packe return Status::OK(); } -Status DeepCpuLstmOp::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) { +static void UseSharedPrePackedBuffersImpl(std::vector& prepacked_buffers, + rnn::detail::PackedWeights& packed_tensor) { + packed_tensor.buffer_ = std::move(prepacked_buffers[0]); +} + +Status DeepCpuLstmOp::PrePack(const Tensor& tensor, int input_idx, + AllocatorPtr alloc, /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) { is_packed = false; if (tensor.IsDataType()) { if (input_idx == 1) { - return TryPackWeights(tensor, packed_W_, is_packed); + ORT_RETURN_IF_ERROR(TryPackWeights(tensor, packed_W_, is_packed, alloc)); + + bool share_prepacked_weights = (prepacked_weights != nullptr); + if (is_packed && share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(packed_W_.buffer_)); + prepacked_weights->buffer_sizes_.push_back(packed_W_.buffer_size_); + } } else if (input_idx == 2) { - return TryPackWeights(tensor, packed_R_, is_packed); + ORT_RETURN_IF_ERROR(TryPackWeights(tensor, packed_R_, is_packed, alloc)); + + bool share_prepacked_weights = (prepacked_weights != nullptr); + if (is_packed && share_prepacked_weights) { + prepacked_weights->buffers_.push_back(std::move(packed_R_.buffer_)); + prepacked_weights->buffer_sizes_.push_back(packed_R_.buffer_size_); + } } } return Status::OK(); } +Status DeepCpuLstmOp::UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) { + used_shared_buffers = false; + + if (input_idx == 1) { + used_shared_buffers = true; + UseSharedPrePackedBuffersImpl(prepacked_buffers, packed_W_); + } else if (input_idx == 2) { + used_shared_buffers = true; + UseSharedPrePackedBuffersImpl(prepacked_buffers, packed_R_); + } + + return Status::OK(); +} + Status DeepCpuLstmOp::Compute(OpKernelContext* context) const { const Tensor& X = *context->Input(0); // inputs. [seq_length, batch_size, input_size] diff --git a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.h b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.h index 9e49adc1d0f6e..9c4c12954022a 100644 --- a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.h +++ b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.h @@ -18,13 +18,21 @@ class DeepCpuLstmOp final : public OpKernel, public LSTMBase { public: DeepCpuLstmOp(const OpKernelInfo& info) : OpKernel(info), LSTMBase(info) {} - Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; + Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, + /*out*/ PrePackedWeights* prepacked_weights) override; + + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) override; + Status Compute(OpKernelContext* context) const override; ~DeepCpuLstmOp() override = default; private: - Status TryPackWeights(const Tensor& weights, rnn::detail::PackedWeights& packed_weights, bool& is_packed); + Status TryPackWeights(const Tensor& weights, rnn::detail::PackedWeights& packed_weights, + bool& is_packed, AllocatorPtr& alloc); template Status ComputeImpl(OpKernelContext& context) const; diff --git a/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h b/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h index 63188c5f3ebc8..711afd9fffed1 100644 --- a/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h +++ b/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h @@ -166,6 +166,7 @@ void ComputeGemm(const int M, struct PackedWeights { BufferUniquePtr buffer_; + size_t buffer_size_; size_t weights_size_; TensorShape shape_; }; diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 550393ea5c198..0556e5c33369a 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -1114,6 +1114,22 @@ static bool ModelHasFP16Inputs(const Graph& graph) { return false; } +common::Status InferenceSession::AddPrePackedWeightsContainer(PrepackedWeightsContainer* prepacked_weights_container) { + if (prepacked_weights_container == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "The provided PrePackedWeightsContainer instance to be added to the session is null"); + } + + if (prepacked_weights_container_ != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "The session already has a PrePackedWeightsContainer instance"); + } + + prepacked_weights_container_ = prepacked_weights_container; + + return Status::OK(); +} + common::Status InferenceSession::Initialize() { Status status = Status::OK(); TimePoint tp; @@ -1191,7 +1207,8 @@ common::Status InferenceSession::Initialize() { *session_logger_, session_profiler_, session_options_.use_deterministic_compute, - session_options_.enable_mem_reuse); + session_options_.enable_mem_reuse, + prepacked_weights_container_); onnxruntime::Graph& graph = model_->MainGraph(); diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index 71d37f8055254..f49900928b538 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -14,6 +14,7 @@ #include "core/framework/framework_common.h" #include "core/framework/iexecutor.h" #include "core/framework/kernel_registry_manager.h" +#include "core/framework/prepacked_weights_container.h" #include "core/framework/session_state.h" #include "core/graph/basic_types.h" #include "core/optimizer/graph_transformer_level.h" @@ -428,6 +429,13 @@ class InferenceSession { return *session_state_; } + /** + * Add a PrepackedWeightsContainer instance to the session so as to store the pre-packed weights + * of shared initializers to be shared across sessions. + * @param prepacked_weights_container PrepackedWeightsContainer instance + */ + Status AddPrePackedWeightsContainer(PrepackedWeightsContainer* prepacked_weights_container); + protected: #if !defined(ORT_MINIMAL_BUILD) /** @@ -704,6 +712,11 @@ class InferenceSession { std::vector ort_format_model_bytes_; std::shared_ptr allocator_manager_; + + // Container to store pre-packed weights to share between sessions. + // The life-cycle of the cache itself is maintained by the user and the user will ensure + // the cache is valid until any session reliant on it is still in scope. + PrepackedWeightsContainer* prepacked_weights_container_ = nullptr; }; struct SessionIOBinding { diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index b4358ec0ce9b8..3201e123d0200 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -410,6 +410,7 @@ static ORT_STATUS_PTR CreateSessionAndLoadModel(_In_ const OrtSessionOptions* op _In_opt_z_ const ORTCHAR_T* model_path, _In_opt_ const void* model_data, size_t model_data_length, + std::unique_ptr& sess) { // quick check here to decide load path. InferenceSession will provide error message for invalid values. // TODO: Could move to a helper @@ -463,7 +464,8 @@ static ORT_STATUS_PTR CreateSessionAndLoadModel(_In_ const OrtSessionOptions* op } static ORT_STATUS_PTR InitializeSession(_In_ const OrtSessionOptions* options, - _In_ std::unique_ptr<::onnxruntime::InferenceSession>& sess) { + _In_ std::unique_ptr<::onnxruntime::InferenceSession>& sess, + _Inout_opt_ OrtPrepackedWeightsContainer* prepacked_weights_container = nullptr) { // we need to disable mem pattern if DML is one of the providers since DML doesn't have the concept of // byte addressable memory std::vector> provider_list; @@ -481,6 +483,11 @@ static ORT_STATUS_PTR InitializeSession(_In_ const OrtSessionOptions* options, } } + if (prepacked_weights_container != nullptr) { + ORT_API_RETURN_IF_STATUS_NOT_OK(sess->AddPrePackedWeightsContainer( + reinterpret_cast(prepacked_weights_container))); + } + ORT_API_RETURN_IF_STATUS_NOT_OK(sess->Initialize()); return nullptr; @@ -1901,6 +1908,68 @@ ORT_API(void, OrtApis::ReleaseArenaCfg, _Frees_ptr_opt_ OrtArenaCfg* ptr) { delete ptr; } +ORT_API_STATUS_IMPL(OrtApis::CreatePrepackedWeightsContainer, _Outptr_ OrtPrepackedWeightsContainer** out) { + API_IMPL_BEGIN + std::unique_ptr container(new PrepackedWeightsContainer()); + *out = reinterpret_cast(container.release()); + return nullptr; + API_IMPL_END +} + +ORT_API(void, OrtApis::ReleasePrepackedWeightsContainer, _Frees_ptr_opt_ OrtPrepackedWeightsContainer* ptr) { + delete reinterpret_cast(ptr); +} + +ORT_API_STATUS_IMPL(OrtApis::CreateSessionWithPrepackedWeightsContainer, _In_ const OrtEnv* env, _In_ const ORTCHAR_T* model_path, + _In_ const OrtSessionOptions* options, _Inout_ OrtPrepackedWeightsContainer* prepacked_weights_container, + _Outptr_ OrtSession** out) { + API_IMPL_BEGIN + std::unique_ptr sess; + OrtStatus* status = nullptr; + *out = nullptr; + + ORT_TRY { + ORT_API_RETURN_IF_ERROR(CreateSessionAndLoadModel(options, env, model_path, nullptr, 0, sess)); + ORT_API_RETURN_IF_ERROR(InitializeSession(options, sess, prepacked_weights_container)); + + *out = reinterpret_cast(sess.release()); + } + ORT_CATCH(const std::exception& e) { + ORT_HANDLE_EXCEPTION([&]() { + status = OrtApis::CreateStatus(ORT_FAIL, e.what()); + }); + } + + return status; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::CreateSessionFromArrayWithPrepackedWeightsContainer, _In_ const OrtEnv* env, + _In_ const void* model_data, size_t model_data_length, + _In_ const OrtSessionOptions* options, _Inout_ OrtPrepackedWeightsContainer* prepacked_weights_container, + _Outptr_ OrtSession** out) { + API_IMPL_BEGIN + std::unique_ptr sess; + OrtStatus* status = nullptr; + *out = nullptr; + + ORT_TRY { + ORT_API_RETURN_IF_ERROR(CreateSessionAndLoadModel(options, env, nullptr, model_data, + model_data_length, sess)); + ORT_API_RETURN_IF_ERROR(InitializeSession(options, sess, prepacked_weights_container)); + + *out = reinterpret_cast(sess.release()); + } + ORT_CATCH(const std::exception& e) { + ORT_HANDLE_EXCEPTION([&]() { + status = OrtApis::CreateStatus(ORT_FAIL, e.what()); + }); + } + + return status; + API_IMPL_END +} + #if defined(ORT_MINIMAL_BUILD) ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider_TensorRT, _In_ OrtSessionOptions* options, _In_ const OrtTensorRTProviderOptions* tensorrt_options) { @@ -2154,6 +2223,10 @@ static constexpr OrtApi ort_api_1_to_8 = { &OrtApis::KernelInfoGetAttributeArray_int64, &OrtApis::CreateArenaCfgV2, &OrtApis::AddRunConfigEntry, + &OrtApis::CreatePrepackedWeightsContainer, + &OrtApis::ReleasePrepackedWeightsContainer, + &OrtApis::CreateSessionWithPrepackedWeightsContainer, + &OrtApis::CreateSessionFromArrayWithPrepackedWeightsContainer, }; // Assert to do a limited check to ensure Version 1 of OrtApi never changes (will detect an addition or deletion but not if they cancel out each other) diff --git a/onnxruntime/core/session/ort_apis.h b/onnxruntime/core/session/ort_apis.h index a7ce1a54c81f8..2d99582724562 100644 --- a/onnxruntime/core/session/ort_apis.h +++ b/onnxruntime/core/session/ort_apis.h @@ -267,4 +267,14 @@ ORT_API_STATUS_IMPL(CreateArenaCfgV2, _In_reads_(num_keys) const char* const* ar _In_ size_t num_keys, _Outptr_ OrtArenaCfg** out); ORT_API_STATUS_IMPL(AddRunConfigEntry, _Inout_ OrtRunOptions* options, _In_z_ const char* config_key, _In_z_ const char* config_value); +ORT_API_STATUS_IMPL(CreatePrepackedWeightsContainer, _Outptr_ OrtPrepackedWeightsContainer** out); +ORT_API(void, ReleasePrepackedWeightsContainer, _Frees_ptr_opt_ OrtPrepackedWeightsContainer*); +ORT_API_STATUS_IMPL(CreateSessionWithPrepackedWeightsContainer, _In_ const OrtEnv* env, _In_ const ORTCHAR_T* model_path, + _In_ const OrtSessionOptions* options, _Inout_ OrtPrepackedWeightsContainer* prepacked_weights_container, + _Outptr_ OrtSession** out); +ORT_API_STATUS_IMPL(CreateSessionFromArrayWithPrepackedWeightsContainer, _In_ const OrtEnv* env, + _In_ const void* model_data, size_t model_data_length, + _In_ const OrtSessionOptions* options, _Inout_ OrtPrepackedWeightsContainer* prepacked_weights_container, + _Outptr_ OrtSession** out); + } // namespace OrtApis diff --git a/onnxruntime/test/contrib_ops/attention_op_test.cc b/onnxruntime/test/contrib_ops/attention_op_test.cc index 97da885edd4f2..2df0b167d707d 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/attention_op_test.cc @@ -15,7 +15,7 @@ enum MaskIndexType { kMaskRaw, kMask3D, kMaskDummy, // Dummy mask with shape [1, 1] or [batch_size, 1] - kMask4D // Megatron GPT2 mask with shape [batch_size, 1, max_sequence_length, max_sequence_length] + kMask4D // Megatron GPT2 mask with shape [batch_size, 1, max_sequence_length, max_sequence_length] }; static void RunAttentionTest( @@ -39,7 +39,7 @@ static void RunAttentionTest( int input_hidden_size = 0, int max_sequence_length = 0, bool only_enable_cuda = false) { - input_hidden_size = (input_hidden_size == 0 ? hidden_size : input_hidden_size); // By default, no pruning. + input_hidden_size = (input_hidden_size == 0 ? hidden_size : input_hidden_size); // By default, no pruning. int min_cuda_architecture = use_float16 ? 530 : 0; bool enable_cuda = HasCudaEnvironment(min_cuda_architecture) && !is_weights_constant; @@ -1410,7 +1410,7 @@ TEST(AttentionTest, Attention4DMask) { int past_sequence_length = 0; int input_hidden_size = 0; int max_sequence_length = 4; - bool only_enable_cuda = true; // only support 4D mask in cuda + bool only_enable_cuda = true; // only support 4D mask in cuda const std::vector* past_data = nullptr; const std::vector* present_data = nullptr; RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data, @@ -1490,7 +1490,6 @@ TEST(AttentionTest, AttentionPastState_dynamic) { test.Run(); } - TEST(AttentionTest, AttentionPrunedModel) { int batch_size = 2; int sequence_length = 2; @@ -1506,13 +1505,79 @@ TEST(AttentionTest, AttentionPrunedModel) { 0.5f, 0.2f, 0.3f, -0.6f, 6.0f, 7.0f}; std::vector weight_data = { - 0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f, - 0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f, - 0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f, - 0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f, - 0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 1.0f, 1.1f, 1.2f, - 1.2f, 1.1f, 1.0f, 0.9f, 0.8f, 0.7f, 0.6f, 0.5f, 0.4f, 0.3f, 0.2f, 0.1f, - }; + 0.1f, + -0.2f, + 0.3f, + 1.0f, + 1.1f, + 0.3f, + 0.5f, + 0.2f, + 0.3f, + -0.6f, + 1.5f, + 2.0f, + 0.5f, + 0.1f, + 0.4f, + 1.6f, + 1.0f, + 2.0f, + 0.4f, + 0.8f, + 0.9f, + 0.1f, + -1.3f, + 0.7f, + 0.3f, + 0.2f, + 4.0f, + 2.2f, + 1.6f, + 1.1f, + 0.7f, + 0.2f, + 0.4f, + 1.0f, + 1.2f, + 0.5f, + 0.2f, + 0.1f, + 0.4f, + 1.6f, + 2.4f, + 3.3f, + 2.1f, + 4.2f, + 8.4f, + 0.0f, + 2.1f, + 3.2f, + 0.1f, + 0.2f, + 0.3f, + 0.4f, + 0.5f, + 0.6f, + 0.7f, + 0.8f, + 0.9f, + 1.0f, + 1.1f, + 1.2f, + 1.2f, + 1.1f, + 1.0f, + 0.9f, + 0.8f, + 0.7f, + 0.6f, + 0.5f, + 0.4f, + 0.3f, + 0.2f, + 0.1f, + }; std::vector bias_data = { -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f}; @@ -1523,8 +1588,7 @@ TEST(AttentionTest, AttentionPrunedModel) { 11.689527f, 2.769937f, 7.05f, 8.350000f, 11.690000f, 2.770000f, 7.05f, 8.350000f, 14.276558f, 5.374159f, 9.650001f, 10.95f, - 14.289073f, 5.370287f, 9.650001f, 10.95f - }; + 14.289073f, 5.370287f, 9.650001f, 10.95f}; bool use_float16 = false; bool is_unidirectional = false; @@ -1536,5 +1600,119 @@ TEST(AttentionTest, AttentionPrunedModel) { batch_size, sequence_length, hidden_size, number_of_heads, use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data, kMaskRaw, input_hidden_size); } + +#ifndef ENABLE_TRAINING // Prepacking is enabled only on non-training builds +TEST(AttentionTest, SharedPrepackedWeights) { + int batch_size = 2; + int sequence_length = 2; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.5f, 0.2f, 0.3f, -0.6f, + 0.8f, -0.5f, 0.0f, 1.f, + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector weight_data = { + 0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f, + 0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f, + 0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f, + 0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f}; + + // Test that all attention masks are zero. + std::vector mask_index_data = {0, 0, 2, 2}; + + std::vector output_data = { + 3.96724534f, 0.07324841f, 4.25f, 5.65f, + 3.14984703f, 0.10842596f, 4.25f, 5.65f, + 3.14984703f, 0.10842596f, 4.25f, 5.65f, + 3.96724534f, 0.07324841f, 4.25f, 5.65f}; + + OpTester tester("Attention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(number_of_heads)); + tester.AddAttribute("unidirectional", static_cast(0)); + + std::vector input_dims = {batch_size, sequence_length, hidden_size}; + std::vector weights_dims = {hidden_size, 3 * hidden_size}; + std::vector bias_dims = {3 * hidden_size}; + + std::vector mask_index_dims = {2 * batch_size}; + std::vector output_dims = {batch_size, sequence_length, hidden_size}; + + tester.AddInput("input", input_dims, input_data); + tester.AddInput("weight", weights_dims, weight_data, true); // Trigger pre-packing + tester.AddInput("bias", bias_dims, bias_data); + tester.AddOutput("output", output_dims, output_data); + tester.AddInput("mask_index", mask_index_dims, mask_index_data); + + auto p_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape(weights_dims), + weight_data.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + OrtValue weight; + + weight.Init(p_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + + SessionOptions so; + + // Set up weight as a shared initializer to be shared between sessions + ASSERT_EQ(so.AddInitializer("weight", &weight), Status::OK()); + + // We want all sessions running using this OpTester to be able to share pre-packed weights if applicable + tester.EnableSharingOfPrePackedWeightsAcrossSessions(); + + // Pre-packing is limited just to the CPU EP for now and we will only test the CPU EP + // and we want to ensure that it is available in this build + auto cpu_ep = []() -> std::vector> { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + return execution_providers; + }; + + size_t number_of_pre_packed_weights_counter_session_1 = 0; + size_t number_of_shared_pre_packed_weights_counter = 0; + + // Session 1 + { + auto ep_vec = cpu_ep(); + tester.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter); + // Assert that no pre-packed weights have been shared thus far + ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0)); + } + + auto number_of_elements_in_shared_prepacked_buffers_container = + tester.GetNumPrePackedWeightsShared(); + // Assert that the number of elements in the shared container + // is the same as the number of weights that have been pre-packed + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_elements_in_shared_prepacked_buffers_container); + + // On some platforms/architectures MLAS may choose to not do any pre-packing and the number of elements + // that have been pre-packed will be zero in which case we do not continue with the testing + // of "sharing" of pre-packed weights as there are no pre-packed weights to be shared at all. + if (number_of_pre_packed_weights_counter_session_1 == 0) + return; + + // Session 2 + { + size_t number_of_pre_packed_weights_counter_session_2 = 0; + auto ep_vec = cpu_ep(); + tester.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter); + + // Assert that the same number of weights were pre-packed in both sessions + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_pre_packed_weights_counter_session_2); + + // Assert that the number of pre-packed weights that were shared equals + // the number of pre-packed weights in the second session + ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, + static_cast(number_of_shared_pre_packed_weights_counter)); + } +} +#endif + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc b/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc index 7074527003980..d3fd22a7da1e2 100644 --- a/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/quantize_attention_op_test.cc @@ -741,7 +741,7 @@ void TestQuantizedAttentionPastState(int64_t batch, constexpr int32_t weight_range = weight_max - weight_min; std::vector weight_zero_point(weight_scale_zp_size); - for(auto& zp : weight_zero_point) { + for (auto& zp : weight_zero_point) { zp = static_cast(random.Uniform({1}, weight_min, weight_max)[0]); } @@ -857,5 +857,120 @@ TEST(QAttentionTest, QAttentionPrunedModel) { input_hidden_size); } +#ifndef ENABLE_TRAINING // Prepacking is enabled only on non-training builds +TEST(QAttentionTest, SharedPrepackedWeights) { + int batch_size = 1; + int sequence_length = 2; + int hidden_size = 4; + int number_of_heads = 2; + + std::vector input_data = { + 0.8f, -0.5f, 0.0f, 1.f, + 0.5f, 0.2f, 0.3f, -0.6f}; + + std::vector weight_data = { + 0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f, + 0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f, + 0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f, + 0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f}; + + std::vector bias_data = { + -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f}; + + std::vector mask_index_data = {2L}; + + std::vector output_data = { + 3.1495983600616455f, 0.10843668878078461f, 4.25f, 5.6499996185302734f, + 3.9696791172027588f, 0.073143675923347473f, 4.2499995231628418f, 5.6499991416931152f}; + + std::vector input_dims = {batch_size, sequence_length, hidden_size}; + std::vector weights_dims = {hidden_size, 3 * hidden_size}; + std::vector bias_dims = {3 * hidden_size}; + std::vector mask_index_dims = {batch_size}; + std::vector output_dims = {batch_size, sequence_length, hidden_size}; + + OpTester tester("QAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(number_of_heads)); + + tester.AddInput("input", input_dims, ToInteger(input_data, 0.1f, 128)); + auto weight_data_converted_to_int = ToInteger(weight_data, 0.1f, 128); + tester.AddInput("weight", weights_dims, weight_data_converted_to_int, true); // Trigger pre-packing + + tester.AddInput("bias", bias_dims, bias_data); + tester.AddInput("input_scale", {1}, {0.1f}); + tester.AddInput("weight_scale", {1}, {0.1f}); + tester.AddOutput("output", output_dims, output_data); + + tester.AddInput("mask_index", mask_index_dims, mask_index_data); + + tester.AddInput("input_zero_point", {1}, {128}); + tester.AddInput("weight_zero_point", {1}, {128}); + + auto p_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape(weights_dims), + weight_data_converted_to_int.data(), + OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + OrtValue weight; + + weight.Init(p_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + + SessionOptions so; + + // Set up weight as a shared initializer to be shared between sessions + ASSERT_EQ(so.AddInitializer("weight", &weight), Status::OK()); + + // We want all sessions running using this OpTester to be able to share pre-packed weights if applicable + tester.EnableSharingOfPrePackedWeightsAcrossSessions(); + + // Pre-packing is limited just to the CPU EP for now and we will only test the CPU EP + // and we want to ensure that it is available in this build + auto cpu_ep = []() -> std::vector> { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + return execution_providers; + }; + + size_t number_of_pre_packed_weights_counter_session_1 = 0; + size_t number_of_shared_pre_packed_weights_counter = 0; + + // Session 1 + { + auto ep_vec = cpu_ep(); + tester.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter); + // Assert that no pre-packed weights have been shared thus far + ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0)); + } + + auto number_of_elements_in_shared_prepacked_buffers_container = + tester.GetNumPrePackedWeightsShared(); + // Assert that the number of elements in the shared container + // is the same as the number of weights that have been pre-packed + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_elements_in_shared_prepacked_buffers_container); + + // On some platforms/architectures MLAS may choose to not do any pre-packing and the number of elements + // that have been pre-packed will be zero in which case we do not continue with the testing + // of "sharing" of pre-packed weights as there are no pre-packed weights to be shared at all. + if (number_of_pre_packed_weights_counter_session_1 == 0) + return; + + // Session 2 + { + size_t number_of_pre_packed_weights_counter_session_2 = 0; + auto ep_vec = cpu_ep(); + tester.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter); + + // Assert that the same number of weights were pre-packed in both sessions + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_pre_packed_weights_counter_session_2); + + // Assert that the number of pre-packed weights that were shared equals + // the number of pre-packed weights in the second session + ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, + static_cast(number_of_shared_pre_packed_weights_counter)); + } +} +#endif + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc index 14b95de9bbf67..f225cf33259cb 100644 --- a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc +++ b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc @@ -362,5 +362,167 @@ TEST(DynamicQuantLSTMTest, LargeSize) { RunQuantLSTM(12, 3, 278); } +#ifndef ENABLE_TRAINING // Prepacking is enabled only on non-training builds +TEST(DynamicQuantLSTMTest, SharedPrepackedWeights) { + OpTester test("DynamicQuantizeLSTM", 1 /*opset_version*/, onnxruntime::kMSDomain /*domain*/); + + int num_directions = 1; + int input_size = 2; + int batch_size = 1; + int hidden_size = 16; + + std::vector activations; + activations = {"sigmoid", "tanh", "tanh"}; + test.AddAttribute>("activations", activations); + + test.AddAttribute("direction", "forward"); + test.AddAttribute("hidden_size", static_cast(hidden_size)); + test.AddAttribute("input_forget", 0); + + RandomValueGenerator rand_gen; + + // X + int64_t seq_len = 1; // only use seq length 1 to model the test + std::vector X_dims = {seq_len, batch_size, input_size}; + std::vector X_data = rand_gen.Gaussian({seq_len, batch_size, input_size}, 0.0f, 0.25f); + test.AddInput("X", X_dims, X_data); + + // W + std::vector W_dims = {num_directions, input_size, 4 * hidden_size}; + std::vector W_data = rand_gen.Gaussian({num_directions, 4 * hidden_size, input_size}, 0.0f, 0.25f); + + std::vector w_scale; + std::vector w_zp; + std::vector w_quant; + QuantizeWeight(w_quant, w_scale, w_zp, W_data, num_directions, 4 * hidden_size, input_size, false); + test.AddInput("W", W_dims, w_quant, true); // Trigger pre-packing + + // R + std::vector R_dims = {num_directions, hidden_size, 4 * hidden_size}; + std::vector R_data = rand_gen.Gaussian({num_directions, 4 * hidden_size, hidden_size}, 0.0f, 0.25f); + + std::vector r_scale; + std::vector r_zp; + std::vector r_quant; + QuantizeWeight(r_quant, r_scale, r_zp, R_data, num_directions, 4 * hidden_size, hidden_size, false); + test.AddInput("R", R_dims, r_quant, true); // Trigger pre-packing + + // B + test.AddMissingOptionalInput(); + + // sequence_lens + test.AddMissingOptionalInput(); + + // initial_h + std::vector initial_h_dims = {num_directions, batch_size, hidden_size}; + std::vector initial_h_data = rand_gen.Gaussian(initial_h_dims, 0.0f, 0.25f); + test.AddInput("initial_h", initial_h_dims, initial_h_data); + + // initial_c + std::vector initial_c_dims = {num_directions, batch_size, hidden_size}; + std::vector initial_c_data = rand_gen.Gaussian(initial_c_dims, 0.0f, 0.25f); + test.AddInput("initial_c", initial_c_dims, initial_c_data); + + test.AddMissingOptionalInput(); + + std::vector per_tensor_dims = {num_directions}; + test.AddInput("W_scale", per_tensor_dims, w_scale); + test.AddInput("W_zero_point", per_tensor_dims, w_zp); + + test.AddInput("R_scale", per_tensor_dims, r_scale); + test.AddInput("R_zero_point", per_tensor_dims, r_zp); + + std::vector Y_data; + std::vector Y_h_data; + std::vector Y_c_data; + ComputeRefOutput(Y_data, Y_h_data, Y_c_data, + input_size, batch_size, hidden_size, + X_data, W_data, R_data, + nullptr, + nullptr, + initial_h_data, initial_c_data, + "forward", activations, false); + + std::vector Y_dims = {seq_len, num_directions, batch_size, hidden_size}; + test.AddOutput("Y", Y_dims, Y_data); + + std::vector Y_h_dims{num_directions, batch_size, hidden_size}; + test.AddOutput("Y_h", Y_h_dims, Y_h_data); + + std::vector Y_c_dims{num_directions, batch_size, hidden_size}; + test.AddOutput("Y_c", Y_c_dims, Y_c_data); + + auto W_quant_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape(W_dims), + w_quant.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + OrtValue W; + + W.Init(W_quant_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + + auto R_quant_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape(R_dims), + r_quant.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + OrtValue R; + + R.Init(R_quant_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + SessionOptions so; + + // Set up weight(s) as a shared initializer to be shared between sessions + ASSERT_EQ(so.AddInitializer("W", &W), Status::OK()); + ASSERT_EQ(so.AddInitializer("R", &R), Status::OK()); + + // We want all sessions running using this OpTester to be able to share pre-packed weights if applicable + test.EnableSharingOfPrePackedWeightsAcrossSessions(); + + // Pre-packing is limited just to the CPU EP for now and we will only test the CPU EP + // and we want to ensure that it is available in this build + auto cpu_ep = []() -> std::vector> { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + return execution_providers; + }; + + size_t number_of_pre_packed_weights_counter_session_1 = 0; + size_t number_of_shared_pre_packed_weights_counter = 0; + + // Session 1 + { + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter); + // Assert that no pre-packed weights have been shared thus far + ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0)); + } + + auto number_of_elements_in_shared_prepacked_buffers_container = + test.GetNumPrePackedWeightsShared(); + // Assert that the number of elements in the shared container + // is the same as the number of weights that have been pre-packed + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_elements_in_shared_prepacked_buffers_container); + + // On some platforms/architectures MLAS may choose to not do any pre-packing and the number of elements + // that have been pre-packed will be zero in which case we do not continue with the testing + // of "sharing" of pre-packed weights as there are no pre-packed weights to be shared at all. + if (number_of_pre_packed_weights_counter_session_1 == 0) + return; + + // Session 2 + { + size_t number_of_pre_packed_weights_counter_session_2 = 0; + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter); + + // Assert that the same number of weights were pre-packed in both sessions + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_pre_packed_weights_counter_session_2); + + // Assert that the number of pre-packed weights that were shared equals + // the number of pre-packed weights in the second session + ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, + static_cast(number_of_shared_pre_packed_weights_counter)); + } +} +#endif + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/framework/session_state_test.cc b/onnxruntime/test/framework/session_state_test.cc index cfe37b2af47c3..2c7647f4fe02a 100644 --- a/onnxruntime/test/framework/session_state_test.cc +++ b/onnxruntime/test/framework/session_state_test.cc @@ -89,7 +89,7 @@ TEST_P(SessionStateAddGetKernelTest, AddGetKernelTest) { ASSERT_STATUS_OK(kernel_registry->Register(KernelCreateInfo( std::move(kernel_def), [](const OpKernelInfo& info) -> OpKernel* { return new TestOpKernel(info); }))); kernel_registry_manager.RegisterKernelRegistry(kernel_registry); - ASSERT_STATUS_OK(s.FinalizeSessionState(ORT_TSTR(""), kernel_registry_manager, SessionOptions())); + ASSERT_STATUS_OK(s.FinalizeSessionState(ORT_TSTR(""), kernel_registry_manager)); auto test_kernel = s.GetKernel(node.Index()); std::cout << "orig: " << orig_num_outputs << " new: " << test_kernel->Node().OutputDefs().size() << std::endl; @@ -292,12 +292,40 @@ class PrePackingTestOpKernel : public OpKernel { return Status::OK(); } - Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override { + Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, + int input_idx, + /*out*/ bool& used_shared_buffers) override { + ORT_UNUSED_PARAMETER(input_idx); + + weight_packed_ = std::move(prepacked_buffers[0]); + used_shared_buffers = true; + ++store_pre_packed_weight_calls_count; + return Status::OK(); + } + + Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, /*out*/ PrePackedWeights* prepacked_weights) override { ORT_UNUSED_PARAMETER(tensor); ORT_UNUSED_PARAMETER(input_idx); + + weight_packed_ = BufferUniquePtr(alloc->Alloc(8), BufferDeleter(alloc)); + float* data_weights_packed = reinterpret_cast(weight_packed_.get()); + data_weights_packed[0] = 1.2345f; + data_weights_packed[1] = data_weights_packed[0] * 2.f; + + if (prepacked_weights != nullptr) { + prepacked_weights->buffers_.push_back(std::move(weight_packed_)); + prepacked_weights->buffer_sizes_.push_back(8); + } + is_packed = true; + ++prepack_calls_count; return Status::OK(); } + + int prepack_calls_count = 0; + int store_pre_packed_weight_calls_count = 0; + BufferUniquePtr weight_packed_; }; static void CreateSimpleGraph(Graph& graph) { @@ -498,6 +526,261 @@ TEST_P(SessionStatePrepackingTest, PrePackingTest) { ASSERT_EQ(const_initialized_tensors.size(), size_t(test_param.test_prepacking ? 0 : 1)); } +TEST(SessionStateTest, SharedInitalizersWithPrePackingTest) { + OrtThreadPoolParams to; + auto tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP); + ONNX_OPERATOR_SCHEMA(PrePackingTest) + .SetDoc("Faking Node for PrePacking") + .Input(0, "Input_0", "input 0", "tensor(float)") + .Input(1, "Input_1", "input 1", "tensor(float)") + .Output(0, "output_0", "docstr for output_0.", "tensor(float)"); + + ExecutionProviders execution_providers; + auto cpu_execution_provider = std::make_unique(CPUExecutionProviderInfo(false)); + execution_providers.Add(kCpuExecutionProvider, std::move(cpu_execution_provider)); + + DataTransferManager dtm; + profiling::Profiler profiler; + + std::unordered_map domain_to_version; + domain_to_version[kOnnxDomain] = 11; + + KernelRegistryManager kernel_registry_manager; + Status status = kernel_registry_manager.RegisterKernels(execution_providers); + ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); + std::shared_ptr kernel_registry = std::make_shared(); + auto kernel_def = KernelDefBuilder().SetName("PrePackingTest").Provider(kCpuExecutionProvider).SinceVersion(1).Build(); + ASSERT_STATUS_OK(kernel_registry->Register( + KernelCreateInfo(std::move(kernel_def), + [](const OpKernelInfo& info) -> OpKernel* { return new PrePackingTestOpKernel(info); }))); + kernel_registry_manager.RegisterKernelRegistry(kernel_registry); + + // Part 1: Pre-packing enabled + no shared initializers = no pre-packed weights caching + { + SessionOptions sess_options; + // Enable pre-packing + sess_options.config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "0"; + + // First session/model + Model model_1("graph_main", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, std::vector(), + DefaultLoggingManager().DefaultLogger()); + + CreateSimpleGraph(model_1.MainGraph()); + PlaceAllNodesToCPUEP(model_1.MainGraph()); + SessionState session_state_1(model_1.MainGraph(), + execution_providers, + true, /*enable_mem_pattern*/ + tp.get(), + nullptr, /*inter_op_thread_pool*/ + dtm, + DefaultLoggingManager().DefaultLogger(), + profiler); + + ASSERT_STATUS_OK(session_state_1.FinalizeSessionState(std::basic_string(), + kernel_registry_manager, + sess_options)); + + const auto* kernel = reinterpret_cast(session_state_1.GetKernel(0)); + + // Assert that a pre-pack call was made and that no mechanism to store weight from shared container was invoked + ASSERT_EQ(session_state_1.GetNumberOfPrepacksCounter(), static_cast(1)); + ASSERT_EQ(kernel->prepack_calls_count, 1); + ASSERT_EQ(kernel->store_pre_packed_weight_calls_count, 0); + + // Second session/model + Model model_2("graph_main", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, std::vector(), + DefaultLoggingManager().DefaultLogger()); + + CreateSimpleGraph(model_2.MainGraph()); + PlaceAllNodesToCPUEP(model_2.MainGraph()); + SessionState session_state_2(model_2.MainGraph(), + execution_providers, + true, /*enable_mem_pattern*/ + tp.get(), + nullptr, /*inter_op_thread_pool*/ + dtm, + DefaultLoggingManager().DefaultLogger(), + profiler); + + ASSERT_STATUS_OK(session_state_2.FinalizeSessionState(std::basic_string(), + kernel_registry_manager, + sess_options)); + + kernel = reinterpret_cast(session_state_2.GetKernel(0)); + + // Assert that a pre-pack call was made and that no mechanism to store weight from shared container was invoked + ASSERT_EQ(session_state_2.GetNumberOfPrepacksCounter(), static_cast(1)); + ASSERT_EQ(kernel->prepack_calls_count, 1); + ASSERT_EQ(kernel->store_pre_packed_weight_calls_count, 0); + } + + // Part 2: Pre-packing enabled + shared initializers + no pre-packed weights container = no pre-packed weights caching + { + SessionOptions sess_options; + // Enable pre-packing + sess_options.config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "0"; + + // Enable shared initializer + OrtMemoryInfo mem_info(CPU, OrtDeviceAllocator); + std::vector float_data(1, 1); + std::unique_ptr tensor = + std::make_unique(DataTypeImpl::GetType(), TensorShape(std::vector{1}), reinterpret_cast(float_data.data()), mem_info, 0); + + auto value = std::make_unique(); + auto ml_tensor = DataTypeImpl::GetType(); + value->Init(tensor.release(), + ml_tensor, + ml_tensor->GetDeleteFunc()); + + sess_options.AddInitializer("node_0_input_1", value.get()); + + // First session/model + Model model_1("graph_main", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, std::vector(), + DefaultLoggingManager().DefaultLogger()); + + CreateSimpleGraph(model_1.MainGraph()); + PlaceAllNodesToCPUEP(model_1.MainGraph()); + SessionState session_state_1(model_1.MainGraph(), + execution_providers, + true, /*enable_mem_pattern*/ + tp.get(), + nullptr, /*inter_op_thread_pool*/ + dtm, + DefaultLoggingManager().DefaultLogger(), + profiler); + + ASSERT_STATUS_OK(session_state_1.FinalizeSessionState(std::basic_string(), + kernel_registry_manager, + sess_options)); + + const auto* kernel = reinterpret_cast(session_state_1.GetKernel(0)); + + // Assert that a pre-pack call was made and that no mechanism to store weight from shared container was invoked + ASSERT_EQ(session_state_1.GetNumberOfPrepacksCounter(), static_cast(1)); + ASSERT_EQ(kernel->prepack_calls_count, 1); + ASSERT_EQ(kernel->store_pre_packed_weight_calls_count, 0); + + // Second session/model + Model model_2("graph_main", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, std::vector(), + DefaultLoggingManager().DefaultLogger()); + + CreateSimpleGraph(model_2.MainGraph()); + PlaceAllNodesToCPUEP(model_2.MainGraph()); + SessionState session_state_2(model_2.MainGraph(), + execution_providers, + true, /*enable_mem_pattern*/ + tp.get(), + nullptr, /*inter_op_thread_pool*/ + dtm, + DefaultLoggingManager().DefaultLogger(), + profiler); + + ASSERT_STATUS_OK(session_state_2.FinalizeSessionState(std::basic_string(), + kernel_registry_manager, + sess_options)); + + kernel = reinterpret_cast(session_state_2.GetKernel(0)); + + // Assert that a pre-pack call was made and that no mechanism to store weight from shared container was invoked + ASSERT_EQ(session_state_2.GetNumberOfPrepacksCounter(), static_cast(1)); + ASSERT_EQ(kernel->prepack_calls_count, 1); + ASSERT_EQ(kernel->store_pre_packed_weight_calls_count, 0); + } + + // Part 3: Pre-packing enabled + shared initializers + pre-packed weights container = pre-packed weights caching enabled + { + SessionOptions sess_options; + // Enable pre-packing + sess_options.config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "0"; + + // Enable shared initializer + OrtMemoryInfo mem_info(CPU, OrtDeviceAllocator); + std::vector float_data(1, 1); + std::unique_ptr tensor = + std::make_unique(DataTypeImpl::GetType(), TensorShape(std::vector{1}), reinterpret_cast(float_data.data()), mem_info, 0); + + auto value = std::make_unique(); + auto ml_tensor = DataTypeImpl::GetType(); + value->Init(tensor.release(), + ml_tensor, + ml_tensor->GetDeleteFunc()); + + sess_options.AddInitializer("node_0_input_1", value.get()); + + // Enable pre-packed weights container + PrepackedWeightsContainer prepacked_weights_container; + + // First session/model + Model model_1("graph_main", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, std::vector(), + DefaultLoggingManager().DefaultLogger()); + + CreateSimpleGraph(model_1.MainGraph()); + PlaceAllNodesToCPUEP(model_1.MainGraph()); + SessionState session_state_1(model_1.MainGraph(), + execution_providers, + true, /*enable_mem_pattern*/ + tp.get(), + nullptr, /*inter_op_thread_pool*/ + dtm, + DefaultLoggingManager().DefaultLogger(), + profiler, + false, true, + &prepacked_weights_container); + + ASSERT_STATUS_OK(session_state_1.FinalizeSessionState(std::basic_string(), + kernel_registry_manager, + sess_options)); + + const auto* kernel = reinterpret_cast(session_state_1.GetKernel(0)); + // Assert that a pre-pack call was made + ASSERT_EQ(session_state_1.GetNumberOfPrepacksCounter(), static_cast(1)); + ASSERT_EQ(kernel->prepack_calls_count, 1); + // Assert that we made a call to store pre-packed weight from a shared container + ASSERT_EQ(kernel->store_pre_packed_weight_calls_count, 1); + // The weight to be "stored" is the same weight that we got by invoking PrePack() in the step above. + // Hence, assert that it wasn't a "cached" pre-packed weight (i.e.) pre-packed weight + // from another instance of the same op_type consuming the same constant initializer. + ASSERT_EQ(session_state_1.GetUsedSharedPrePackedWeightCounter(), static_cast(0)); + + // Second session/model + Model model_2("graph_main", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, std::vector(), + DefaultLoggingManager().DefaultLogger()); + + CreateSimpleGraph(model_2.MainGraph()); + PlaceAllNodesToCPUEP(model_2.MainGraph()); + SessionState session_state_2(model_2.MainGraph(), + execution_providers, + true, /*enable_mem_pattern*/ + tp.get(), + nullptr, /*inter_op_thread_pool*/ + dtm, + DefaultLoggingManager().DefaultLogger(), + profiler, + false, true, + &prepacked_weights_container); + + ASSERT_STATUS_OK(session_state_2.FinalizeSessionState(std::basic_string(), + kernel_registry_manager, + sess_options)); + + // Assert that a pre-pack call was made + ASSERT_EQ(session_state_2.GetNumberOfPrepacksCounter(), static_cast(1)); + ASSERT_EQ(kernel->prepack_calls_count, 1); + // Assert that we made a call to store pre-packed weight from a shared container + ASSERT_EQ(kernel->store_pre_packed_weight_calls_count, 1); + // The weight to be "stored" is a "cached" weight (i.e.) a pre-packed weight + // from another instance of the same op_type consuming the same constant initializer. + // Assert this. + ASSERT_EQ(session_state_2.GetUsedSharedPrePackedWeightCounter(), static_cast(1)); + } +} + INSTANTIATE_TEST_SUITE_P(SessionStateTests, SessionStatePrepackingTest, testing::Values(PrepackingTestParam{false, false}, diff --git a/onnxruntime/test/providers/cpu/math/gemm_test.cc b/onnxruntime/test/providers/cpu/math/gemm_test.cc index c3eef92126b2d..c99b7af5bb657 100644 --- a/onnxruntime/test/providers/cpu/math/gemm_test.cc +++ b/onnxruntime/test/providers/cpu/math/gemm_test.cc @@ -448,5 +448,91 @@ TEST(GemmOpTest, GemmWithAlphaOpset11) { TestGemmWithAlphaOpset11(); } +#ifndef ENABLE_TRAINING // Prepacking is enabled only on non-training builds +TEST(GemmOpTest, SharedPrepackedWeights) { + OpTester test("Gemm"); + + test.AddAttribute("transA", (int64_t)0); + test.AddAttribute("transB", (int64_t)0); + test.AddAttribute("alpha", 1.0f); + test.AddAttribute("beta", 1.0f); + + std::vector b_init_values(12, 1.0f); + test.AddInput("A", {2, 4}, + {1.0f, 2.0f, 3.0f, 4.0f, + -1.0f, -2.0f, -3.0f, -4.0f}); + // B is to be an initializer for triggering pre-packing + test.AddInput("B", {4, 3}, b_init_values, true); + test.AddInput("C", {2, 3}, std::vector(6, 1.0f)); + test.AddOutput("Y", {2, 3}, + {11.0f, 11.0f, 11.0f, + -9.0f, -9.0f, -9.0f}); + + auto p_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape({4, 3}), + b_init_values.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + + OrtValue b; + + b.Init(p_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + + SessionOptions so; + + // Set up B as a shared initializer to be shared between sessions + ASSERT_EQ(so.AddInitializer("B", &b), Status::OK()); + + // We want all sessions running using this OpTester to be able to share pre-packed weights if applicable + test.EnableSharingOfPrePackedWeightsAcrossSessions(); + + // Pre-packing is limited just to the CPU EP for now and we will only test the CPU EP + // and we want to ensure that it is available in this build + auto cpu_ep = []() -> std::vector> { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + return execution_providers; + }; + + size_t number_of_pre_packed_weights_counter_session_1 = 0; + size_t number_of_shared_pre_packed_weights_counter = 0; + + // Session 1 + { + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter); + // Assert that no pre-packed weights have been shared thus far + ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0)); + } + + auto number_of_elements_in_shared_prepacked_buffers_container = + test.GetNumPrePackedWeightsShared(); + // Assert that the number of elements in the shared container + // is the same as the number of weights that have been pre-packed + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_elements_in_shared_prepacked_buffers_container); + + // On some platforms/architectures MLAS may choose to not do any pre-packing and the number of elements + // that have been pre-packed will be zero in which case we do not continue with the testing + // of "sharing" of pre-packed weights as there are no pre-packed weights to be shared at all. + if (number_of_pre_packed_weights_counter_session_1 == 0) + return; + + // Session 2 + { + size_t number_of_pre_packed_weights_counter_session_2 = 0; + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter); + + // Assert that the same number of weights were pre-packed in both sessions + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_pre_packed_weights_counter_session_2); + + // Assert that the number of pre-packed weights that were shared equals + // the number of pre-packed weights in the second session + ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, + static_cast(number_of_shared_pre_packed_weights_counter)); + } +} +#endif + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc index 877ba9295e22c..57722628ef652 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc @@ -399,5 +399,80 @@ TEST(MatmulIntegerOpTest, MatMulInteger_Uint8_Int8_GEMM) { RunMatMulIntegerU8X8TestBatch(4, 8, 68); } +#ifndef ENABLE_TRAINING // Prepacking is enabled only on non-training builds +TEST(MatmulIntegerOpTest, SharedPrepackedWeights) { + OpTester test("MatMulInteger", 10); + test.AddInput("T1", {1, 1}, {11}); + test.AddInput("T2", {1, 1}, {13}, true); // Trigger pre-packing + test.AddInput("a_zero_point", {}, {12}); + test.AddInput("b_zero_point", {}, {12}); + test.AddOutput("T3", {1, 1}, {-1}); + + std::vector t2_init_values(1, 13); + + auto p_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape({1, 1}), + t2_init_values.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + OrtValue t2; + + t2.Init(p_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + + SessionOptions so; + // Set up T2 as a shared initializer to be shared between sessions + ASSERT_EQ(so.AddInitializer("T2", &t2), Status::OK()); + + // We want all sessions running using this OpTester to be able to share pre-packed weights if applicable + test.EnableSharingOfPrePackedWeightsAcrossSessions(); + + // Pre-packing is limited just to the CPU EP for now and we will only test the CPU EP + // and we want to ensure that it is available in this build + auto cpu_ep = []() -> std::vector> { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + return execution_providers; + }; + + size_t number_of_pre_packed_weights_counter_session_1 = 0; + size_t number_of_shared_pre_packed_weights_counter = 0; + + // Session 1 + { + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter); + // Assert that no pre-packed weights have been shared thus far + ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0)); + } + + auto number_of_elements_in_shared_prepacked_buffers_container = + test.GetNumPrePackedWeightsShared(); + // Assert that the number of elements in the shared container + // is the same as the number of weights that have been pre-packed + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_elements_in_shared_prepacked_buffers_container); + + // On some platforms/architectures MLAS may choose to not do any pre-packing and the number of elements + // that have been pre-packed will be zero in which case we do not continue with the testing + // of "sharing" of pre-packed weights as there are no pre-packed weights to be shared at all. + if (number_of_pre_packed_weights_counter_session_1 == 0) + return; + + // Session 2 + { + size_t number_of_pre_packed_weights_counter_session_2 = 0; + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter); + + // Assert that the same number of weights were pre-packed in both sessions + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_pre_packed_weights_counter_session_2); + + // Assert that the number of pre-packed weights that were shared equals + // the number of pre-packed weights in the second session + ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, + static_cast(number_of_shared_pre_packed_weights_counter)); + } +} +#endif + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/math/matmul_test.cc b/onnxruntime/test/providers/cpu/math/matmul_test.cc index 045d703b78e42..b0e21d9d5db8f 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_test.cc @@ -3,6 +3,7 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "default_providers.h" namespace onnxruntime { namespace test { @@ -145,5 +146,85 @@ TEST(MathOpTest, MatMulUint64Type) { RunMatMulTest(9); } +#ifndef ENABLE_TRAINING // Prepacking is enabled only on non-training builds +TEST(MathOpTest, MatMulSharedPrepackedWeights) { + OpTester test("MatMul"); + + std::vector b_init_values(12, 1.0f); + test.AddInput("A", {2, 4}, + {1.0f, 2.0f, 3.0f, 4.0f, + -1.0f, -2.0f, -3.0f, -4.0f}); + // B is to be an initializer for triggering pre-packing + test.AddInput("B", {4, 3}, b_init_values, true); + + test.AddOutput("Y", {2, 3}, + {10.0f, 10.0f, 10.0f, + -10.0f, -10.0f, -10.0f}); + + auto p_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape({4, 3}), + b_init_values.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + OrtValue b; + + b.Init(p_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + + SessionOptions so; + // Set up B as a shared initializer to be shared between sessions + ASSERT_EQ(so.AddInitializer("B", &b), Status::OK()); + + // We want all sessions running using this OpTester to be able to share pre-packed weights if applicable + test.EnableSharingOfPrePackedWeightsAcrossSessions(); + + // Pre-packing is limited just to the CPU EP for now and we will only test the CPU EP + // and we want to ensure that it is available in this build + auto cpu_ep = []() -> std::vector> { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + return execution_providers; + }; + + size_t number_of_pre_packed_weights_counter_session_1 = 0; + size_t number_of_shared_pre_packed_weights_counter = 0; + + // Session 1 + { + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter); + // Assert that no pre-packed weights have been shared thus far + ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0)); + } + + auto number_of_elements_in_shared_prepacked_buffers_container = + test.GetNumPrePackedWeightsShared(); + // Assert that the number of elements in the shared container + // is the same as the number of weights that have been pre-packed + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_elements_in_shared_prepacked_buffers_container); + + // On some platforms/architectures MLAS may choose to not do any pre-packing and the number of elements + // that have been pre-packed will be zero in which case we do not continue with the testing + // of "sharing" of pre-packed weights as there are no pre-packed weights to be shared at all. + if (number_of_pre_packed_weights_counter_session_1 == 0) + return; + + // Session 2 + { + size_t number_of_pre_packed_weights_counter_session_2 = 0; + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter); + + // Assert that the same number of weights were pre-packed in both sessions + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_pre_packed_weights_counter_session_2); + + // Assert that the number of pre-packed weights that were shared equals + // the number of pre-packed weights in the second session + ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, + static_cast(number_of_shared_pre_packed_weights_counter)); + } +} + +#endif + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc b/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc index e774011591211..7b2618efe2799 100644 --- a/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc @@ -3,6 +3,8 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "default_providers.h" + using namespace std; namespace onnxruntime { namespace test { @@ -282,30 +284,104 @@ TEST(ConvTransposeTest, ConvTranspose_2D_OutputShape_1_group_2_for_tranpose_path vector W_shape = {6, 3, 3, 3}; vector Y_shape = {1, 6, 4, 4}; - auto expected_vals = {12.0f, 18.0f, 18.0f, 12.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 12.0f, 18.0f, 18.0f, 12.0f, - 12.0f, 18.0f, 18.0f, 12.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 12.0f, 18.0f, 18.0f, 12.0f, - 12.0f, 18.0f, 18.0f, 12.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 12.0f, 18.0f, 18.0f, 12.0f, // duplicate below - 12.0f, 18.0f, 18.0f, 12.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 12.0f, 18.0f, 18.0f, 12.0f, - 12.0f, 18.0f, 18.0f, 12.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 12.0f, 18.0f, 18.0f, 12.0f, - 12.0f, 18.0f, 18.0f, 12.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 18.0f, 27.0f, 27.0f, 18.0f, - 12.0f, 18.0f, 18.0f, 12.0f,}; + auto expected_vals = { + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, // duplicate below + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + }; TestConvTransposeOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape); } @@ -920,5 +996,191 @@ TEST(ConvTransposeTest, ConvTranspose_AutoPad_with_non_default_strides) { OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); } +#ifndef ENABLE_TRAINING // Prepacking is enabled only on non-training builds +TEST(ConvTransposeTest, SharedPrepackedWeights) { + OpTester test("ConvTranspose", 11); + test.AddAttribute("kernel_shape", vector{3, 3}); + test.AddAttribute("group", static_cast(2)); + test.AddAttribute("pads", vector{0, 0, 0, 0}); + test.AddAttribute("output_shape", vector{1, 6, 4, 4}); + + int image_size = 4 * 4; + int input_channels = 3 * 2; + int output_channels = 3; + std::vector X; + for (int i = 0; i < input_channels * image_size; i++) + X.push_back(1.0f); + test.AddInput("X", {1, 6, 4, 4}, X, false); + + std::vector W; + int kernel_size = output_channels * input_channels * 3 * 3; + for (int i = 0; i < kernel_size; i++) + W.push_back(1.0f); + test.AddInput("W", {6, 3, 3, 3}, W, true); // Trigger pre-packing + + auto expected_vals = { + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, // duplicate below + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 18.0f, + 27.0f, + 27.0f, + 18.0f, + 12.0f, + 18.0f, + 18.0f, + 12.0f, + }; + test.AddOutput("Y", {1, 6, 4, 4}, expected_vals); + + auto p_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape({6, 3, 3, 3}), + W.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + OrtValue w; + + w.Init(p_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + + SessionOptions so; + // Set up W as a shared initializer to be shared between sessions + ASSERT_EQ(so.AddInitializer("W", &w), Status::OK()); + + // We want all sessions running using this OpTester to be able to share pre-packed weights if applicable + test.EnableSharingOfPrePackedWeightsAcrossSessions(); + + // Pre-packing is limited just to the CPU EP for now and we will only test the CPU EP + // and we want to ensure that it is available in this build + auto cpu_ep = []() -> std::vector> { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + return execution_providers; + }; + + size_t number_of_pre_packed_weights_counter_session_1 = 0; + size_t number_of_shared_pre_packed_weights_counter = 0; + + // Session 1 + { + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter); + // Assert that no pre-packed weights have been shared thus far + ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0)); + } + + auto number_of_elements_in_shared_prepacked_buffers_container = + test.GetNumPrePackedWeightsShared(); + // Assert that the number of elements in the shared container + // is the same as the number of weights that have been pre-packed + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_elements_in_shared_prepacked_buffers_container); + + // On some platforms/architectures MLAS may choose to not do any pre-packing and the number of elements + // that have been pre-packed will be zero in which case we do not continue with the testing + // of "sharing" of pre-packed weights as there are no pre-packed weights to be shared at all. + if (number_of_pre_packed_weights_counter_session_1 == 0) + return; + + // Session 2 + { + size_t number_of_pre_packed_weights_counter_session_2 = 0; + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter); + + // Assert that the same number of weights were pre-packed in both sessions + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_pre_packed_weights_counter_session_2); + + // Assert that the number of pre-packed weights that were shared equals + // the number of pre-packed weights in the second session + ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, + static_cast(number_of_shared_pre_packed_weights_counter)); + } +} +#endif + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc b/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc index f3d65b76c8db6..ea1869aec8d96 100644 --- a/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/qlinearconv_op_test.cc @@ -6,6 +6,7 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" #include +#include "default_providers.h" namespace onnxruntime { namespace test { @@ -386,7 +387,9 @@ class QLinearConvOpTester { Y_shape.push_back(output_channels); for (size_t n = 0; n < kernel_rank; n++) { Y_shape.push_back(((input_shape[n] + pads[n] + pads[kernel_rank + n]) - - (dilations[n] * (kernel_shape[n] - 1) + 1)) / strides[n] + 1); + (dilations[n] * (kernel_shape[n] - 1) + 1)) / + strides[n] + + 1); } const int64_t* output_shape = Y_shape.data() + 2; Y_data.resize(ShapeSize(Y_shape)); @@ -834,6 +837,118 @@ TEST(QLinearConvTest, Conv2D_U8S8_Requantize_Bias_PerChannel) { } } +#ifndef ENABLE_TRAINING // Prepacking is enabled only on non-training builds +TEST(QLinearConvTest, SharedPrepackedWeights) { + QuantizedTensor X({0.45246148109436035f, 0.15498268604278564f, 0.11199361085891724f, -0.39421093463897705f, + 0.2626858949661255f, 0.13414543867111206f, -0.27184486389160156f, -0.43028733134269714f, + -0.26825493574142456f, 0.3893144130706787f, -0.13631996512413025f, -0.009590476751327515f, + -0.48771554231643677f, -0.25256502628326416f, -0.2812897562980652f, 0.4043201804161072f, + 0.07795023918151855f, 0.326981782913208f, 0.13114392757415771f, -0.4416425824165344f, + 0.12446999549865723f, 0.36739975214004517f, 0.1698915958404541f, 0.2008744478225708f, + 0.23339951038360596f, 0.38613730669021606f, 0.11117297410964966f, 0.3877097964286804f, + 0.20812749862670898f, -0.34297940135002136f, -0.029246658086776733f, -0.20483523607254028f, + -0.19244328141212463f, -0.11104947328567505f, -0.32830488681793213f, -0.01800677180290222f, + 0.3618946671485901f, -0.40949052572250366f, -0.18248388171195984f, -0.3349453806877136f, + -0.34091079235076904f, 0.006497859954833984f, 0.4537564516067505f, 0.08006560802459717f, + -0.14788749814033508f, 0.034442365169525146f, -0.33322954177856445f, 0.06049239635467529f, + 0.42619407176971436f}); + QuantizedTensor W({-0.4406261742115021f}); + QuantizedTensor Y({-0.19936637580394745f, -0.06828942894935608f, -0.04934731498360634f, 0.17369966208934784f, + -0.11574628204107285f, -0.05910799279808998f, 0.1197819635272026f, 0.18959586322307587f, + 0.1182001456618309f, -0.17154212296009064f, 0.06006614491343498f, 0.0042258151806890965f, + 0.21490024030208588f, 0.11128675937652588f, 0.12394362688064575f, -0.17815405130386353f, + -0.034346915781497955f, -0.14407673478126526f, -0.05778544768691063f, 0.19459928572177887f, + -0.05484473705291748f, -0.16188594698905945f, -0.07485868036746979f, -0.08851054310798645f, + -0.10284193605184555f, -0.17014220356941223f, -0.04898572340607643f, -0.17083507776260376f, + -0.09170642495155334f, 0.1511256992816925f, 0.012886842712759972f, 0.09025576710700989f, + 0.08479554951190948f, 0.0489313043653965f, 0.14465972781181335f, 0.007934254594147205f, + -0.15946026146411896f, 0.1804322451353073f, 0.08040717244148254f, 0.1475857049226761f, + 0.15021422505378723f, -0.0028631272725760937f, -0.19993697106838226f, -0.03527900204062462f, + 0.06516310572624207f, -0.015176207758486271f, 0.14682966470718384f, -0.02665453404188156f, + -0.18779225647449493f}); + + OpTester test("QLinearConv", 10); + + test.AddInput("x", {1, 1, 7, 7}, X.quantized_); + test.AddInput("x_scale", {}, {X.scale_}, true); + test.AddInput("x_zero_point", {}, {X.zero_point_}, true); + + test.AddInput("w", {1, 1, 1, 1}, W.quantized_, true); + test.AddInput("w_scale", {}, {W.scale_}, true); + test.AddInput("w_zero_point", {}, {W.zero_point_}, true); + + test.AddInput("y_scale", {}, {Y.scale_}, true); + test.AddInput("y_zero_point", {}, {Y.zero_point_}, true); + + test.AddOutput("y", {1, 1, 7, 7}, Y.quantized_); + + // W + auto W_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape({1, 1, 1, 1}), + W.quantized_.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + + OrtValue W_ortvalue; + + W_ortvalue.Init(W_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + + SessionOptions so; + + // Set up weight(s) as a shared initializer to be shared between sessions + ASSERT_EQ(so.AddInitializer("w", &W_ortvalue), Status::OK()); + + // We want all sessions running using this OpTester to be able to share pre-packed weights if applicable + test.EnableSharingOfPrePackedWeightsAcrossSessions(); + + // Pre-packing is limited just to the CPU EP for now and we will only test the CPU EP + // and we want to ensure that it is available in this build + auto cpu_ep = []() -> std::vector> { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + return execution_providers; + }; + + size_t number_of_pre_packed_weights_counter_session_1 = 0; + size_t number_of_shared_pre_packed_weights_counter = 0; + + // Session 1 + { + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter); + // Assert that no pre-packed weights have been shared thus far + ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0)); + } + + auto number_of_elements_in_shared_prepacked_buffers_container = + test.GetNumPrePackedWeightsShared(); + // Assert that the number of elements in the shared container + // is the same as the number of weights that have been pre-packed + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_elements_in_shared_prepacked_buffers_container); + + // On some platforms/architectures MLAS may choose to not do any pre-packing and the number of elements + // that have been pre-packed will be zero in which case we do not continue with the testing + // of "sharing" of pre-packed weights as there are no pre-packed weights to be shared at all. + if (number_of_pre_packed_weights_counter_session_1 == 0) + return; + + // Session 2 + { + size_t number_of_pre_packed_weights_counter_session_2 = 0; + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter); + + // Assert that the same number of weights were pre-packed in both sessions + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_pre_packed_weights_counter_session_2); + + // Assert that the number of pre-packed weights that were shared equals + // the number of pre-packed weights in the second session + ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, + static_cast(number_of_shared_pre_packed_weights_counter)); + } +} +#endif + } // namespace } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc b/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc index fdcc9df886cd5..e2a4c9ffaa5b5 100644 --- a/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc +++ b/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc @@ -8,6 +8,8 @@ #include "core/providers/cpu/rnn/deep_cpu_lstm.h" #include "test/providers/provider_test_utils.h" +#include "default_providers.h" + using namespace std; namespace onnxruntime { namespace test { @@ -1203,5 +1205,149 @@ TEST(LSTMTest, ONNXRuntime_TestLSTMZeroSeqInMiddle) { &sequence_length, use_bias, use_peepholes, 0.0f, false, false); } +#ifndef ENABLE_TRAINING // Prepacking is enabled only on non-training builds +TEST(LSTMTest, SharedPrepackedWeights) { + int64_t seq_length = 2; + int batch_size = 2; + int64_t input_size = 1; + int64_t hidden_size = 3; + int num_directions = 1; + + std::vector X_data{1.f, 2.f, 10.f, 11.f}; + + std::vector W_data{ + 0.1f, 0.2f, 0.3f, 0.4f, + 1.f, 2.f, 3.f, 4.f, + 10.f, 11.f, 12.f, 13.f}; + + std::vector R_data(num_directions * 4 * hidden_size * hidden_size, 0.1f); + + std::vector Y_data{ + 0.28828835f, 0.36581863f, 0.45679406f, + 0.34526032f, 0.47220859f, 0.55850911f, + + 0.84196719f, 0.89402526f, 0.91073048f, + 0.85882828f, 0.90703777f, 0.92382453f}; + + OpTester test("LSTM"); + + std::vector activations = {"sigmoid", "tanh", "tanh"}; + + test.AddAttribute>("activations", activations); + + test.AddAttribute("direction", "forward"); + test.AddAttribute("hidden_size", hidden_size); + test.AddAttribute("input_forget", false); + test.AddAttribute("clip", 9999.f); + + std::vector X_dims = {seq_length, batch_size, input_size}; + std::vector W_dims = {num_directions, 4 * hidden_size, input_size}; + std::vector R_dims = {num_directions, 4 * hidden_size, hidden_size}; + + test.AddInput("X", X_dims, X_data); + test.AddInput("W", W_dims, W_data, true); //Trigger pre-packing + test.AddInput("R", R_dims, R_data, true); // Trigger pre-packing + + // B data + test.AddMissingOptionalInput(); + + // sequence + test.AddMissingOptionalInput(); + + // initial_h + test.AddMissingOptionalInput(); + + // initial_c + test.AddMissingOptionalInput(); + + // P_data + test.AddMissingOptionalInput(); + + std::vector Y_dims = {seq_length, num_directions, batch_size, hidden_size}; + test.AddOutput("Y", Y_dims, Y_data); + + // Y_h + test.AddMissingOptionalOutput(); + + // Y_c + test.AddMissingOptionalOutput(); + + // W + auto W_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape(W_dims), + W_data.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + + OrtValue W; + + W.Init(W_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + + // R + auto R_tensor = std::make_unique(DataTypeImpl::GetType(), TensorShape(R_dims), + R_data.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); + + OrtValue R; + + R.Init(R_tensor.release(), DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + + SessionOptions so; + + // Set up weight(s) as a shared initializer to be shared between sessions + ASSERT_EQ(so.AddInitializer("W", &W), Status::OK()); + ASSERT_EQ(so.AddInitializer("R", &R), Status::OK()); + + // We want all sessions running using this OpTester to be able to share pre-packed weights if applicable + test.EnableSharingOfPrePackedWeightsAcrossSessions(); + + // Pre-packing is limited just to the CPU EP for now and we will only test the CPU EP + // and we want to ensure that it is available in this build + auto cpu_ep = []() -> std::vector> { + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + return execution_providers; + }; + + size_t number_of_pre_packed_weights_counter_session_1 = 0; + size_t number_of_shared_pre_packed_weights_counter = 0; + + // Session 1 + { + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_1, &number_of_shared_pre_packed_weights_counter); + // Assert that no pre-packed weights have been shared thus far + ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0)); + } + + auto number_of_elements_in_shared_prepacked_buffers_container = + test.GetNumPrePackedWeightsShared(); + // Assert that the number of elements in the shared container + // is the same as the number of weights that have been pre-packed + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_elements_in_shared_prepacked_buffers_container); + + // On some platforms/architectures MLAS may choose to not do any pre-packing and the number of elements + // that have been pre-packed will be zero in which case we do not continue with the testing + // of "sharing" of pre-packed weights as there are no pre-packed weights to be shared at all. + if (number_of_pre_packed_weights_counter_session_1 == 0) + return; + + // Session 2 + { + size_t number_of_pre_packed_weights_counter_session_2 = 0; + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, + &ep_vec, {}, &number_of_pre_packed_weights_counter_session_2, &number_of_shared_pre_packed_weights_counter); + + // Assert that the same number of weights were pre-packed in both sessions + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, number_of_pre_packed_weights_counter_session_2); + + // Assert that the number of pre-packed weights that were shared equals + // the number of pre-packed weights in the second session + ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, + static_cast(number_of_shared_pre_packed_weights_counter)); + } +} +#endif + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/provider_test_utils.cc b/onnxruntime/test/providers/provider_test_utils.cc index 19499121d9913..52c6393b59b58 100644 --- a/onnxruntime/test/providers/provider_test_utils.cc +++ b/onnxruntime/test/providers/provider_test_utils.cc @@ -26,7 +26,7 @@ using namespace ::onnxruntime::logging; namespace onnxruntime { namespace test { -template +template Tensor copy_sort(const Tensor& src, const AllocatorPtr& allocator) { Tensor result(src.DataType(), src.Shape(), allocator); memcpy(result.MutableDataRaw(), src.DataRaw(), src.SizeInBytes()); @@ -39,7 +39,6 @@ Tensor copy_sort(const Tensor& src, const AllocatorPtr& allocator) { template void sort_expected_and_actual_buffers(const Tensor& expected, Tensor& expected_sorted, const Tensor& actual, Tensor& actual_sorted) { - auto allocator = TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault); expected_sorted = copy_sort(expected, allocator); actual_sorted = copy_sort(actual, allocator); @@ -71,7 +70,6 @@ template struct TensorCheck { void operator()(const Tensor& expected_tensor, const Tensor& output_tensor, const std::string& provider_type, const CheckParams& params) const { - Tensor expected_sorted, output_sorted; const T* expected; const T* output; @@ -103,7 +101,6 @@ struct TensorCheck { void operator()(const Tensor& expected_tensor, const Tensor& output_tensor, const std::string& provider_type, const CheckParams& params) const { - const bool has_abs_err = params.absolute_error_.has_value(); const bool has_rel_err = params.relative_error_.has_value(); @@ -217,7 +214,6 @@ void InternalNumericalCheck(const Tensor& expected_tensor, const Tensor& output_tensor, const std::string& provider_type, const CheckParams& params) { - const bool has_abs_err = params.absolute_error_.has_value(); const bool has_rel_err = params.relative_error_.has_value(); @@ -238,7 +234,7 @@ void InternalNumericalCheck(const Tensor& expected_tensor, #if defined(USE_CUDA) || defined(USE_ROCM) constexpr float threshold = 0.005f; -#else +#else constexpr float threshold = 0.0001f; #endif @@ -622,6 +618,7 @@ std::vector OpTester::ExecuteModel( } status = session_object.Initialize(); + if (!status.IsOK()) { if (expect_result == ExpectResult::kExpectFailure) { EXPECT_TRUE(!status.IsOK()); @@ -758,7 +755,9 @@ void OpTester::Run( const std::unordered_set& excluded_provider_types, const RunOptions* run_options, std::vector>* execution_providers, - const Graph::ResolveOptions& options) { + const Graph::ResolveOptions& options, + /*out*/ size_t* number_of_pre_packed_weights_counter, + /*out*/ size_t* number_of_shared_pre_packed_weights_counter) { std::string cur_provider = "not set"; ORT_TRY { #ifndef NDEBUG @@ -852,6 +851,10 @@ void OpTester::Run( InferenceSession session_object{so, GetEnvironment()}; + if (add_prepacked_shared_container_to_sessions_) { + session_object.AddPrePackedWeightsContainer(&prepacked_weights_container_); + } + ASSERT_TRUE(!execution_providers->empty()) << "Empty execution providers vector."; std::string provider_types; @@ -865,6 +868,21 @@ void OpTester::Run( *p_model, session_object, expect_result, expected_failure_string, run_options, feeds, output_names, provider_types); + // After the model has initialized (happens in ExecuteModel), + // we should be able to tell how many constant initializers were pre-packed + // and out of these pre-packed ones how many of them used a "cached" version + // from the shared container. + // Populate these value if the user has requested this information. + if (number_of_pre_packed_weights_counter) { + *number_of_pre_packed_weights_counter = + session_object.GetSessionState().GetNumberOfPrepacksCounter(); + } + + if (number_of_shared_pre_packed_weights_counter) { + *number_of_shared_pre_packed_weights_counter = + session_object.GetSessionState().GetUsedSharedPrePackedWeightCounter(); + } + } else { for (const std::string& provider_type : all_provider_types) { if (excluded_provider_types.count(provider_type) > 0) @@ -878,6 +896,10 @@ void OpTester::Run( } InferenceSession session_object{so, GetEnvironment()}; + if (add_prepacked_shared_container_to_sessions_) { + session_object.AddPrePackedWeightsContainer(&prepacked_weights_container_); + } + for (auto& custom_session_registry : custom_session_registries_) ASSERT_PROVIDER_STATUS_OK(session_object.RegisterCustomRegistry(custom_session_registry)); @@ -955,6 +977,21 @@ void OpTester::Run( *p_model, session_object, expect_result, expected_failure_string, run_options, feeds, output_names, provider_type); + // After the model has initialized (happens in ExecuteModel), + // we should be able to tell how many constant initializers were pre-packed + // and out of these pre-packed ones how many of them used a "cached" version + // from the shared container. + // Populate these value if the user has requested this information. + if (number_of_pre_packed_weights_counter) { + *number_of_pre_packed_weights_counter = + session_object.GetSessionState().GetNumberOfPrepacksCounter(); + } + + if (number_of_shared_pre_packed_weights_counter) { + *number_of_shared_pre_packed_weights_counter = + session_object.GetSessionState().GetUsedSharedPrePackedWeightCounter(); + } + cur_provider = "not set"; } diff --git a/onnxruntime/test/providers/provider_test_utils.h b/onnxruntime/test/providers/provider_test_utils.h index 47831e0909c18..f14941a95d7e9 100644 --- a/onnxruntime/test/providers/provider_test_utils.h +++ b/onnxruntime/test/providers/provider_test_utils.h @@ -12,6 +12,7 @@ #include "core/framework/run_options.h" #include "core/framework/session_state.h" #include "core/framework/tensor.h" +#include "core/framework/prepacked_weights_container.h" #include "core/graph/graph_viewer.h" #include "core/graph/model.h" #include "core/framework/data_types.h" @@ -436,9 +437,12 @@ class OpTester { const std::unordered_set& excluded_provider_types = {}, const RunOptions* run_options = nullptr, std::vector>* execution_providers = nullptr, - const Graph::ResolveOptions& resolve_options = {}); + const Graph::ResolveOptions& resolve_options = {}, + /*out*/ size_t* number_of_pre_packed_weights_counter = nullptr, + /*out*/ size_t* number_of_shared_pre_packed_weights_counter = nullptr); - std::vector GetFetches() { return fetches_; } + std::vector + GetFetches() { return fetches_; } std::unique_ptr BuildGraph(const std::unordered_map& extra_domain_to_version = {}); @@ -487,6 +491,14 @@ class OpTester { use_determinism_ = use_determinism; } + void EnableSharingOfPrePackedWeightsAcrossSessions() { + add_prepacked_shared_container_to_sessions_ = true; + } + + size_t GetNumPrePackedWeightsShared() const { + return prepacked_weights_container_.GetNumberOfElements(); + } + protected: virtual void AddNodes(onnxruntime::Graph& graph, std::vector& graph_input_defs, std::vector& graph_output_defs, @@ -645,6 +657,10 @@ class OpTester { bool use_determinism_ = false; CustomOutputVerifierFn custom_output_verifier_; + + bool add_prepacked_shared_container_to_sessions_ = false; + + onnxruntime::PrepackedWeightsContainer prepacked_weights_container_; }; template diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index 169bbca5ac5dc..ad116f75b0afa 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -161,6 +161,7 @@ static void TestInference(Ort::Env& env, const std::basic_string& mod } static constexpr PATH_TYPE MODEL_URI = TSTR("testdata/mul_1.onnx"); +static constexpr PATH_TYPE MATMUL_MODEL_URI = TSTR("testdata/matmul_1.onnx"); static constexpr PATH_TYPE SEQUENCE_MODEL_URI = TSTR("testdata/sequence_length.onnx"); static constexpr PATH_TYPE CUSTOM_OP_MODEL_URI = TSTR("testdata/foo_1.onnx"); static constexpr PATH_TYPE CUSTOM_OP_LIBRARY_TEST_MODEL_URI = TSTR("testdata/custom_op_library/custom_op_test.onnx"); @@ -1326,7 +1327,7 @@ TEST(CApiTest, TestSharedAllocatorUsingCreateAndRegisterAllocator) { nullptr); } -TEST(CApiTest, TestSharingOfInitializer) { +TEST(CApiTest, TestSharingOfInitializerWithPrepackedWeightsCaching) { // simple inference test // prepare inputs std::vector inputs(1); @@ -1336,22 +1337,30 @@ TEST(CApiTest, TestSharingOfInitializer) { input.values = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; // prepare expected inputs and outputs - std::vector expected_dims_y = {3, 2}; - std::vector expected_values_y = {2.0f, 2.0f, 12.0f, 12.0f, 30.0f, 30.0f}; + std::vector expected_dims_y = {3, 1}; + std::vector expected_values_y = {4.0f, 10.0f, 16.0f}; Ort::SessionOptions session_options; Ort::MemoryInfo mem_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); // These values are different from the actual initializer values in the model - float data[] = {2., 1., 4., 3., 6., 5.}; + float data[] = {2.0f, 1.0f}; const int data_len = sizeof(data) / sizeof(data[0]); - const int64_t shape[] = {3, 2}; + const int64_t shape[] = {2, 1}; const size_t shape_len = sizeof(shape) / sizeof(shape[0]); Ort::Value val = Ort::Value::CreateTensor(mem_info, data, data_len, shape, shape_len); session_options.AddInitializer("W", val); + const auto& api = Ort::GetApi(); + + OrtPrepackedWeightsContainer* prepacked_weights_container = nullptr; + ASSERT_TRUE(api.CreatePrepackedWeightsContainer(&prepacked_weights_container) == nullptr); + std::unique_ptr + rel_prepacked_weights_container(prepacked_weights_container, api.ReleasePrepackedWeightsContainer); + auto default_allocator = std::make_unique(); + // create session 1 - Ort::Session session1(*ort_env, MODEL_URI, session_options); + Ort::Session session1(*ort_env, MATMUL_MODEL_URI, session_options, prepacked_weights_container); RunSession(default_allocator.get(), session1, inputs, @@ -1361,7 +1370,7 @@ TEST(CApiTest, TestSharingOfInitializer) { nullptr); // create session 2 - Ort::Session session2(*ort_env, MODEL_URI, session_options); + Ort::Session session2(*ort_env, MATMUL_MODEL_URI, session_options, prepacked_weights_container); RunSession(default_allocator.get(), session2, inputs,