diff --git a/target-s390x/cpu.h b/target-s390x/cpu.h
index a2c077bdcd..68b5ab7056 100644
--- a/target-s390x/cpu.h
+++ b/target-s390x/cpu.h
@@ -352,6 +352,9 @@ static inline hwaddr decode_basedisp_s(CPUS390XState *env, uint32_t ipb)
     return addr;
 }
 
+/* Base/displacement are at the same locations. */
+#define decode_basedisp_rs decode_basedisp_s
+
 void s390x_tod_timer(void *opaque);
 void s390x_cpu_timer(void *opaque);
 
diff --git a/target-s390x/kvm.c b/target-s390x/kvm.c
index 02ac4ba995..b00a6617a3 100644
--- a/target-s390x/kvm.c
+++ b/target-s390x/kvm.c
@@ -562,11 +562,19 @@ static void kvm_handle_diag_308(S390CPU *cpu, struct kvm_run *run)
     handle_diag_308(&cpu->env, r1, r3);
 }
 
-static int handle_diag(S390CPU *cpu, struct kvm_run *run, int ipb_code)
+#define DIAG_KVM_CODE_MASK 0x000000000000ffff
+
+static int handle_diag(S390CPU *cpu, struct kvm_run *run, uint32_t ipb)
 {
     int r = 0;
+    uint16_t func_code;
 
-    switch (ipb_code) {
+    /*
+     * For any diagnose call we support, bits 48-63 of the resulting
+     * address specify the function code; the remainder is ignored.
+     */
+    func_code = decode_basedisp_rs(&cpu->env, ipb) & DIAG_KVM_CODE_MASK;
+    switch (func_code) {
     case DIAG_IPL:
         kvm_handle_diag_308(cpu, run);
         break;
@@ -577,7 +585,7 @@ static int handle_diag(S390CPU *cpu, struct kvm_run *run, int ipb_code)
         sleep(10);
         break;
     default:
-        DPRINTF("KVM: unknown DIAG: 0x%x\n", ipb_code);
+        DPRINTF("KVM: unknown DIAG: 0x%x\n", func_code);
         r = -1;
         break;
     }
@@ -684,7 +692,6 @@ static void handle_instruction(S390CPU *cpu, struct kvm_run *run)
 {
     unsigned int ipa0 = (run->s390_sieic.ipa & 0xff00);
     uint8_t ipa1 = run->s390_sieic.ipa & 0x00ff;
-    int ipb_code = (run->s390_sieic.ipb & 0x0fff0000) >> 16;
     int r = -1;
 
     DPRINTF("handle_instruction 0x%x 0x%x\n",
@@ -696,7 +703,7 @@ static void handle_instruction(S390CPU *cpu, struct kvm_run *run)
         r = handle_priv(cpu, run, ipa0 >> 8, ipa1);
         break;
     case IPA0_DIAG:
-        r = handle_diag(cpu, run, ipb_code);
+        r = handle_diag(cpu, run, run->s390_sieic.ipb);
         break;
     case IPA0_SIGP:
         r = handle_sigp(cpu, run, ipa1);