Add vp2intersect to AVX512 dialect.

[springerm]

Summary

Adds vp2intersect to the AVX512 dialect and defines a lowering to the
LLVM dialect.

Author: Matthias Springer <[e-mail]>

[nicolasvasilache]

Looks great, thanks @springerm !

Let's address the nits and land this, please lmk if you need help landing.
If you do need help landing, please also read https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access.

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

106

In LLVM OSS we just write // TODO: stufff.

107

I wouldn't leave this as a TODO but more as a simple note.

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

49

Note that this op and the one above are candidates to just disappearing and mlir-translate just converting the ops from AVX512/AVX512.td directly.
This requires machinery so don't worry about it for now.

However, note that this is not the case for the ops you just added as the struct result type does not exist in std.

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

29

nit: missing newline

mlir/test/Dialect/AVX512/roundtrip.mlir

32

nit: missing newline

[ftynse]

mlir/include/mlir/Dialect/AVX512/AVX512Dialect.h

19

We shouldn't be needing the OpImplementation header in another _header_, as its name indicates. LLVM in general prefers not to include headers in other headers unless strictly necessary to reduce compilation time.

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

49

However, note that this is not the case for the ops you just added as the struct result type does not exist in std.

We can use multi-return ops and convert them to struct-return + extract-element...

56–57

Nit: We have a preference for op def names (which will end up being C++ class names) to match the mnemonic of the operation, it's easier to understand what in the code produces which IR this way. So epi64 -> q.

mlir/lib/Conversion/AVX512ToLLVM/ConvertAVX512ToLLVM.cpp

81–85

: public ConvertOpToLLVMPattern<Vp2IntersectOp> will likely save some boilerplate.

mlir/test/Dialect/AVX512/roundtrip.mlir

23

Generally, there is no need to test generated code such as the syntax obtained from let assemblyFormat. We checked the generator instead.

[springerm]

Minor changes.

[springerm]

mlir/include/mlir/Dialect/AVX512/AVX512Dialect.h

19

Without this, the generated Dialect/AVX512/AVX512.h.inc file does not compile:

[...] AVX512IncGen/mlir/Dialect/AVX512/AVX512.h.inc:162:240: error: no member named 'OpAsmOpInterface' in namespace 'mlir'
class Vp2IntersectOp : public ::mlir::Op<Vp2IntersectOp, ::mlir::OpTrait::ZeroRegion, ::mlir::OpTrait::NResults<2>::Impl, ::mlir::OpTrait::ZeroSuccessor, ::mlir::OpTrait::NOperands<2>::Impl, ::mlir::MemoryEffectOpInterface::Trait, ::mlir::OpAsmOpInterface::Trait> {
[...]

Not sure what the OpAsmOpInterface is doing, but it is added only when specifying a second output in the td file.

Is there a better way to fix this?

mlir/test/Dialect/AVX512/roundtrip.mlir

23

So you mean

// CHECK: avx512.vp2intersect {{.*}} : vector<16xi32>

could be simplified to just

// CHECK: avx512.vp2intersect {{.*}}

and I can just return a single value such as %0 for simplicity?

[ftynse]

mlir/include/mlir/Dialect/AVX512/AVX512Dialect.h

19

Hmm, seems like layering was broken in some of these headers. Okay to keep the include for now.

mlir/test/Dialect/AVX512/roundtrip.mlir

23

I mean the entire test in roundtrip.mlir here is useless, it does not exercise C++ you wrote, only the code that was generated from ODS and we have already tested the generator.

[nicolasvasilache]

mlir/test/Dialect/AVX512/roundtrip.mlir

23

Hmm, I generally find those useful as general documentation that runs (whatever people add in the .td has no guarantee of executing and not getting stale).
If we drop this, should we start a more general cleanup effort to hunt down ops all such ops that come from an assembly format?

[ftynse]

mlir/test/Dialect/AVX512/roundtrip.mlir

23

It's fine to keep some minimal amount of these IMO, but having a lot is just extra review/maintenance load.

[aartbik]

Cool to see new AVX512 stuff, thanks!

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

31

super minor nit (since it is only local), but we tend to use "unsigned" instead of int for such situations

mlir/test/Dialect/AVX512/roundtrip.mlir

23

+1 on having such tests since it acts as some documentation

[ftynse]

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

31

There is no unsigned in tablegen AFAIK.

[aartbik]

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

31

Ah, in that case sorry for the noise.

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

let arguments = (ins VectorOfLengthAndType<[16, 8], [F32, F64]>:$src,
                    I32:$rounding
             );
   let results = (outs VectorOfLengthAndType<[16, 8], [F32, F64]>:$dst);
   // Fully specified by traits.
   let assemblyFormat =
     "$src `,` $a `,` $b `,` $k `,` $rounding attr-dict `:` type($dst)";
 }
 
+def Vp2IntersectOp : AVX512_Op<"vp2intersect", [NoSideEffect,
+  AllTypesMatch<["a", "b"]>,
+  TypesMatchWith<"k1 has the same number of bits as elements in a",
+                 "a", "k1",
+                 "IntegerType::get($_self.getContext(), "
+                 "($_self.cast<VectorType>().getShape()[0]))">,
+  TypesMatchWith<"k2 has the same number of bits as elements in b",
+                 // Should use `b` instead of `a`, but that would require
+                 // adding `type($b)` to assemblyFormat.
+                 "a", "k2",
+                 "IntegerType::get($_self.getContext(), "
+                 "($_self.cast<VectorType>().getShape()[0]))">]> {
+  let summary = "Vp2Intersect op";
+  let description = [{
+    The `vp2intersect` op is an AVX512 specific op that can lower to the proper
+    LLVMAVX512 operation: `llvm.vp2intersect.d.512` or
+    `llvm.vp2intersect.q.512` depending on the type of MLIR vectors it is
+    applied to.
+
+    #### From the Intel Intrinsics Guide:
+
+    Compute intersection of packed integer vectors `a` and `b`, and store
+    indication of match in the corresponding bit of two mask registers
+    specified by `k1` and `k2`. A match in corresponding elements of `a` and
+    `b` is indicated by a set bit in the corresponding bit of the mask
+    registers.
+  }];
+  let arguments = (ins VectorOfLengthAndType<[16, 8], [I32, I64]>:$a,
+                   VectorOfLengthAndType<[16, 8], [I32, I64]>:$b
+                   );
+  let results = (outs AnyTypeOf<[I16, I8]>:$k1,
+                 AnyTypeOf<[I16, I8]>:$k2
+                 );
+  let assemblyFormat =
+    "$a `,` $b attr-dict `:` type($a)";
+}
+
 #endif // AVX512_OPS

mlir/include/mlir/Dialect/AVX512/AVX512Dialect.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_AVX512_AVX512DIALECT_H_
 #define MLIR_DIALECT_AVX512_AVX512DIALECT_H_
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #include "mlir/Dialect/AVX512/AVX512Dialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/AVX512/AVX512.h.inc"
 
 #endif // MLIR_DIALECT_AVX512_AVX512DIALECT_H_

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

def LLVMAVX512_Dialect : Dialect {
   let name = "llvm_avx512";
   let cppNamespace = "::mlir::LLVM";
 }
 
 //----------------------------------------------------------------------------//
 // MLIR LLVM AVX512 intrinsics using the MLIR LLVM Dialect type system
 //----------------------------------------------------------------------------//
 
-class LLVMAVX512_IntrOp<string mnemonic, list<OpTrait> traits = []> :
+class LLVMAVX512_IntrOp<string mnemonic, int numResults, list<OpTrait> traits = []> :
   LLVM_IntrOpBase<LLVMAVX512_Dialect, mnemonic,
                   "x86_avx512_" # !subst(".", "_", mnemonic),
-                  [], [], traits, 1>;
+                  [], [], traits, numResults>;
 
 def LLVM_x86_avx512_mask_rndscale_ps_512 :
-  LLVMAVX512_IntrOp<"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">,
+  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">,
+  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">,
+  LLVMAVX512_IntrOp<"mask.scalef.pd.512", 1>,
   Arguments<(ins LLVM_Type, LLVM_Type, 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

if (elementType.isF32())
           *getTypeConverter(), rewriter);
     if (elementType.isF64())
       return LLVM::detail::oneToOneRewrite(
           op, LLVM::x86_avx512_mask_scalef_pd_512::getOperationName(), operands,
           *getTypeConverter(), rewriter);
     return failure();
   }
 };
+
+struct Vp2IntersectOp512Conversion
+    : public ConvertOpToLLVMPattern<Vp2IntersectOp> {
+  explicit Vp2IntersectOp512Conversion(MLIRContext *context,
+                                       LLVMTypeConverter &typeConverter)
+      : ConvertOpToLLVMPattern<Vp2IntersectOp>(typeConverter) {}
+
+  LogicalResult
+  matchAndRewrite(Vp2IntersectOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type elementType =
+        op.a().getType().template cast<VectorType>().getElementType();
+    if (elementType.isInteger(32))
+      return LLVM::detail::oneToOneRewrite(
+          op, LLVM::x86_avx512_vp2intersect_d_512::getOperationName(), operands,
+          *getTypeConverter(), rewriter);
+    if (elementType.isInteger(64))
+      return LLVM::detail::oneToOneRewrite(
+          op, LLVM::x86_avx512_vp2intersect_q_512::getOperationName(), operands,
+          *getTypeConverter(), rewriter);
+    return failure();
+  }
+};
 } // namespace
 
 /// 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>(&converter.getContext(), converter);
+                  ScaleFOp512Conversion,
+                  Vp2IntersectOp512Conversion>(&converter.getContext(),
+                                               converter);
   // clang-format on
 }

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

func @avx512_mask_rndscale(%a: vector<16xf32>, %b: vector<8xf64>, %i32: i32, %i16: i16, %i8: i8)
   // CHECK: llvm_avx512.mask.scalef.ps.512
   %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_vp2intersect(%a: vector<16xi32>, %b: vector<8xi64>)
+  -> (i16, i16, i8, i8)
+{
+  // 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 : i16, i16, i8, i8
+}

mlir/test/Dialect/AVX512/roundtrip.mlir

func @avx512_scalef(%a: vector<16xf32>, %b: vector<8xf64>, %i32: i32, %i16: i16, %i8: i8)
   -> (vector<16xf32>, vector<8xf64>)
 {
   // CHECK: avx512.mask.scalef {{.*}}: vector<16xf32>
   %0 = avx512.mask.scalef %a, %a, %a, %i16, %i32: vector<16xf32>
   // CHECK: avx512.mask.scalef {{.*}}: vector<8xf64>
   %1 = avx512.mask.scalef %b, %b, %b, %i8, %i32 : vector<8xf64>
   return %0, %1: vector<16xf32>, vector<8xf64>
 }
+
+func @avx512_vp2intersect(%a: vector<16xi32>, %b: vector<8xi64>)
+  -> (i16, i16, i8, i8)
+{
+  // 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 : i16, i16, i8, i8
+}

mlir/test/Target/avx512.mlir

llvm.func @LLVM_x86_avx512_mask_pd_512(%a: vector<8xf64>,
   // CHECK: call <8 x double> @llvm.x86.avx512.mask.rndscale.pd.512(<8 x double>
   %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 <{ i16, i16 }> @LLVM_x86_vp2intersect_d_512
+llvm.func @LLVM_x86_vp2intersect_d_512(%a: vector<16xi32>, %b: vector<16xi32>)
+  -> !llvm.struct<packed (i16, i16)>
+{
+  // 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<packed (i16, i16)>
+  llvm.return %0 : !llvm.struct<packed (i16, i16)>
+}
+
+// CHECK-LABEL: define <{ i8, i8 }> @LLVM_x86_vp2intersect_q_512
+llvm.func @LLVM_x86_vp2intersect_q_512(%a: vector<8xi64>, %b: vector<8xi64>)
+  -> !llvm.struct<packed (i8, i8)>
+{
+  // CHECK: call { <8 x i1>, <8 x i1> } @llvm.x86.avx512.vp2intersect.q.512(<8 x i64>
+  %0 = "llvm_avx512.vp2intersect.q.512"(%a, %b) :
+    (vector<8xi64>, vector<8xi64>) -> !llvm.struct<packed (i8, i8)>
+  llvm.return %0 : !llvm.struct<packed (i8, i8)>
+}