[mlir][ArmSME] Fix crash on func decls in 'arm_za' legality checks

Reviewed By: dcaballe, Dinistro

Differential Revision: https://reviews.llvm.org/D153750

diff --git a/mlir/lib/Dialect/ArmSME/Transforms/LegalizeForLLVMExport.cpp b/mlir/lib/Dialect/ArmSME/Transforms/LegalizeForLLVMExport.cpp
index 3fe9e78..a946126 100644 (file)
--- a/mlir/lib/Dialect/ArmSME/Transforms/LegalizeForLLVMExport.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/LegalizeForLLVMExport.cpp
@@ -59,6 +59,8 @@ void mlir::configureArmSMELegalizeForExportTarget(
   //   2. the 'arm_za' function attribute is present and the first op in the
   //      function is an 'arm_sme::aarch64_sme_za_enable' intrinsic.
   target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp funcOp) {
+    if (funcOp.isDeclaration())
+      return true;
     auto firstOp = funcOp.getBody().front().begin();
     return !funcOp->hasAttr("arm_za") ||
            isa<arm_sme::aarch64_sme_za_enable>(firstOp);

diff --git a/mlir/test/Dialect/ArmSME/enable-arm-za.mlir b/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
index ae0bbdc..d415b19 100644 (file)
--- a/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
+++ b/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
@@ -2,6 +2,9 @@
 // RUN: mlir-opt %s -enable-arm-streaming -convert-vector-to-llvm="enable-arm-sme" | FileCheck %s -check-prefix=DISABLE-ZA
 // RUN: mlir-opt %s -convert-vector-to-llvm="enable-arm-sme" | FileCheck %s -check-prefix=NO-ARM-STREAMING
 
+// CHECK-LABEL: @declaration
+func.func private @declaration()
+
 // CHECK-LABEL: @arm_za
 func.func @arm_za() {
   // ENABLE-ZA: arm_sme.intr.za.enable