#define ECA_AIV                0x00200000
 #define ECA_VX         0x00020000
 #define ECA_PROTEXCI   0x00002000
+#define ECA_APIE       0x00000008
 #define ECA_SII                0x00000001
        __u32   eca;                    /* 0x004c */
 #define ICPT_INST      0x04
        __u8    reservede4[4];          /* 0x00e4 */
        __u64   tecmc;                  /* 0x00e8 */
        __u8    reservedf0[12];         /* 0x00f0 */
+#define CRYCB_FORMAT_MASK 0x00000003
 #define CRYCB_FORMAT1 0x00000001
 #define CRYCB_FORMAT2 0x00000003
        __u32   crycbd;                 /* 0x00fc */
        __u32 crycbd;
        __u8 aes_kw;
        __u8 dea_kw;
+       __u8 apie;
 };
 
 #define APCB0_MASK_SIZE 1
 
 #include <asm/sclp.h>
 #include <asm/cpacf.h>
 #include <asm/timex.h>
+#include <asm/ap.h>
 #include "kvm-s390.h"
 #include "gaccess.h"
 
        return r;
 }
 
-static int kvm_s390_query_ap_config(u8 *config)
-{
-       u32 fcn_code = 0x04000000UL;
-       u32 cc = 0;
-
-       memset(config, 0, 128);
-       asm volatile(
-               "lgr 0,%1\n"
-               "lgr 2,%2\n"
-               ".long 0xb2af0000\n"            /* PQAP(QCI) */
-               "0: ipm %0\n"
-               "srl %0,28\n"
-               "1:\n"
-               EX_TABLE(0b, 1b)
-               : "+r" (cc)
-               : "r" (fcn_code), "r" (config)
-               : "cc", "0", "2", "memory"
-       );
-
-       return cc;
-}
-
 static int kvm_s390_apxa_installed(void)
 {
-       u8 config[128];
-       int cc;
+       struct ap_config_info info;
 
-       if (test_facility(12)) {
-               cc = kvm_s390_query_ap_config(config);
-
-               if (cc)
-                       pr_err("PQAP(QCI) failed with cc=%d", cc);
-               else
-                       return config[0] & 0x40;
+       if (ap_instructions_available()) {
+               if (ap_qci(&info) == 0)
+                       return info.apxa;
        }
 
        return 0;
 }
 
+/*
+ * The format of the crypto control block (CRYCB) is specified in the 3 low
+ * order bits of the CRYCB designation (CRYCBD) field as follows:
+ * Format 0: Neither the message security assist extension 3 (MSAX3) nor the
+ *          AP extended addressing (APXA) facility are installed.
+ * Format 1: The APXA facility is not installed but the MSAX3 facility is.
+ * Format 2: Both the APXA and MSAX3 facilities are installed
+ */
 static void kvm_s390_set_crycb_format(struct kvm *kvm)
 {
        kvm->arch.crypto.crycbd = (__u32)(unsigned long) kvm->arch.crypto.crycb;
 
+       /* Clear the CRYCB format bits - i.e., set format 0 by default */
+       kvm->arch.crypto.crycbd &= ~(CRYCB_FORMAT_MASK);
+
+       /* Check whether MSAX3 is installed */
+       if (!test_kvm_facility(kvm, 76))
+               return;
+
        if (kvm_s390_apxa_installed())
                kvm->arch.crypto.crycbd |= CRYCB_FORMAT2;
        else
 
 static void kvm_s390_crypto_init(struct kvm *kvm)
 {
-       if (!test_kvm_facility(kvm, 76))
-               return;
-
        kvm->arch.crypto.crycb = &kvm->arch.sie_page2->crycb;
        kvm_s390_set_crycb_format(kvm);
 
+       if (!test_kvm_facility(kvm, 76))
+               return;
+
        /* Enable AES/DEA protected key functions by default */
        kvm->arch.crypto.aes_kw = 1;
        kvm->arch.crypto.dea_kw = 1;
 
 static void kvm_s390_vcpu_crypto_setup(struct kvm_vcpu *vcpu)
 {
-       if (!test_kvm_facility(vcpu->kvm, 76))
+       /*
+        * If the AP instructions are not being interpreted and the MSAX3
+        * facility is not configured for the guest, there is nothing to set up.
+        */
+       if (!vcpu->kvm->arch.crypto.apie && !test_kvm_facility(vcpu->kvm, 76))
                return;
 
+       vcpu->arch.sie_block->crycbd = vcpu->kvm->arch.crypto.crycbd;
        vcpu->arch.sie_block->ecb3 &= ~(ECB3_AES | ECB3_DEA);
 
+       if (vcpu->kvm->arch.crypto.apie)
+               vcpu->arch.sie_block->eca |= ECA_APIE;
+
+       /* Set up protected key support */
        if (vcpu->kvm->arch.crypto.aes_kw)
                vcpu->arch.sie_block->ecb3 |= ECB3_AES;
        if (vcpu->kvm->arch.crypto.dea_kw)
                vcpu->arch.sie_block->ecb3 |= ECB3_DEA;
-
-       vcpu->arch.sie_block->crycbd = vcpu->kvm->arch.crypto.crycbd;
 }
 
 void kvm_s390_vcpu_unsetup_cmma(struct kvm_vcpu *vcpu)