[llvm] [llvm] Change int_aarch64_sme_in_streaming_mode to use IntrReadMem (PR #123205)

[Nicholas Guy via llvm-commits]

https://github.com/NickGuy-Arm created https://github.com/llvm/llvm-project/pull/123205

As the int_aarch64_sme_in_streaming_mode intrinsic depends on modifiable state, the use of IntrNoMem was not the correct attribute to use. Switch it to use IntrReadMem and fix up the codegen.

>From 8983b9d921bff17f3ce7a43e310de6beef55c9f7 Mon Sep 17 00:00:00 2001
From: Nick Guy <nicholas.guy at arm.com>
Date: Thu, 16 Jan 2025 14:18:45 +0000
Subject: [PATCH] [llvm] Change int_aarch64_sme_in_streaming_mode to use
 IntrReadMem

As the int_aarch64_sme_in_streaming_mode intrinsic depends on modifiable state, the use of IntrNoMem was not the correct attribute to use. Switch it to use IntrReadMem and fix up the codegen
---
 llvm/include/llvm/IR/IntrinsicsAArch64.td     |  2 +-
 .../Target/AArch64/AArch64ISelLowering.cpp    | 29 ++++++++++++-------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/llvm/include/llvm/IR/IntrinsicsAArch64.td b/llvm/include/llvm/IR/IntrinsicsAArch64.td
index cc7a81e15f6609..1ecefb26cfe13e 100644
--- a/llvm/include/llvm/IR/IntrinsicsAArch64.td
+++ b/llvm/include/llvm/IR/IntrinsicsAArch64.td
@@ -2974,7 +2974,7 @@ let TargetPrefix = "aarch64" in {
 
 
   def int_aarch64_sme_zero : DefaultAttrsIntrinsic<[], [llvm_i32_ty], [ImmArg<ArgIndex<0>>]>;
-  def int_aarch64_sme_in_streaming_mode : DefaultAttrsIntrinsic<[llvm_i1_ty], [], [IntrNoMem]>, ClangBuiltin<"__builtin_arm_in_streaming_mode">;
+  def int_aarch64_sme_in_streaming_mode : DefaultAttrsIntrinsic<[llvm_i1_ty], [], [IntrReadMem]>, ClangBuiltin<"__builtin_arm_in_streaming_mode">;
 
   class SME_OuterProduct_Intrinsic
       : DefaultAttrsIntrinsic<[],
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index d4a114c275fb76..3982a35c7ccdfd 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -1190,7 +1190,7 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
 
   setOperationAction(ISD::INTRINSIC_WO_CHAIN, MVT::Other, Custom);
   if (Subtarget->hasSME())
-    setOperationAction(ISD::INTRINSIC_WO_CHAIN, MVT::i1, Custom);
+    setOperationAction(ISD::INTRINSIC_W_CHAIN, MVT::i1, Custom);
 
   if (Subtarget->isNeonAvailable()) {
     // FIXME: v1f64 shouldn't be legal if we can avoid it, because it leads to
@@ -27631,15 +27631,6 @@ void AArch64TargetLowering::ReplaceNodeResults(
       Results.push_back(DAG.getNode(ISD::TRUNCATE, DL, VT, V));
       return;
     }
-    case Intrinsic::aarch64_sme_in_streaming_mode: {
-      SDLoc DL(N);
-      SDValue Chain = DAG.getEntryNode();
-      SDValue RuntimePStateSM =
-          getRuntimePStateSM(DAG, Chain, DL, N->getValueType(0));
-      Results.push_back(
-          DAG.getNode(ISD::TRUNCATE, DL, MVT::i1, RuntimePStateSM));
-      return;
-    }
     case Intrinsic::experimental_vector_match:
     case Intrinsic::get_active_lane_mask: {
       if (!VT.isFixedLengthVector() || VT.getVectorElementType() != MVT::i1)
@@ -27657,6 +27648,24 @@ void AArch64TargetLowering::ReplaceNodeResults(
     }
     }
   }
+  case ISD::INTRINSIC_W_CHAIN: {
+    Intrinsic::ID IntID =
+        static_cast<Intrinsic::ID>(N->getConstantOperandVal(1));
+    switch (IntID) {
+    default:
+      return;
+    case Intrinsic::aarch64_sme_in_streaming_mode: {
+      SDValue Chain = N->getOperand(0);
+      SDLoc DL(N);
+      SDValue RuntimePStateSM =
+          getRuntimePStateSM(DAG, Chain, DL, N->getValueType(0));
+      Results.push_back(
+          DAG.getNode(ISD::TRUNCATE, DL, MVT::i1, RuntimePStateSM));
+      Results.push_back(Chain);
+      return;
+    }
+    }
+  }
   case ISD::READ_REGISTER: {
     SDLoc DL(N);
     assert(N->getValueType(0) == MVT::i128 &&