[Operator] Adding CPU support for matrix multiplication

python/hidet/backend/build.py

@@ -120,14 +120,14 @@ def compile(self, src_path: str, out_lib_path: str, options: Optional[Dict[str,
             *['-L{}'.format(library_dir) for library_dir in self.library_dirs],
             # optimize host side code via -O3
             '-O3',
-            # enable openmp support for cpu kernels
-            '-Xcompiler -fopenmp',
+            # host compiler options: enable openmp, avx2, unroll loops and fast math
+            '-Xcompiler -fopenmp,-fPIC,-m64,-mavx2,-march=native,-O3,-funroll-loops,-ffast-math',
             # the target PTX and SASS version.
             '-gencode arch=compute_{cc},code=sm_{cc}'.format(cc=cc_code),
             # allow ptxas (PTX assembler) to output information like register/smem usage.
             '--ptxas-options=-v',
             # compile into position independent code.
-            '--compiler-options -fPIC',
+            # '--compiler-options -fPIC,-m64,-mavx2,-march=native, -O3',
             # embed the line information into the binary, allow Nsight Compute to get the source code for profiling.
             '-lineinfo',
             # ftz=true and prec-div=false for fast math
@@ -184,6 +184,10 @@ def compile(self, src_path: str, out_lib_path: str, options: Optional[Dict[str,
             *['-L{}'.format(library_dir) for library_dir in self.library_dirs],
             # apply -O3 optimization.
             '-O3',
+            # support avx intrinsics
+            '-mavx2',
+            '-m64',
+            '-march=native',
             # compile into position independent code.
             '-fPIC',
             # enable OpenMP.

python/hidet/backend/codegen.py

@@ -441,7 +441,9 @@ def visit_ForStmt(self, stmt: ForStmt):
                 doc += NewLine() + '#pragma unroll'
         elif stmt.attr.parallel:
             if stmt.attr.parallel_threads:
-                doc += NewLine() + '#pragma omp parallel for num_threads({})'.format(stmt.attr.parallel_threads)
+                doc += NewLine() + '#pragma omp parallel for schedule(dynamic) num_threads({})'.format(
+                    stmt.attr.parallel_threads
+                )
             else:
                 doc += NewLine() + '#pragma omp parallel for'
         doc += NewLine() + Text('for (') + init_doc + '; ' + cond_doc + '; ' + update_doc + ') '
@@ -555,6 +557,8 @@ def visit_DataType(self, t: DataType):
             'tfloat32': 'tfloat32_t',
             'complex64': 'complex64_t',
             'complex128': 'complex128_t',
+            'float32x4': '__m128',
+            'float32x8': '__m256',
         }
         return Text(scalar_type_map[t.name])
 
@@ -613,6 +617,8 @@ def require_headers(self) -> Doc:
         doc += Text('#include <hidet/runtime/cuda/complex.h>') + NewLine()
         doc += Text('#include <hidet/runtime/cuda/context.h>') + NewLine()
 
+        doc += Text('#include <immintrin.h>') + NewLine()
+
         # nvcc use float to 'store' tfloat32 data
         doc += Text('typedef float tfloat32_t;') + NewLine()
         doc += Text('typedef __nv_bfloat16 bfloat16_t;') + NewLine()
@@ -684,6 +690,7 @@ def require_headers(self) -> Doc:
         doc += Text('#include <hidet/runtime/cpu/float16.h>') + NewLine()
         doc += Text('#include <hidet/runtime/cpu/bfloat16.h>') + NewLine()
         doc += Text('#include <hidet/runtime/cpu/complex.h>') + NewLine()
+        doc += Text('#include <immintrin.h>')
         doc += NewLine()
         return doc
 

python/hidet/graph/ops/__init__.py

@@ -18,7 +18,7 @@
 from .definitions.conv2d_transpose import conv2d_transpose, conv2d_transpose_gemm
 from .definitions.conv3d import conv3d, conv3d_gemm
 from .definitions.conv3d_transpose import conv3d_transpose
-from .definitions.matmul import batch_matmul, matmul
+from .definitions.matmul import batch_matmul, matmul, matmul_x86
 from .definitions.pool import avg_pool2d, avg_pool3d, adaptive_avg_pool1d, adaptive_avg_pool2d, adaptive_avg_pool3d
 from .definitions.pool import max_pool2d, max_pool3d, adaptive_max_pool1d, adaptive_max_pool2d, adaptive_max_pool3d
 from .definitions.activation import relu, leaky_relu, sigmoid, hardsigmoid, clip, relu6, prelu, gelu, silu, hardswish

python/hidet/graph/ops/definitions/__init__.py

@@ -34,7 +34,7 @@
 from .conv3d_transpose import conv3d_transpose
 from .matmul import batch_matmul, matmul
 
-from .matmul import BatchMatmulOp, MatmulOp
+from .matmul import BatchMatmulOp, MatmulOp, Matmulx86Op
 from .conv2d import Conv2dOp
 from .arithmetic import ErfOp, PowOp, AddOp, SubtractOp, MultiplyOp, DivideOp, WhereOp
 from .compare import EqualOp

python/hidet/graph/ops/definitions/matmul/__init__.py

@@ -12,3 +12,7 @@
 from .matmul import matmul, MatmulOp, MatmulTask
 from .batch_matmul import batch_matmul, BatchMatmulOp, BatchMatmulTask
 from . import resolve
+
+from .matmul_f32_x86 import matmul_x86
+
+from .matmul_f32_x86 import MatmulF32Taskx86, Matmulx86Op