[mlir][vector][avx] add AVX dot product to X86Vector dialect with lowering

[aartbik]

Summary

In the long run, we want to unify the dot product codegen solutions between
all target architectures, but this intrinsic enables experimenting with AVX
specific implementations in the meantime.

[nicolasvasilache]

mlir/include/mlir/Dialect/X86Vector/X86Vector.td

313

Like we argued/settled in other dialects (i.e. neon IIRC), this should have a
"codegen public" form x86.avx.dot on vector<2x4xf32> and an x86.avx.intrin.dot form on vector<8xf32>.

325

Then this can become:

%0 = x86vector.avx.dot %a, %b : vector<2x4xf32>
%1 = vector.extract %0[0, 0]: vector<2x4xf32>
%2 = vector.extract %0[1, 0]: vector<2x4xf32>
%d = addf %1, %2 : f32

[aartbik]

mlir/include/mlir/Dialect/X86Vector/X86Vector.td

313

Ah, I suppose you are referring to https://reviews.llvm.org/D98198

I somehow missed the second part of the discussion during the review. Let me have a careful read. Are you okay with this CL + later improvement or do you want me to make those changes first?

[nicolasvasilache]

I'd like to make sure we are on the same page globally on this.
As far as simple step forward, assuming you'd be in agreement, you could just rename to x86vector.avx.intr.dot and not change anything else.
The "codegen exposed" version would then be a new op and the lowering to x86vector would be a simple reshape.

The big thing then will be the iterative canonicalizations of reshapes all the way to vector.transfer ops which we have to do in a portable fashion.

[aartbik]

In D100593#2692703, @nicolasvasilache wrote:

I'd like to make sure we are on the same page globally on this.

Yes, I am! The lower/higher lanes semantics are always a bit confusing, and I like the solution of exploiting the higher dimensionality of the vector dialect to represent that better.
I will make both intrinsics, and then add a new op in next revision.

[aartbik]

see if you like what I have done to the place....

mlir/include/mlir/Dialect/X86Vector/X86Vector.td

def X86Vector_Dialect : Dialect {
   let name = "x86vector";
   let cppNamespace = "::mlir::x86vector";
 }
 
 //===----------------------------------------------------------------------===//
 // AVX512 op definitions
 //===----------------------------------------------------------------------===//
 
+// Operation that is part of the input dialect.
 class AVX512_Op<string mnemonic, list<OpTrait> traits = []> :
   Op<X86Vector_Dialect, "avx512." # mnemonic, traits> {}
 
+// Intrinsic operation used during lowering to LLVM IR.
 class AVX512_IntrOp<string mnemonic, int numResults, list<OpTrait> traits = []> :
   LLVM_IntrOpBase<X86Vector_Dialect, "avx512.intr." # mnemonic,
                   "x86_avx512_" # !subst(".", "_", mnemonic),
                   [], [], traits, numResults>;
 
 // Defined by first result overload. May have to be extended for other
 // instructions in the future.
 class AVX512_IntrOverloadedOp<string mnemonic,
                               list<OpTrait> traits = []> :
   LLVM_IntrOpBase<X86Vector_Dialect, "avx512.intr." # mnemonic,
                   "x86_avx512_" # !subst(".", "_", mnemonic),
                   /*list<int> overloadedResults=*/[0],
                   /*list<int> overloadedOperands=*/[],
                   traits, /*numResults=*/1>;
+
 //----------------------------------------------------------------------------//
 // MaskCompressOp
 //----------------------------------------------------------------------------//
 
 def MaskCompressOp : AVX512_Op<"mask.compress", [NoSideEffect,
   // TODO: Support optional arguments in `AllTypesMatch`. "type($src)" could
   // then be removed from assemblyFormat.
   AllTypesMatch<["a", "dst"]>,
def Vp2IntersectQIntrOp : AVX512_IntrOp<"vp2intersect.q.512", 2, [
   let arguments = (ins VectorOfLengthAndType<[8], [I64]>:$a,
                    VectorOfLengthAndType<[8], [I64]>:$b);
 }
 
 //===----------------------------------------------------------------------===//
 // AVX op definitions
 //===----------------------------------------------------------------------===//
 
+// Operation that is part of the input dialect.
 class AVX_Op<string mnemonic, list<OpTrait> traits = []> :
   Op<X86Vector_Dialect, "avx." # mnemonic, traits> {}
 
+// Operation that may be part of the input dialect, but whose
+// form is somewhere between the user view of the operation
+// and the actual lower level intrinsic in LLVM IR.
+class AVX_LowOp<string mnemonic, list<OpTrait> traits = []> :
+  Op<X86Vector_Dialect, "avx.intr." # mnemonic, traits> {}
+
+// Intrinsic operation used during lowering to LLVM IR.
 class AVX_IntrOp<string mnemonic, int numResults, list<OpTrait> traits = []> :
   LLVM_IntrOpBase<X86Vector_Dialect, "avx.intr." # mnemonic,
                   "x86_avx_" # !subst(".", "_", mnemonic),
                   [], [], traits, numResults>;
 
 //----------------------------------------------------------------------------//
 // AVX Rsqrt
 //----------------------------------------------------------------------------//
 
 def RsqrtOp : AVX_Op<"rsqrt", [NoSideEffect, SameOperandsAndResultType]> {
   let summary = "Rsqrt";
   let arguments = (ins VectorOfLengthAndType<[8], [F32]>:$a);
   let results = (outs VectorOfLengthAndType<[8], [F32]>:$b);
   let assemblyFormat = "$a attr-dict `:` type($a)";
 }
 
 def RsqrtIntrOp : AVX_IntrOp<"rsqrt.ps.256", 1, [NoSideEffect,
   SameOperandsAndResultType]> {
   let arguments = (ins VectorOfLengthAndType<[8], [F32]>:$a);
 }
 
+//----------------------------------------------------------------------------//
+// AVX Dot
+//----------------------------------------------------------------------------//
+
+def DotOp : AVX_LowOp<"dot", [NoSideEffect, SameOperandsAndResultType]> {
+  let summary = "Dot";
+  let description = [{
+    Computes the 4-way dot products of the lower and higher parts of the source
+    vectors and broadcasts the two results to the lower and higher elements of
+    the destination vector, respectively. Adding one element of the lower part
+    to one element of the higher part in the destination vector yields the full
+    dot product of the two source vectors.
+
+    Example:
+
+    ```mlir
+    %0 = x86vector.avx.intr.dot %a, %b : vector<8xf32>
+    %1 = vector.extractelement %0[%i0 : i32]: vector<8xf32>
+    %2 = vector.extractelement %0[%i4 : i32]: vector<8xf32>
+    %d = addf %1, %2 : f32
+    ```
+  }];
+  let arguments = (ins VectorOfLengthAndType<[8], [F32]>:$a,
+                       VectorOfLengthAndType<[8], [F32]>:$b);
+  let results = (outs VectorOfLengthAndType<[8], [F32]>:$res);
+  let assemblyFormat = "$a `,` $b attr-dict `:` type($res)";
+}
+
+def DotIntrOp : AVX_IntrOp<"dp.ps.256", 1, [NoSideEffect,
+    AllTypesMatch<["a", "b", "res"]>]> {
+  let arguments = (ins VectorOfLengthAndType<[8], [F32]>:$a,
+                       VectorOfLengthAndType<[8], [F32]>:$b, I8:$c);
+  let results = (outs VectorOfLengthAndType<[8], [F32]>:$res);
+}
+
 #endif // X86VECTOR_OPS

mlir/lib/Dialect/X86Vector/Transforms/LegalizeForLLVMExport.cpp

matchAndRewrite(RsqrtOp op, ArrayRef<Value> operands,
     RsqrtOp::Adaptor adaptor(operands);
 
     auto opType = adaptor.a().getType();
     rewriter.replaceOpWithNewOp<RsqrtIntrOp>(op, opType, adaptor.a());
     return success();
   }
 };
 
+struct DotOpConversion : public ConvertOpToLLVMPattern<DotOp> {
+  using ConvertOpToLLVMPattern<DotOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(DotOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    DotOp::Adaptor adaptor(operands);
+    auto opType = adaptor.a().getType();
+    Type llvmIntType = IntegerType::get(&getTypeConverter()->getContext(), 8);
+    // Dot product of all elements, broadcasted to all elements.
+    auto attr = rewriter.getI8IntegerAttr(0xff);
+    Value scale =
+        rewriter.create<LLVM::ConstantOp>(op.getLoc(), llvmIntType, attr);
+    rewriter.replaceOpWithNewOp<DotIntrOp>(op, opType, adaptor.a(), adaptor.b(),
+                                           scale);
+    return success();
+  }
+};
+
 /// An entry associating the "main" AVX512 op with its instantiations for
 /// vectors of 32-bit and 64-bit elements.
 template <typename OpTy, typename Intr32OpTy, typename Intr64OpTy>
 struct RegEntry {
   using MainOp = OpTy;
   using Intr32Op = Intr32OpTy;
   using Intr64Op = Intr64OpTy;
 };
using Registry = RegistryImpl<
     RegEntry<Vp2IntersectOp, Vp2IntersectDIntrOp, Vp2IntersectQIntrOp>>;
 
 } // namespace
 
 /// Populate the given list with patterns that convert from X86Vector to LLVM.
 void mlir::populateX86VectorLegalizeForLLVMExportPatterns(
     LLVMTypeConverter &converter, RewritePatternSet &patterns) {
   Registry::registerPatterns(converter, patterns);
-  patterns.add<MaskCompressOpConversion, RsqrtOpConversion>(converter);
+  patterns.add<MaskCompressOpConversion, RsqrtOpConversion, DotOpConversion>(
+      converter);
 }
 
 void mlir::configureX86VectorLegalizeForExportTarget(
     LLVMConversionTarget &target) {
   Registry::configureTarget(target);
   target.addLegalOp<MaskCompressIntrOp>();
   target.addIllegalOp<MaskCompressOp>();
   target.addLegalOp<RsqrtIntrOp>();
   target.addIllegalOp<RsqrtOp>();
+  target.addLegalOp<DotIntrOp>();
+  target.addIllegalOp<DotOp>();
 }

mlir/test/Dialect/X86Vector/legalize-for-llvm.mlir

 
 // CHECK-LABEL: func @avx_rsqrt
 func @avx_rsqrt(%a: vector<8xf32>) -> (vector<8xf32>)
 {
   // CHECK: x86vector.avx.intr.rsqrt.ps.256
   %0 = x86vector.avx.rsqrt %a : vector<8xf32>
   return %0 : vector<8xf32>
 }
+
+// CHECK-LABEL: func @avx_dot
+func @avx_dot(%a: vector<8xf32>, %b: vector<8xf32>) -> (vector<8xf32>)
+{
+  // CHECK: x86vector.avx.intr.dp.ps.256
+  %0 = x86vector.avx.intr.dot %a, %b : vector<8xf32>
+  return %0 : vector<8xf32>
+}

mlir/test/Dialect/X86Vector/roundtrip.mlir

 
 // CHECK-LABEL: func @avx_rsqrt
 func @avx_rsqrt(%a: vector<8xf32>) -> (vector<8xf32>)
 {
   // CHECK: x86vector.avx.rsqrt {{.*}} : vector<8xf32>
   %0 = x86vector.avx.rsqrt %a : vector<8xf32>
   return %0 : vector<8xf32>
 }
+
+// CHECK-LABEL: func @avx_dot
+func @avx_dot(%a: vector<8xf32>, %b: vector<8xf32>) -> (vector<8xf32>)
+{
+  // CHECK: x86vector.avx.intr.dot {{.*}} : vector<8xf32>
+  %0 = x86vector.avx.intr.dot %a, %b : vector<8xf32>
+  return %0 : vector<8xf32>
+}

mlir/test/Integration/Dialect/Vector/CPU/X86Vector/test-dot.mlir

+// RUN: mlir-opt %s -convert-scf-to-std -convert-vector-to-llvm="enable-x86vector" -convert-std-to-llvm | \
+// RUN: mlir-translate --mlir-to-llvmir | \
+// RUN: %lli --entry-function=entry --mattr="avx" --dlopen=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+func @entry() -> i32 {
+  %i0 = constant 0 : i32
+  %i4 = constant 4 : i32
+
+  %a = std.constant dense<[1.0,  2.0,  3.0,  4.0,  5.0,  6.0,  7.0,  8.0]> : vector<8xf32>
+  %b = std.constant dense<[9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0]> : vector<8xf32>
+  %r = x86vector.avx.intr.dot %a, %b : vector<8xf32>
+
+  %1 = vector.extractelement %r[%i0 : i32]: vector<8xf32>
+  %2 = vector.extractelement %r[%i4 : i32]: vector<8xf32>
+  %d = addf %1, %2 : f32
+
+  // CHECK: ( 110, 110, 110, 110, 382, 382, 382, 382 )
+  // CHECK: 492
+  vector.print %r : vector<8xf32>
+  vector.print %d : f32
+
+  return %i0 : i32
+}