-
Notifications
You must be signed in to change notification settings - Fork 10.8k
/
VectorOps.td
2682 lines (2334 loc) · 102 KB
/
VectorOps.td
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
//===- VectorOps.td - Vector op definitions ---------------*- tablegen -*-====//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Defines MLIR vector operations.
//
//===----------------------------------------------------------------------===//
#ifndef VECTOR_OPS
#define VECTOR_OPS
include "mlir/IR/EnumAttr.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Interfaces/VectorInterfaces.td"
include "mlir/Interfaces/ViewLikeInterface.td"
def Vector_Dialect : Dialect {
let name = "vector";
let cppNamespace = "::mlir::vector";
let useDefaultAttributePrinterParser = 1;
let hasConstantMaterializer = 1;
let dependentDialects = ["arith::ArithmeticDialect"];
let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
}
// Base class for Vector dialect ops.
class Vector_Op<string mnemonic, list<Trait> traits = []> :
Op<Vector_Dialect, mnemonic, traits>;
// The "kind" of combining function for contractions and reductions.
def COMBINING_KIND_ADD : I32BitEnumAttrCaseBit<"ADD", 0, "add">;
def COMBINING_KIND_MUL : I32BitEnumAttrCaseBit<"MUL", 1, "mul">;
def COMBINING_KIND_MINUI : I32BitEnumAttrCaseBit<"MINUI", 2, "minui">;
def COMBINING_KIND_MINSI : I32BitEnumAttrCaseBit<"MINSI", 3, "minsi">;
def COMBINING_KIND_MINF : I32BitEnumAttrCaseBit<"MINF", 4, "minf">;
def COMBINING_KIND_MAXUI : I32BitEnumAttrCaseBit<"MAXUI", 5, "maxui">;
def COMBINING_KIND_MAXSI : I32BitEnumAttrCaseBit<"MAXSI", 6, "maxsi">;
def COMBINING_KIND_MAXF : I32BitEnumAttrCaseBit<"MAXF", 7, "maxf">;
def COMBINING_KIND_AND : I32BitEnumAttrCaseBit<"AND", 8, "and">;
def COMBINING_KIND_OR : I32BitEnumAttrCaseBit<"OR", 9, "or">;
def COMBINING_KIND_XOR : I32BitEnumAttrCaseBit<"XOR", 10, "xor">;
def CombiningKind : I32BitEnumAttr<
"CombiningKind",
"Kind of combining function for contractions and reductions",
[COMBINING_KIND_ADD, COMBINING_KIND_MUL, COMBINING_KIND_MINUI,
COMBINING_KIND_MINSI, COMBINING_KIND_MINF, COMBINING_KIND_MAXUI,
COMBINING_KIND_MAXSI, COMBINING_KIND_MAXF, COMBINING_KIND_AND,
COMBINING_KIND_OR, COMBINING_KIND_XOR]> {
let cppNamespace = "::mlir::vector";
let genSpecializedAttr = 0;
}
def Vector_CombiningKindAttr : DialectAttr<
Vector_Dialect,
CPred<"$_self.isa<::mlir::vector::CombiningKindAttr>()">,
"Kind of combining function for contractions and reductions"> {
let storageType = "::mlir::vector::CombiningKindAttr";
let returnType = "::mlir::vector::CombiningKind";
let convertFromStorage = "$_self.getKind()";
let constBuilderCall =
"::mlir::vector::CombiningKindAttr::get($0, $_builder.getContext())";
}
def Vector_AffineMapArrayAttr : TypedArrayAttrBase<AffineMapAttr,
"AffineMap array attribute"> {
let returnType = [{ ::llvm::SmallVector<::mlir::AffineMap, 4> }];
let convertFromStorage = [{
llvm::to_vector<4>($_self.getAsValueRange<::mlir::AffineMapAttr>());
}];
let constBuilderCall = "$_builder.getAffineMapArrayAttr($0)";
}
// TODO: Add an attribute to specify a different algebra with operators other
// than the current set: {*, +}.
def Vector_ContractionOp :
Vector_Op<"contract", [
NoSideEffect,
PredOpTrait<"lhs and rhs have same element type", TCopVTEtIsSameAs<0, 1>>,
PredOpTrait<"third operand acc and result have same element type",
TCresVTEtIsSameAsOpBase<0, 2>>,
DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
]>,
Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, AnyType:$acc,
Variadic<VectorOf<[I1]>>:$masks,
Vector_AffineMapArrayAttr:$indexing_maps,
ArrayAttr:$iterator_types,
DefaultValuedAttr<Vector_CombiningKindAttr,
"CombiningKind::ADD">:$kind)>,
Results<(outs AnyType)> {
let summary = "vector contraction operation";
let description = [{
Computes the sum of products of vector elements along contracting
dimension pairs from 2 vectors of rank M and N respectively, adds this
intermediate result to the accumulator argument of rank K, and returns a
vector result of rank K (where K = num_lhs_free_dims + num_rhs_free_dims +
num_batch_dims (see dimension type descriptions below)). For K = 0 (no
free or batch dimensions), the accumulator and output are a scalar.
Optional vector mask arguments (produced by CreateMaskOp or ConstantMaskOp)
specify the dynamic dimension sizes of valid data within the lhs/rhs vector
arguments.
An iterator type attribute list must be specified, where each element of
the list represents an iterator with one of the following types:
*) "reduction": reduction dimensions are present in the lhs and rhs
arguments but not in the output (and accumulator
argument). These are the dimensions along which the vector
contraction op computes the sum of products, and
contracting dimension pair dimension sizes must match
between lhs/rhs.
*) "parallel": Batch dimensions are iterator type "parallel", and
are non-contracting dimensions present in the lhs, rhs and
output. The lhs/rhs co-iterate along the batch dimensions,
which should be expressed in their indexing maps.
Free dimensions are iterator type "parallel", and are
non-contraction, non-batch dimensions accessed by either the
lhs or rhs (but not both). The lhs and rhs free dimensions
are unrelated to each other and do not co-iterate, which
should be expressed in their indexing maps.
An indexing map attribute list must be specified with an entry for lhs, rhs
and acc arguments. An indexing map attribute specifies a mapping from each
iterator in the iterator type list, to each dimension of an N-D vector.
An optional kind attribute may be used to specify the combining function
between the intermediate result and accumulator argument of rank K. This
attribute can take the values add/mul/min/max for int/fp, and/or/xor for
int only. The default is "add".
Example:
```mlir
// Simple DOT product (K = 0).
#contraction_accesses = [
affine_map<(i) -> (i)>,
affine_map<(i) -> (i)>,
affine_map<(i) -> ()>
]
#contraction_trait = {
indexing_maps = #contraction_accesses,
iterator_types = ["reduction"]
}
%3 = vector.contract #contraction_trait %0, %1, %2
: vector<10xf32>, vector<10xf32> into f32
// 2D vector contraction with one contracting dimension (matmul, K = 2).
#contraction_accesses = [
affine_map<(i, j, k) -> (i, k)>,
affine_map<(i, j, k) -> (k, j)>,
affine_map<(i, j, k) -> (i, j)>
]
#contraction_trait = {
indexing_maps = #contraction_accesses,
iterator_types = ["parallel", "parallel", "reduction"]
}
%3 = vector.contract #contraction_trait %0, %1, %2
: vector<4x3xf32>, vector<3x7xf32> into vector<4x7xf32>
// 4D to 3D vector contraction with two contracting dimensions and
// one batch dimension (K = 3).
#contraction_accesses = [
affine_map<(b0, f0, f1, c0, c1) -> (c0, b0, c1, f0)>,
affine_map<(b0, f0, f1, c0, c1) -> (b0, c1, c0, f1)>,
affine_map<(b0, f0, f1, c0, c1) -> (b0, f0, f1)>
]
#contraction_trait = {
indexing_maps = #contraction_accesses,
iterator_types = ["parallel", "parallel", "parallel",
"reduction", "reduction"]
}
%4 = vector.contract #contraction_trait %0, %1, %2
: vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
// 4D vector contraction with two contracting dimensions and optional
// vector mask arguments.
%lhs_mask = vector.constant_mask [7, 8, 16, 15] : vector<7x8x16x15xi1>
%rhs_mask = vector.constant_mask [8, 16, 7, 5] : vector<8x16x7x5xi1>
%5 = vector.contract #contraction_trait %0, %1, %2, %lhs_mask, %rhs_mask
: vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x8x5xf32>
// Vector contraction with mixed typed. lhs/rhs have different element
// types than accumulator/result.
%6 = vector.contract #contraction_trait %0, %1, %2
: vector<10xf16>, vector<10xf16> into f32
// Contract with max (K = 0).
#contraction_accesses = [
affine_map<(i) -> (i)>,
affine_map<(i) -> (i)>,
affine_map<(i) -> ()>
]
#contraction_trait = {
indexing_maps = #contraction_accesses,
iterator_types = ["reduction"],
kind = #vector.kind<max>
}
%7 = vector.contract #contraction_trait %0, %1, %2
: vector<10xf32>, vector<10xf32> into f32
```
}];
let builders = [
OpBuilder<(ins "Value":$lhs, "Value":$rhs, "Value":$acc,
"ArrayAttr":$indexingMaps, "ArrayAttr":$iteratorTypes)>,
OpBuilder<(ins "Value":$lhs, "Value":$rhs, "Value":$acc,
"ArrayRef<ArrayRef<AffineExpr>>":$indexingExprs,
"ArrayRef<StringRef>":$iteratorTypes)>,
OpBuilder<(ins "Value":$lhs, "Value":$rhs, "Value":$acc,
"ArrayAttr":$indexingMaps, "ArrayAttr":$iteratorTypes,
"CombiningKind":$kind)>
];
let extraClassDeclaration = [{
VectorType getLhsType() {
return getLhs().getType().cast<VectorType>();
}
VectorType getRhsType() {
return getRhs().getType().cast<VectorType>();
}
Type getAccType() { return getAcc().getType(); }
VectorType getLHSVectorMaskType() {
if (llvm::size(getMasks()) != 2) return VectorType();
return getOperand(3).getType().cast<VectorType>();
}
VectorType getRHSVectorMaskType() {
if (llvm::size(getMasks()) != 2) return VectorType();
return getOperand(4).getType().cast<VectorType>();
}
Type getResultType() { return getResult().getType(); }
ArrayRef<StringRef> getTraitAttrNames();
static unsigned getAccOperandIndex() { return 2; }
// Returns the bounds of each dimension in the iteration space spanned
// by the iterator types of this operation.
void getIterationBounds(SmallVectorImpl<int64_t> &iterationBounds);
// Returns a list of index maps, where there is a list entry for each
// op indexing map attribute (i.e. one for each input and output, with
// the output listed last). Each index map, maps from this operations
// iteration space, to vector dimensions of the maps input/output.
void getIterationIndexMap(
std::vector<DenseMap<int64_t, int64_t>> &iterationIndexMap);
std::vector<std::pair<int64_t, int64_t>> getContractingDimMap();
std::vector<std::pair<int64_t, int64_t>> getBatchDimMap();
static constexpr StringRef getKindAttrStrName() { return "kind"; }
static CombiningKind getDefaultKind() {
return CombiningKind::ADD;
}
}];
let hasCanonicalizer = 1;
let hasCustomAssemblyFormat = 1;
let hasVerifier = 1;
}
def Vector_ReductionOp :
Vector_Op<"reduction", [NoSideEffect,
PredOpTrait<"source operand and result have same element type",
TCresVTEtIsSameAsOpBase<0, 0>>,
DeclareOpInterfaceMethods<VectorUnrollOpInterface,
["getShapeForUnroll"]>]>,
Arguments<(ins Vector_CombiningKindAttr:$kind, AnyVector:$vector,
Optional<AnyType>:$acc)>,
Results<(outs AnyType:$dest)> {
let summary = "reduction operation";
let description = [{
Reduces an 1-D vector "horizontally" into a scalar using the given
operation (add/mul/min/max for int/fp and and/or/xor for int only).
Reductions also allow an optional fused accumulator.
Note that these operations are restricted to 1-D vectors to remain
close to the corresponding LLVM intrinsics:
http://llvm.org/docs/LangRef.html#vector-reduction-intrinsics
Example:
```mlir
%1 = vector.reduction <add>, %0 : vector<16xf32> into f32
%3 = vector.reduction <xor>, %2 : vector<4xi32> into i32
%4 = vector.reduction <mul>, %0, %1 : vector<16xf32> into f32
```
}];
let extraClassDeclaration = [{
VectorType getVectorType() {
return getVector().getType().cast<VectorType>();
}
}];
let builders = [
// Builder that infers the type of `dest`.
OpBuilder<(ins "CombiningKind":$kind, "Value":$vector, "Value":$acc)>,
// Builder that infers the type of `dest` and has no accumulator.
OpBuilder<(ins "CombiningKind":$kind, "Value":$vector)>
];
// TODO: Migrate to assemblyFormat once `AllTypesMatch` supports optional
// operands.
let hasCustomAssemblyFormat = 1;
let hasCanonicalizer = 1;
let hasVerifier = 1;
}
def Vector_MultiDimReductionOp :
Vector_Op<"multi_reduction", [NoSideEffect,
PredOpTrait<"source operand and result have same element type",
TCresVTEtIsSameAsOpBase<0, 0>>,
DeclareOpInterfaceMethods<InferTypeOpInterface>,
DeclareOpInterfaceMethods<VectorUnrollOpInterface,
["getShapeForUnroll"]>]>,
Arguments<(ins Vector_CombiningKindAttr:$kind,
AnyVector:$source,
I64ArrayAttr:$reduction_dims)>,
Results<(outs AnyType:$dest)> {
let summary = "Multi-dimensional reduction operation";
let description = [{
Reduces an n-D vector into an (n-k)-D vector (or a scalar when k == n)
using the given operation (add/mul/min/max for int/fp and and/or/xor for
int only).
Example:
```mlir
%1 = vector.multi_reduction <add>, %0 [1, 3] :
vector<4x8x16x32xf32> into vector<4x16xf32>
%2 = vector.multi_reduction <add>, %1 [0, 1] :
vector<4x16xf32> into f32
```
}];
let builders = [
OpBuilder<(ins "Value":$source, "ArrayRef<bool>":$reductionMask,
"CombiningKind":$kind)>
];
let extraClassDeclaration = [{
static StringRef getKindAttrStrName() { return "kind"; }
static StringRef getReductionDimsAttrStrName() { return "reduction_dims"; }
VectorType getSourceVectorType() {
return getSource().getType().cast<VectorType>();
}
Type getDestType() {
return getDest().getType();
}
bool isReducedDim(int64_t d) {
assert(d >= 0 && d < static_cast<int64_t>(getReductionMask().size()) &&
"d overflows the number of dims");
return getReductionMask()[d];
}
SmallVector<bool> getReductionMask() {
SmallVector<bool> res(getSourceVectorType().getRank(), false);
for (auto ia : getReductionDims().getAsRange<IntegerAttr>())
res[ia.getInt()] = true;
return res;
}
static SmallVector<bool> getReductionMask(
ArrayRef<int64_t> reductionDims, unsigned sourceRank) {
SmallVector<bool> res(sourceRank, false);
for (auto idx : reductionDims)
res[idx] = true;
return res;
}
}];
let assemblyFormat =
"$kind `,` $source attr-dict $reduction_dims `:` type($source) `to` type($dest)";
let hasFolder = 1;
}
def Vector_BroadcastOp :
Vector_Op<"broadcast", [NoSideEffect,
PredOpTrait<"source operand and result have same element type",
TCresVTEtIsSameAsOpBase<0, 0>>]>,
Arguments<(ins AnyType:$source)>,
Results<(outs AnyVectorOfAnyRank:$vector)> {
let summary = "broadcast operation";
let description = [{
Broadcasts the scalar or k-D vector value in the source operand
to a n-D result vector such that the broadcast makes sense, i.e.,
the source operand is duplicated to match the given rank and sizes
in the result vector. The legality rules are:
* the source operand must have the same element type as the result type
* a k-D vector <s_1 x .. x s_k x type> can be broadcast to
a n-D vector <t_1 x .. x t_n x type> if
* k <= n, and
* the sizes in the trailing dimensions n-k < i <= n with j=i+k-n
match exactly as s_j = t_i or s_j = 1:
```
t_1 x .. t_n-k x t_n-k+1 x .. x t_i x .. x t_n
s_1 x .. x s_j x .. x s_k
<duplication> <potential stretch>
```
The source operand is duplicated over all the missing leading dimensions
and stretched over the trailing dimensions where the source has a non-equal
dimension of 1. These rules imply that any scalar broadcast (k=0) to any
shaped vector with the same element type is always legal.
Example:
```mlir
%0 = arith.constant 0.0 : f32
%1 = vector.broadcast %0 : f32 to vector<16xf32>
%2 = vector.broadcast %1 : vector<16xf32> to vector<4x16xf32>
```
}];
let extraClassDeclaration = [{
Type getSourceType() { return getSource().getType(); }
VectorType getVectorType() {
return getVector().getType().cast<VectorType>();
}
}];
let assemblyFormat = "$source attr-dict `:` type($source) `to` type($vector)";
let hasFolder = 1;
let hasCanonicalizer = 1;
let hasVerifier = 1;
}
def Vector_ShuffleOp :
Vector_Op<"shuffle", [NoSideEffect,
PredOpTrait<"first operand v1 and result have same element type",
TCresVTEtIsSameAsOpBase<0, 0>>,
PredOpTrait<"second operand v2 and result have same element type",
TCresVTEtIsSameAsOpBase<0, 1>>,
DeclareOpInterfaceMethods<InferTypeOpInterface>]>,
Arguments<(ins AnyVector:$v1, AnyVector:$v2, I64ArrayAttr:$mask)>,
Results<(outs AnyVector:$vector)> {
let summary = "shuffle operation";
let description = [{
The shuffle operation constructs a permutation (or duplication) of elements
from two input vectors, returning a vector with the same element type as
the input and a length that is the same as the shuffle mask. The two input
vectors must have the same element type, rank, and trailing dimension sizes
and shuffles their values in the leading dimension (which may differ in size)
according to the given mask. The legality rules are:
* the two operands must have the same element type as the result
* the two operands and the result must have the same rank and trailing
dimension sizes, viz. given two k-D operands
v1 : <s_1 x s_2 x .. x s_k x type> and
v2 : <t_1 x t_2 x .. x t_k x type>
we have s_i = t_i for all 1 < i <= k
* the mask length equals the leading dimension size of the result
* numbering the input vector indices left to right across the operands, all
mask values must be within range, viz. given two k-D operands v1 and v2
above, all mask values are in the range [0,s_1+t_1)
Example:
```mlir
%0 = vector.shuffle %a, %b[0, 3]
: vector<2xf32>, vector<2xf32> ; yields vector<2xf32>
%1 = vector.shuffle %c, %b[0, 1, 2]
: vector<2x16xf32>, vector<1x16xf32> ; yields vector<3x16xf32>
%2 = vector.shuffle %a, %b[3, 2, 1, 0]
: vector<2xf32>, vector<2xf32> ; yields vector<4xf32>
```
}];
let builders = [
OpBuilder<(ins "Value":$v1, "Value":$v2, "ArrayRef<int64_t>")>
];
let hasFolder = 1;
let extraClassDeclaration = [{
static StringRef getMaskAttrStrName() { return "mask"; }
VectorType getV1VectorType() {
return getV1().getType().cast<VectorType>();
}
VectorType getV2VectorType() {
return getV2().getType().cast<VectorType>();
}
VectorType getVectorType() {
return getVector().getType().cast<VectorType>();
}
}];
let assemblyFormat = "operands $mask attr-dict `:` type(operands)";
let hasVerifier = 1;
}
def Vector_ExtractElementOp :
Vector_Op<"extractelement", [NoSideEffect,
TypesMatchWith<"result type matches element type of vector operand",
"vector", "result",
"$_self.cast<ShapedType>().getElementType()">]>,
Arguments<(ins AnyVectorOfAnyRank:$vector,
Optional<AnySignlessIntegerOrIndex>:$position)>,
Results<(outs AnyType:$result)> {
let summary = "extractelement operation";
let description = [{
Takes a 0-D or 1-D vector and a optional dynamic index position and
extracts the scalar at that position.
Note that this instruction resembles vector.extract, but is restricted to
0-D and 1-D vectors and relaxed to dynamic indices.
If the vector is 0-D, the position must be llvm::None.
It is meant to be closer to LLVM's version:
https://llvm.org/docs/LangRef.html#extractelement-instruction
Example:
```mlir
%c = arith.constant 15 : i32
%1 = vector.extractelement %0[%c : i32]: vector<16xf32>
%2 = vector.extractelement %z[]: vector<f32>
```
}];
let assemblyFormat = [{
$vector `[` ($position^ `:` type($position))? `]` attr-dict `:` type($vector)
}];
let builders = [
// 0-D builder.
OpBuilder<(ins "Value":$source)>,
// 1-D + position builder.
OpBuilder<(ins "Value":$source, "Value":$position)>,
];
let extraClassDeclaration = [{
VectorType getVectorType() {
return getVector().getType().cast<VectorType>();
}
}];
let hasVerifier = 1;
let hasFolder = 1;
}
def Vector_ExtractOp :
Vector_Op<"extract", [NoSideEffect,
PredOpTrait<"operand and result have same element type",
TCresVTEtIsSameAsOpBase<0, 0>>,
DeclareOpInterfaceMethods<InferTypeOpInterface>]>,
Arguments<(ins AnyVector:$vector, I64ArrayAttr:$position)>,
Results<(outs AnyType)> {
let summary = "extract operation";
let description = [{
Takes an n-D vector and a k-D position and extracts the (n-k)-D vector at
the proper position. Degenerates to an element type in the 0-D case.
Example:
```mlir
%1 = vector.extract %0[3]: vector<4x8x16xf32>
%2 = vector.extract %0[3, 3, 3]: vector<4x8x16xf32>
```
}];
let builders = [
OpBuilder<(ins "Value":$source, "ArrayRef<int64_t>":$position)>,
// Convenience builder which assumes the values in `position` are defined by
// ConstantIndexOp.
OpBuilder<(ins "Value":$source, "ValueRange":$position)>
];
let extraClassDeclaration = [{
static StringRef getPositionAttrStrName() { return "position"; }
VectorType getVectorType() {
return getVector().getType().cast<VectorType>();
}
static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
}];
let assemblyFormat = "$vector `` $position attr-dict `:` type($vector)";
let hasCanonicalizer = 1;
let hasFolder = 1;
let hasVerifier = 1;
}
def Vector_ExtractMapOp :
Vector_Op<"extract_map", [NoSideEffect]>,
Arguments<(ins AnyVector:$vector, Variadic<Index>:$ids)>,
Results<(outs AnyVector)> {
let summary = "vector extract map operation";
let description = [{
Takes an N-D vector and extracts a sub-part of the vector starting at id
along each dimension.
The dimension associated to each element of `ids` used to extract are
implicitly deduced from the destination type. For each dimension the
multiplicity is the destination dimension size divided by the source
dimension size, each dimension with a multiplicity greater than 1 is
associated to the next id, following ids order.
For example if the source type is `vector<64x4x32xf32>` and the destination
type is `vector<4x4x2xf32>`, the first id maps to dimension 0 and the second
id to dimension 2.
Similarly to vector.tuple_get, this operation is used for progressive
lowering and should be folded away before converting to LLVM.
It is different than `vector.extract_slice` and
`vector.extract_strided_slice` as it takes a Value as index instead of an
attribute. Also in the future it is meant to support extracting along any
dimensions and not only the most major ones.
For instance:
```
// dynamic computation producing the value 0 of index type
%idx0 = ... : index
// dynamic computation producing the value 1 of index type
%idx1 = ... : index
%0 = arith.constant dense<0, 1, 2, 3>: vector<4xi32>
// extracts values [0, 1]
%1 = vector.extract_map %0[%idx0] : vector<4xi32> to vector<2xi32>
// extracts values [1, 2]
%2 = vector.extract_map %0[%idx1] : vector<4xi32> to vector<2xi32>
```
Example:
```mlir
%ev = vector.extract_map %v[%id] : vector<32xf32> to vector<1xf32>
%ev1 = vector.extract_map %v1[%id1, %id2] : vector<64x4x32xf32>
to vector<4x4x2xf32>
```
}];
let builders = [
OpBuilder<(ins "Value":$vector, "ValueRange":$ids,
"ArrayRef<int64_t>":$multiplicity,
"AffineMap":$map)>];
let extraClassDeclaration = [{
VectorType getSourceVectorType() {
return getVector().getType().cast<VectorType>();
}
VectorType getResultType() {
return getResult().getType().cast<VectorType>();
}
void getMultiplicity(SmallVectorImpl<int64_t> &multiplicity);
AffineMap map();
}];
let assemblyFormat = [{
$vector `[` $ids `]` attr-dict `:` type($vector) `to` type(results)
}];
let hasFolder = 1;
let hasVerifier = 1;
}
def Vector_FMAOp :
Op<Vector_Dialect, "fma", [
NoSideEffect, AllTypesMatch<["lhs", "rhs", "acc", "result"]>,
DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
] # ElementwiseMappable.traits>,
Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, AnyVector:$acc)>,
Results<(outs AnyVector:$result)> {
let summary = "vector fused multiply-add";
let description = [{
Multiply-add expressions operate on n-D vectors and compute a fused
pointwise multiply-and-accumulate: `$result = `$lhs * $rhs + $acc`.
All operands and result have the same vector type. The semantics
of the operation correspond to those of the `llvm.fma`
[intrinsic](https://llvm.org/docs/LangRef.html#int-fma). In the
particular case of lowering to LLVM, this is guaranteed to lower
to the `llvm.fma.*` intrinsic.
Example:
```mlir
%3 = vector.fma %0, %1, %2: vector<8x16xf32>
```
}];
let assemblyFormat = "$lhs `,` $rhs `,` $acc attr-dict `:` type($lhs)";
let extraClassDeclaration = [{
VectorType getVectorType() { return getLhs().getType().cast<VectorType>(); }
}];
}
def Vector_InsertElementOp :
Vector_Op<"insertelement", [NoSideEffect,
TypesMatchWith<"source operand type matches element type of result",
"result", "source",
"$_self.cast<ShapedType>().getElementType()">,
AllTypesMatch<["dest", "result"]>]>,
Arguments<(ins AnyType:$source, AnyVectorOfAnyRank:$dest,
Optional<AnySignlessIntegerOrIndex>:$position)>,
Results<(outs AnyVectorOfAnyRank:$result)> {
let summary = "insertelement operation";
let description = [{
Takes a scalar source, a 0-D or 1-D destination vector and a dynamic index
position and inserts the source into the destination at the proper position.
Note that this instruction resembles vector.insert, but is restricted to 0-D
and 1-D vectors and relaxed to dynamic indices.
It is meant to be closer to LLVM's version:
https://llvm.org/docs/LangRef.html#insertelement-instruction
Example:
```mlir
%c = arith.constant 15 : i32
%f = arith.constant 0.0f : f32
%1 = vector.insertelement %f, %0[%c : i32]: vector<16xf32>
%2 = vector.insertelement %f, %z[]: vector<f32>
```
}];
let assemblyFormat = [{
$source `,` $dest `[` ($position^ `:` type($position))? `]` attr-dict `:`
type($result)
}];
let builders = [
// 0-D builder.
OpBuilder<(ins "Value":$source, "Value":$dest)>,
];
let extraClassDeclaration = [{
Type getSourceType() { return getSource().getType(); }
VectorType getDestVectorType() {
return getDest().getType().cast<VectorType>();
}
}];
let hasVerifier = 1;
let hasFolder = 1;
}
def Vector_InsertOp :
Vector_Op<"insert", [NoSideEffect,
PredOpTrait<"source operand and result have same element type",
TCresVTEtIsSameAsOpBase<0, 0>>,
AllTypesMatch<["dest", "res"]>]>,
Arguments<(ins AnyType:$source, AnyVector:$dest, I64ArrayAttr:$position)>,
Results<(outs AnyVector:$res)> {
let summary = "insert operation";
let description = [{
Takes an n-D source vector, an (n+k)-D destination vector and a k-D position
and inserts the n-D source into the (n+k)-D destination at the proper
position. Degenerates to a scalar source type when n = 0.
Example:
```mlir
%2 = vector.insert %0, %1[3] : vector<8x16xf32> into vector<4x8x16xf32>
%5 = vector.insert %3, %4[3, 3, 3] : f32 into vector<4x8x16xf32>
```
}];
let assemblyFormat = [{
$source `,` $dest $position attr-dict `:` type($source) `into` type($dest)
}];
let builders = [
OpBuilder<(ins "Value":$source, "Value":$dest,
"ArrayRef<int64_t>":$position)>,
// Convenience builder which assumes all values are constant indices.
OpBuilder<(ins "Value":$source, "Value":$dest, "ValueRange":$position)>
];
let extraClassDeclaration = [{
static StringRef getPositionAttrStrName() { return "position"; }
Type getSourceType() { return getSource().getType(); }
VectorType getDestVectorType() {
return getDest().getType().cast<VectorType>();
}
}];
let hasCanonicalizer = 1;
let hasFolder = 1;
let hasVerifier = 1;
}
def Vector_InsertMapOp :
Vector_Op<"insert_map", [NoSideEffect, AllTypesMatch<["dest", "result"]>]>,
Arguments<(ins AnyVector:$vector, AnyVector:$dest, Variadic<Index>:$ids)>,
Results<(outs AnyVector:$result)> {
let summary = "vector insert map operation";
let description = [{
Inserts a N-D vector and within a larger vector starting at id. The new
vector created will have the same size as the destination operand vector.
The dimension associated to each element of `ids` used to insert is
implicitly deduced from the source type (see `ExtractMapOp` for details).
For example if source type is `vector<4x4x2xf32>` and the destination type
is `vector<64x4x32xf32>`, the first id maps to dimension 0 and the second id
to dimension 2.
Similarly to vector.tuple_get, this operation is used for progressive
lowering and should be folded away before converting to LLVM.
It is different than `vector.insert` and `vector.insert_strided_slice` as it
takes a Value as index instead of an attribute. Also in the future it is
meant to support inserting along any dimensions and not only the most major
ones.
This operations is meant to be used in combination with vector.extract_map.
For instance:
```
// dynamic computation producing the value 0 of index type
%idx0 = ... : index
// dynamic computation producing the value 1 of index type
%idx1 = ... : index /
%0 = arith.constant dense<0, 1, 2, 3>: vector<4xi32>
// extracts values [0, 1]
%1 = vector.extract_map %0[%idx0] : vector<4xi32> to vector<2xi32>
// extracts values [1, 2]
%2 = vector.extract_map %0[%idx1] : vector<4xi32> to vector<2xi32>
// insert [0, 1] into [x, x, x, x] and produce [0, 1, x, x]
%3 = vector.insert_map %1, %0[%idx0] : vector<2xi32> into vector<4xi32>
// insert [1, 2] into [x, x, x, x] and produce [x, 1, 2, x]
%4 = vector.insert_map %2, %0[%idx1] : vector<2xi32> into vector<4xi32>
```
Example:
```mlir
%v = vector.insert_map %ev %v[%id] : vector<1xf32> into vector<32xf32>
%v1 = vector.insert_map %ev1, %v1[%arg0, %arg1] : vector<2x4x1xf32>
into vector<64x4x32xf32>
```
}];
let extraClassDeclaration = [{
VectorType getSourceVectorType() {
return getVector().getType().cast<VectorType>();
}
VectorType getResultType() {
return getResult().getType().cast<VectorType>();
}
// Return a map indicating the dimension mapping to the given ids.
AffineMap map();
}];
let assemblyFormat = [{
$vector `,` $dest `[` $ids `]` attr-dict
`:` type($vector) `into` type($result)
}];
let hasVerifier = 1;
}
def Vector_InsertStridedSliceOp :
Vector_Op<"insert_strided_slice", [NoSideEffect,
PredOpTrait<"operand #0 and result have same element type",
TCresVTEtIsSameAsOpBase<0, 0>>,
AllTypesMatch<["dest", "res"]>]>,
Arguments<(ins AnyVector:$source, AnyVector:$dest, I64ArrayAttr:$offsets,
I64ArrayAttr:$strides)>,
Results<(outs AnyVector:$res)> {
let summary = "strided_slice operation";
let description = [{
Takes a k-D source vector, an n-D destination vector (n >= k), n-sized
`offsets` integer array attribute, a k-sized `strides` integer array attribute
and inserts the k-D source vector as a strided subvector at the proper offset
into the n-D destination vector.
At the moment strides must contain only 1s.
Returns an n-D vector that is a copy of the n-D destination vector in which
the last k-D dimensions contain the k-D source vector elements strided at
the proper location as specified by the offsets.
Example:
```mlir
%2 = vector.insert_strided_slice %0, %1
{offsets = [0, 0, 2], strides = [1, 1]}:
vector<2x4xf32> into vector<16x4x8xf32>
```
}];
let assemblyFormat = [{
$source `,` $dest attr-dict `:` type($source) `into` type($dest)
}];
let builders = [
OpBuilder<(ins "Value":$source, "Value":$dest,
"ArrayRef<int64_t>":$offsets, "ArrayRef<int64_t>":$strides)>
];
let extraClassDeclaration = [{
static StringRef getOffsetsAttrStrName() { return "offsets"; }
static StringRef getStridesAttrStrName() { return "strides"; }
VectorType getSourceVectorType() {
return getSource().getType().cast<VectorType>();
}
VectorType getDestVectorType() {
return getDest().getType().cast<VectorType>();
}
bool hasNonUnitStrides() {
return llvm::any_of(getStrides(), [](Attribute attr) {
return attr.cast<IntegerAttr>().getInt() != 1;
});
}
}];
let hasFolder = 1;
let hasVerifier = 1;
let hasCanonicalizer = 1;
}
def Vector_OuterProductOp :
Vector_Op<"outerproduct", [NoSideEffect,
PredOpTrait<"lhs operand and result have same element type",
TCresVTEtIsSameAsOpBase<0, 0>>,
PredOpTrait<"rhs operand and result have same element type",
TCresVTEtIsSameAsOpBase<0, 1>>]>,
Arguments<(ins AnyVector:$lhs, AnyType:$rhs,
Variadic<AnyVector>:$acc,
DefaultValuedAttr<Vector_CombiningKindAttr, "CombiningKind::ADD">:$kind)>,
Results<(outs AnyVector)> {
let summary = "vector outerproduct with optional fused add";
let description = [{
Takes 2 1-D vectors and returns the 2-D vector containing the outer-product,
as illustrated below:
```
outer | [c, d]
------+------------
[a, | [ [a*c, a*d],
b] | [b*c, b*d] ]
```
This operation also accepts a 1-D vector lhs and a scalar rhs. In this
case a simple AXPY operation is performed, which returns a 1-D vector.
```
[a, b] * c = [a*c, b*c]
```
An optional extra vector argument with the same shape as the output
vector may be specified in which case the operation returns the sum of
the outer-product and the extra vector. In this multiply-accumulate
scenario for floating-point arguments, the rounding mode is enforced
by guaranteeing that a fused-multiply add operation is emitted. When
lowered to the LLVMIR dialect, this form emits `llvm.intr.fma`, which
is guaranteed to lower to actual `fma` instructions on x86.
An optional kind attribute may be specified to be add/mul/min/max
for int/fp, and and/or/xor for int only. The default is "add", in which
case the operation returns a fused multiply-add. In other cases it returns
a multiply followed by the appropriate operation (for example, a compare and
select for "max").
Example:
```
%2 = vector.outerproduct %0, %1: vector<4xf32>, vector<8xf32>
return %2: vector<4x8xf32>
%3 = vector.outerproduct %0, %1, %2:
vector<4xf32>, vector<8xf32>, vector<4x8xf32>
return %3: vector<4x8xf32>
%4 = vector.outerproduct %0, %1, %2 {kind = #vector.kind<max>}:
vector<4xf32>, vector<8xf32>, vector<4x8xf32>
return %3: vector<4x8xf32>
%6 = vector.outerproduct %4, %5: vector<10xf32>, f32
return %6: vector<10xf32>
```
}];
let builders = [
// Build an op without mask, use the type of `acc` as the return type.
OpBuilder<(ins "Value":$lhs, "Value":$rhs, "Value":$acc)>
];
let extraClassDeclaration = [{
VectorType getOperandVectorTypeLHS() {
return getLhs().getType().cast<VectorType>();
}
Type getOperandTypeRHS() {
return getRhs().getType();
}
VectorType getOperandVectorTypeACC() {
return (llvm::size(getAcc()) == 0)
? VectorType()
: (*getAcc().begin()).getType().cast<VectorType>();
}
VectorType getVectorType() {
return getResult().getType().cast<VectorType>();
}
static constexpr StringRef getKindAttrStrName() {
return "kind";
}
static CombiningKind getDefaultKind() {
return CombiningKind::ADD;
}
}];
let hasCustomAssemblyFormat = 1;
let hasVerifier = 1;
}
// TODO: Add transformation which decomposes ReshapeOp into an optimized
// sequence of vector rotate/shuffle/select operations.
def Vector_ReshapeOp :
Vector_Op<"reshape", [AttrSizedOperandSegments, NoSideEffect]>,
Arguments<(ins AnyVector:$vector, Variadic<Index>:$input_shape,
Variadic<Index>:$output_shape,
I64ArrayAttr:$fixed_vector_sizes)>,
Results<(outs AnyVector:$result)> {
let summary = "vector reshape operation";
let description = [{
Reshapes its vector operand from 'input_shape' to 'output_shape' maintaining
fixed vector dimension 'fixed_vector_sizes' on the innermost vector
dimensions.
The parameters 'input_shape' and 'output_shape' represent valid data shapes
across fixed vector shapes. For example, if a vector has a valid data
shape [6] with fixed vector size [8], then the valid data elements are
assumed to be stored at the beginning of the vector with the remaining
vector elements undefined.