[WIP] amx conv im2col fp16s

Signed-off-by: Molly Sophia <[email protected]>
Tencent · Oct 10, 2024 · 7340660 · 7340660
1 parent bb58f2e
commit 7340660
Show file tree

Hide file tree

Showing 2 changed files with 1,738 additions and 2 deletions.
diff --git a/src/layer/arm/amx_usability.h b/src/layer/arm/amx_usability.h
@@ -103,16 +103,48 @@ void amx_stz(bool pair, unsigned int z_row, const void * ptr)
     AMX_STZ(oprand);
 }
 
-void amx_fma32(bool vector, unsigned int x_offset, unsigned int y_offset, int z_row)
+void amx_fma16_masked(bool vector, unsigned int x_offset, unsigned int y_offset, int z_row, uint8_t x_mode, uint8_t x_mask, uint8_t y_mode, uint8_t y_mask)
+{
+    uint64_t oprand = 0;
+    if (vector)
+        oprand |= 1ULL << 63;
+
+    oprand |= (uint64_t)y_offset & 0x1FF;
+    oprand |= ((uint64_t)x_offset & 0x1FF) << 10;
+    oprand |= ((uint64_t)z_row & 0x3F) << 20;
+    oprand |= ((uint64_t)y_mask & 0x1F) << 32;
+    oprand |= ((uint64_t)y_mode & 0x3) << 37;
+    oprand |= ((uint64_t)x_mask & 0x1F) << 41;
+    oprand |= ((uint64_t)x_mode & 0x3) << 46;
+
+    AMX_FMA16(oprand);
+}
+
+void amx_fma16(bool vector, unsigned int x_offset, unsigned int y_offset, int z_row)
+{
+    amx_fma16_masked(vector, x_offset, y_offset, z_row, 0, 0, 0, 0);
+}
+
+void amx_fma32_masked(bool vector, unsigned int x_offset, unsigned int y_offset, int z_row, uint8_t x_mode, uint8_t x_mask, uint8_t y_mode, uint8_t y_mask)
 {
     uint64_t oprand = 0;
     if (vector)
         oprand |= 1ULL << 63;
 
     oprand |= (uint64_t)y_offset & 0x1FF;
     oprand |= ((uint64_t)x_offset & 0x1FF) << 10;
+    oprand |= ((uint64_t)z_row & 0x3F) << 20;
+    oprand |= ((uint64_t)y_mask & 0x1F) << 32;
+    oprand |= ((uint64_t)y_mode & 0x3) << 37;
+    oprand |= ((uint64_t)x_mask & 0x1F) << 41;
+    oprand |= ((uint64_t)x_mode & 0x3) << 46;
 
     AMX_FMA32(oprand);
 }
 
+void amx_fma32(bool vector, unsigned int x_offset, unsigned int y_offset, int z_row)
+{
+    amx_fma32_masked(vector, x_offset, y_offset, z_row, 0, 0, 0, 0);
+}
+
 #endif // AMX_USABILITY_H