-
Notifications
You must be signed in to change notification settings - Fork 2.8k
/
mlasi.h
1679 lines (1492 loc) · 43.2 KB
/
mlasi.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
mlasi.h
Abstract:
This module contains the private data structures and procedure prototypes
for the Microsoft Machine Learning algebra subprogram library.
--*/
#pragma once
#include <mlas.h>
#include <memory.h>
#include <algorithm>
#include <limits>
#include <cmath>
#include <type_traits>
#if defined(_WIN32)
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#include <intrin.h>
#else
#if defined(__arm__) || defined(__aarch64__)
#include <arm_neon.h>
#endif
#if defined(__x86_64__) || defined(__i386__)
#include <cpuid.h>
#include <immintrin.h>
#endif
#if defined(__VSX__)
#include <altivec.h>
// Undefine unwanted aliases from altivec.h.
#undef vector
#undef pixel
#undef bool
#endif
#endif
//
// Macro to place variables at a specified alignment.
//
#ifdef _WIN32
#define MLAS_DECLSPEC_ALIGN(variable, alignment) DECLSPEC_ALIGN(alignment) variable
#else
#define MLAS_DECLSPEC_ALIGN(variable, alignment) variable __attribute__ ((aligned(alignment)))
#endif
//
// Macro to force inline expansion of a function.
//
#if defined(_MSC_VER)
#define MLAS_FORCEINLINE __forceinline
#else
#define MLAS_FORCEINLINE __attribute__ ((always_inline)) inline
#endif
//
// Macro to tag globals as internal data shared with kernels written in
// assembly. These globals are marked with having hidden visibility to avoid
// needing to access the data through the global object table.
//
#if defined(_MSC_VER)
#define MLAS_INTERNAL_DATA extern "C"
#else
#define MLAS_INTERNAL_DATA extern "C" __attribute ((visibility("hidden")))
#endif
//
// Macro to suppress unreferenced parameter warnings.
//
#define MLAS_UNREFERENCED_PARAMETER(parameter) ((void)(parameter))
//
// Select the target architecture.
//
#if defined(_M_AMD64) || defined(__x86_64__)
#define MLAS_TARGET_AMD64
#endif
#if defined(_M_IX86) || defined(__i386__)
#define MLAS_TARGET_IX86
#endif
#if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_IX86)
#define MLAS_TARGET_AMD64_IX86
#endif
#if defined(_M_ARM64) || defined(__aarch64__)
#define MLAS_TARGET_ARM64
#endif
#if defined(_M_ARM) || defined(__arm__)
#define MLAS_TARGET_ARM
#endif
#if defined(__VSX__)
#define MLAS_TARGET_POWER
#endif
//
// Select the threading model.
//
// N.B. MLAS_NO_ONNXRUNTIME_THREADPOOL is used to build MLAS test code outside
// of the ONNX Runtime source tree. OpenMP may or may not be enabled in this
// configuration.
//
#if !defined(MLAS_NO_ONNXRUNTIME_THREADPOOL)
#include "core/platform/threadpool.h"
#endif
#if defined(_OPENMP)
#include <omp.h>
#endif
//
// Define the maximum number of threads supported by this implementation.
//
#define MLAS_MAXIMUM_THREAD_COUNT 16
//
// Define the default strides to step through slices of the input matrices.
//
#define MLAS_SGEMM_STRIDEN 128
#define MLAS_SGEMM_STRIDEK 128
#define MLAS_SGEMM_PACKED_STRIDEN 128
#define MLAS_SGEMM_PACKED_STRIDEK 256
#define MLAS_DGEMM_STRIDEN 64
#define MLAS_DGEMM_STRIDEK 128
//
// Define the alignment for segmenting a GEMM operation across multiple
// threads.
//
// All of the SGEMM kernels can efficiently handle 16 elements. AVX512F can
// efficiently handle 32 elements, but making this value dynamic is not worth
// the effort at this time.
//
#define MLAS_SGEMM_STRIDEN_THREAD_ALIGN 16
#define MLAS_DGEMM_STRIDEN_THREAD_ALIGN 8
#define MLAS_QGEMM_STRIDEN_THREAD_ALIGN 16
//
// Define the prototypes of the platform optimized routines.
//
#if defined(MLAS_TARGET_AMD64_IX86) || defined(MLAS_TARGET_POWER)
typedef
size_t
(MLASCALL MLAS_GEMM_FLOAT_KERNEL)(
const float* A,
const float* B,
float* C,
size_t CountK,
size_t CountM,
size_t CountN,
size_t lda,
size_t ldc,
float alpha,
bool ZeroMode
);
typedef
size_t
(MLASCALL MLAS_GEMM_DOUBLE_KERNEL)(
const double* A,
const double* B,
double* C,
size_t CountK,
size_t CountM,
size_t CountN,
size_t lda,
size_t ldc,
double alpha,
bool ZeroMode
);
#else
typedef
size_t
(MLASCALL MLAS_GEMM_FLOAT_KERNEL)(
const float* A,
const float* B,
float* C,
size_t CountK,
size_t CountM,
size_t CountN,
size_t lda,
size_t ldc,
float alpha
);
typedef
size_t
(MLASCALL MLAS_GEMM_DOUBLE_KERNEL)(
const double* A,
const double* B,
double* C,
size_t CountK,
size_t CountM,
size_t CountN,
size_t lda,
size_t ldc,
double alpha
);
#endif
typedef MLAS_GEMM_FLOAT_KERNEL* PMLAS_GEMM_FLOAT_KERNEL;
typedef MLAS_GEMM_DOUBLE_KERNEL* PMLAS_GEMM_DOUBLE_KERNEL;
typedef
size_t
(MLASCALL MLAS_GEMV_FLOAT_KERNEL)(
const float* A,
const float* B,
float* C,
size_t CountK,
size_t CountN,
size_t ldb,
bool ZeroMode
);
typedef MLAS_GEMV_FLOAT_KERNEL* PMLAS_GEMV_FLOAT_KERNEL;
typedef
void
(MLASCALL MLAS_SGEMM_KERNEL_M1_ROUTINE)(
const float* A,
const float* B,
float* C,
size_t CountK,
size_t CountN,
size_t ldb,
float beta
);
typedef MLAS_SGEMM_KERNEL_M1_ROUTINE* PMLAS_SGEMM_KERNEL_M1_ROUTINE;
typedef
void
(MLASCALL MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE)(
float* D,
const float* B,
size_t ldb
);
typedef MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE* PMLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE;
typedef
void
(MLASCALL MLAS_GEMM_U8X8_OPERATION)(
const struct MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock
);
typedef MLAS_GEMM_U8X8_OPERATION* PMLAS_GEMM_U8X8_OPERATION;
typedef
size_t
(MLASCALL MLAS_GEMM_U8S8_KERNEL)(
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,
int32_t DepthValue,
bool ZeroMode
);
typedef MLAS_GEMM_U8S8_KERNEL* PMLAS_GEMM_U8S8_KERNEL;
typedef
size_t
(MLASCALL MLAS_GEMV_U8S8_KERNEL)(
const uint8_t* A,
const uint8_t* B,
int32_t* C,
size_t CountK,
size_t CountN,
size_t ldb
);
typedef MLAS_GEMV_U8S8_KERNEL* PMLAS_GEMV_U8S8_KERNEL;
typedef
size_t
(MLASCALL MLAS_GEMM_U8U8_KERNEL)(
const int16_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,
int32_t DepthValue,
bool ZeroMode
);
typedef MLAS_GEMM_U8U8_KERNEL* PMLAS_GEMM_U8U8_KERNEL;
typedef
void
(MLASCALL MLAS_CONV_FLOAT_KERNEL)(
const float* Input,
const float* Filter,
float* Output,
size_t StrideWidth,
size_t DilationWidth,
size_t FilterCount,
size_t InputStride,
size_t FilterStride,
size_t OutputStride,
size_t KernelHeight,
size_t KernelWidth,
const float* InputBase,
size_t InputWidth,
size_t DilatedInputWidth,
size_t OutputCountLeftPad,
size_t OutputCount,
size_t OutputCountRightPad,
const float* Bias,
unsigned KernelFlags
);
typedef MLAS_CONV_FLOAT_KERNEL* PMLAS_CONV_FLOAT_KERNEL;
typedef
void
(MLASCALL MLAS_CONV_DEPTHWISE_FLOAT_KERNEL)(
const float* Input,
const float* Filter,
float* Output,
size_t StrideWidth,
size_t DilationWidth,
size_t InputStride,
size_t KernelHeight,
size_t KernelWidth,
const float* InputBase,
size_t InputWidth,
size_t DilatedInputWidth,
size_t OutputCountLeftPad,
size_t OutputCount,
size_t OutputCountRightPad,
const float* Bias,
unsigned KernelFlags
);
typedef MLAS_CONV_DEPTHWISE_FLOAT_KERNEL* PMLAS_CONV_DEPTHWISE_FLOAT_KERNEL;
typedef
void
(MLASCALL MLAS_CONV_POINTWISE_FLOAT_KERNEL)(
const float* Input,
const float* Filter,
float* Output,
size_t StrideWidth,
size_t InputChannels,
size_t FilterCount,
size_t InputStride,
size_t FilterStride,
size_t OutputStride,
size_t OutputCount,
const float* Bias,
unsigned KernelFlags
);
typedef MLAS_CONV_POINTWISE_FLOAT_KERNEL* PMLAS_CONV_POINTWISE_FLOAT_KERNEL;
typedef
void
(MLASCALL MLAS_POOL_FLOAT_KERNEL)(
const float* Input,
float* Output,
size_t StrideWidth,
size_t DilationWidth,
size_t InputStride,
size_t ActualKernelSize,
size_t KernelHeight,
size_t KernelWidth,
const float* InputBase,
size_t InputWidth,
size_t DilatedInputWidth,
size_t OutputCountLeftPad,
size_t OutputCount,
size_t OutputCountRightPad
);
typedef MLAS_POOL_FLOAT_KERNEL* PMLAS_POOL_FLOAT_KERNEL;
typedef
void
(MLASCALL MLAS_COMPUTE_UNARY_FLOAT_KERNEL)(
const float* Input,
float* Output,
size_t N
);
typedef MLAS_COMPUTE_UNARY_FLOAT_KERNEL* PMLAS_COMPUTE_UNARY_FLOAT_KERNEL;
typedef
float
(MLASCALL MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL)(
const float* Input,
float* Output,
size_t N,
const float* NegativeMaximum
);
typedef MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL* PMLAS_COMPUTE_SUMEXP_FLOAT_KERNEL;
typedef
void
(MLASCALL MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL)(
float* Output,
size_t N,
const float* Parameters
);
typedef MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL* PMLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL;
typedef
void
(MLASCALL MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL)(
const float* Input,
float* Output,
size_t N,
const float* Parameters
);
typedef MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL* PMLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL;
typedef
float
(MLASCALL MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL)(
const float* Input,
size_t N
);
typedef MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL* PMLAS_REDUCE_MAXIMUM_FLOAT_KERNEL;
typedef
void
(MLASCALL MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL)(
const float* Input,
float* Min,
float* Max,
size_t N
);
typedef MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL* PMLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL;
typedef
void
(MLASCALL MLAS_QLINEAR_BINARY_OP_S8_KERNEL)(
const int8_t* InputA,
float ScaleA,
int32_t ZeroPointA,
const int8_t* InputB,
float ScaleB,
int32_t ZeroPointB,
float ScaleC,
int32_t ZeroPointC,
int8_t* OutputC,
size_t N,
bool IsScalarB
);
typedef MLAS_QLINEAR_BINARY_OP_S8_KERNEL* PMLAS_QLINEAR_BINARY_OP_S8_KERNEL;
typedef
void
(MLASCALL MLAS_QLINEAR_BINARY_OP_U8_KERNEL)(
const uint8_t* InputA,
float ScaleA,
int32_t ZeroPointA,
const uint8_t* InputB,
float ScaleB,
int32_t ZeroPointB,
float ScaleC,
int32_t ZeroPointC,
uint8_t* OutputC,
size_t N,
bool IsScalarB
);
typedef MLAS_QLINEAR_BINARY_OP_U8_KERNEL* PMLAS_QLINEAR_BINARY_OP_U8_KERNEL;
extern "C" {
#if defined(MLAS_TARGET_AMD64_IX86)
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelSse;
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelAvx;
#if defined(MLAS_TARGET_AMD64)
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelFma3;
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelAvx512F;
MLAS_GEMM_DOUBLE_KERNEL MlasGemmDoubleKernelSse;
MLAS_GEMM_DOUBLE_KERNEL MlasGemmDoubleKernelAvx;
MLAS_GEMM_DOUBLE_KERNEL MlasGemmDoubleKernelFma3;
MLAS_GEMM_DOUBLE_KERNEL MlasGemmDoubleKernelAvx512F;
#endif
#elif defined(MLAS_TARGET_POWER)
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernel;
#else
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelZero;
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelAdd;
MLAS_GEMM_DOUBLE_KERNEL MlasDgemmKernelZero;
MLAS_GEMM_DOUBLE_KERNEL MlasDgemmKernelAdd;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_SGEMM_KERNEL_M1_ROUTINE MlasSgemmKernelM1Avx;
MLAS_SGEMM_KERNEL_M1_ROUTINE MlasSgemmKernelM1TransposeBAvx;
#elif defined(MLAS_TARGET_ARM64)
MLAS_GEMV_FLOAT_KERNEL MlasGemvFloatKernel;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE MlasSgemmTransposePackB16x4Sse;
MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE MlasSgemmTransposePackB16x4Avx;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx2;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx2;
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx512Core;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx512Core;
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx512Vnni;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx512Vnni;
MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx2;
MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx512Core;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelSse;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelSse;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelSse;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelSse;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelAvx;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelAvx;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelAvx;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelAvx;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelFma3;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelFma3;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelFma3;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelFma3;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelAvx512F;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelAvx512F;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelAvx512F;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelAvx512F;
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernelSse;
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernelAvx;
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernelAvx512F;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernelSse;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernelAvx;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernelAvx512F;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernelSse;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernelAvx;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernelAvx512F;
#else
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernel;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernel;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernel;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernel;
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernel;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernel;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernel;
#endif
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasErfKernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeExpF32Kernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasLogisticKernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasTanhKernel;
MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL MlasComputeSumExpF32Kernel;
MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeSoftmaxOutputF32Kernel;
MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeLogSoftmaxOutputF32Kernel;
MLAS_QLINEAR_BINARY_OP_S8_KERNEL MlasQLinearAddS8Kernel;
MLAS_QLINEAR_BINARY_OP_U8_KERNEL MlasQLinearAddU8Kernel;
#if defined(MLAS_TARGET_AMD64)
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasErfKernelFma3;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeExpF32KernelFma3;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeExpF32KernelAvx512F;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasLogisticKernelFma3;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasTanhKernelFma3;
MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL MlasComputeSumExpF32KernelFma3;
MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL MlasComputeSumExpF32KernelAvx512F;
MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeSoftmaxOutputF32KernelAvx;
MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeLogSoftmaxOutputF32KernelAvx;
MLAS_QLINEAR_BINARY_OP_S8_KERNEL MlasQLinearAddS8KernelAvx2;
MLAS_QLINEAR_BINARY_OP_U8_KERNEL MlasQLinearAddU8KernelAvx2;
#endif
MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32Kernel;
MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL MlasReduceMinimumMaximumF32Kernel;
#if defined(MLAS_TARGET_AMD64)
MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32KernelAvx;
MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL MlasReduceMinimumMaximumF32KernelAvx;
#endif
}
//
// Define the default preferred byte alignment for buffers.
//
// MLAS_TARGET_AMD64_IX86: The typical architecture uses AVX instructions
// accessing 256-bit vectors. MLAS_TARGET_AMD64 returns a larger value if the
// platform supports 512-bit vectors to ensure that vectors are not split.
//
// MLAS_TARGET_ARM64: The kernels use "load pair" instructions to access 128-bit
// vectors, so this value keeps both vectors in the same cache line.
//
// MLAS_TARGET_ARM: Using 16 for a single 128-bit vector may be sufficient for
// this architecture, but the ONNX Runtime has historically used this larger
// value.
//
#define MLAS_DEFAULT_PREFERRED_BUFFER_ALIGNMENT 32
//
// Define the target number of per-thread multiplies before using another
// thread to perform additional work.
//
#define MLAS_SGEMM_THREAD_COMPLEXITY (64 * 1024)
#define MLAS_DGEMM_THREAD_COMPLEXITY (64 * 1024)
#define MLAS_QGEMM_THREAD_COMPLEXITY (64 * 1024)
//
// Single-threaded single precision matrix/matrix multiply operation.
//
void
MlasSgemmOperation(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
float alpha,
const float* A,
size_t lda,
const float* B,
size_t ldb,
float beta,
float* C,
size_t ldc
);
//
// Quantized integer matrix/matrix multiply operation.
//
struct MLAS_GEMM_U8X8_KERNEL_SSE;
struct MLAS_GEMM_U8S8_KERNEL_AVX2;
struct MLAS_GEMM_U8U8_KERNEL_AVX2;
template<typename KernelType>
void
MLASCALL
MlasGemmU8X8Operation(
const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock
);
template<typename KernelType>
void
MLASCALL
MlasGemmU8X8PackedOperation(
const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock
);
//
// Environment information class.
//
struct MLAS_PLATFORM {
MLAS_PLATFORM(void);
#if defined(MLAS_TARGET_AMD64_IX86)
PMLAS_GEMM_FLOAT_KERNEL GemmFloatKernel;
#endif
#if defined(MLAS_TARGET_AMD64)
PMLAS_SGEMM_KERNEL_M1_ROUTINE KernelM1Routine;
PMLAS_SGEMM_KERNEL_M1_ROUTINE KernelM1TransposeBRoutine;
PMLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE TransposePackB16x4Routine;
PMLAS_GEMM_DOUBLE_KERNEL GemmDoubleKernel;
PMLAS_GEMM_U8X8_OPERATION GemmU8S8Operation;
PMLAS_GEMM_U8X8_OPERATION GemmU8S8PackedOperation;
PMLAS_GEMM_U8S8_KERNEL GemmU8S8Kernel;
PMLAS_GEMV_U8S8_KERNEL GemvU8S8Kernel;
PMLAS_GEMM_U8X8_OPERATION GemmU8U8Operation;
PMLAS_GEMM_U8X8_OPERATION GemmU8U8PackedOperation;
PMLAS_GEMM_U8U8_KERNEL GemmU8U8Kernel;
PMLAS_CONV_FLOAT_KERNEL ConvNchwFloatKernel;
PMLAS_CONV_FLOAT_KERNEL ConvNchwcFloatKernel;
PMLAS_CONV_DEPTHWISE_FLOAT_KERNEL ConvDepthwiseFloatKernel;
PMLAS_CONV_POINTWISE_FLOAT_KERNEL ConvPointwiseFloatKernel;
PMLAS_POOL_FLOAT_KERNEL PoolFloatKernel[MlasPoolingKindCount];
PMLAS_COMPUTE_UNARY_FLOAT_KERNEL ErfKernelRoutine;
PMLAS_QLINEAR_BINARY_OP_S8_KERNEL QLinearAddS8Kernel;
PMLAS_QLINEAR_BINARY_OP_U8_KERNEL QLinearAddU8Kernel;
PMLAS_COMPUTE_UNARY_FLOAT_KERNEL ComputeExpF32Kernel;
PMLAS_COMPUTE_UNARY_FLOAT_KERNEL LogisticKernelRoutine;
PMLAS_COMPUTE_UNARY_FLOAT_KERNEL TanhKernelRoutine;
PMLAS_COMPUTE_SUMEXP_FLOAT_KERNEL ComputeSumExpF32Kernel;
PMLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL ComputeSoftmaxOutputF32Kernel;
PMLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL ComputeLogSoftmaxOutputF32Kernel;
PMLAS_REDUCE_MAXIMUM_FLOAT_KERNEL ReduceMaximumF32Kernel;
PMLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL ReduceMinimumMaximumF32Kernel;
uint32_t NchwcBlockSize;
uint32_t PreferredBufferAlignment;
#endif
};
extern MLAS_PLATFORM MlasPlatform;
//
// Threading support.
//
typedef
void
(MLAS_THREADED_ROUTINE)(
void* Context,
int32_t Index
);
typedef MLAS_THREADED_ROUTINE* PMLAS_THREADED_ROUTINE;
void
MlasExecuteThreaded(
PMLAS_THREADED_ROUTINE ThreadedRoutine,
void* Context,
int32_t Iterations,
MLAS_THREADPOOL* ThreadPool
);
inline
int32_t
MlasGetMaximumThreadCount(
MLAS_THREADPOOL* ThreadPool
)
{
#if defined(MLAS_NO_ONNXRUNTIME_THREADPOOL)
MLAS_UNREFERENCED_PARAMETER(ThreadPool);
#if defined(_OPENMP)
return (omp_get_num_threads() == 1) ? omp_get_max_threads() : 1;
#else
return 1;
#endif
#else
return onnxruntime::concurrency::ThreadPool::DegreeOfParallelism(ThreadPool);
#endif
}
inline
void
MlasPartitionWork(
int32_t ThreadId,
int32_t ThreadCount,
size_t TotalWork,
size_t* WorkIndex,
size_t* WorkRemaining
)
{
const size_t WorkPerThread = TotalWork / ThreadCount;
const size_t WorkPerThreadExtra = TotalWork % ThreadCount;
if (uint32_t(ThreadId) < WorkPerThreadExtra) {
*WorkIndex = (WorkPerThread + 1) * ThreadId;
*WorkRemaining = WorkPerThread + 1;
} else {
*WorkIndex = WorkPerThread * ThreadId + WorkPerThreadExtra;
*WorkRemaining = WorkPerThread;
}
}
//
// Define the missing ARM64 NEON intrinsic macros from arm64_neon.h that enable
// cross-compiler support.
//
// Also define additional standard NEON intrinsics using the MSVC aliases.
//
#if defined(_M_ARM64)
#ifndef vmaxvq_f32
#define vmaxvq_f32(src) neon_fmaxv(src)
#endif
#ifndef vminvq_f32
#define vminvq_f32(src) neon_fminv(src)
#endif
#endif
//
// Cross-platform wrappers for 32-bit vector intrinsics.
//
#if defined(MLAS_TARGET_ARM)
#define MLAS_NEON_INTRINSICS
#define MLAS_NEON32_INTRINSICS
#elif defined(MLAS_TARGET_ARM64)
#define MLAS_NEON_INTRINSICS
#define MLAS_NEON64_INTRINSICS
#elif defined(MLAS_TARGET_POWER)
#define MLAS_VSX_INTRINSICS
#elif defined(MLAS_TARGET_AMD64_IX86)
#define MLAS_SSE2_INTRINSICS
#if defined(__SSE4_1__) || (defined(_MSC_VER) && defined(__AVX__))
#define MLAS_SSE41_INTRINSICS
#endif
#if defined(__AVX__)
#define MLAS_AVX_INTRINSICS
#endif
#if defined(__AVX2__)
#define MLAS_AVX2_INTRINSICS
#endif
#if defined(__FMA__) || (defined(_MSC_VER) && defined(__AVX2__))
#define MLAS_FMA3_INTRINSICS
#endif
#endif
#if defined(MLAS_NEON_INTRINSICS)
typedef float32x4_t MLAS_FLOAT32X4;
typedef int32x4_t MLAS_INT32X4;
#elif defined(MLAS_SSE2_INTRINSICS)
typedef __m128 MLAS_FLOAT32X4;
typedef __m128i MLAS_INT32X4;
#elif defined(MLAS_VSX_INTRINSICS)
typedef __vector float MLAS_FLOAT32X4;
typedef __vector int MLAS_INT32X4;
typedef __vector unsigned MLAS_UINT32X4;
#else
typedef float MLAS_FLOAT32X4 __attribute__ ((vector_size(16)));
typedef int32_t MLAS_INT32X4 __attribute__ ((vector_size(16)));
#endif
MLAS_FORCEINLINE
MLAS_INT32X4
MlasReinterpretAsInt32x4(MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
return vreinterpretq_s32_f32(Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_castps_si128(Vector);
#else
return MLAS_INT32X4(Vector);
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasCastToInt32x4(MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
return vcvtq_s32_f32(Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_cvttps_epi32(Vector);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_cts(Vector, 0);
#else
return MLAS_INT32X4{int32_t(Vector[0]), int32_t(Vector[1]), int32_t(Vector[2]), int32_t(Vector[3])};
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasBroadcastInt32x4(int32_t Value)
{
#if defined(MLAS_NEON_INTRINSICS)
return vdupq_n_s32(Value);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_set1_epi32(Value);
#else
return MLAS_INT32X4{Value, Value, Value, Value};
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasLoadInt32x4(const int32_t* Buffer)
{
#if defined(MLAS_NEON_INTRINSICS)
return vld1q_s32(Buffer);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_loadu_si128((const __m128i*)Buffer);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_vsx_ld(0, Buffer);
#else
return *((MLAS_INT32X4*)Buffer);
#endif
}
MLAS_FORCEINLINE
void
MlasStoreInt32x4(int32_t* Buffer, MLAS_INT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
vst1q_s32(Buffer, Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
_mm_storeu_si128((__m128i*)Buffer, Vector);
#elif defined(MLAS_VSX_INTRINSICS)
vec_vsx_st(Vector, 0, Buffer);
#else
*((MLAS_INT32X4*)Buffer) = Vector;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasAddInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vaddq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_add_epi32(Vector1, Vector2);
#else
return Vector1 + Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasSubtractInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vsubq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_sub_epi32(Vector1, Vector2);
#else
return Vector1 - Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasAndInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vandq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_and_si128(Vector1, Vector2);
#else
return Vector1 & Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasOrInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vorrq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_or_si128(Vector1, Vector2);
#else
return Vector1 | Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasAndNotInt32x4(MLAS_INT32X4 VectorNot, MLAS_INT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
return vandq_s32(vmvnq_s32(VectorNot), Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_andnot_si128(VectorNot, Vector);
#else
return (~VectorNot) & Vector;
#endif
}
MLAS_FORCEINLINE