[llvm] [AArch64] Add immediate range checks for more MTE instructions (PR #119216)

[Oliver Stannard via llvm-commits]

https://github.com/ostannard created https://github.com/llvm/llvm-project/pull/119216

This would have turned the bug fixed in #117146 from a miscompilation into an assertion failure.

>From 65b7001c2e80fb8786a170f2bb738d2895a52d0e Mon Sep 17 00:00:00 2001
From: Oliver Stannard <oliver.stannard at arm.com>
Date: Mon, 9 Dec 2024 14:06:45 +0000
Subject: [PATCH] [AArch64] Add immediate range checks for more MTE
 instructions

This would have turned the bug fixed in #117146 from a miscompilation
into an assertion failure.
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 8a3ed10b8e0bd2..fd24e49f948a20 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -2555,6 +2555,14 @@ unsigned AArch64InstrInfo::getLoadStoreImmIdx(unsigned Opc) {
   case AArch64::ST4H_IMM:
   case AArch64::ST4W_IMM:
   case AArch64::STGPi:
+  case AArch64::STGPreIndex:
+  case AArch64::STZGPreIndex:
+  case AArch64::ST2GPreIndex:
+  case AArch64::STZ2GPreIndex:
+  case AArch64::STGPostIndex:
+  case AArch64::STZGPostIndex:
+  case AArch64::ST2GPostIndex:
+  case AArch64::STZ2GPostIndex:
   case AArch64::STNPDi:
   case AArch64::STNPQi:
   case AArch64::STNPSi:
@@ -2598,6 +2606,8 @@ unsigned AArch64InstrInfo::getLoadStoreImmIdx(unsigned Opc) {
   case AArch64::LDPWpre:
   case AArch64::LDPXpost:
   case AArch64::LDPXpre:
+  case AArch64::STGPpre:
+  case AArch64::STGPpost:
   case AArch64::STPDpost:
   case AArch64::STPDpre:
   case AArch64::STPQpost:
@@ -4167,7 +4177,11 @@ bool AArch64InstrInfo::getMemOpInfo(unsigned Opcode, TypeSize &Scale,
     break;
   case AArch64::LDG:
   case AArch64::STGi:
+  case AArch64::STGPreIndex:
+  case AArch64::STGPostIndex:
   case AArch64::STZGi:
+  case AArch64::STZGPreIndex:
+  case AArch64::STZGPostIndex:
     Scale = TypeSize::getFixed(16);
     Width = TypeSize::getFixed(16);
     MinOffset = -256;
@@ -4334,13 +4348,19 @@ bool AArch64InstrInfo::getMemOpInfo(unsigned Opcode, TypeSize &Scale,
     MaxOffset = 7;
     break;
   case AArch64::ST2Gi:
+  case AArch64::ST2GPreIndex:
+  case AArch64::ST2GPostIndex:
   case AArch64::STZ2Gi:
+  case AArch64::STZ2GPreIndex:
+  case AArch64::STZ2GPostIndex:
     Scale = TypeSize::getFixed(16);
     Width = TypeSize::getFixed(32);
     MinOffset = -256;
     MaxOffset = 255;
     break;
   case AArch64::STGPi:
+  case AArch64::STGPpost:
+  case AArch64::STGPpre:
     Scale = TypeSize::getFixed(16);
     Width = TypeSize::getFixed(16);
     MinOffset = -64;