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

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

80 cols plz

24

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

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

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

31 ↗

(On Diff #345541)

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.

41 ↗

(On Diff #345541)

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

32 ↗

(On Diff #350973)

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.

38 ↗

(On Diff #350973)

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

32 ↗

(On Diff #350973)

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

38 ↗

(On Diff #350973)

ah ok, thanks for the note.

[nicolasvasilache]

mlir/lib/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.cpp

32 ↗

(On Diff #350973)

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/ArmNeonStructuredToIntr/ArmNeonStructuredToIntr.h

+#ifndef MLIR_CONVERSION_ARMNEONSTRUCTUREDTOINTR_ARMNEONSTRUCTUREDTOINTR_H_
+#define MLIR_CONVERSION_ARMNEONSTRUCTUREDTOINTR_ARMNEONSTRUCTUREDTOINTR_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+class FuncOp;
+template <typename T>
+class OperationPass;
+
+void populateConvertArmNeonStructuredToIntrPatterns(
+    RewritePatternSet &patterns);
+std::unique_ptr<OperationPass<FuncOp>>
+createConvertArmNeonStructuredToIntrPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_ARMNEONSTRUCTUREDTOINTR_ARMNEONSTRUCTUREDTOINTR_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/ArmNeonStructuredToIntr/ArmNeonStructuredToIntr.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"];
 }
 
+//===----------------------------------------------------------------------===//
+// ArmNeonStructuredToIntr
+//===----------------------------------------------------------------------===//
+
+def ConvertArmNeonStructuredToIntr : Pass<"arm-neon-structured-to-intr", "FuncOp"> {
+  let summary = "Convert Arm NEON structured ops to intrinsics";
+  let constructor = "mlir::createConvertArmNeonStructuredToIntrPass()";
+  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_StructuredOp<string mnemonic,
+                     list<OpTrait> traits = []>
+    : Op</*dialect=*/ArmNeon_Dialect,
+                      /*opName=*/"structured." # mnemonic,
+                      /*traits=*/traits>;
+
+def StructuredSdotOp : ArmNeon_StructuredOp<"sdot", [
+      NoSideEffect,
+      AllTypesMatch<["b", "c"]>,
+      AllTypesMatch<["a", "res"]>,
+      TypesMatchWith<"res has the same number of elements as operand b",
+                     "b", "res",
+                     "VectorType::get({$_self.cast<VectorType>().getShape()[0]},"
+                     "IntegerType::get($_self.getContext(), 32))">]
+      // TODO: how do we express the 2d input shape requirement here, and the
+      // requirement that the inner dimension of b and c is 4?
+  > {
+  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)";
+  }
+
 #endif // ARMNEON_OPS

mlir/lib/Conversion/ArmNeonStructuredToIntr/ArmNeonStructuredToIntr.cpp

+//===- ArmNeonStructuredToIntr.cpp - conversion from Arm Neon structured 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/ArmNeonStructuredToIntr/ArmNeonStructuredToIntr.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 StructuredSdotLoweringPattern
+    : public OpRewritePattern<StructuredSdotOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  /// Converts the type of the result to an LLVM type, pass operands as is,
+  /// preserve attributes.
+  LogicalResult matchAndRewrite(StructuredSdotOp op,
+                                PatternRewriter &rewriter) const override {
+    auto elemType = op.b().getType().cast<VectorType>().getElementType();
+    int length = op.b().getType().cast<VectorType>().getShape()[0] * 4;
+    auto flattenedVectorType = VectorType::get({length}, elemType);
+    auto structured_b = op.b();
+    auto structured_c = op.c();
+    auto loc = op.getLoc();
+    auto flat_b = rewriter.create<vector::ShapeCastOp>(loc, flattenedVectorType,
+                                                       structured_b);
+    auto flat_c = rewriter.create<vector::ShapeCastOp>(loc, flattenedVectorType,
+                                                       structured_c);
+    Value newOp = rewriter.create<SdotOp>(loc, op->getResult(0).getType(),
+                                          op.a(), flat_b, flat_c);
+    rewriter.replaceOp(op, {newOp});
+    return success();
+  }
+};
+
+class ConvertArmNeonStructuredToIntr
+    : public ConvertArmNeonStructuredToIntrBase<
+          ConvertArmNeonStructuredToIntr> {
+  void runOnOperation() override {
+    auto func = getOperation();
+    auto *context = &getContext();
+
+    RewritePatternSet patterns(context);
+    populateConvertArmNeonStructuredToIntrPatterns(patterns);
+
+    if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns))))
+      return signalPassFailure();
+  }
+};
+
+} // namespace
+
+namespace mlir {
+
+void populateConvertArmNeonStructuredToIntrPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<StructuredSdotLoweringPattern>(patterns.getContext());
+}
+
+std::unique_ptr<OperationPass<FuncOp>>
+createConvertArmNeonStructuredToIntrPass() {
+  return std::make_unique<ConvertArmNeonStructuredToIntr>();
+}
+
+} // namespace mlir

mlir/lib/Conversion/ArmNeonStructuredToIntr/CMakeLists.txt

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

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
+add_subdirectory(ArmNeonStructuredToIntr)
 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/Target/LLVMIR/arm-neon-structured.mlir

+// RUN: mlir-opt -arm-neon-structured-to-intr %s | FileCheck %s
+
+// CHECK-LABEL: arm_neon_structured_sdot_4x4_i8i8
+func @arm_neon_structured_sdot_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.structured.sdot %a, %b, %c : vector<4x4xi8>, vector<4x4xi8> to vector<4xi32>
+  return %0 : vector<4xi32>
+}
+
+// CHECK-LABEL: arm_neon_structured_sdot_2x4_i8i8
+func @arm_neon_structured_sdot_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.structured.sdot %a, %b, %c : vector<2x4xi8>, vector<2x4xi8> to vector<2xi32>
+  return %0 : vector<2xi32>
+}