2d Arm Neon sdot op, and lowering to the intrinsic.

[Benoit]

Summary

This adds Sdot2d op, which is similar to the usual Neon
intrinsic except that it takes 2d vector operands, reflecting the
structure of the arithmetic that it's performing: 4 separate
4-dimensional dot products, whence the vector<4x4xi8> shape.

This also adds a new pass, arm-neon-2d-to-intr, lowering
this new 2d op to the 1d intrinsic.

[nicolasvasilache]

Cool ! Accepting conditioned on these minor changes.

mlir/include/mlir/Conversion/ArmNeonStructuredToIntr/ArmNeonStructuredToIntr.h

11 ↗

(On Diff #345453)

some minor comments + spacing prefixed with /// please

mlir/include/mlir/Dialect/ArmNeon/ArmNeon.td

127

So structured op has a specific meaning in a bunch of MLIR places related to Linalg and Vector.
Could we call this MamtulSomethingOp ?

135

Not sure there is tablegen support for that.
Just add it in C++ with a custom verifier ? e.g. https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/X86Vector/X86Vector.td#L78

mlir/lib/Conversion/ArmNeonStructuredToIntr/ArmNeonStructuredToIntr.cpp

2 ↗

(On Diff #345453)

80 cols plz

24 ↗

(On Diff #345453)

Note that one of the remaining challenges is to properly propagate the shape casts you introduce here all the way to memory operations and ideally have them fold into reindexings.
Add a TODO along those lines ?

29 ↗

(On Diff #345453)

There is non LLVM type here, just 1-D vector type that map 1-1 with LLVM (unlike the 2+D vectors that don't).
Can you please rephrase a bit? I'd just say "convert to 1-D vector type + attribute to match the neon.intr.xxx operation requirements".

Those intr. ops are the separation of concern between LLVM and MLIR.

36 ↗

(On Diff #345453)

camelcase everywhere plz

[Benoit]

Applied review comments

[Benoit]

Please take a look - nontrivial enough changes that I'd like your opinion again.

mlir/include/mlir/Dialect/ArmNeon/ArmNeon.td

127

I renamed 'Structured' to '2d' everywhere.
This is the best descriptive term that I could think of.
'Matmul' would be both overly specific to certain Neon intrinsic, and not even well descriptive of the present intrinsic: there are 2 forms of ARM SDOT instruction, only one of them is a 4x4x1 matmul, and the one that LLVM has as an intrinsic (which we are dealing with here) is the other one. It's not a 4x4x1 matmul, it's 4 completely independent dot products of separate 4d vectors.

[Benoit]

There is one of your comments which I haven't applied yet. It's the one about writing a TODO about the handling of these reshapes. Part of the problem is I didn't understand 100% of what you were saying, and part is that this topic is going to get more subtle still in the following commit, adding the lowering from vector.contract, which will involve a broadcast to map the contract to the flavor of SDOT that we have here, which (as said in the other comment) is actually 4 separate vector dot products. So if you're OK, I would suggest to do nothing about this in this CL, discuss with you next week, decide something for the next commit.

[nicolasvasilache]

please add an invalid.mlir in the armneon dialect tests, one per verifier error that is not automatically produced by tablegen (we test those separately), look at other invalid.mlir to see how they are structured.

mlir/include/mlir/Dialect/ArmNeon/ArmNeon.td

127

sgtm, thanks!

mlir/lib/Dialect/ArmNeon/IR/ArmNeonDialect.cpp

28 ↗

(On Diff #345541)

with custom verifiers we also add invalid tests in an invalid.mlir that exercise each of the expected errors.

[ftynse]

mlir/include/mlir/Dialect/ArmNeon/ArmNeon.td

123–125

mlir/lib/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp

32

Please expand auto here and below. MLIR uses auto when the type is either obvious from context (e.g., there's a cast on the RHS) or prohibitively long/impossible to express.

42

Nit: consider op.res() instead of op->getResult(0) to avoid the generic API with magic numbers.

mlir/lib/Dialect/ArmNeon/IR/ArmNeonDialect.cpp

32–36 ↗

(On Diff #345541)

These seem to be already covered by the type constraint in ODS so no need to check them manually. If you follow Nicolas' advice from above, you'll notice which messages are never emitted.

33 ↗

(On Diff #345541)

Nit: op.emitOpError() / op.emitError() spare you the manual fetching of location and produce more semantically-meaningful errors.

[Benoit]

Addressed review comments. Predicates now implemented using CPred.

[Benoit]

addressed ftynse's review comments

[Benoit]

one last review comment

[Benoit]

one more whitespace fix

[nicolasvasilache]

mlir/lib/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp

33

You could give the magic constant a meaningful name and expose it as a static extra method of the op.
You could also use that static method in the .td and then everything gets normalized.

39

just a warning, a lot of the canonicalization work to push these ShapeCastOp all the way to address computations is missing.
So we'll prob. need to iterate when you start looking at a real kernel from memory and the quality of the assembly that gets generated.

[Benoit]

mlir/lib/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp

33

Sorry, I don't know how to add such a static method or constant to a class defined in tablegen?

39

ah ok, thanks for the note.

[nicolasvasilache]

mlir/lib/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp

33

You could do something like this: https://sourcegraph.com/github.com/llvm/llvm-project/-/blob/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td#L61

let extraClassDeclarations = [] ...

[Benoit]

one more review comment

[nicolasvasilache]

Great, let's land this!

[Benoit]

I don't think that I have write access. I read https://llvm.org/docs/Phabricator.html#committing-a-change but it's not very clear to me. As far as I can see all I can do here is write this message asking if someone can submit for me? I don't see something else than this free-text field to input to (so is this process dependent on contributors freely browsing open reviews for ones that need submitting?)

mlir/include/mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h

+//===- ArmNeon2dToIntr.h - convert Arm Neon 2d ops to intrinsics ----------===//
+//
+// 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_ARMNEON2DTOINTR_ARMNEON2DTOINTR_H_
+#define MLIR_CONVERSION_ARMNEON2DTOINTR_ARMNEON2DTOINTR_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+class FuncOp;
+template <typename T>
+class OperationPass;
+
+/// Populates patterns for the lowering of Arm NEON 2D ops to intrinsics.
+/// See createConvertArmNeon2dToIntrPass.
+void populateConvertArmNeon2dToIntrPatterns(RewritePatternSet &patterns);
+
+/// Creates a pass to lower Arm NEON 2D ops to intrinsics, i.e.
+/// equivalent ops operating on flattened 1D vectors and mapping more
+/// directly to the corresponding Arm NEON instruction.
+std::unique_ptr<OperationPass<FuncOp>> createConvertArmNeon2dToIntrPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_ARMNEON2DTOINTR_ARMNEON2DTOINTR_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/AffineToStandard/AffineToStandard.h"
+#include "mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h"
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h"
 #include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"

mlir/include/mlir/Conversion/Passes.td

 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToSPIRV : Pass<"convert-vector-to-spirv", "ModuleOp"> {
   let summary = "Convert Vector dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertVectorToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];
 }
 
+//===----------------------------------------------------------------------===//
+// ArmNeon2dToIntr
+//===----------------------------------------------------------------------===//
+
+def ConvertArmNeon2dToIntr : Pass<"arm-neon-2d-to-intr", "FuncOp"> {
+  let summary = "Convert Arm NEON structured ops to intrinsics";
+  let constructor = "mlir::createConvertArmNeon2dToIntrPass()";
+  let dependentDialects = ["arm_neon::ArmNeonDialect", "vector::VectorDialect"];
+}
+
+
 #endif // MLIR_CONVERSION_PASSES

mlir/include/mlir/Dialect/ArmNeon/ArmNeon.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef ARMNEON_OPS
 #define ARMNEON_OPS
 
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/OpBase.td"
 
 //===----------------------------------------------------------------------===//
 // ArmNeon dialect definition
 //===----------------------------------------------------------------------===//
 
 def ArmNeon_Dialect : Dialect {
   let name = "arm_neon";
   let cppNamespace = "::mlir::arm_neon";
def SdotOp : ArmNeon_OverloadedOperandsWithOneResultIntrOp<"sdot",[1], [
   let arguments = (ins VectorOfLengthAndType<[4, 2], [I32]>:$a,
                        VectorOfLengthAndType<[16, 8], [I8]>:$b,
                        VectorOfLengthAndType<[16, 8], [I8]>:$c);
   let results = (outs VectorOfLengthAndType<[4, 2], [I32]>:$res);
   let assemblyFormat =
     "$a `,` $b `,` $c attr-dict `:` type($b) `,` type($c) `to` type($res)";
   }
 
+class ArmNeon_2dOp<string mnemonic, list<OpTrait> traits = []>
+    : Op</*dialect=*/ArmNeon_Dialect,
+         /*opName=*/"2d." # mnemonic,
+         /*traits=*/traits>;
+
+def Sdot2dOp : ArmNeon_2dOp<"sdot", [
+      NoSideEffect,
+      AllTypesMatch<["b", "c"]>,
+      AllTypesMatch<["a", "res"]>,
+      PredOpTrait<
+        "operand `a` should be 1-dimensional",
+        CPred<"a().getType().cast<VectorType>().getShape().size() == 1">
+      >,
+      PredOpTrait<
+        "operand `b` should be 2-dimensional",
+        CPred<"b().getType().cast<VectorType>().getShape().size() == 2">
+      >,
+      PredOpTrait<
+        "operand `b` should have 4 columns",
+        CPred<"b().getType().cast<VectorType>().getShape()[1] == 4">
+      >,
+      PredOpTrait<
+        "operand `b` should have as many rows as the size of operand `a`",
+        CPred<"b().getType().cast<VectorType>().getShape()[0] == a().getType().cast<VectorType>().getShape()[0]">
+      >,
+      ]
+  > {
+  let summary = "sdot op";
+  let description = [{
+    The two input vectors `b` and `c` have a 2D shape, consisting of either 2
+    or 4 rows, each row having length 4. This operation computes the pair-wise
+    dot-products of the rows of `b` and `c` and accumulates them with the
+    corresponding entry of `a`:
+
+    ```
+    res[i] := a[i] + dot_product(b[i, ...], c[i, ...])
+    ```
+
+  }];
+  // Supports either:
+  //   (vector<2xi32>, vector<2x4xi8>, vector<2x4xi8>) -> vector<2xi32>
+  //   (vector<4xi32>, vector<4x4xi8>, vector<4x4xi8>) -> vector<4xi32>
+  // TODO: how do we express 2D shape requirements here?
+  let arguments = (ins VectorOfLengthAndType<[4, 2], [I32]>:$a,
+                       VectorOfLengthAndType<[16, 8], [I8]>:$b,
+                       VectorOfLengthAndType<[16, 8], [I8]>:$c);
+  let results = (outs VectorOfLengthAndType<[4, 2], [I32]>:$res);
+  let assemblyFormat =
+    "$a `,` $b `,` $c attr-dict `:` type($b) `,` type($c) `to` type($res)";
+  let extraClassDeclaration = [{
+    static constexpr int kReductionSize = 4;
+  }];
+}
+
 #endif // ARMNEON_OPS

mlir/lib/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp

+//===- ArmNeon2dToIntr.cpp - convert Arm Neon 2d ops to intrinsics --------===//
+//
+// 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/ArmNeon2dToIntr/ArmNeon2dToIntr.h"
+#include "../PassDetail.h"
+#include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassRegistry.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::arm_neon;
+
+namespace {
+
+class Sdot2dLoweringPattern : public OpRewritePattern<Sdot2dOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  /// Convert to 1-dimensional vector type to match the requirements of
+  /// arm.neon.intr.sdot
+  LogicalResult matchAndRewrite(Sdot2dOp op,
+                                PatternRewriter &rewriter) const override {
+    Type elemType = op.b().getType().cast<VectorType>().getElementType();
+    int length = op.b().getType().cast<VectorType>().getShape()[0] *
+                 Sdot2dOp::kReductionSize;
+    VectorType flattenedVectorType = VectorType::get({length}, elemType);
+    Value b2d = op.b();
+    Value c2d = op.c();
+    Location loc = op.getLoc();
+    Value b1d =
+        rewriter.create<vector::ShapeCastOp>(loc, flattenedVectorType, b2d);
+    Value c1d =
+        rewriter.create<vector::ShapeCastOp>(loc, flattenedVectorType, c2d);
+    Value newOp =
+        rewriter.create<SdotOp>(loc, op.res().getType(), op.a(), b1d, c1d);
+    rewriter.replaceOp(op, {newOp});
+    return success();
+  }
+};
+
+class ConvertArmNeon2dToIntr
+    : public ConvertArmNeon2dToIntrBase<ConvertArmNeon2dToIntr> {
+  void runOnOperation() override {
+    auto func = getOperation();
+    auto *context = &getContext();
+
+    RewritePatternSet patterns(context);
+    populateConvertArmNeon2dToIntrPatterns(patterns);
+
+    if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns))))
+      return signalPassFailure();
+  }
+};
+
+} // namespace
+
+namespace mlir {
+
+void populateConvertArmNeon2dToIntrPatterns(RewritePatternSet &patterns) {
+  patterns.add<Sdot2dLoweringPattern>(patterns.getContext());
+}
+
+std::unique_ptr<OperationPass<FuncOp>> createConvertArmNeon2dToIntrPass() {
+  return std::make_unique<ConvertArmNeon2dToIntr>();
+}
+
+} // namespace mlir

mlir/lib/Conversion/ArmNeon2dToIntr/CMakeLists.txt

+add_mlir_conversion_library(MLIRArmNeon2dToIntr
+  ArmNeon2dToIntr.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ArmNeon2dToIntr
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRArmNeon
+  MLIRPass
+  MLIRTransforms
+  MLIRIR
+  )

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
+add_subdirectory(ArmNeon2dToIntr)
 add_subdirectory(AsyncToLLVM)
 add_subdirectory(ComplexToLLVM)
 add_subdirectory(ComplexToStandard)
 add_subdirectory(GPUCommon)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)

mlir/lib/Conversion/PassDetail.h

 namespace tosa {
 class TosaDialect;
 } // end namespace tosa
 
 namespace vector {
 class VectorDialect;
 } // end namespace vector
 
+namespace arm_neon {
+class ArmNeonDialect;
+} // end namespace arm_neon
+
 #define GEN_PASS_CLASSES
 #include "mlir/Conversion/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // CONVERSION_PASSDETAIL_H_

mlir/test/Dialect/ArmNeon/invalid.mlir

+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// -----
+
+func @a_is_2d(%a : vector<2x2xi32>, %b : vector<4x4xi8>) -> vector<2x2xi32> {
+    // expected-error@+1 {{operand `a` should be 1-dimensional}}
+    %0 = arm_neon.2d.sdot %a, %b, %b : vector<4x4xi8>, vector<4x4xi8> to vector<2x2xi32>
+    return %0 : vector<2x2xi32>
+}
+
+// -----
+
+func @b_is_3d(%a : vector<4xi32>, %b : vector<1x4x4xi8>) -> vector<4xi32> {
+    // expected-error@+1 {{operand `b` should be 2-dimensional}}
+    %0 = arm_neon.2d.sdot %a, %b, %b : vector<1x4x4xi8>, vector<1x4x4xi8> to vector<4xi32>
+    return %0 : vector<4xi32>
+}
+
+// -----
+
+func @b_has_2_columns(%a : vector<4xi32>, %b : vector<4x2xi8>) -> vector<4xi32> {
+    // expected-error@+1 {{operand `b` should have 4 columns}}
+    %0 = arm_neon.2d.sdot %a, %b, %b : vector<4x2xi8>, vector<4x2xi8> to vector<4xi32>
+    return %0 : vector<4xi32>
+}
+
+// -----
+
+func @b_has_2_rows_but_a_has_length_4(%a : vector<4xi32>, %b : vector<2x4xi8>) -> vector<4xi32> {
+    // expected-error@+1 {{operand `b` should have as many rows as the size of operand `a`}}
+    %0 = arm_neon.2d.sdot %a, %b, %b : vector<2x4xi8>, vector<2x4xi8> to vector<4xi32>
+    return %0 : vector<4xi32>
+}

mlir/test/Target/LLVMIR/arm-neon-2d.mlir

+// RUN: mlir-opt -arm-neon-2d-to-intr %s | FileCheck %s
+
+// CHECK-LABEL: arm_neon_sdot2d_4x4_i8i8
+func @arm_neon_sdot2d_4x4_i8i8(%a: vector<4xi32>, %b: vector<4x4xi8>, %c: vector<4x4xi8>) -> vector<4xi32> {
+  // CHECK: arm_neon.intr.sdot %{{.*}}, %{{.*}}, %{{.*}} : vector<16xi8>, vector<16xi8> to vector<4xi32>
+  // CHECK-NEXT: return %{{.*}} : vector<4xi32>
+  %0 = arm_neon.2d.sdot %a, %b, %c : vector<4x4xi8>, vector<4x4xi8> to vector<4xi32>
+  return %0 : vector<4xi32>
+}
+
+// CHECK-LABEL: arm_neon_sdot2d_2x4_i8i8
+func @arm_neon_sdot2d_2x4_i8i8(%a: vector<2xi32>, %b: vector<2x4xi8>, %c: vector<2x4xi8>) -> vector<2xi32> {
+  // CHECK: arm_neon.intr.sdot %{{.*}}, %{{.*}}, %{{.*}} : vector<8xi8>, vector<8xi8> to vector<2xi32>
+  // CHECK-NEXT: return %{{.*}} : vector<2xi32>
+  %0 = arm_neon.2d.sdot %a, %b, %c : vector<2x4xi8>, vector<2x4xi8> to vector<2xi32>
+  return %0 : vector<2xi32>
+}