[mlir][vector][avx512] move avx512 lowering pass into general vector lowering

[aartbik]

Summary

A separate AVX512 lowering pass does not compose well with the regular
vector lowering pass. As such, it is at risk of code duplication and
lowering inconsistencies. This change removes the separate AVX512 lowering
pass and makes it an "option" in the regular vector lowering pass
(viz. vector dialect "augmented" with AVX512 dialect).

[rriddle]

mlir/include/mlir/Conversion/Passes.td

403

nit: Add a space after the comma

410–414

nit: Wrap these at 80 characters.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

1590

Can we move this out to a new file ConvertVectorToLLVMPass.cpp so that this file doesn't need to start depending on other dialects?

[rriddle]

mlir/include/mlir/Conversion/Passes.td

390

Could you(in a followup is fine) add a proper description for this pass?

[aartbik]

mlir/include/mlir/Conversion/Passes.td

390

Good idea. Might as well make a start with that here :-)

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

1590

Yes, excellent idea. That allows for a cleaner separation of that tasks. Thanks.

[aartbik]

addressed reviewer's feedback

[nicolasvasilache]

Great, thanks for refactoring this!

[bondhugula]

mlir/include/mlir/Conversion/Passes.td

414

Nit: Augments the vector dialect ... isn't really accurate here. Instead:

Enables the use of AVX512 dialect while lowering the vector dialect.

mlir/include/mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h

31

setEnable -> enable ?

35

setEnable -> enable?

[aartbik]

Thanks!

mlir/include/mlir/Conversion/Passes.td

414

I thought this phrasing was okay, but I will send a follow up CL with your suggestion

mlir/include/mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h

31

Agreed, this is perhaps becoming a bit too verbose
(I used the "set" since it sets it to either true/false, but using enable verb wise is indeed better)

mlir/include/mlir/Conversion/AVX512ToLLVM/ConvertAVX512ToLLVM.h

 //===- ConvertAVX512ToLLVM.h - Conversion Patterns from AVX512 to LLVM ----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_AVX512TOLLVM_CONVERTAVX512TOLLVM_H_
 #define MLIR_CONVERSION_AVX512TOLLVM_CONVERTAVX512TOLLVM_H_
 
-#include <memory>
-
 namespace mlir {
+
 class LLVMTypeConverter;
-class ModuleOp;
-template <typename T> class OperationPass;
 class OwningRewritePatternList;
 
 /// Collect a set of patterns to convert from the AVX512 dialect to LLVM.
 void populateAVX512ToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                             OwningRewritePatternList &patterns);
 
-/// Create a pass to convert AVX512 operations to the LLVMIR dialect.
-std::unique_ptr<OperationPass<ModuleOp>> createConvertAVX512ToLLVMPass();
-
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_AVX512TOLLVM_CONVERTAVX512TOLLVM_H_

mlir/include/mlir/Conversion/Passes.h

 //===- Passes.h - Conversion Pass Construction and Registration -----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_PASSES_H
 #define MLIR_CONVERSION_PASSES_H
 
-#include "mlir/Conversion/AVX512ToLLVM/ConvertAVX512ToLLVM.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"

mlir/include/mlir/Conversion/Passes.td

def ConvertAffineToStandard : Pass<"lower-affine"> {
   let dependentDialects = [
     "scf::SCFDialect",
     "StandardOpsDialect",
     "vector::VectorDialect"
   ];
 }
 
 //===----------------------------------------------------------------------===//
-// AVX512ToLLVM
-//===----------------------------------------------------------------------===//
-
-def ConvertAVX512ToLLVM : Pass<"convert-avx512-to-llvm", "ModuleOp"> {
-  let summary = "Convert the operations from the avx512 dialect into the LLVM "
-                "dialect";
-  let constructor = "mlir::createConvertAVX512ToLLVMPass()";
-  let dependentDialects = ["LLVM::LLVMDialect", "LLVM::LLVMAVX512Dialect"];
-}
-
-//===----------------------------------------------------------------------===//
 // AsyncToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertAsyncToLLVM : Pass<"convert-async-to-llvm", "ModuleOp"> {
   let summary = "Convert the operations from the async dialect into the LLVM "
                 "dialect";
   let description = [{
     Convert `async.execute` operations to LLVM coroutines and use async runtime
Option<"fullUnroll", "full-unroll", "bool", /*default=*/"false",
            "Perform full unrolling when converting vector transfers to SCF">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // VectorToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToLLVM : Pass<"convert-vector-to-llvm", "ModuleOp"> {
   let summary = "Lower the operations from the vector dialect into the LLVM "
                 "dialect";
+  let description = [{
+
+    Convert operations from the vector dialect into the LLVM IR dialect
+    operations. The lowering pass provides several options to control
+    the kind of optimizations that are allowed. It also provides options
+    that augment the architectural-neutral vector dialect with
+    architectural-specific dialects (AVX512, Neon, etc.).
+
+  }];
   let constructor = "mlir::createConvertVectorToLLVMPass()";
-  let dependentDialects = ["LLVM::LLVMDialect"];
+  let dependentDialects = ["LLVM::LLVMDialect", "LLVM::LLVMAVX512Dialect"];
   let options = [
     Option<"reassociateFPReductions", "reassociate-fp-reductions",
            "bool", /*default=*/"false",
            "Allows llvm to reassociate floating-point reductions for speed">,
     Option<"enableIndexOptimizations", "enable-index-optimizations",
            "bool", /*default=*/"true",
-           "Allows compiler to assume indices fit in 32-bit if that yields faster code">
+           "Allows compiler to assume indices fit in 32-bit if that yields "
+     "faster code">,
+    Option<"enableAVX512", "enable-avx512",
+           "bool", /*default=*/"false",
+           "Augments the vector dialect with the AVX512 dialect during lowering">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // VectorToROCDL
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToROCDL : Pass<"convert-vector-to-rocdl", "ModuleOp"> {

mlir/include/mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h

 
 /// Options to control Vector to LLVM lowering.
 ///
 /// This should kept in sync with VectorToLLVM options defined for the
 /// ConvertVectorToLLVM pass in include/mlir/Conversion/Passes.td
 struct LowerVectorToLLVMOptions {
   bool reassociateFPReductions = false;
   bool enableIndexOptimizations = true;
+  bool enableAVX512 = false;
   LowerVectorToLLVMOptions &setReassociateFPReductions(bool b) {
     reassociateFPReductions = b;
     return *this;
   }
   LowerVectorToLLVMOptions &setEnableIndexOptimizations(bool b) {
     enableIndexOptimizations = b;
     return *this;
   }
+  LowerVectorToLLVMOptions &setEnableAVX512(bool b) {
+    enableAVX512 = b;
+    return *this;
+  }
 };
 
 /// Collect a set of patterns to convert from Vector contractions to LLVM Matrix
 /// Intrinsics. To lower to assembly, the LLVM flag -lower-matrix-intrinsics
 /// will be needed when invoking LLVM.
 void populateVectorToLLVMMatrixConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns);
 

mlir/lib/Conversion/AVX512ToLLVM/ConvertAVX512ToLLVM.cpp

 //===- ConvertAVX512ToLLVM.cpp - Convert AVX512 to the LLVM dialect -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/AVX512ToLLVM/ConvertAVX512ToLLVM.h"
 
-#include "../PassDetail.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
-#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
-#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Dialect/AVX512/AVX512Dialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMAVX512Dialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/PatternMatch.h"
-#include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 using namespace mlir::vector;
 using namespace mlir::avx512;
 
 template <typename OpTy>
 static Type getSrcVectorElementType(OpTy op) {
   return op.src().getType().template cast<VectorType>().getElementType();
void mlir::populateAVX512ToLLVMConversionPatterns(
   MLIRContext *ctx = converter.getDialect()->getContext();
   // clang-format off
   patterns.insert<MaskRndScaleOpPS512Conversion,
                   MaskRndScaleOpPD512Conversion,
                   ScaleFOpPS512Conversion,
                   ScaleFOpPD512Conversion>(ctx, converter);
   // clang-format on
 }
-
-namespace {
-struct ConvertAVX512ToLLVMPass
-    : public ConvertAVX512ToLLVMBase<ConvertAVX512ToLLVMPass> {
-  void runOnOperation() override;
-};
-} // namespace
-
-void ConvertAVX512ToLLVMPass::runOnOperation() {
-  // Convert to the LLVM IR dialect.
-  OwningRewritePatternList patterns;
-  LLVMTypeConverter converter(&getContext());
-  populateAVX512ToLLVMConversionPatterns(converter, patterns);
-  populateVectorToLLVMConversionPatterns(converter, patterns);
-  populateStdToLLVMConversionPatterns(converter, patterns);
-
-  ConversionTarget target(getContext());
-  target.addLegalDialect<LLVM::LLVMDialect>();
-  target.addLegalDialect<LLVM::LLVMAVX512Dialect>();
-  target.addIllegalDialect<avx512::AVX512Dialect>();
-  if (failed(applyPartialConversion(getOperation(), target,
-                                    std::move(patterns)))) {
-    signalPassFailure();
-  }
-}
-
-std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertAVX512ToLLVMPass() {
-  return std::make_unique<ConvertAVX512ToLLVMPass>();
-}

mlir/lib/Conversion/VectorToLLVM/CMakeLists.txt

 add_mlir_conversion_library(MLIRVectorToLLVM
   ConvertVectorToLLVM.cpp
+  ConvertVectorToLLVMPass.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/VectorToLLVM
 
   DEPENDS
   MLIRConversionPassIncGen
   intrinsics_gen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
+  MLIRAVX512
+  MLIRAVX512ToLLVM
+  MLIRLLVMAVX512
   MLIRLLVMIR
   MLIRStandardToLLVM
   MLIRTargetLLVMIRModuleTranslation
   MLIRTransforms
   MLIRVector
   )

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 //===- VectorToLLVM.cpp - Conversion from Vector to the LLVM dialect ------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 
-#include "../PassDetail.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
-#include "mlir/IR/AffineMap.h"
-#include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Target/LLVMIR/TypeTranslation.h"
 #include "mlir/Transforms/DialectConversion.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "mlir/Transforms/Passes.h"
-#include "llvm/IR/DerivedTypes.h"
-#include "llvm/IR/Module.h"
-#include "llvm/IR/Type.h"
-#include "llvm/Support/Allocator.h"
-#include "llvm/Support/ErrorHandling.h"
 
 using namespace mlir;
 using namespace mlir::vector;
 
 // Helper to reduce vector type by one rank at front.
 static VectorType reducedVectorTypeFront(VectorType tp) {
   assert((tp.getRank() > 1) && "unlowerable vector type");
   return VectorType::get(tp.getShape().drop_front(), tp.getElementType());
void mlir::populateVectorToLLVMConversionPatterns(
   // clang-format on
 }
 
 void mlir::populateVectorToLLVMMatrixConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   MLIRContext *ctx = converter.getDialect()->getContext();
   patterns.insert<VectorMatmulOpConversion>(ctx, converter);
   patterns.insert<VectorFlatTransposeOpConversion>(ctx, converter);
 }
-
-namespace {
-struct LowerVectorToLLVMPass
-    : public ConvertVectorToLLVMBase<LowerVectorToLLVMPass> {
-  LowerVectorToLLVMPass(const LowerVectorToLLVMOptions &options) {
-    this->reassociateFPReductions = options.reassociateFPReductions;
-    this->enableIndexOptimizations = options.enableIndexOptimizations;
-  }
-  void runOnOperation() override;
-};
-} // namespace
-
-void LowerVectorToLLVMPass::runOnOperation() {
-  // Perform progressive lowering of operations on slices and
-  // all contraction operations. Also applies folding and DCE.
-  {
-    OwningRewritePatternList patterns;
-    populateVectorToVectorCanonicalizationPatterns(patterns, &getContext());
-    populateVectorSlicesLoweringPatterns(patterns, &getContext());
-    populateVectorContractLoweringPatterns(patterns, &getContext());
-    applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
-  }
-
-  // Convert to the LLVM IR dialect.
-  LLVMTypeConverter converter(&getContext());
-  OwningRewritePatternList patterns;
-  populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
-  populateVectorToLLVMConversionPatterns(
-      converter, patterns, reassociateFPReductions, enableIndexOptimizations);
-  populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
-  populateStdToLLVMConversionPatterns(converter, patterns);
-
-  LLVMConversionTarget target(getContext());
-  if (failed(
-          applyPartialConversion(getOperation(), target, std::move(patterns))))
-    signalPassFailure();
-}
-
-std::unique_ptr<OperationPass<ModuleOp>>
-mlir::createConvertVectorToLLVMPass(const LowerVectorToLLVMOptions &options) {
-  return std::make_unique<LowerVectorToLLVMPass>(options);
-}

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp

+//===- VectorToLLVM.cpp - Conversion from Vector to the LLVM dialect ------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
+
+#include "../PassDetail.h"
+
+#include "mlir/Conversion/AVX512ToLLVM/ConvertAVX512ToLLVM.h"
+#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
+#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Dialect/AVX512/AVX512Dialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMAVX512Dialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::vector;
+
+namespace {
+struct LowerVectorToLLVMPass
+    : public ConvertVectorToLLVMBase<LowerVectorToLLVMPass> {
+  LowerVectorToLLVMPass(const LowerVectorToLLVMOptions &options) {
+    this->reassociateFPReductions = options.reassociateFPReductions;
+    this->enableIndexOptimizations = options.enableIndexOptimizations;
+    this->enableAVX512 = options.enableAVX512;
+  }
+  void runOnOperation() override;
+};
+} // namespace
+
+void LowerVectorToLLVMPass::runOnOperation() {
+  // Perform progressive lowering of operations on slices and
+  // all contraction operations. Also applies folding and DCE.
+  {
+    OwningRewritePatternList patterns;
+    populateVectorToVectorCanonicalizationPatterns(patterns, &getContext());
+    populateVectorSlicesLoweringPatterns(patterns, &getContext());
+    populateVectorContractLoweringPatterns(patterns, &getContext());
+    applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+  }
+
+  // Convert to the LLVM IR dialect.
+  LLVMTypeConverter converter(&getContext());
+  OwningRewritePatternList patterns;
+  populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
+  populateVectorToLLVMConversionPatterns(
+      converter, patterns, reassociateFPReductions, enableIndexOptimizations);
+  populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
+  populateStdToLLVMConversionPatterns(converter, patterns);
+
+  // Architecture specific augmentations.
+  LLVMConversionTarget target(getContext());
+  if (enableAVX512) {
+    target.addLegalDialect<LLVM::LLVMAVX512Dialect>();
+    target.addIllegalDialect<avx512::AVX512Dialect>();
+    populateAVX512ToLLVMConversionPatterns(converter, patterns);
+  }
+
+  if (failed(
+          applyPartialConversion(getOperation(), target, std::move(patterns))))
+    signalPassFailure();
+}
+
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::createConvertVectorToLLVMPass(const LowerVectorToLLVMOptions &options) {
+  return std::make_unique<LowerVectorToLLVMPass>(options);
+}

mlir/lib/Dialect/AVX512/CMakeLists.txt

 add_mlir_dialect_library(MLIRAVX512
   IR/AVX512Dialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/AVX512
 
   DEPENDS
   MLIRAVX512IncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRSideEffectInterfaces
-  MLIRVectorToLLVM
   )

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

-// RUN: mlir-opt %s -convert-avx512-to-llvm | mlir-opt | FileCheck %s
+// RUN: mlir-opt %s -convert-vector-to-llvm="enable-avx512" | mlir-opt | FileCheck %s
 
 func @avx512_mask_rndscale(%a: vector<16xf32>, %b: vector<8xf64>, %i32: i32, %i16: i16, %i8: i8)
-  -> (vector<16xf32>, vector<8xf64>)
+  -> (vector<16xf32>, vector<8xf64>, vector<16xf32>, vector<8xf64>)
 {
   // CHECK: llvm_avx512.mask.rndscale.ps.512
   %0 = avx512.mask.rndscale %a, %i32, %a, %i16, %i32: vector<16xf32>
   // CHECK: llvm_avx512.mask.rndscale.pd.512
   %1 = avx512.mask.rndscale %b, %i32, %b, %i8, %i32: vector<8xf64>
 
   // CHECK: llvm_avx512.mask.scalef.ps.512
-  %a0 = avx512.mask.scalef %a, %a, %a, %i16, %i32: vector<16xf32>
+  %2 = avx512.mask.scalef %a, %a, %a, %i16, %i32: vector<16xf32>
   // CHECK: llvm_avx512.mask.scalef.pd.512
-  %a1 = avx512.mask.scalef %b, %b, %b, %i8, %i32: vector<8xf64>
+  %3 = avx512.mask.scalef %b, %b, %b, %i8, %i32: vector<8xf64>
 
-  return %a0, %a1: vector<16xf32>, vector<8xf64>
+  // Keep results alive.
+  return %0, %1, %2, %3 : vector<16xf32>, vector<8xf64>, vector<16xf32>, vector<8xf64>
 }