Remove the size restriction of self_attention_fuse pass (#56909)

abenmao · web-flow · commit 4998c32f75d4 · 2023-09-11T16:28:17.000+08:00
diff --git a/paddle/fluid/operators/fused/scaled_dp_attention.h b/paddle/fluid/operators/fused/scaled_dp_attention.h
@@ -227,7 +227,6 @@ void softmax_sum_max(float* AB,
                      float refac,
                      int m,
                      int k) {
-  assert(k % 16 == 0);
   float max_val = std::numeric_limits<float>::lowest();
   __m512 vrefac = _mm512_set1_ps(refac);
   for (int i = 0; i < m; ++i) {
@@ -290,18 +289,19 @@ void update_out_blk(float* output,
                     float* max,
                     int m,
                     int n) {
-  assert(n % 16 == 0);
   for (int i = 0; i < m; ++i) {
     const float* buf = exp_ABC + i * n;
     float* outbuf = output + i * n;
     __m512 merr = _mm512_set1_ps(pre_max[i] - max[i]);
     merr = vexp(merr);
     __m512 vfac = _mm512_set1_ps(pre_sum[i] / sum[i]);
     for (int off = 0; off < n; off += 16) {
-      __m512 vout = _mm512_loadu_ps(outbuf + off);
-      __m512 vabc = _mm512_loadu_ps(buf + off);
+      int remain = n - off;
+      __mmask16 mask = (remain >= 16 ? 0xffff : (1 << remain) - 1);
+      __m512 vout = _mm512_maskz_loadu_ps(mask, outbuf + off);
+      __m512 vabc = _mm512_maskz_loadu_ps(mask, buf + off);
       __m512 vupt = vout * merr * vfac + vabc;
-      _mm512_storeu_ps(outbuf + off, vupt);
+      _mm512_mask_storeu_ps(outbuf + off, mask, vupt);
     }
     pre_sum[i] = sum[i];
     pre_max[i] = max[i];
@@ -348,8 +348,6 @@ void scaled_dp_attention(const float* query,
   int iblk = std::min(512, itsize / 1);
   int oblk = std::min(512, otsize / 1);
   float refac = scale;
-  assert(itsize % iblk == 0);
-  assert(otsize % oblk == 0);
 
 #ifdef PADDLE_WITH_MKLML
   int nth = omp_get_max_threads();
diff --git a/paddle/fluid/operators/fused/self_dp_attention_op.cc b/paddle/fluid/operators/fused/self_dp_attention_op.cc
@@ -30,13 +30,6 @@ void SelfDPAttenOp::InferShape(framework::InferShapeContext* ctx) const {
                         "[batchsize, tokensize, 3, nhead, headsize] "
                         ", but now Input X dim is:[%s] ",
                         dim_input));
-  PADDLE_ENFORCE_EQ(dim_input[4] % 16,
-                    0,
-                    platform::errors::InvalidArgument(
-                        "The last dim of input X should be a multiple of 16, "
-                        ", but now the dim is:[%d] "
-                        "Please remove self_attention_fuse_pass from the lists",
-                        dim_input[4]));
   framework::DDim out_dims(
       {dim_input[0], dim_input[1], dim_input[3], dim_input[4]});
   ctx->SetOutputDim("Out", out_dims);
diff --git a/test/mkldnn/test_fused_vit_attention.py b/test/mkldnn/test_fused_vit_attention.py
@@ -24,50 +24,53 @@
 
 
 def test_fuse_resenet_unit():
-    place = paddle.CPUPlace()
-    program = paddle.static.Program()
-    startup_program = paddle.static.Program()
-    batch_size = 1
-    token_size = 4097
-    hidden_size = 768
-    num_heads = 12
-    dtype = np.float32
-    with paddle.static.program_guard(program, startup_program):
-        x = paddle.static.data(
-            "x", [batch_size, token_size, hidden_size * 3], dtype=dtype
-        )
-        qkv = x.reshape(
-            (batch_size, token_size, 3, num_heads, hidden_size // num_heads)
-        ).transpose((2, 0, 3, 1, 4))
+    tests = [[1, 4096, 768, 12], [10, 4097, 756, 12], [10, 4097, 756, 12]]
+    for test in tests:
+        batch_size = test[0]
+        token_size = test[1]
+        hidden_size = test[2]
+        num_heads = test[3]
+        dtype = np.float32
 
-        q, k, v = qkv[0], qkv[1], qkv[2]
+        place = paddle.CPUPlace()
+        program = paddle.static.Program()
+        startup_program = paddle.static.Program()
+        with paddle.static.program_guard(program, startup_program):
+            x = paddle.static.data(
+                "x", [batch_size, token_size, hidden_size * 3], dtype=dtype
+            )
+            qkv = x.reshape(
+                (batch_size, token_size, 3, num_heads, hidden_size // num_heads)
+            ).transpose((2, 0, 3, 1, 4))
 
-        attn = q.matmul(k.transpose((0, 1, 3, 2)))
+            q, k, v = qkv[0], qkv[1], qkv[2]
 
-        attn = paddle.nn.functional.softmax(attn, axis=-1)
+            attn = q.matmul(k.transpose((0, 1, 3, 2)))
 
-        out = (
-            (attn.matmul(v))
-            .transpose((0, 2, 1, 3))
-            .reshape((-1, token_size, hidden_size))
-        )
+            attn = paddle.nn.functional.softmax(attn, axis=-1)
+
+            out = (
+                (attn.matmul(v))
+                .transpose((0, 2, 1, 3))
+                .reshape((-1, token_size, hidden_size))
+            )
 
-    graph = core.Graph(program.desc)
-    core.get_pass("self_attention_fuse_pass").apply(graph)
-    after_program = paddle.base.framework.IrGraph(graph).to_program()
-    exe = paddle.static.Executor(place)
-    exe.run(startup_program)
+        graph = core.Graph(program.desc)
+        core.get_pass("self_attention_fuse_pass").apply(graph)
+        after_program = paddle.framework.IrGraph(graph).to_program()
+        exe = paddle.static.Executor(place)
+        exe.run(startup_program)
 
-    feed = {
-        "x": np.random.randn(batch_size, token_size, hidden_size * 3).astype(
-            dtype
+        feed = {
+            "x": np.random.randn(
+                batch_size, token_size, hidden_size * 3
+            ).astype(dtype)
+        }
+        before_out = exe.run(program, feed=feed, fetch_list=[out.name])
+        after_out = exe.run(after_program, feed=feed, fetch_list=[out.name])
+        np.testing.assert_allclose(
+            before_out[0], after_out[0], rtol=1e-05, atol=0.005
         )
-    }
-    before_out = exe.run(program, feed=feed, fetch_list=[out.name])
-    after_out = exe.run(after_program, feed=feed, fetch_list=[out.name])
-    np.testing.assert_allclose(
-        before_out[0], after_out[0], rtol=1e-05, atol=0.005
-    )
 
 
 if __name__ == '__main__':
