[mlir][AVX512] Add mask.compress to AVX512 dialect.

[springerm]

Summary

Adds mask.compress to the AVX512 dialect and defines a lowering to the LLVM dialect.

[springerm]

mlir/include/mlir/Dialect/AVX512/AVX512.td

34

There are _mm512_maskz_compress_ps etc. intrinsics for zeroed compress, but they use the same AVX instruction as _mm512_mask_compress_ps: mask.compress. So technically we don't really need the zeroed variant as a separate operation, and can instead set src := vector of all zeros.

Should I keep MaskzCompressOp or remove it from this commit? (In most cases (e.g., sparse dot product), we probably use the zeroed variant.)

[aartbik]

mlir/include/mlir/Dialect/AVX512/AVX512.td

36

Maybe say "Register k ...." just so it does not look like a typo to people less familiar with the mask registers

47

put quotes around the names "a" so it does not look like a dangling article

62

same as above

[ftynse]

mlir/include/mlir/Dialect/AVX512/AVX512.td

34

I don't see the value proposition of the zeroed variant, especially since it lowers to the same LLVM IR intrinsic as the non-zeroed one. If you want to avoid looking up the src operand and checking if it is a constant zero vector, you can either have a unit attribute zero on the "main" op + make the src argument optional + add a verifier that either the operand or the attribute is present or have a constant_src attribute with _any_ value (not only zero), that can be used instead of the src argument + the same scheme as above. The latter is closer in spirit to Linalg, which is the dialect that will target this I presume.

47

Surrounding variable names with quotes or backticks is a generally good practice. Doing so for k in error messages in TypesMatchWith would also address the issue of misreading them as a typo.

mlir/include/mlir/Dialect/LLVMIR/LLVMAVX512.td

36

Naming nit: this is a specific kind of overloading that is fully defined by the first result, there may be other kinds so I wouldn't use IntrOverloadedOp here only to find out later that it cannot be used for, e.g., an op that is overloaded by its second operand.

42

Nit: add a comment for 1 for consistency with other literal arguments.

[nicolasvasilache]

mlir/include/mlir/Dialect/AVX512/AVX512.td

34

Hmm I thought I had replied but apparently did not commit ..

I'm be curious to see if using one vs the other can result in significant difference in the final assembly generated by LLVM in interesting cases (likely a few commits down the line until we have a full dot/matvec example). Technically there may be a difference for LLVM's optimizations.
Not sure yet whether this matters in practice but I'd be interested in a confirmation based on data.

Re. where the funneling should go, we could indeed drop 1 op and use an attribute.
However, I think we would then need to make some changes to the autogenerated LLVMIR translation in non-trivial ways (I forget)?
If so, is it worth it?

[springerm]

Remove maskz.compress. Make src optional.

[springerm]

mlir/include/mlir/Dialect/AVX512/AVX512.td

34

I decided to add a constant_src attribute. If neither constant_src nor src is specified, a vector of all zeros is used (same as maskz.compress).

mlir/include/mlir/Dialect/LLVMIR/LLVMAVX512.td

36

Added a comment. I can change this if necessary when adding additional ops.

mlir/include/mlir/Dialect/AVX512/AVX512.td

 
 //===----------------------------------------------------------------------===//
 // AVX512 op definitions
 //===----------------------------------------------------------------------===//
 
 class AVX512_Op<string mnemonic, list<OpTrait> traits = []> :
   Op<AVX512_Dialect, mnemonic, traits> {}
 
+def MaskCompressOp : AVX512_Op<"mask.compress", [NoSideEffect,
+  // TODO: Support optional arguments in `AllTypesMatch`. "type($src)" could
+  // then be removed from assemblyFormat.
+  AllTypesMatch<["a", "dst"]>,
+  TypesMatchWith<"`k` has the same number of bits as elements in `dst`",
+                 "dst", "k",
+                 "VectorType::get({$_self.cast<VectorType>().getShape()[0]}, "
+                 "IntegerType::get($_self.getContext(), 1))">]> {
+  let summary = "Masked compress op";
+  let description = [{
+  The mask.compress op is an AVX512 specific op that can lower to the
+  `llvm.mask.compress` instruction. Instead of `src`, a constant vector
+  vector attribute `constant_src` may be specified. If neither `src` nor
+  `constant_src` is specified, the remaining elements in the result vector are
+  set to zero.
+
+  #### From the Intel Intrinsics Guide:
+
+  Contiguously store the active integer/floating-point elements in `a` (those
+  with their respective bit set in writemask `k`) to `dst`, and pass through the
+  remaining elements from `src`.
+  }];
+  let verifier = [{ return ::verify(*this); }];
+  let arguments = (ins VectorOfLengthAndType<[16, 16, 8, 8],
+                                             [I1, I1, I1, I1]>:$k,
+                   VectorOfLengthAndType<[16, 16, 8, 8],
+                                         [F32, I32, F64, I64]>:$a,
+                   Optional<VectorOfLengthAndType<[16, 16, 8, 8],
+                                                  [F32, I32, F64, I64]>>:$src,
+                   OptionalAttr<ElementsAttr>:$constant_src);
+  let results = (outs VectorOfLengthAndType<[16, 16, 8, 8],
+                                            [F32, I32, F64, I64]>:$dst);
+  let assemblyFormat = "$k `,` $a (`,` $src^)? attr-dict"
+                       " `:` type($dst) (`,` type($src)^)?";
+}
+
 def MaskRndScaleOp : AVX512_Op<"mask.rndscale", [NoSideEffect,
   AllTypesMatch<["src", "a", "dst"]>,
   TypesMatchWith<"imm has the same number of bits as elements in dst",
                  "dst", "imm",
                  "IntegerType::get($_self.getContext(), "
                  "($_self.cast<VectorType>().getShape()[0]))">]> {
   let summary = "Masked roundscale op";
   let description = [{

mlir/include/mlir/Dialect/LLVMIR/LLVMAVX512.td

 // MLIR LLVM AVX512 intrinsics using the MLIR LLVM Dialect type system
 //----------------------------------------------------------------------------//
 
 class LLVMAVX512_IntrOp<string mnemonic, int numResults, list<OpTrait> traits = []> :
   LLVM_IntrOpBase<LLVMAVX512_Dialect, mnemonic,
                   "x86_avx512_" # !subst(".", "_", mnemonic),
                   [], [], traits, numResults>;
 
+// Defined by first result overload. May have to be extended for other
+// instructions in the future.
+class LLVMAVX512_IntrOverloadedOp<string mnemonic,
+                                  list<OpTrait> traits = []> :
+  LLVM_IntrOpBase<LLVMAVX512_Dialect, mnemonic,
+                  "x86_avx512_" # !subst(".", "_", mnemonic),
+                  /*list<int> overloadedResults=*/[0],
+                  /*list<int> overloadedOperands=*/[],
+                  traits, /*numResults=*/1>;
+
 def LLVM_x86_avx512_mask_rndscale_ps_512 :
   LLVMAVX512_IntrOp<"mask.rndscale.ps.512", 1>,
   Arguments<(ins LLVM_Type, LLVM_Type, LLVM_Type, LLVM_Type, LLVM_Type)>;
 
 def LLVM_x86_avx512_mask_rndscale_pd_512 :
   LLVMAVX512_IntrOp<"mask.rndscale.pd.512", 1>,
   Arguments<(ins LLVM_Type, LLVM_Type, LLVM_Type, LLVM_Type, LLVM_Type)>;
 
 def LLVM_x86_avx512_mask_scalef_ps_512 :
   LLVMAVX512_IntrOp<"mask.scalef.ps.512", 1>,
   Arguments<(ins LLVM_Type, LLVM_Type, LLVM_Type, LLVM_Type, LLVM_Type)>;
 
 def LLVM_x86_avx512_mask_scalef_pd_512 :
   LLVMAVX512_IntrOp<"mask.scalef.pd.512", 1>,
   Arguments<(ins LLVM_Type, LLVM_Type, LLVM_Type, LLVM_Type, LLVM_Type)>;
 
+def LLVM_x86_avx512_mask_compress :
+  LLVMAVX512_IntrOverloadedOp<"mask.compress">,
+  Arguments<(ins LLVM_Type, LLVM_Type, LLVM_Type)>;
+
 def LLVM_x86_avx512_vp2intersect_d_512 :
   LLVMAVX512_IntrOp<"vp2intersect.d.512", 2>,
   Arguments<(ins LLVM_Type, LLVM_Type)>;
 
 def LLVM_x86_avx512_vp2intersect_q_512 :
   LLVMAVX512_IntrOp<"vp2intersect.q.512", 2>,
   Arguments<(ins LLVM_Type, LLVM_Type)>;
 
 #endif // AVX512_OPS

mlir/lib/Conversion/AVX512ToLLVM/ConvertAVX512ToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     if (elementType.isF64())
       return LLVM::detail::oneToOneRewrite(
           op, LLVM::x86_avx512_mask_rndscale_pd_512::getOperationName(),
           operands, *getTypeConverter(), rewriter);
     return failure();
   }
 };
 
+struct MaskCompressOpConversion
+    : public ConvertOpToLLVMPattern<MaskCompressOp> {
+  using ConvertOpToLLVMPattern<MaskCompressOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(MaskCompressOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    MaskCompressOp::Adaptor adaptor(operands);
+    auto opType = adaptor.a().getType();
+
+    Value src;
+    if (op.src()) {
+      src = adaptor.src();
+    } else if (op.constant_src()) {
+      src = rewriter.create<ConstantOp>(op.getLoc(), opType,
+                                        op.constant_srcAttr());
+    } else {
+      Attribute zeroAttr = rewriter.getZeroAttr(opType);
+      src = rewriter.create<ConstantOp>(op->getLoc(), opType, zeroAttr);
+    }
+
+    rewriter.replaceOpWithNewOp<LLVM::x86_avx512_mask_compress>(
+        op, opType, adaptor.a(), src, adaptor.k());
+
+    return success();
+  }
+};
+
 struct ScaleFOp512Conversion : public ConvertToLLVMPattern {
   explicit ScaleFOp512Conversion(MLIRContext *context,
                                  LLVMTypeConverter &typeConverter)
       : ConvertToLLVMPattern(MaskScaleFOp::getOperationName(), context,
                              typeConverter) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,

 /// Populate the given list with patterns that convert from AVX512 to LLVM.
 void mlir::populateAVX512ToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   // clang-format off
   patterns.insert<MaskRndScaleOp512Conversion,
                   ScaleFOp512Conversion,
                   Vp2IntersectOp512Conversion>(&converter.getContext(),
                                                converter);
+  patterns.insert<MaskCompressOpConversion>(converter);
   // clang-format on
 }

mlir/lib/Dialect/AVX512/IR/AVX512Dialect.cpp

 
 void avx512::AVX512Dialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/AVX512/AVX512.cpp.inc"
       >();
 }
 
+static LogicalResult verify(avx512::MaskCompressOp op) {
+  if (op.src() && op.constant_src())
+    return emitError(op.getLoc(), "cannot use both src and constant_src");
+
+  if (op.src() && (op.src().getType() != op.dst().getType()))
+    return emitError(op.getLoc(),
+                     "failed to verify that src and dst have same type");
+
+  if (op.constant_src() && (op.constant_src()->getType() != op.dst().getType()))
+    return emitError(
+        op.getLoc(),
+        "failed to verify that constant_src and dst have same type");
+
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/AVX512/AVX512.cpp.inc"

mlir/test/Conversion/AVX512ToLLVM/convert-to-llvm.mlir

func @avx512_mask_rndscale(%a: vector<16xf32>, %b: vector<8xf64>, %i32: i32, %i16: i16, %i8: i8)
   %2 = avx512.mask.scalef %a, %a, %a, %i16, %i32: vector<16xf32>
   // CHECK: llvm_avx512.mask.scalef.pd.512
   %3 = avx512.mask.scalef %b, %b, %b, %i8, %i32: vector<8xf64>
 
   // Keep results alive.
   return %0, %1, %2, %3 : vector<16xf32>, vector<8xf64>, vector<16xf32>, vector<8xf64>
 }
 
+func @avx512_mask_compress(%k1: vector<16xi1>, %a1: vector<16xf32>,
+                           %k2: vector<8xi1>, %a2: vector<8xi64>)
+  -> (vector<16xf32>, vector<16xf32>, vector<8xi64>)
+{
+  // CHECK: llvm_avx512.mask.compress
+  %0 = avx512.mask.compress %k1, %a1 : vector<16xf32>
+  // CHECK: llvm_avx512.mask.compress
+  %1 = avx512.mask.compress %k1, %a1 {constant_src = dense<5.0> : vector<16xf32>} : vector<16xf32>
+  // CHECK: llvm_avx512.mask.compress
+  %2 = avx512.mask.compress %k2, %a2, %a2 : vector<8xi64>, vector<8xi64>
+  return %0, %1, %2 : vector<16xf32>, vector<16xf32>, vector<8xi64>
+}
+
 func @avx512_vp2intersect(%a: vector<16xi32>, %b: vector<8xi64>)
   -> (vector<16xi1>, vector<16xi1>, vector<8xi1>, vector<8xi1>)
 {
   // CHECK: llvm_avx512.vp2intersect.d.512
   %0, %1 = avx512.vp2intersect %a, %a : vector<16xi32>
   // CHECK: llvm_avx512.vp2intersect.q.512
   %2, %3 = avx512.vp2intersect %b, %b : vector<8xi64>
   return %0, %1, %2, %3 : vector<16xi1>, vector<16xi1>, vector<8xi1>, vector<8xi1>
 }

mlir/test/Dialect/AVX512/roundtrip.mlir

func @avx512_vp2intersect(%a: vector<16xi32>, %b: vector<8xi64>)
   -> (vector<16xi1>, vector<16xi1>, vector<8xi1>, vector<8xi1>)
 {
   // CHECK: avx512.vp2intersect {{.*}} : vector<16xi32>
   %0, %1 = avx512.vp2intersect %a, %a : vector<16xi32>
   // CHECK: avx512.vp2intersect {{.*}} : vector<8xi64>
   %2, %3 = avx512.vp2intersect %b, %b : vector<8xi64>
   return %0, %1, %2, %3 : vector<16xi1>, vector<16xi1>, vector<8xi1>, vector<8xi1>
 }
+
+func @avx512_mask_compress(%k1: vector<16xi1>, %a1: vector<16xf32>,
+                           %k2: vector<8xi1>, %a2: vector<8xi64>)
+  -> (vector<16xf32>, vector<16xf32>, vector<8xi64>)
+{
+  // CHECK: avx512.mask.compress {{.*}} : vector<16xf32>
+  %0 = avx512.mask.compress %k1, %a1 : vector<16xf32>
+  // CHECK: avx512.mask.compress {{.*}} : vector<16xf32>
+  %1 = avx512.mask.compress %k1, %a1 {constant_src = dense<5.0> : vector<16xf32>} : vector<16xf32>
+  // CHECK: avx512.mask.compress {{.*}} : vector<8xi64>
+  %2 = avx512.mask.compress %k2, %a2, %a2 : vector<8xi64>, vector<8xi64>
+  return %0, %1, %2 : vector<16xf32>, vector<16xf32>, vector<8xi64>
+}

mlir/test/Integration/Dialect/Vector/CPU/AVX512/test-mask-compress.mlir

+// RUN: mlir-opt %s -convert-scf-to-std -convert-vector-to-llvm="enable-avx512" -convert-std-to-llvm  | \
+// RUN: mlir-translate  --mlir-to-llvmir | \
+// RUN: %lli --entry-function=entry --mattr="avx512bw" --dlopen=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+func @entry() -> i32 {
+  %i0 = constant 0 : i32
+
+  %a = std.constant dense<[1., 0., 0., 2., 4., 3., 5., 7., 8., 1., 5., 5., 3., 1., 0., 7.]> : vector<16xf32>
+  %k = std.constant dense<[1,  0,  1,  1,  1,  0,  0,  0,  0,  0,  1,  1,  1,  1,  1,  0]> : vector<16xi1>
+  %r1 = avx512.mask.compress %k, %a : vector<16xf32>
+  %r2 = avx512.mask.compress %k, %a {constant_src = dense<5.0> : vector<16xf32>} : vector<16xf32>
+
+  vector.print %r1 : vector<16xf32>
+  // CHECK: ( 1, 0, 2, 4, 5, 5, 3, 1, 0, 0, 0, 0, 0, 0, 0, 0 )
+
+  vector.print %r2 : vector<16xf32>
+  // CHECK: ( 1, 0, 2, 4, 5, 5, 3, 1, 0, 5, 5, 5, 5, 5, 5, 5 )
+
+  %src = std.constant dense<[0., 2., 1., 8., 6., 4., 4., 3., 2., 8., 5., 6., 3., 7., 6., 9.]> : vector<16xf32>
+  %r3 = avx512.mask.compress %k, %a, %src : vector<16xf32>, vector<16xf32>
+
+  vector.print %r3 : vector<16xf32>
+  // CHECK: ( 1, 0, 2, 4, 5, 5, 3, 1, 0, 8, 5, 6, 3, 7, 6, 9 )
+
+  return %i0 : i32
+}

mlir/test/Integration/Dialect/Vector/CPU/AVX512/test-vp2intersect-i32.mlir

 // RUN: mlir-opt %s -convert-scf-to-std -convert-vector-to-llvm="enable-avx512" -convert-std-to-llvm  | \
-// RUN: mlir-translate  --avx512-mlir-to-llvmir | \
+// RUN: mlir-translate  --mlir-to-llvmir | \
 // RUN: %lli --entry-function=entry --mattr="avx512bw,avx512vp2intersect" --dlopen=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 // Note: To run this test, your CPU must support AVX512 vp2intersect.
 
 func @entry() -> i32 {
   %i0 = constant 0 : i32
   %i1 = constant 1: i32

mlir/test/Target/avx512.mlir

llvm.func @LLVM_x86_avx512_mask_pd_512(%a: vector<8xf64>,
   %0 = "llvm_avx512.mask.rndscale.pd.512"(%a, %b, %a, %c, %b) :
     (vector<8xf64>, i32, vector<8xf64>, i8, i32) -> vector<8xf64>
   // CHECK: call <8 x double> @llvm.x86.avx512.mask.scalef.pd.512(<8 x double>
   %1 = "llvm_avx512.mask.scalef.pd.512"(%a, %a, %a, %c, %b) :
     (vector<8xf64>, vector<8xf64>, vector<8xf64>, i8, i32) -> vector<8xf64>
   llvm.return %1: vector<8xf64>
 }
 
+// CHECK-LABEL: define <16 x float> @LLVM_x86_mask_compress
+llvm.func @LLVM_x86_mask_compress(%k: vector<16xi1>, %a: vector<16xf32>)
+  -> vector<16xf32>
+{
+  // CHECK: call <16 x float> @llvm.x86.avx512.mask.compress.v16f32(
+  %0 = "llvm_avx512.mask.compress"(%a, %a, %k) :
+    (vector<16xf32>, vector<16xf32>, vector<16xi1>) -> vector<16xf32>
+  llvm.return %0 : vector<16xf32>
+}
+
 // CHECK-LABEL: define { <16 x i1>, <16 x i1> } @LLVM_x86_vp2intersect_d_512
 llvm.func @LLVM_x86_vp2intersect_d_512(%a: vector<16xi32>, %b: vector<16xi32>)
   -> !llvm.struct<(vector<16 x i1>, vector<16 x i1>)>
 {
   // CHECK: call { <16 x i1>, <16 x i1> } @llvm.x86.avx512.vp2intersect.d.512(<16 x i32>
   %0 = "llvm_avx512.vp2intersect.d.512"(%a, %b) :
     (vector<16xi32>, vector<16xi32>) -> !llvm.struct<(vector<16 x i1>, vector<16 x i1>)>
   llvm.return %0 : !llvm.struct<(vector<16 x i1>, vector<16 x i1>)>