diff --git a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp --- a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp +++ b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp @@ -219,6 +219,11 @@ 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: @@ -231,6 +236,20 @@ } } +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) { @@ -1334,10 +1353,15 @@ if (!PairIsValidLdStrOpc) return false; - // FIXME: We don't support merging narrow stores with mixed scaled/unscaled - // offsets. - if (isNarrowStore(OpcA) || isNarrowStore(OpcB)) + bool IsNarrowStoreA = isNarrowStore(OpcA); + bool IsNarrowStoreB = isNarrowStore(OpcB); + // skip if A and B aren't in the same store class. Narrow and wide stores + // can't be merged. + if (IsNarrowStoreA != IsNarrowStoreB) return false; + // Check if we can merge two narrow stores. + if (IsNarrowStoreA && IsNarrowStoreB) + return getNarrowStoreType(OpcA) == getNarrowStoreType(OpcB); // The STRpre - STRui and // LDRpre-LDRui diff --git a/llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir b/llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir --- a/llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir +++ b/llvm/test/CodeGen/AArch64/str-narrow-zero-merge.mir @@ -29,8 +29,7 @@ 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)) @@ -41,8 +40,7 @@ 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)) @@ -75,8 +73,7 @@ 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)) @@ -87,8 +84,7 @@ 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))