Enable Transpose-free BF16 MHA based on the mha_calc fusions (#992) (#1048)

* Enable Transpose-free BF16 MHA based on the mha_calc fusions (#992)

* Enable Transpose-free BF16 MHA based on the mha_calc fusions

* Enable transfree MHA & BF16 Matmul

* Add test cases for Transfree MHA and OutTransfree BF16 BMM

* fall back to BF16 MHA based on the mha_calc fusions

* merge cpu-device llga

* Combine the Matmul kernel with contiguous checks

* Revise the descriptions for some kernels

Co-authored-by: Wang Weihan <[email protected]>

* remove aligned restriction using loadu/storeu and add UT (#938)

Co-authored-by: Wei Lin <[email protected]>
Co-authored-by: jianan-gu <[email protected]>

Author: 3 people

intel_extension_for_pytorch/csrc/aten/cpu/AddLayerNorm.cpp

@@ -48,12 +48,9 @@ at::Tensor dil_add_layernorm(
       break;
     }
   }
-  // Only support 64byte aligned
-  bool aligned_64_bytes = a.size(a.ndimension() - 1) % 16 == 0 &&
-      b.size(b.ndimension() - 1) % 16 == 0;
   // Only support contiguous tensor
   bool is_contiguous = a.is_contiguous() && b.is_contiguous();
-  if (no_broadcast && aligned_64_bytes && is_contiguous && alpha == 1.0f) {
+  if (no_broadcast && is_contiguous && alpha == 1.0f) {
     return AddLayerNorm(
         a, b, alpha, normalized_shape, weight_opt, bias_opt, eps);
   } else {

intel_extension_for_pytorch/csrc/aten/cpu/kernels/DivSoftmaxKrnl.cpp

@@ -18,7 +18,6 @@ using namespace torch_ipex::cpu::kernel;
  * There are some assumptions for this operator.
  * - The reduce dimension for softmax is the last dimension
  * - The reduce dimension for softmax is the leading dimension
- * - The elements number of the reduce dimension for softmax is n*16
  * - The input tensors are contiguous
  * - The number of the input tensor dimension should be >=2
  * - The mask b can be expand_as a with the mask_reshape (bs :: seq_length),

intel_extension_for_pytorch/csrc/cpu/ideep/ideep/operators/matmul.hpp

@@ -315,6 +315,8 @@ struct matmul_forward : public dnnl::matmul,
 
     dst_data_type = dst_type == data_type::undef ? dst_data_type : dst_type;
     tensor::desc dst_desc(dst_dims, dst_data_type, tag::any);
+    if (!dst.is_empty())
+      dst_desc = dst.get_desc().to_type(dst_data_type);
     auto key = utils::create_key(
         src_desc,
         weights_desc,

intel_extension_for_pytorch/csrc/cpu/vec/vec512/perf_kernel/add_layernorm.h

@@ -21,32 +21,33 @@ std::pair<float, float> _add_and_compute_mean_var(
     float* out) {
   // compute add and mean/var of the value after add
   // we should firstly store add value
-  auto vec_a = _load_f32_data(a_ptr);
-  auto vec_b = _load_f32_data(b_ptr);
+  auto vec_a = _loadu(a_ptr);
+  auto vec_b = _loadu(b_ptr);
   auto vec_add = _mm512_add_ps(vec_a, vec_b);
   auto vec_acc_mean = vec_add;
   auto vec_acc_pow = _mm512_mul_ps(vec_add, vec_add);
-  _mm512_store_ps(out, vec_add);
+  _mm512_storeu_ps(out, vec_add);
 
   int i = 16;
   for (; i <= size - 16; i += 16) {
-    vec_a = _load_f32_data(a_ptr + i);
-    vec_b = _load_f32_data(b_ptr + i);
+    vec_a = _loadu(a_ptr + i);
+    vec_b = _loadu(b_ptr + i);
     vec_add = _mm512_add_ps(vec_a, vec_b);
     vec_acc_mean = _mm512_add_ps(vec_add, vec_acc_mean);
-    _mm512_store_ps(out + i, vec_add);
+    _mm512_storeu_ps(out + i, vec_add);
     vec_acc_pow = _mm512_fmadd_ps(vec_add, vec_add, vec_acc_pow);
   }
 
   if (i < size) {
     __mmask16 mask = (1 << (size - i)) - 1;
-    vec_a = _maskz_load_f32_data(a_ptr + i, mask);
-    vec_b = _maskz_load_f32_data(b_ptr + i, mask);
+    vec_a = _maskz_loadu(a_ptr + i, mask);
+    vec_b = _maskz_loadu(b_ptr + i, mask);
     vec_add = _mm512_add_ps(vec_a, vec_b);
     auto vec_zero = _mm512_set1_ps(0);
-    _mm512_mask_store_ps(out + i, mask, vec_add);
-    vec_acc_mean = _mm512_maskz_add_ps(mask, vec_add, vec_acc_mean);
-    vec_acc_pow = _mm512_maskz_fmadd_ps(mask, vec_add, vec_add, vec_acc_pow);
+
+    vec_acc_mean = _mm512_add_ps(vec_add, vec_acc_mean);
+    _mm512_mask_storeu_ps(out + i, mask, vec_add);
+    vec_acc_pow = _mm512_fmadd_ps(vec_add, vec_add, vec_acc_pow);
   }
   float mean_var = _mm512_reduce_add_ps(vec_acc_mean) / float(size);
   float var_val = _mm512_reduce_add_ps(vec_acc_pow);
@@ -68,35 +69,35 @@ void _normalize_kernel(
   auto vec_bias = _mm512_set1_ps(bias);
   int i = 0;
   for (; i <= size - 16; i += 16) {
-    auto vec_input = _load_f32_data(input_ptr + i);
+    auto vec_input = _loadu(input_ptr + i);
     auto vec_gamma = vec_one;
     auto vec_beta = vec_zero;
     if (gamma_ptr) {
-      vec_gamma = _load_f32_data(gamma_ptr + i);
+      vec_gamma = _loadu(gamma_ptr + i);
     }
     if (beta_ptr) {
-      vec_beta = _load_f32_data(beta_ptr + i);
+      vec_beta = _loadu(beta_ptr + i);
     }
     //(a_ptr[i] * scale + bias) * gamma + beta;
     auto vec_norm = _mm512_fmadd_ps(vec_input, vec_scale, vec_bias);
     auto vec_res = _mm512_fmadd_ps(vec_norm, vec_gamma, vec_beta);
-    _store_data(out_ptr + i, vec_res);
+    _storeu(out_ptr + i, vec_res);
   }
   if (i < size) {
     __mmask16 mask = (1 << (size - i)) - 1;
-    auto vec_input = _maskz_load_f32_data(input_ptr + i, mask);
+    auto vec_input = _maskz_loadu(input_ptr + i, mask);
     auto vec_gamma = vec_one;
     auto vec_beta = vec_zero;
-    if (!gamma_ptr) {
-      vec_gamma = _maskz_load_f32_data(gamma_ptr + i, mask);
+    if (gamma_ptr) {
+      vec_gamma = _maskz_loadu(gamma_ptr + i, mask);
     }
-    if (!beta_ptr) {
-      vec_beta = _maskz_load_f32_data(beta_ptr + i, mask);
+    if (beta_ptr) {
+      vec_beta = _maskz_loadu(beta_ptr + i, mask);
     }
     //(a_ptr[i] * scale + bias) * gamma + beta;
-    auto vec_norm = _mm512_maskz_fmadd_ps(mask, vec_input, vec_scale, vec_bias);
-    auto vec_res = _mm512_maskz_fmadd_ps(mask, vec_norm, vec_gamma, vec_beta);
-    _mask_store_data(out_ptr + i, vec_res, mask);
+    auto vec_norm = _mm512_fmadd_ps(vec_input, vec_scale, vec_bias);
+    auto vec_res = _mm512_fmadd_ps(vec_norm, vec_gamma, vec_beta);
+    _mask_storeu(out_ptr + i, vec_res, mask);
   }
 }
 

intel_extension_for_pytorch/csrc/cpu/vec/vec512/perf_kernel/add_softmax.h

@@ -98,7 +98,7 @@ inline void _dil_div_add_reduce_max_fusion_kernel(
     vec_b = _loadu(b + i);
     vec_out = _mm512_fmadd_ps(vec_a, vec_r_dim_per_head, vec_b);
     vec_ps_min = _mm512_max_ps(vec_ps_min, vec_out);
-    _mm512_store_ps(out + i, vec_out);
+    _mm512_storeu_ps(out + i, vec_out);
   }
 
   if (i < size) {
@@ -107,7 +107,7 @@ inline void _dil_div_add_reduce_max_fusion_kernel(
     vec_b = _maskz_loadu(b + i, mask);
     vec_out = _mm512_fmadd_ps(vec_a, vec_r_dim_per_head, vec_b);
     vec_ps_min = _mm512_mask_max_ps(vec_ps_min, mask, vec_out, vec_ps_min);
-    _mm512_mask_store_ps(out + i, mask, vec_out);
+    _mm512_mask_storeu_ps(out + i, mask, vec_out);
   }
 
   // NOTE: _mm512_reduce_max_ps is sequence instruction
@@ -134,22 +134,22 @@ inline void _dil_maskedfill_div_max_fusion_kernel(
 
   int i = 0;
   for (; i <= size - 16; i += 16) {
-    vec_a = _load_f32_data(a + i);
-    vec_b = _load_f32_data(b + i);
+    vec_a = _loadu(a + i);
+    vec_b = _loadu(b + i);
     __mmask16 fill_mask = _mm512_cmp_ps_mask(vec_b, mask_c, 12);
     vec_out = _mm512_mask_div_ps(vec_fill, fill_mask, vec_a, vec_dim_per_head);
     vec_ps_min = _mm512_max_ps(vec_ps_min, vec_out);
-    _mm512_store_ps(out + i, vec_out);
+    _mm512_storeu_ps(out + i, vec_out);
   }
 
   if (i < size) {
     __mmask16 mask = (1 << (size - i)) - 1;
-    vec_a = _maskz_load_f32_data(a + i, mask);
-    vec_b = _maskz_load_f32_data(b + i, mask);
+    vec_a = _maskz_loadu(a + i, mask);
+    vec_b = _maskz_loadu(b + i, mask);
     __mmask16 fill_mask = _mm512_cmp_ps_mask(vec_b, mask_c, 12);
     vec_out = _mm512_mask_div_ps(vec_fill, fill_mask, vec_a, vec_dim_per_head);
     vec_ps_min = _mm512_max_ps(vec_ps_min, vec_out);
-    _mm512_mask_store_ps(out + i, mask, vec_out);
+    _mm512_mask_storeu_ps(out + i, mask, vec_out);
   }
 
   // NOTE: _mm512_reduce_max_ps is sequence instruction

intel_extension_for_pytorch/csrc/jit/cpu/kernels/Matmul.cpp

@@ -23,16 +23,42 @@ namespace cpu {
  * @param out Optinal output provided by user for matmul
  * @attr Attribute for matmul oneDNN primitive
  * @return output Tensor.
- */
+ * Since oneDNN 2.6.0, AMX and AVX512 brgemm are enabled for the DNNL MATMUL
+ * primitive if the input tensors are with the following tags:
+ * 3-dim - abc, acb; 4-dim - abcd, acbd, adbc, abdc.
+ * If the input tensor has one of the above layouts, the contiguous should NOT
+ * be applied to avoid unnecessary transpose (copy).
+ **/
 at::Tensor bmm_impl(
     const at::Tensor& tensor1,
     const at::Tensor& tensor2,
     at::Tensor out,
     const ideep::attr_t& attr,
     const std::vector<ideep::tensor>& postop_tensors,
     const float dst_coeff = 1.0f) {
-  auto tensor1_ = tensor1.is_contiguous() ? tensor1 : tensor1.contiguous();
-  auto tensor2_ = tensor2.is_contiguous() ? tensor2 : tensor2.contiguous();
+  // The following conditions are strict to exclude some extreme cases when the
+  // tensors have the undefined stride values. For the sake of reliability of
+  // transpose-free Matmul kernel, contiguous will be applied to these tensors.
+  auto check_tensor_layout = [](at::Tensor tensor) {
+    // Check if the Tensor is 3-dim or 4-dim
+    if (tensor.dim() != 3 && tensor.dim() != 4)
+      return false;
+    // Check if 'a' is the first dim
+    for (int64_t i = 1; i < tensor.dim(); ++i) {
+      if (tensor.stride(0) < tensor.stride(i))
+        return false;
+    }
+    // Check if the tensor has one of the above memory tags:
+    // The strides of the tensor should not be out of the tensor's ranges.
+    // 4-dim: 'b' should not be the last dim.
+    if (tensor.stride(0) * tensor.size(0) != tensor.numel() ||
+        (tensor.dim() == 4 && tensor.stride(1) == 1))
+      return false;
+    return true;
+  };
+  auto tensor1_ = check_tensor_layout(tensor1) ? tensor1 : tensor1.contiguous();
+  auto tensor2_ = check_tensor_layout(tensor2) ? tensor2 : tensor2.contiguous();
+
   const int64_t dim = tensor1.dim();
   const ideep::tensor mkldnn_input = itensor_view_from_dense(tensor1_);
   const ideep::tensor mkldnn_tensor2 = itensor_view_from_dense(tensor2_);