[mlir][ArmSME] Introduce custom TypeConverter for ArmSME

[awarzynski]

Summary

At the moment, SME-to-LLVM lowerings relies entirely on
LLVMTypeConverter. This patch introduces a dedicated TypeConverter
that inherits from LLVMTypeConverter (it will also be used when
lowering ArmSME Ops to LLVM).

The new type converter merely disables lowerings for VectorType to
prevent 2-d scalable vectors (common in the context of ArmSME), e.g.

`vector<[16]x[16]xi8>`,

entering the LLVM Type converter. LLVM does not support arrays of
scalable vectors and hence the need for specialisation. In the case of
SME such types are effectively eliminated when emitting LLVM IR
intrinsics for SME.

Depends on D154867

[c-rhodes]

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp

100

nit: this could be moved to if (armSME) {

mlir/lib/Dialect/ArmSME/Transforms/ArmSMETypeConverter.cpp

2

nit: update filename

16

does this disable all vector type conversion? Just wondering what happens for 1-d

[awarzynski]

Cheers for taking a look!

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp

100

That would be nice, but then it gets destroyed before it is used. I've not tried to investigate how to "move" ownership for these classes.

mlir/lib/Dialect/ArmSME/Transforms/ArmSMETypeConverter.cpp

16

does this disable all vector type conversion?

Yes, but only for SME Ops.

Just wondering what happens for 1-d

Once we hit a case like that, we could just copy LLVMTypeConverter::convertVectorType. I am not too worried about that just yet.

[awarzynski]

Update filename

[c-rhodes]

LGTM cheers

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp

100

That would be nice, but then it gets destroyed before it is used. I've not tried to investigate how to "move" ownership for these classes.

No problem thanks for clarifying didnt consider that

mlir/lib/Dialect/ArmSME/Transforms/ArmSMETypeConverter.cpp

16

does this disable all vector type conversion?

Yes, but only for SME Ops.

Just wondering what happens for 1-d

Once we hit a case like that, we could just copy LLVMTypeConverter::convertVectorType. I am not too worried about that just yet.

We use 1-d scalable vectors in mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector-ops.mlir so I was wondering how this still worked, but of course it's because LLVMTypeConverter::convertVectorType doesn't actually do anything for 1-d 😄

[dcaballe]

Glad that this worked! Great!

mlir/lib/Conversion/VectorToLLVM/CMakeLists.txt

20

sort

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.h

 //===- Passes.h - Pass Entrypoints ------------------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_ARMSME_TRANSFORMS_PASSES_H
 #define MLIR_DIALECT_ARMSME_TRANSFORMS_PASSES_H
 
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 class RewritePatternSet;
 
 namespace arm_sme {
+//===----------------------------------------------------------------------===//
+// The EnableArmStreaming pass.
+//===----------------------------------------------------------------------===//
 // Options for Armv9 Streaming SVE mode. By default, streaming-mode is part of
 // the function interface (ABI) and the caller manages PSTATE.SM on entry/exit.
 // In a locally streaming function PSTATE.SM is kept internal and the callee
 // manages it on entry/exit.
 enum class ArmStreaming { Default = 0, Locally = 1 };
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/ArmSME/Transforms/Passes.h.inc"
 
 /// Pass to enable Armv9 Streaming SVE mode.
 std::unique_ptr<Pass>
 createEnableArmStreamingPass(const ArmStreaming mode = ArmStreaming::Default,
                              const bool enableZA = false);
 
 /// Pass that replaces 'arm_sme.get_tile_id' ops with actual tiles.
 std::unique_ptr<Pass> createTileAllocationPass();
 
 //===----------------------------------------------------------------------===//
+// Type ArmSMETypeConverter pass.
+//===----------------------------------------------------------------------===//
+class ArmSMETypeConverter : public LLVMTypeConverter {
+public:
+  ArmSMETypeConverter(MLIRContext *ctx, const LowerToLLVMOptions &options);
+};
+
+//===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/ArmSME/Transforms/Passes.h.inc"
 
 } // namespace arm_sme
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARMSME_TRANSFORMS_PASSES_H

mlir/lib/Conversion/VectorToLLVM/CMakeLists.txt

add_mlir_conversion_library(MLIRVectorToLLVM
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRArmNeonDialect
   MLIRArmSMEDialect
   MLIRArmSMETransforms
+  MLIRVectorToArmSME
   MLIRArmSVEDialect
   MLIRArmSVETransforms
   MLIRAMXDialect
   MLIRAMXTransforms
   MLIRLLVMCommonConversion
   MLIRLLVMDialect
   MLIRMemRefDialect
   MLIRTargetLLVMIRExport
   MLIRTransforms
   MLIRVectorDialect
   MLIRVectorTransforms
   MLIRX86VectorDialect
   MLIRX86VectorTransforms
   )

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp

 
 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/AMX/Transforms.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
 #include "mlir/Dialect/ArmSME/IR/ArmSME.h"
+#include "mlir/Dialect/ArmSME/Transforms/Passes.h"
 #include "mlir/Dialect/ArmSME/Transforms/Transforms.h"
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/ArmSVE/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
populateVectorToLLVMConversionPatterns(
       converter, patterns, reassociateFPReductions, force32BitVectorIndices);
   populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
 
   // Architecture specific augmentations.
   LLVMConversionTarget target(getContext());
   target.addLegalDialect<arith::ArithDialect>();
   target.addLegalDialect<memref::MemRefDialect>();
   target.addLegalOp<UnrealizedConversionCastOp>();
+  arm_sme::ArmSMETypeConverter armSMEConverter(&getContext(), options);
+
   if (armNeon) {
     // TODO: we may or may not want to include in-dialect lowering to
     // LLVM-compatible operations here. So far, all operations in the dialect
     // can be translated to LLVM IR so there is no conversion necessary.
     target.addLegalDialect<arm_neon::ArmNeonDialect>();
   }
   if (armSVE) {
     configureArmSVELegalizeForExportTarget(target);
     populateArmSVELegalizeForLLVMExportPatterns(converter, patterns);
   }
   if (armSME) {
     configureArmSMELegalizeForExportTarget(target);
-    populateArmSMELegalizeForLLVMExportPatterns(converter, patterns);
+    populateArmSMELegalizeForLLVMExportPatterns(armSMEConverter, patterns);
   }
   if (amx) {
     configureAMXLegalizeForExportTarget(target);
     populateAMXLegalizeForLLVMExportPatterns(converter, patterns);
   }
   if (x86Vector) {
     configureX86VectorLegalizeForExportTarget(target);
     populateX86VectorLegalizeForLLVMExportPatterns(converter, patterns);
   }
 
   if (failed(
           applyPartialConversion(getOperation(), target, std::move(patterns))))
     signalPassFailure();
 }

mlir/lib/Dialect/ArmSME/Transforms/ArmSMETypeConverter.cpp

+//===- ArmSMETypeConverter.cpp - Convert builtin to LLVM dialect types ----===//
+//
+// 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/Dialect/ArmSME/Transforms/Passes.h"
+
+using namespace mlir;
+arm_sme::ArmSMETypeConverter::ArmSMETypeConverter(
+    MLIRContext *ctx, const LowerToLLVMOptions &options)
+    : LLVMTypeConverter(ctx, options) {
+  // Disable LLVM type conversion for vectors. This is to prevent 2-d scalable
+  // vectors (common in the context of ArmSME), e.g.
+  //    `vector<[16]x[16]xi8>`,
+  // entering the LLVM Type converter. LLVM does not support arrays of scalable
+  // vectors, but in the case of SME such types are effectively eliminated when
+  // emitting ArmSME LLVM IR intrinsics.
+  addConversion([&](VectorType type) { return type; });
+}

mlir/lib/Dialect/ArmSME/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRArmSMETransforms
+  ArmSMETypeConverter.cpp
   EnableArmStreaming.cpp
   LegalizeForLLVMExport.cpp
   TileAllocation.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ArmSME/Transforms
 
   DEPENDS