[AArch64] merge scaled and unscaled 8-bit and 18-bit zero narrow stores.

[zjaffal]

Summary

Previous patch (https://reviews.llvm.org/D150963) merged scaled and unscaled 32-bit stores. This patch adds the support for smaller size types.

[zjaffal]

Fix crashing tests

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

 
 char AArch64LoadStoreOpt::ID = 0;
 
 } // end anonymous namespace
 
 INITIALIZE_PASS(AArch64LoadStoreOpt, "aarch64-ldst-opt",
                 AARCH64_LOAD_STORE_OPT_NAME, false, false)
 
+enum NarrowStoreTy {
+  N8_STORE,
+  N16_STORE,
+};
+
 static bool isNarrowStore(unsigned Opc) {
   switch (Opc) {
   default:
     return false;
   case AArch64::STRBBui:
   case AArch64::STURBBi:
   case AArch64::STRHHui:
   case AArch64::STURHHi:
     return true;
   }
 }
 
+static NarrowStoreTy getNarrowStoreType(unsigned Opc) {
+  switch (Opc) {
+  default:
+    llvm_unreachable("Opcode is not a narrow store");
+
+  case AArch64::STRBBui:
+  case AArch64::STURBBi:
+    return NarrowStoreTy::N8_STORE;
+  case AArch64::STRHHui:
+  case AArch64::STURHHi:
+    return NarrowStoreTy::N16_STORE;
+  }
+}
+
 // These instruction set memory tag and either keep memory contents unchanged or
 // set it to zero, ignoring the address part of the source register.
 static bool isTagStore(const MachineInstr &MI) {
   switch (MI.getOpcode()) {
   default:
     return false;
   case AArch64::STGi:
   case AArch64::STZGi:
if (NonSExtOpc == getMatchingNonSExtOpcode(OpcB, &PairIsValidLdStrOpc)) {
     return true;
   }
 
   // If the second instruction isn't even a mergable/pairable load/store, bail
   // out.
   if (!PairIsValidLdStrOpc)
     return false;
 
-  // FIXME: We don't support merging narrow stores with mixed scaled/unscaled
-  // offsets.
-  if (isNarrowStore(OpcA) || isNarrowStore(OpcB))
-    return false;
+  // Check if we can merge two narrow stores.
+  if (isNarrowStore(OpcA) && isNarrowStore(OpcB))
+    return getNarrowStoreType(OpcA) == getNarrowStoreType(OpcB);
 
   // The STR<S,D,Q,W,X>pre - STR<S,D,Q,W,X>ui and
   // LDR<S,D,Q,W,X>pre-LDR<S,D,Q,W,X>ui
   // are candidate pairs that can be merged.
   if (isPreLdStPairCandidate(FirstMI, MI))
     return true;
 
   // Try to match an unscaled load/store with a scaled load/store.

llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir

bb.0.entry:
     STRBBui $wzr, $x0, 5 :: (store (s8))
     RET undef $lr
 ...
 ---
 name:  merge_scaled_str_with_unscaled_8
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: merge_scaled_str_with_unscaled_8
-    ; CHECK: STRBBui $wzr, $x0, 4 :: (store (s8))
-    ; CHECK-NEXT: STURBBi $wzr, $x0, 5 :: (store (s8))
+    ; CHECK: STRHHui $wzr, $x0, 2 :: (store (s8))
     ; CHECK-NEXT: RET undef $lr
     STRBBui $wzr, $x0, 4 :: (store (s8))
     STURBBi $wzr, $x0, 5 :: (store (s8))
     RET undef $lr
 ...
 ---
 name:  merge_unscaled_str_with_scaled_8
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: merge_unscaled_str_with_scaled_8
-    ; CHECK: STURBBi $wzr, $x0, 4 :: (store (s8))
-    ; CHECK-NEXT: STRBBui $wzr, $x0, 5 :: (store (s8))
+    ; CHECK: STURHHi $wzr, $x0, 4 :: (store (s8))
     ; CHECK-NEXT: RET undef $lr
     STURBBi $wzr, $x0, 4 :: (store (s8))
     STRBBui $wzr, $x0, 5 :: (store (s8))
     RET undef $lr
 ...
 ---
 name:  merge_unscaled_str_with_unscaled_str_16
 body:             |
bb.0.entry:
     STRHHui $wzr, $x0, 5 :: (store (s16))
     RET undef $lr
 ...
 ---
 name:  merge_scaled_str_with_unscaled_16
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: merge_scaled_str_with_unscaled_16
-    ; CHECK: STRHHui $wzr, $x0, 2 :: (store (s16))
-    ; CHECK-NEXT: STURHHi $wzr, $x0, 6 :: (store (s16))
+    ; CHECK: STRWui $wzr, $x0, 1 :: (store (s16))
     ; CHECK-NEXT: RET undef $lr
     STRHHui $wzr, $x0, 2 :: (store (s16))
     STURHHi $wzr, $x0, 6 :: (store (s16))
     RET undef $lr
 ...
 ---
 name:  merge_unscaled_str_with_scaled_16
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: merge_unscaled_str_with_scaled_16
-    ; CHECK: STURHHi $wzr, $x0, 4 :: (store (s16))
-    ; CHECK-NEXT: STRHHui $wzr, $x0, 3 :: (store (s16))
+    ; CHECK: STURWi $wzr, $x0, 4 :: (store (s16))
     ; CHECK-NEXT: RET undef $lr
     STURHHi $wzr, $x0, 4 :: (store (s16))
     STRHHui $wzr, $x0, 3 :: (store (s16))
     RET undef $lr
 ...
 ---
 name:  merge_unscaled_str_with_unscaled_32
 body:             |