[mlir][ArmSME] Insert intrinsics to enable/disable ZA

[c-rhodes]

Summary

This patch adds two LLVM intrinsics to the ArmSME dialect:

• llvm.aarch64.sme.za.enable
• llvm.aarch64.sme.za.disable

for enabling the ZA storage array [1], as well as patterns for inserting
them during legalization to LLVM at the start and end of functions if
the function has the 'arm_za' attribute (D152695).

In the future ZA should probably be automatically enabled/disabled when
lowering from vector to SME, but this should be sufficient for now at
least until we have patterns lowering to SME instructions that use ZA.

N.B. The backend function attribute 'aarch64_pstate_za_new' can be used
manage ZA state (as was originally tried in D152694), but it emits calls
to the following SME support routines [2] for the lazy-save mechanism
[3]:

• __arm_tpidr2_restore
• __arm_tpidr2_save

These will soon be added to compiler-rt but there's currently no public
implementation, and using this attribute would introduce an MLIR
dependency on compiler-rt. Furthermore, this mechanism is for routines
with ZA enabled calling other routines with it also enabled. We can
choose not to enable ZA in the compiler when this is case.

Depends on D152695

[1] https://developer.arm.com/documentation/ddi0616/aa
[2] https://github.com/ARM-software/abi-aa/blob/main/aapcs64/aapcs64.rst#sme-support-routines
[3] https://github.com/ARM-software/abi-aa/blob/main/aapcs64/aapcs64.rst#the-za-lazy-saving-scheme

[dcaballe]

Thanks!

[awarzynski]

Thanks Cullen, just a couple of minor suggestions.

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

57–58

This comment suggests that target.addDynamicallyLegalOp() would enable ZA, but that's not quite true, is it? Instead, it marks funcOp as dynamically legal and provides a custom legality check:

• if arm_za attribute is present then legal if the first function Op is arm_sme::aarch64_sme_za_enable,
• if arm_za is not present, mark as legal.

I guess that the comment refers to the overall legalisation process rather then specifically to this hook?

mlir/test/Dialect/ArmSME/enable-arm-za.mlir

1

Could you add another RUN line without enable-za and verify that there's no asm_sme.intr.za.{enable|disable}?

It would also be good to test without the -enable-arm-streaming pass.

[c-rhodes]

Address comments

[c-rhodes]

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

57–58

This comment suggests that target.addDynamicallyLegalOp() would enable ZA, but that's not quite true, is it? Instead, it marks funcOp as dynamically legal and provides a custom legality check:

• if arm_za attribute is present then legal if the first function Op is arm_sme::aarch64_sme_za_enable,
• if arm_za is not present, mark as legal.

I guess that the comment refers to the overall legalisation process rather then specifically to this hook?

Good point I've updated the comment hopefully makes more sense, cheers

mlir/test/Dialect/ArmSME/enable-arm-za.mlir

1

Could you add another RUN line without enable-za and verify that there's no asm_sme.intr.za.{enable|disable}?

Done

It would also be good to test without the -enable-arm-streaming pass.

I misunderstood what you mean here somehow, will update again to add this 😄

[c-rhodes]

Add test checking intrinsics aren't emitted when -enable-arm-streaming pass isn't specified.

[awarzynski]

Cheers for addressing my comments, LGTM!

mlir/test/Dialect/ArmSME/enable-arm-za.mlir

2–4

Personal opinion, feel free to ignore.

1. Custom prefix everywhere - this way you use the prefixes to "document" the test cases (e.g. "ENABLE-ZA" clearly suggests that in this case ZA is explicitly enabled)
2. Drop CHECK from custom prefixes. I think that most people know that these are CHECK ... prefixes. And shorter prefix means less noise/repetition.
3. Match the order of checks with the RUN lines (so, right now, it would be CHECK, then CHECK-ENABLE-ZA and CHECK-NO-ARM-STREAMING).

Like I said, person opinion, aka nit :)

[awarzynski]

mlir/test/Dialect/ArmSME/enable-arm-za.mlir

1

I misunderstood what you mean here somehow, will update again to add this 😄

Ah, I meant a test in which the input already contains arm_za. Added by hand rather then a compiler :)

[c-rhodes]

mlir/test/Dialect/ArmSME/enable-arm-za.mlir

2–4

Personal opinion, feel free to ignore.

1. Custom prefix everywhere - this way you use the prefixes to "document" the test cases (e.g. "ENABLE-ZA" clearly suggests that in this case ZA is explicitly enabled)
2. Drop CHECK from custom prefixes. I think that most people know that these are CHECK ... prefixes. And shorter prefix means less noise/repetition.
3. Match the order of checks with the RUN lines (so, right now, it would be CHECK, then CHECK-ENABLE-ZA and CHECK-NO-ARM-STREAMING).

Like I said, person opinion, aka nit :)

Done, cheers

mlir/include/mlir/Conversion/Passes.td

let options = [
     Option<"armNeon", "enable-arm-neon",
            "bool", /*default=*/"false",
            "Enables the use of ArmNeon dialect while lowering the vector "
      "dialect.">,
     Option<"armSVE", "enable-arm-sve",
            "bool", /*default=*/"false",
            "Enables the use of ArmSVE dialect while lowering the vector "
        "dialect.">,
+    Option<"armSME", "enable-arm-sme",
+           "bool", /*default=*/"false",
+           "Enables the use of ArmSME dialect while lowering the vector "
+       "dialect.">,
     Option<"x86Vector", "enable-x86vector",
            "bool", /*default=*/"false",
            "Enables the use of X86Vector dialect while lowering the vector "
      "dialect.">,
     Option<"useOpaquePointers", "use-opaque-pointers", "bool",
                /*default=*/"true", "Generate LLVM IR using opaque pointers "
                "instead of typed pointers">
   ];

mlir/include/mlir/Dialect/ArmSME/IR/ArmSME.td

 def LLVM_aarch64_sme_st1d_horiz : ArmSME_IntrStoreOp<"st1d.horiz">;
 def LLVM_aarch64_sme_st1q_horiz : ArmSME_IntrStoreOp<"st1q.horiz">;
 def LLVM_aarch64_sme_st1b_vert : ArmSME_IntrStoreOp<"st1b.vert">;
 def LLVM_aarch64_sme_st1h_vert : ArmSME_IntrStoreOp<"st1h.vert">;
 def LLVM_aarch64_sme_st1w_vert : ArmSME_IntrStoreOp<"st1w.vert">;
 def LLVM_aarch64_sme_st1d_vert : ArmSME_IntrStoreOp<"st1d.vert">;
 def LLVM_aarch64_sme_st1q_vert : ArmSME_IntrStoreOp<"st1q.vert">;
 
+def LLVM_aarch64_sme_za_enable : ArmSME_IntrOp<"za.enable">;
+def LLVM_aarch64_sme_za_disable : ArmSME_IntrOp<"za.disable">;
+
 #endif // ARMSME_OPS

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

+//===- Transforms.h - ArmSME Dialect Transformation 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_H
+#define MLIR_DIALECT_ARMSME_TRANSFORMS_H
+
+namespace mlir {
+
+class LLVMConversionTarget;
+class LLVMTypeConverter;
+class RewritePatternSet;
+
+/// Collect a set of patterns to lower ArmSME ops to ops that map to LLVM
+/// intrinsics.
+void populateArmSMELegalizeForLLVMExportPatterns(LLVMTypeConverter &converter,
+                                                 RewritePatternSet &patterns);
+
+/// Configure the target to support lowering ArmSME ops to ops that map to LLVM
+/// intrinsics.
+void configureArmSMELegalizeForExportTarget(LLVMConversionTarget &target);
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_ARMSME_TRANSFORMS_H

mlir/lib/Conversion/VectorToLLVM/CMakeLists.txt

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

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp

 //===- VectorToLLVM.cpp - Conversion from Vector to the LLVM dialect ------===//
 //
 // 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/VectorToLLVM/ConvertVectorToLLVM.h"
 
 #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/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"
 #include "mlir/Dialect/X86Vector/Transforms.h"
struct LowerVectorToLLVMPass
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<LLVM::LLVMDialect>();
     registry.insert<arith::ArithDialect>();
     registry.insert<memref::MemRefDialect>();
     if (armNeon)
       registry.insert<arm_neon::ArmNeonDialect>();
     if (armSVE)
       registry.insert<arm_sve::ArmSVEDialect>();
+    if (armSME)
+      registry.insert<arm_sme::ArmSMEDialect>();
     if (amx)
       registry.insert<amx::AMXDialect>();
     if (x86Vector)
       registry.insert<x86vector::X86VectorDialect>();
   }
   void runOnOperation() override;
 };
 } // namespace
if (armNeon) {
     // 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);
+  }
   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/CMakeLists.txt

 add_mlir_dialect_library(MLIRArmSMETransforms
   EnableArmStreaming.cpp
+  LegalizeForLLVMExport.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ArmSME/Transforms
 
   DEPENDS
   MLIRArmSMETransformsIncGen
 
   LINK_LIBS PUBLIC
+  MLIRArmSMEDialect
   MLIRFuncDialect
+  MLIRLLVMCommonConversion
   MLIRPass
   )

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

+//===- LegalizeForLLVMExport.cpp - Prepare ArmSME for LLVM translation ----===//
+//
+// 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/LLVMCommon/ConversionTarget.h"
+#include "mlir/Dialect/ArmSME/IR/ArmSME.h"
+#include "mlir/Dialect/ArmSME/Transforms/Transforms.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+
+using namespace mlir;
+using namespace mlir::arm_sme;
+
+namespace {
+/// Insert 'llvm.aarch64.sme.za.enable' intrinsic at the start of 'func.func'
+/// ops to enable the ZA storage array.
+struct EnableZAPattern : public OpRewritePattern<func::FuncOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(func::FuncOp op,
+                                PatternRewriter &rewriter) const final {
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPointToStart(&op.front());
+    rewriter.create<arm_sme::aarch64_sme_za_enable>(op->getLoc());
+    rewriter.updateRootInPlace(op, [] {});
+    return success();
+  }
+};
+
+/// Insert 'llvm.aarch64.sme.za.disable' intrinsic before 'func.return' ops to
+/// disable the ZA storage array.
+struct DisableZAPattern : public OpRewritePattern<func::ReturnOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(func::ReturnOp op,
+                                PatternRewriter &rewriter) const final {
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPoint(op);
+    rewriter.create<arm_sme::aarch64_sme_za_disable>(op->getLoc());
+    rewriter.updateRootInPlace(op, [] {});
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateArmSMELegalizeForLLVMExportPatterns(
+    LLVMTypeConverter &converter, RewritePatternSet &patterns) {
+  patterns.add<EnableZAPattern, DisableZAPattern>(patterns.getContext());
+}
+
+void mlir::configureArmSMELegalizeForExportTarget(
+    LLVMConversionTarget &target) {
+  target.addLegalOp<arm_sme::aarch64_sme_za_enable,
+                    arm_sme::aarch64_sme_za_disable>();
+
+  // Enable ZA if the function has the 'arm_za' attribute and it hasn't already
+  // been enabled.
+  target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp funcOp) {
+    auto firstOp = funcOp.getBody().front().begin();
+    return !funcOp->hasAttr("arm_za") ||
+           isa<arm_sme::aarch64_sme_za_enable>(firstOp);
+  });
+
+  // Disable ZA if the function has the 'arm_za' attribute and it hasn't
+  // already been disabled.
+  target.addDynamicallyLegalOp<func::ReturnOp>([&](func::ReturnOp returnOp) {
+    bool hasDisableZA = false;
+    auto funcOp = returnOp->getParentOp();
+    funcOp->walk<WalkOrder::PreOrder>(
+        [&](arm_sme::aarch64_sme_za_disable op) { hasDisableZA = true; });
+    return !funcOp->hasAttr("arm_za") || hasDisableZA;
+  });
+}

mlir/test/Dialect/ArmSME/enable-arm-za.mlir

+// RUN: mlir-opt %s -enable-arm-streaming=enable-za -convert-vector-to-llvm="enable-arm-sme" | FileCheck %s
+
+// CHECK-LABEL: @arm_za
+func.func @arm_za() {
+  // CHECK: arm_sme.intr.za.enable
+  // CHECK-NEXT: arm_sme.intr.za.disable
+  // CHECK-NEXT: return
+  return
+}

mlir/test/Target/LLVMIR/arm-sme.mlir

llvm.func @arm_sme_store(%nxv1i1  : vector<[1]xi1>,
   // CHECK: call void @llvm.aarch64.sme.st1h.vert
   "arm_sme.intr.st1h.vert"(%nxv8i1, %p16, %c0, %c0) :
               (vector<[8]xi1>, !llvm.ptr<i16>, i32, i32) -> ()
   // CHECK: call void @llvm.aarch64.sme.st1b.vert
   "arm_sme.intr.st1b.vert"(%nxv16i1, %p8, %c0, %c0) :
               (vector<[16]xi1>, !llvm.ptr<i8>, i32, i32) -> ()
   llvm.return
 }
+
+// -----
+
+// CHECK-LABEL: @arm_sme_toggle_za
+llvm.func @arm_sme_toggle_za() {
+  // CHECK: call void @llvm.aarch64.sme.za.enable()
+  "arm_sme.intr.za.enable"() : () -> ()
+  // CHECK: call void @llvm.aarch64.sme.za.disable()
+  "arm_sme.intr.za.disable"() : () -> ()
+  llvm.return
+}