Comment multiplier plus efficacement des matrices A*B^T ou A^T*B^T (T pour la transposition) en utilisant SSE ?

Question

Je n'arrête pas de me prendre la tête avec ça. J'ai un algorithme basé sur l'ESS pour multiplier la matrice A par matrice B . Je dois également mettre en œuvre les opérations pour lesquelles A, B ou les deux sont transposés. J'ai fait une implémentation naïve, le code de la matrice 4x4 représenté ci-dessous (qui est une opération SSE assez standard, je pense), mais le code de la matrice 4x4 n'a pas été implémenté. A*B^T L'opération prend environ deux fois plus de temps que l'opération A*B . L'implémentation d'ATLAS renvoie des valeurs similaires pour A*B et des résultats presque identiques pour la multiplication par une transposition, ce qui me laisse penser qu'il existe un moyen efficace de procéder.

MM-Multiplication :

m1 = (mat1.m_>>2)<<2;
n2 = (mat2.n_>>2)<<2;
n  = (mat1.n_>>2)<<2;

for (k=0; k<n; k+=4) {
  for (i=0; i<m1; i+=4) {
    // fetch: get 4x4 matrix from mat1
    // row-major storage, so get 4 rows
    Float* a0 = mat1.el_[i]+k;
    Float* a1 = mat1.el_[i+1]+k;
    Float* a2 = mat1.el_[i+2]+k;
    Float* a3 = mat1.el_[i+3]+k;

    for (j=0; j<n2; j+=4) {
      // fetch: get 4x4 matrix from mat2
      // row-major storage, so get 4 rows
      Float* b0 = mat2.el_[k]+j;
      Float* b1 = mat2.el_[k+1]+j;
      Float* b2 = mat2.el_[k+2]+j;
      Float* b3 = mat2.el_[k+3]+j;

      __m128 b0r = _mm_loadu_ps(b0);
      __m128 b1r = _mm_loadu_ps(b1);
      __m128 b2r = _mm_loadu_ps(b2);
      __m128 b3r = _mm_loadu_ps(b3);

      {  // first row of result += first row of mat1 * 4x4 of mat2
        __m128 cX1 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a0+0), b0r), _mm_mul_ps(_mm_load_ps1(a0+1), b1r));
        __m128 cX2 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a0+2), b2r), _mm_mul_ps(_mm_load_ps1(a0+3), b3r));
        Float* c0 = this->el_[i]+j;
        _mm_storeu_ps(c0, _mm_add_ps(_mm_add_ps(cX1, cX2), _mm_loadu_ps(c0)));
      }

      { // second row of result += second row of mat1 * 4x4 of mat2
        __m128 cX1 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a1+0), b0r), _mm_mul_ps(_mm_load_ps1(a1+1), b1r));
        __m128 cX2 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a1+2), b2r), _mm_mul_ps(_mm_load_ps1(a1+3), b3r));
        Float* c1 = this->el_[i+1]+j;
        _mm_storeu_ps(c1, _mm_add_ps(_mm_add_ps(cX1, cX2), _mm_loadu_ps(c1)));
      }

      { // third row of result += third row of mat1 * 4x4 of mat2
        __m128 cX1 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a2+0), b0r), _mm_mul_ps(_mm_load_ps1(a2+1), b1r));
        __m128 cX2 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a2+2), b2r), _mm_mul_ps(_mm_load_ps1(a2+3), b3r));
        Float* c2 = this->el_[i+2]+j;
        _mm_storeu_ps(c2, _mm_add_ps(_mm_add_ps(cX1, cX2), _mm_loadu_ps(c2)));
      }

      { // fourth row of result += fourth row of mat1 * 4x4 of mat2
        __m128 cX1 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a3+0), b0r), _mm_mul_ps(_mm_load_ps1(a3+1), b1r));
        __m128 cX2 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a3+2), b2r), _mm_mul_ps(_mm_load_ps1(a3+3), b3r));
        Float* c3 = this->el_[i+3]+j;
        _mm_storeu_ps(c3, _mm_add_ps(_mm_add_ps(cX1, cX2), _mm_loadu_ps(c3)));
      }
  }
// Code omitted to handle remaining rows and columns
}

Pour la multiplication MT (matrice multipliée par la matrice transposée), j'ai stocké b0r à b3r avec les commandes suivantes et j'ai modifié les variables de la boucle en conséquence :

__m128 b0r = _mm_set_ps(b3[0], b2[0], b1[0], b0[0]);
__m128 b1r = _mm_set_ps(b3[1], b2[1], b1[1], b0[1]);
__m128 b2r = _mm_set_ps(b3[2], b2[2], b1[2], b0[2]);
__m128 b3r = _mm_set_ps(b3[3], b2[3], b1[3], b0[3]);

Je soupçonne que le ralentissement est en partie dû à la différence entre l'extraction d'une ligne à la fois et le fait de devoir stocker 4 valeurs à chaque fois pour obtenir la colonne, mais j'ai l'impression que l'autre façon de procéder, l'extraction de lignes de B et la multiplication par la colonne de As, ne fera que déplacer le coût vers le stockage de 4 colonnes de résultats.

J'ai également essayé d'insérer les lignes de B en tant que lignes et d'utiliser ensuite la fonction _MM_TRANSPOSE4_PS(b0r, b1r, b2r, b3r); pour effectuer la transposition (je pensais qu'il y avait des optimisations supplémentaires dans cette macro), mais il n'y a pas d'amélioration réelle.

A priori, je pense que cela devrait être plus rapide... les produits points impliqués seraient une ligne par une ligne, ce qui semble intrinsèquement plus efficace, mais en essayant de faire les produits points directement, on se retrouve à devoir faire la même chose pour stocker les résultats.

Qu'est-ce qui m'échappe ?

Ajouté : Pour clarifier, j'essaie de ne pas transposer les matrices. Je préfère itérer le long des matrices. Le problème, d'après ce que je sais, est que la commande _mm_set_ps est beaucoup plus lente que _mm_load_ps.

J'ai également essayé une variante dans laquelle j'ai stocké les 4 lignes de la matrice A, puis j'ai remplacé les 4 segments entre crochets contenant 1 chargement, 4 multiplications et 2 additions par 4 instructions de multiplication et 3 instructions d'addition. hadds mais sans grand résultat. Le timing reste le même (et oui, j'ai essayé avec une instruction de débogage pour vérifier que le code avait changé dans ma compilation de test. Cette déclaration de débogage a été supprimée avant le profilage, bien sûr) :

    {  // first row of result += first row of mat1 * 4x4 of mat2
      __m128 cX1 = _mm_hadd_ps(_mm_mul_ps(a0r, b0r), _mm_mul_ps(a0r, b1r));
      __m128 cX2 = _mm_hadd_ps(_mm_mul_ps(a0r, b2r), _mm_mul_ps(a0r, b3r));
      Float* c0 = this->el_[i]+j;
      _mm_storeu_ps(c0, _mm_add_ps(_mm_hadd_ps(cX1, cX2), _mm_loadu_ps(c0)));
    }

    { // second row of result += second row of mat1 * 4x4 of mat2
      __m128 cX1 = _mm_hadd_ps(_mm_mul_ps(a1r, b0r), _mm_mul_ps(a1r, b1r));
      __m128 cX2 = _mm_hadd_ps(_mm_mul_ps(a1r, b2r), _mm_mul_ps(a1r, b3r));
      Float* c0 = this->el_[i+1]+j;
      _mm_storeu_ps(c0, _mm_add_ps(_mm_hadd_ps(cX1, cX2), _mm_loadu_ps(c0)));
    }

    { // third row of result += third row of mat1 * 4x4 of mat2
      __m128 cX1 = _mm_hadd_ps(_mm_mul_ps(a2r, b0r), _mm_mul_ps(a2r, b1r));
      __m128 cX2 = _mm_hadd_ps(_mm_mul_ps(a2r, b2r), _mm_mul_ps(a2r, b3r));
      Float* c0 = this->el_[i+2]+j;
      _mm_storeu_ps(c0, _mm_add_ps(_mm_hadd_ps(cX1, cX2), _mm_loadu_ps(c0)));
    }

    { // fourth row of result += fourth row of mat1 * 4x4 of mat2
      __m128 cX1 = _mm_hadd_ps(_mm_mul_ps(a3r, b0r), _mm_mul_ps(a3r, b1r));
      __m128 cX2 = _mm_hadd_ps(_mm_mul_ps(a3r, b2r), _mm_mul_ps(a3r, b3r));
      Float* c0 = this->el_[i+3]+j;
      _mm_storeu_ps(c0, _mm_add_ps(_mm_hadd_ps(cX1, cX2), _mm_loadu_ps(c0)));
    }

Mise à jour : droite, et en déplaçant le chargement des lignes de a0r a a3r dans les accolades, afin d'éviter les perturbations du registre, a également échoué.

Answer 1

Quelques recommandations qui peuvent vous aider :

• N'utilisez pas de mémoire non alignée (ces _mm_loadu* sont lents).
• Vous n'accédez pas à la mémoire de manière séquentielle, ce qui tue le cache. Essayez de transposer la matrice avant l'accès réel à cette mémoire, ce qui permettra à l'unité centrale de rechercher et d'utiliser le cache autant que possible. De cette manière __m128 b0r = _mm_set_ps(b3[0], b2[0], b1[0], b0[0]); // and b1r, etc.. n'est pas nécessaire. L'idée est de saisir les 4 composants de manière séquentielle. Si vous devez réorganiser votre mémoire avant que le code SSE ne soit appelé, faites-le.
• Vous chargez dans la boucle intérieure ceci : _mm_load_ps1(a0+0) (idem pour a1, a2 et a3) mais elle est constante pour toutes les itérations de la boucle intérieure. Vous pouvez charger ces valeurs à l'extérieur et économiser quelques cycles. Gardez un œil sur ce que vous pouvez réutiliser de l'itération précédente.
• Profil. Utilisez Intel VTune ou quelque chose de similaire, il vous dira où se trouve le goulot d'étranglement.

Answer 2

Je pense que c'est l'un des rares cas où l'addition horizontale est utile. Vous voulez C = A B^T mais B n'est pas stocké en mémoire en tant que transposition. C'est là le problème. Il est stocké comme un AoS au lieu d'un SoA. Dans ce cas, prendre la transposition de B et faire une addition verticale est plus lent que d'utiliser une addition horizontale, je pense. C'est au moins vrai pour Matrix vecteur Multiplication vectorielle efficace de matrices 4x4 avec SSE : addition horizontale et produit de points - quel est l'intérêt ? . Dans le code ci-dessous, la fonction m4x4 est un produit matriciel 4x4 sans SSE, m4x4_vec utilise le SSE, m4x4T fait C=A _B^T sans SSE, et m4x4T_vec fait C=A_ B^T usisg SSE. La dernière est celle que vous voulez, je pense.

Note : pour les matrices plus grandes, je n'utiliserais pas cette méthode. Dans ce cas, il est plus rapide de prendre d'abord la transposition et d'utiliser l'addition verticale (avec SSE/AVX, vous faites quelque chose de plus compliqué, vous transposez les bandes avec la largeur de SSE/AVX). En effet, la transposition est égale à O(n^2) et le produit matriciel est égal à O(n^3), de sorte que pour les grandes matrices, la transposition est insignifiante. Cependant, pour les matrices 4x4, la transposition est significative, de sorte que l'addition horizontale l'emporte.

Modifier : J'ai mal compris ce que vous vouliez. Vous voulez C = (A B)^T. Cela devrait être aussi rapide que (A B) et que le code est presque le même, il suffit d'intervertir les rôles de A et B.
Nous pouvons écrire le calcul comme suit :

C = A*B in Einstein notation is C_i,j = A_i,k * B_k,j.  
Since (A*B)^T = B^T*A^T we can write 
C = (A*B)^T in Einstein notation is C_i,j = B^T_i,k * A^T_k,j = A_j,k * B_k,i

Si vous comparez les deux, la seule chose qui change est que nous échangeons les rôles de j et i. J'ai mis du code pour faire cela à la fin de cette réponse.

#include "stdio.h"
#include <nmmintrin.h>    

void m4x4(const float *A, const float *B, float *C) {
    for(int i=0; i<4; i++) {
        for(int j=0; j<4; j++) {
            float sum = 0.0f;
            for(int k=0; k<4; k++) {
                sum += A[i*4+k]*B[k*4+j];
            }
            C[i*4 + j] = sum;
        }
    }
}

void m4x4T(const float *A, const float *B, float *C) {
    for(int i=0; i<4; i++) {
        for(int j=0; j<4; j++) {
            float sum = 0.0f;
            for(int k=0; k<4; k++) {
                sum += A[i*4+k]*B[j*4+k];
            }
            C[i*4 + j] = sum;
        }
    }
}

void m4x4_vec(const float *A, const float *B, float *C) {
    __m128 Brow[4], Mrow[4];
    for(int i=0; i<4; i++) {
        Brow[i] = _mm_load_ps(&B[4*i]);
    }

    for(int i=0; i<4; i++) {
        Mrow[i] = _mm_set1_ps(0.0f);
        for(int j=0; j<4; j++) {
            __m128 a = _mm_set1_ps(A[4*i +j]);
            Mrow[i] = _mm_add_ps(Mrow[i], _mm_mul_ps(a, Brow[j]));
        }
    }
    for(int i=0; i<4; i++) {
        _mm_store_ps(&C[4*i], Mrow[i]);
    }
}

void m4x4T_vec(const float *A, const float *B, float *C) {
    __m128 Arow[4], Brow[4], Mrow[4];
    for(int i=0; i<4; i++) {
        Arow[i] = _mm_load_ps(&A[4*i]);
        Brow[i] = _mm_load_ps(&B[4*i]);
    }

    for(int i=0; i<4; i++) {
        __m128 prod[4];
        for(int j=0; j<4; j++) {
            prod[j] =  _mm_mul_ps(Arow[i], Brow[j]);
        }
        Mrow[i] = _mm_hadd_ps(_mm_hadd_ps(prod[0], prod[1]), _mm_hadd_ps(prod[2], prod[3]));    
    }
    for(int i=0; i<4; i++) {
        _mm_store_ps(&C[4*i], Mrow[i]);
    }

}

float compare_4x4(const float* A, const float*B) {
    float diff = 0.0f;
    for(int i=0; i<4; i++) {
        for(int j=0; j<4; j++) {
            diff += A[i*4 +j] - B[i*4+j];
            printf("A %f, B %f\n", A[i*4 +j], B[i*4 +j]);
        }
    }
    return diff;    
}

int main() {
    float *A = (float*)_mm_malloc(sizeof(float)*16,16);
    float *B = (float*)_mm_malloc(sizeof(float)*16,16);
    float *C1 = (float*)_mm_malloc(sizeof(float)*16,16);
    float *C2 = (float*)_mm_malloc(sizeof(float)*16,16);

    for(int i=0; i<4; i++) {
        for(int j=0; j<4; j++) {
            A[i*4 +j] = i*4+j;
            B[i*4 +j] = i*4+j;
            C1[i*4 +j] = 0.0f;
            C2[i*4 +j] = 0.0f;
        }
    }
    m4x4T(A, B, C1);
    m4x4T_vec(A, B, C2);
    printf("compare %f\n", compare_4x4(C1,C2));

}

Edita:

Voici les fonctions scalaires et SSE qui font C = (A B)^T. Ils devraient être aussi rapides que leurs homologues A Versions B.

void m4x4TT(const float *A, const float *B, float *C) {
    for(int i=0; i<4; i++) {
        for(int j=0; j<4; j++) {
            float sum = 0.0f;
            for(int k=0; k<4; k++) {
                sum += A[j*4+k]*B[k*4+i];
            }
            C[i*4 + j] = sum;
        }
    }
}

void m4x4TT_vec(const float *A, const float *B, float *C) {
    __m128 Arow[4], Crow[4];
    for(int i=0; i<4; i++) {
        Arow[i] = _mm_load_ps(&A[4*i]);
    }

    for(int i=0; i<4; i++) {
        Crow[i] = _mm_set1_ps(0.0f);
        for(int j=0; j<4; j++) {
            __m128 a = _mm_set1_ps(B[4*i +j]);
            Crow[i] = _mm_add_ps(Crow[i], _mm_mul_ps(a, Arow[j]));
        }
    }

    for(int i=0; i<4; i++) {
        _mm_store_ps(&C[4*i], Crow[i]);
    }
}