[mlir][ArmSME] Remove ArmSMETypeConverter (and configure LLVM one instead) (#73639)

This patch removes the ArmSMETypeConverter, and instead updates
`populateArmSMEToLLVMConversionPatterns()` to add an ArmSME vector type
conversion to the existing LLVMTypeConverter. This makes it easier to
add these patterns to an existing `-to-llvm` lowering pass.

Author: MacDue

mlir/include/mlir/Conversion/ArmSMEToLLVM/ArmSMEToLLVM.h

@@ -20,8 +20,6 @@ class RewritePatternSet;
 #define GEN_PASS_DECL_CONVERTARMSMETOLLVM
 #include "mlir/Conversion/Passes.h.inc"
 
-using arm_sme::ArmSMETypeConverter;
-
 /// Create a pass to convert from the ArmSME dialect to LLVM intrinsics.
 std::unique_ptr<Pass> createConvertArmSMEToLLVMPass();
 
@@ -30,7 +28,7 @@ void configureArmSMEToLLVMConversionLegality(ConversionTarget &target);
 
 /// Populate the given list with patterns that convert from the ArmSME dialect
 /// to LLVM intrinsics.
-void populateArmSMEToLLVMConversionPatterns(ArmSMETypeConverter &converter,
+void populateArmSMEToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                             RewritePatternSet &patterns);
 
 } // namespace mlir

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

@@ -32,14 +32,6 @@ std::unique_ptr<Pass> createEnableArmStreamingPass(
 /// Pass that allocates tile IDs to ArmSME operations.
 std::unique_ptr<Pass> createTileAllocationPass();
 
-//===----------------------------------------------------------------------===//
-// Type ArmSMETypeConverter pass.
-//===----------------------------------------------------------------------===//
-class ArmSMETypeConverter : public LLVMTypeConverter {
-public:
-  ArmSMETypeConverter(MLIRContext *ctx, const LowerToLLVMOptions &options);
-};
-
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//

mlir/lib/Conversion/ArmSMEToLLVM/ArmSMEToLLVM.cpp

@@ -538,9 +538,7 @@ struct ConvertArmSMEToLLVMPass
   void runOnOperation() override {
     LLVMConversionTarget target(getContext());
     RewritePatternSet patterns(&getContext());
-    ArmSMETypeConverter converter(&getContext(),
-                                  LowerToLLVMOptions(&getContext()));
-
+    LLVMTypeConverter converter(&getContext());
     configureArmSMEToLLVMConversionLegality(target);
     populateArmSMEToLLVMConversionPatterns(converter, patterns);
 
@@ -573,8 +571,16 @@ void mlir::configureArmSMEToLLVMConversionLegality(ConversionTarget &target) {
   target.addLegalOp<UnrealizedConversionCastOp>();
 }
 
-void mlir::populateArmSMEToLLVMConversionPatterns(
-    ArmSMETypeConverter &converter, RewritePatternSet &patterns) {
+void mlir::populateArmSMEToLLVMConversionPatterns(LLVMTypeConverter &converter,
+                                                  RewritePatternSet &patterns) {
+  converter.addConversion([&](VectorType type) -> std::optional<Type> {
+    // There's no LLVM type for SME tiles, but after lowering to intrinsics all
+    // SME vector types should be eliminated.
+    if (arm_sme::isValidSMETileVectorType(type))
+      return type;
+    return std::nullopt;
+  });
+
   patterns.add<LoadTileSliceConversion, MoveTileSliceToVectorConversion,
                MoveVectorToTileSliceConversion, StoreTileSliceConversion,
                OuterProductOpConversion, ZeroOpConversion, GetTileConversion>(

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

@@ -1,5 +1,4 @@
 add_mlir_dialect_library(MLIRArmSMETransforms
-  ArmSMETypeConverter.cpp
   EnableArmStreaming.cpp
   TileAllocation.cpp
 

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

@@ -0,0 +1,5 @@
+add_mlir_unittest(MLIRArmSMETests
+  TileTypeConversionTest.cpp)
+target_link_libraries(MLIRArmSMETests
+  PRIVATE
+  MLIRArmSMEToLLVM)

mlir/unittests/Dialect/ArmSME/CMakeLists.txt

@@ -0,0 +1,51 @@
+//===- TileTypeConversionTest.cpp - Tests ArmSME tile type conversion -----===//
+//
+// 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/ArmSMEToLLVM/ArmSMEToLLVM.h"
+#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
+#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Dialect/ArmSME/IR/ArmSME.h"
+
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+class ArmSMETest : public ::testing::Test {
+protected:
+  ArmSMETest() { context.getOrLoadDialect<mlir::arm_sme::ArmSMEDialect>(); }
+
+  mlir::MLIRContext context;
+};
+
+TEST_F(ArmSMETest, TestTileTypeConversion) {
+  LLVMTypeConverter llvmConverter(&context);
+  LLVMTypeConverter llvmConverterWithArmSMEConversion(&context);
+
+  RewritePatternSet patterns(&context);
+  populateArmSMEToLLVMConversionPatterns(llvmConverterWithArmSMEConversion,
+                                         patterns);
+
+  Type i32 = IntegerType::get(&context, 32);
+  auto smeTileType = VectorType::get({4, 4}, i32, {true, true});
+
+  // An unmodified LLVMTypeConverer should fail to convert an ArmSME tile type.
+  {
+    SmallVector<Type> convertedType;
+    ASSERT_TRUE(failed(llvmConverter.convertType(smeTileType, convertedType)));
+  }
+
+  // An updated LLVMTypeConverer should return the ArmSME tile vector type
+  // unchanged.
+  {
+    SmallVector<Type> convertedType;
+    ASSERT_TRUE(succeeded(llvmConverterWithArmSMEConversion.convertType(
+        smeTileType, convertedType)));
+    ASSERT_EQ(ArrayRef<Type>(convertedType), ArrayRef<Type>{smeTileType});
+  }
+}

mlir/unittests/Dialect/ArmSME/TileTypeConversionTest.cpp

@@ -6,6 +6,7 @@ target_link_libraries(MLIRDialectTests
   MLIRIR
   MLIRDialect)
 
+add_subdirectory(ArmSME)
 add_subdirectory(Index)
 add_subdirectory(LLVMIR)
 add_subdirectory(MemRef)