[mlir][AVX512] Start an AVX512 dialect

[nicolasvasilache]

Summary

The [Vector Dialect]((https://mlir.llvm.org/docs/Dialects/Vector/)) implements a retargetable n-D vector abstraction.
One of the layer currently missing in OSS atm is the Hardware Vector Ops (HWV) level.

This revision introduces a AVX512-specific dialect for the purpose of implementing portions of XNNPack in MLIR and benchmarking them.
This lives under Dialect/Targets/AVX512 and will target useful intrinsics to implement XNNPack.
The first such function is xnn_math_f32_exp__avx512f_p5_scalef.

At the moment, MLIR relies too much on LLVM’s peephole optimizer to do a good job from small insertelement/extractelement/shufflevector.
Some intrinsics defined and used in the LLVMDialect but these are all “portable” intrinsics.
This revision starts defining layering the layering to attack the right instructions in avx512 directly.

Experience with avx512 is also expected to be useful for a future ARM SVE dialect.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[flaub]

To be consistent with LLVM proper, I suggest using Target instead of Targets for the directory name.

[jsetoain]

Brilliant! Thanks, Nicholas. This will indeed be a helpful reference for "close-to-ISA" dialects in general :-)

[rriddle]

(Blocking for now to allow for time to look/discussion)

[aartbik]

mlir/include/mlir/Dialect/Targets/AVX512/AVX512Ops.td

28 ↗

(On Diff #242630)

Bit new to this syntax, so please forgive this question ;-)

I see the 16xf32, but where is the 8xf64?

[nicolasvasilache]

Move to Dialect/Target.

[nicolasvasilache]

@flaub thanks! addressed.

mlir/include/mlir/Dialect/Targets/AVX512/AVX512Ops.td

28 ↗

(On Diff #242630)

Dropped it because XNNPack does not target that atm (in that function at least).
Once I have all the pieces connected end-to-end we can think about generalizing.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[nicolasvasilache]

The RFC in its current form has been abandoned.
I will start a separate lower level dialect following up on the comments and suggestions I received at a future time.

mlir/include/mlir/Dialect/CMakeLists.txt

 add_subdirectory(AffineOps)
 add_subdirectory(FxpMathOps)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(LoopOps)
 add_subdirectory(OpenMP)
 add_subdirectory(QuantOps)
 add_subdirectory(SPIRV)
 add_subdirectory(StandardOps)
+add_subdirectory(Target)
 add_subdirectory(VectorOps)

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

+//===- AVX512Dialect.h - MLIR Dialect for OpenMP ----------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the Target dialect for AVX512 in MLIR.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TARGETS_AVX512_AVX512DIALECT_H_
+#define MLIR_DIALECT_TARGETS_AVX512_AVX512DIALECT_H_
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+namespace avx512 {
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Target/AVX512/AVX512Ops.h.inc"
+
+class AVX512Dialect : public Dialect {
+public:
+  explicit AVX512Dialect(MLIRContext *context);
+
+  static StringRef getDialectNamespace() { return "avx512"; }
+};
+
+} // namespace avx512
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TARGETS_AVX512_AVX512DIALECT_H_

mlir/include/mlir/Dialect/Target/AVX512/AVX512Ops.td

+//===-- AVX512Ops.td - AVX512 dialect operation definitions *- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the basic operations for the AVX512 dialect.
+//
+//===----------------------------------------------------------------------===//
+
+
+#ifndef AVX512_OPS
+#define AVX512_OPS
+
+include "mlir/IR/OpBase.td"
+
+def AVX512_Dialect : Dialect {
+  let name = "avx512";
+}
+
+class AVX512_Op<string mnemonic, list<OpTrait> traits = []> :
+      Op<AVX512_Dialect, mnemonic, traits> {}
+
+def ExpP5ScaleOp : AVX512_Op<"exp_p5_scale", [NoSideEffect,
+    AllTypesMatch<["source", "result"]>]>,
+    // Supports vector<16xf32>.
+    Arguments<(ins VectorOfLengthAndType<[16], [F32]>:$source)>,
+    Results<(outs VectorOfLengthAndType<[16], [F32]>:$result)>  {
+  let summary = "exp_p5_scale Op";
+  let description = [{
+    The exp_p5_scale constructs a degree-5 polynomial approximation for exp(t)
+    on [-log(2)/2, log(2)/2]. It corresponds to the XNNPack function
+    [xnn_math_f32_exp__avx512f_p5_scalef](https://github.com/google/XNNPACK/blob/master/src/math/exp-avx512f-p5-scalef.c).
+  }];
+  // Fully specified by traits.
+  let verifier = ?;
+  let assemblyFormat = "$source attr-dict `:` type($source)";
+}
+
+#endif // AVX512_OPS

mlir/include/mlir/Dialect/Target/AVX512/CMakeLists.txt

+add_mlir_dialect(AVX512Ops AVX512Ops)

mlir/include/mlir/Dialect/Target/CMakeLists.txt

+add_subdirectory(AVX512)

mlir/lib/Dialect/CMakeLists.txt

 add_subdirectory(AffineOps)
 add_subdirectory(FxpMathOps)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(LoopOps)
 add_subdirectory(OpenMP)
 add_subdirectory(QuantOps)
 add_subdirectory(SDBM)
 add_subdirectory(SPIRV)
 add_subdirectory(StandardOps)
+add_subdirectory(Target)
 add_subdirectory(VectorOps)
 
 add_llvm_library(MLIRDialect
   Traits.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect
   )
 target_link_libraries(MLIRDialect MLIRIR)

mlir/lib/Dialect/Target/AVX512/AVX512Ops.cpp

+//===- AVX512Dialect.cpp - MLIR Dialect for AVX512 implementation ---------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the AVX512 dialect and its operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Target/AVX512/AVX512Dialect.h"
+#include "mlir/Dialect/VectorOps/VectorOps.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/TypeUtilities.h"
+
+using namespace mlir;
+using namespace mlir::avx512;
+
+AVX512Dialect::AVX512Dialect(MLIRContext *context)
+    : Dialect(getDialectNamespace(), context) {
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/Target/AVX512/AVX512Ops.cpp.inc"
+      >();
+}
+
+namespace mlir {
+namespace avx512 {
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Target/AVX512/AVX512Ops.cpp.inc"
+} // namespace avx512
+} // namespace mlir
+
+// Static initialization for AVX512Ops dialect registration.
+static DialectRegistration<AVX512Dialect> avx512Ops;

mlir/lib/Dialect/Target/AVX512/CMakeLists.txt

+add_llvm_library(MLIRAVX512Ops
+  AVX512Ops.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Targets/AVX512
+  )
+
+add_dependencies(MLIRAVX512Ops MLIRAVX512OpsIncGen)

mlir/lib/Dialect/Target/CMakeLists.txt

+add_subdirectory(AVX512)

mlir/test/Dialect/Targets/AVX512/roundtrip.mlir

+// RUN: mlir-opt -verify-diagnostics %s | mlir-opt | FileCheck %s
+
+func @avx512.exp_p5_scale(%a: vector<16xf32>) -> (vector<16xf32>) {
+  // CHECK: avx512.exp_p5_scale {{.*}}: vector<16xf32>
+  %0 = avx512.exp_p5_scale %a : vector<16xf32>
+  return %0: vector<16xf32>
+}

mlir/tools/mlir-opt/CMakeLists.txt

 )
 target_link_libraries(MLIRMlirOptLib
   ${LIB_LIBS}
 )
 
 set(LIBS
   MLIRAnalysis
   MLIRAffineOps
   MLIRAffineToStandard
+  MLIRAVX512Ops
   MLIRDialect
   MLIRLoopsToGPU
   MLIRLinalgToLLVM
 
   MLIRLoopToStandard
   MLIREDSC
   MLIRFxpMathOps
   MLIRGPU