diff --git a/llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp b/llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp
--- a/llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp
+++ b/llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp
@@ -317,6 +317,7 @@
       .addImm(Imm0)
       .addImm(12);
 
+  MRI->constrainRegClass(NewDstReg, MRI->getRegClass(DstReg));
   BuildMI(*MBB, MI, DL, TII->get(Opcode), NewDstReg)
       .addReg(NewTmpReg)
       .addImm(Imm1)
diff --git a/llvm/test/CodeGen/AArch64/addsub-24bit-imm.mir b/llvm/test/CodeGen/AArch64/addsub-24bit-imm.mir
new file mode 100644
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/addsub-24bit-imm.mir
@@ -0,0 +1,63 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -run-pass=aarch64-mi-peephole-opt -o - -mtriple=aarch64-unknown-linux -verify-machineinstrs %s | FileCheck %s
+
+# Main intention is to verify machine instructions have valid register classes.
+# Use of UBFM[W|X]ri is used as an arbitrary instruction that requires GPR[32|64]RegClass.
+# If the ADD/SUB optimization generates invalid register classes, this test will fail.
+---
+name: addi
+body: |
+  bb.0.entry:
+    liveins: $w0
+    ; CHECK-LABEL: name: addi
+    ; CHECK: [[COPY:%[0-9]+]]:gpr32common = COPY $w0
+    ; CHECK-NEXT: [[ADDWri:%[0-9]+]]:gpr32sp = ADDWri [[COPY]], 273, 12
+    ; CHECK-NEXT: [[ADDWri1:%[0-9]+]]:gpr32common = ADDWri [[ADDWri]], 3549, 0
+    ; CHECK-NEXT: [[UBFMWri:%[0-9]+]]:gpr32 = UBFMWri [[ADDWri1]], 28, 31
+    ; CHECK-NEXT: $w0 = COPY [[UBFMWri]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %0:gpr32 = COPY $w0
+    %1:gpr32 = MOVi32imm 1121757
+    %2:gpr32 = ADDWrr %0, %1
+    %3:gpr32 = UBFMWri %2, 28, 31
+    $w0 = COPY %3
+    RET_ReallyLR implicit $w0
+...
+---
+name: addl
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: addl
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64common = COPY $x0
+    ; CHECK-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY]], 273, 12
+    ; CHECK-NEXT: [[ADDXri1:%[0-9]+]]:gpr64common = ADDXri [[ADDXri]], 3549, 0
+    ; CHECK-NEXT: [[UBFMXri:%[0-9]+]]:gpr64 = UBFMXri [[ADDXri1]], 28, 31
+    ; CHECK-NEXT: $x0 = COPY [[UBFMXri]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %0:gpr64 = COPY $x0
+    %1:gpr32 = MOVi32imm 1121757
+    %2:gpr64 = SUBREG_TO_REG 0, %1, %subreg.sub_32
+    %3:gpr64 = ADDXrr %0, killed %2
+    %4:gpr64 = UBFMXri %3, 28, 31
+    $x0 = COPY %4
+    RET_ReallyLR implicit $x0
+...
+---
+name: addl_negate
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: addl_negate
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64common = COPY $x0
+    ; CHECK-NEXT: [[SUBXri:%[0-9]+]]:gpr64sp = SUBXri [[COPY]], 273, 12
+    ; CHECK-NEXT: [[SUBXri1:%[0-9]+]]:gpr64common = SUBXri [[SUBXri]], 3549, 0
+    ; CHECK-NEXT: [[UBFMXri:%[0-9]+]]:gpr64 = UBFMXri [[SUBXri1]], 28, 31
+    ; CHECK-NEXT: $x0 = COPY [[UBFMXri]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %0:gpr64 = COPY $x0
+    %1:gpr64 = MOVi64imm -1121757
+    %2:gpr64 = ADDXrr %0, killed %1
+    %3:gpr64 = UBFMXri %2, 28, 31
+    $x0 = COPY %3
+    RET_ReallyLR implicit $x0