@@ -49,6 +49,12 @@ def SVBoolMask : VectorWithTrailingDimScalableOfSizeAndType<
4949def SVEPredicateMask : VectorWithTrailingDimScalableOfSizeAndType<
5050 [16, 8, 4, 2, 1], [I1]>;
5151
52+ // A constraint for a 1-D scalable vector of `length`.
53+ class Scalable1DVectorOfLength<int length, list<Type> elementTypes> : ShapedContainerType<
54+ elementTypes, And<[IsVectorOfShape<[length]>, IsVectorTypeWithAnyDimScalablePred]>,
55+ "a 1-D scalable vector with length " # length,
56+ "::mlir::VectorType">;
57+
5258//===----------------------------------------------------------------------===//
5359// ArmSVE op definitions
5460//===----------------------------------------------------------------------===//
@@ -321,6 +327,121 @@ def ConvertToSvboolOp : ArmSVE_Op<"convert_to_svbool",
321327 let assemblyFormat = "$source attr-dict `:` type($source)";
322328}
323329
330+ // Inputs valid for the multi-vector zips (not including the 128-bit element zipqs)
331+ def ZipInputVectorType : AnyTypeOf<[
332+ Scalable1DVectorOfLength<2, [I64, F64]>,
333+ Scalable1DVectorOfLength<4, [I32, F32]>,
334+ Scalable1DVectorOfLength<8, [I16, F16, BF16]>,
335+ Scalable1DVectorOfLength<16, [I8]>],
336+ "an SVE vector with element size <= 64-bit">;
337+
338+ def ZipX2Op : ArmSVE_Op<"zip.x2", [
339+ Pure,
340+ AllTypesMatch<["sourceV1", "sourceV2", "resultV1", "resultV2"]>]
341+ > {
342+ let summary = "Multi-vector two-way zip op";
343+
344+ let description = [{
345+ This operation interleaves elements from two input SVE vectors, returning
346+ two new SVE vectors (`resultV1` and `resultV2`), which contain the low and
347+ high halves of the result respectively.
348+
349+ Example:
350+ ```mlir
351+ // sourceV1 = [ A1, A2, A3, ... An ]
352+ // sourceV2 = [ B1, B2, B3, ... Bn ]
353+ // (resultV1, resultV2) = [ A1, B1, A2, B2, A3, B3, ... An, Bn ]
354+ %resultV1, %resultV2 = arm_sve.zip.x2 %sourceV1, %sourceV2 : vector<[16]xi8>
355+ ```
356+
357+ Note: This requires SME 2 (`+sme2` in LLVM target features)
358+
359+ [Source](https://developer.arm.com/documentation/ddi0602/2023-12/SME-Instructions/ZIP--two-registers---Interleave-elements-from-two-vectors-?lang=en)
360+ }];
361+
362+ let arguments = (ins ZipInputVectorType:$sourceV1,
363+ ZipInputVectorType:$sourceV2);
364+
365+ let results = (outs ZipInputVectorType:$resultV1,
366+ ZipInputVectorType:$resultV2);
367+
368+ let builders = [
369+ OpBuilder<(ins "Value":$v1, "Value":$v2), [{
370+ build($_builder, $_state, v1.getType(), v1.getType(), v1, v2);
371+ }]>];
372+
373+ let assemblyFormat = "$sourceV1 `,` $sourceV2 attr-dict `:` type($sourceV1)";
374+
375+ let extraClassDeclaration = [{
376+ VectorType getVectorType() {
377+ return ::llvm::cast<VectorType>(getSourceV1().getType());
378+ }
379+ }];
380+ }
381+
382+ def ZipX4Op : ArmSVE_Op<"zip.x4", [
383+ Pure,
384+ AllTypesMatch<[
385+ "sourceV1", "sourceV2", "sourceV3", "sourceV4",
386+ "resultV1", "resultV2", "resultV3", "resultV4"]>]
387+ > {
388+ let summary = "Multi-vector four-way zip op";
389+
390+ let description = [{
391+ This operation interleaves elements from four input SVE vectors, returning
392+ four new SVE vectors, each of which contain a quarter of the result. The
393+ first quarter will be in `resultV1`, second in `resultV2`, third in
394+ `resultV3`, and fourth in `resultV4`.
395+
396+ ```mlir
397+ // sourceV1 = [ A1, A2, ... An ]
398+ // sourceV2 = [ B1, B2, ... Bn ]
399+ // sourceV3 = [ C1, C2, ... Cn ]
400+ // sourceV4 = [ D1, D2, ... Dn ]
401+ // (resultV1, resultV2, resultV3, resultV4)
402+ // = [ A1, B1, C1, D1, A2, B2, C2, D2, ... An, Bn, Cn, Dn ]
403+ %resultV1, %resultV2, %resultV3, %resultV4 = arm_sve.zip.x4
404+ %sourceV1, %sourceV2, %sourceV3, %sourceV4 : vector<[16]xi8>
405+ ```
406+
407+ **Warning:** The result of this op is undefined for 64-bit elements on
408+ hardware with less than 256-bit vectors!
409+
410+ Note: This requires SME 2 (`+sme2` in LLVM target features)
411+
412+ [Source](https://developer.arm.com/documentation/ddi0602/2023-12/SME-Instructions/ZIP--four-registers---Interleave-elements-from-four-vectors-?lang=en)
413+ }];
414+
415+ let arguments = (ins ZipInputVectorType:$sourceV1,
416+ ZipInputVectorType:$sourceV2,
417+ ZipInputVectorType:$sourceV3,
418+ ZipInputVectorType:$sourceV4);
419+
420+ let results = (outs ZipInputVectorType:$resultV1,
421+ ZipInputVectorType:$resultV2,
422+ ZipInputVectorType:$resultV3,
423+ ZipInputVectorType:$resultV4);
424+
425+ let builders = [
426+ OpBuilder<(ins "Value":$v1, "Value":$v2, "Value":$v3, "Value":$v4), [{
427+ build($_builder, $_state,
428+ v1.getType(), v1.getType(),
429+ v1.getType(), v1.getType(),
430+ v1, v2, v3, v4);
431+ }]>];
432+
433+ let assemblyFormat = [{
434+ $sourceV1 `,` $sourceV2 `,` $sourceV3 `,` $sourceV4 attr-dict
435+ `:` type($sourceV1)
436+ }];
437+
438+ let extraClassDeclaration = [{
439+ VectorType getVectorType() {
440+ return ::llvm::cast<VectorType>(getSourceV1().getType());
441+ }
442+ }];
443+ }
444+
324445def ScalableMaskedAddIOp : ScalableMaskedIOp<"masked.addi", "addition",
325446 [Commutative]>;
326447
0 commit comments