Begin adding SIMD support

Author: davschneller

tensorforge/backend/instructions/compute/primitives/amd.py

@@ -374,7 +374,7 @@ def mfma_emu_int8(writer: Writer, C, B, A, c, a, b):
     for x in const:
         constM *= x
     constI = [pow(constM // const[i], -1, const[i]) for i in range(len(const))]
-    const2 = [(constM // const[i]) * constI[i] for i in range(len(const))]
+    const2 = [float((constM // const[i]) * constI[i]) for i in range(len(const))]
     acc = len(const)
 
     Aa = writer.varalloc()

tensorforge/backend/instructions/compute/primitives/intel.py

@@ -1,16 +1,19 @@
 
-def dpas(C, B, A, rc, sd):
+def dpas(writer, C, B, A, rc, sd):
     # cf. https://github.com/intel/intel-graphics-compiler/blob/master/documentation/visa/instructions/DPAS.md
     # sd == depth == k * elemsIn32Bit
-    # rc == m
-    writer(f'asm("DPAS.tf32.tf32.{sd}.{rc} (16) %[D], %[C], %[B], %[A]" : [D]"=f"({C}) : [C]"f"({C}), [B]"d"({B}), [A]"d"({A}) :);')
+    # rc == m [1,2,4,8]
+    writer(f'tensorforge::intel_esimd::simd<tensorforge::TF32, 32> {A};')
+    writer(f'tensorforge::intel_esimd::simd<tensorforge::TF32, 32> {B};')
+    writer(f'tensorforge::intel_esimd::simd<float, 32> {C};')
+    writer(f'{C} = tensorforge::intel_xmx::dpas<{sd}, {rc}, float>({C}, {B}, {A});')
 
-def matmul(writer, C, A, B, M, N, K, kx, threads, dtype, sparse, ctx):
+def fmadpp(writer, C, B, A, size, offset, lane):
+    writer(f'{C}.select<{size}, 1>({offset}) += {A}[{lane}] * {B}.select<{size}, 1>({offset});')
 
+def load(writer, C):
+    writer(f'{C}')
 
+def matmul(writer, C, A, B, M, N, K, kx, threads, dtype, sparse, ctx):
     rc = 8
     sd = 8
-
-    dpas(C, A, B, rc, sd)
-
-    # TODO

tensorforge/backend/instructions/compute/primitives/nvidia.py

@@ -222,11 +222,14 @@ def asmcall(self, writer, D, A, B, C):
 : {arggrp2(A, f"{typeidx}")}, {arggrp2(B, f"{typeidx}")}, {arggrp2(C, f"{typeid}")}
 );""")
 
+    def epilogue(self):
+        pass
+
     def generate(self, writer, context, A, B, C):
-        Cstr = ','.join(f'{c}' for c in C)
         with writer.Scope():
             if self.mode == MMAMode.I8:
-                raise NotImplementedError()
+
+                pass
             if self.mode == MMAMode.TF32:
                 Atf32 = tfconvert(writer, A)
                 Btf32 = tfconvert(writer, B)
@@ -297,11 +300,15 @@ def generate(self, writer, context, A, B, C):
 ]
 
 INSTRS = [
-    MMAInstr(16,8,4,1,Datatype.F32,'mma.sync.aligned.m16n8k4.row.col.f32.tf32.tf32.f32', MMAMode.TF32),
-    MMAInstr(16,8,8,1,Datatype.F32,'mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32', MMAMode.TF32),
-    MMAInstr(8,8,4,1,Datatype.F64,'mma.sync.aligned.m8n8k4.row.col.f64.f64.f64.f64', MMAMode.DIRECT),
-    MMAInstr(16,8,4,1,Datatype.F64,'mma.sync.aligned.m16n8k4.row.col.f64.f64.f64.f64', MMAMode.DIRECT),
-    MMAInstr(16,8,8,1,Datatype.F64,'mma.sync.aligned.m16n8k8.row.col.f64.f64.f64.f64', MMAMode.DIRECT),
+    MMAInstr(16,8,4,1,Datatype.F32,'mma.sync.aligned.m16n8k4.row.col.f32.tf32.tf32.f32', MMAMode.TF32), # SM_80
+    MMAInstr(16,8,8,1,Datatype.F32,'mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32', MMAMode.TF32), # SM_80
+    MMAInstr(8,8,4,1,Datatype.F64,'mma.sync.aligned.m8n8k4.row.col.f64.f64.f64.f64', MMAMode.DIRECT), # SM_80
+    MMAInstr(16,8,4,1,Datatype.F64,'mma.sync.aligned.m16n8k4.row.col.f64.f64.f64.f64', MMAMode.DIRECT), # SM_90
+    MMAInstr(16,8,8,1,Datatype.F64,'mma.sync.aligned.m16n8k8.row.col.f64.f64.f64.f64', MMAMode.DIRECT), # SM_90
+    MMAInstr(16,8,16,1,Datatype.F64,'mma.sync.aligned.m16n8k16.row.col.f64.f64.f64.f64', MMAMode.DIRECT), # SM_90
+    MMAInstr(8,8,16,1,Datatype.F64,'mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32', MMAMode.I8), # SM_75
+    MMAInstr(16,8,16,1,Datatype.F64,'mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32', MMAMode.I8), # SM_80
+    MMAInstr(16,8,32,1,Datatype.F64,'mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32', MMAMode.I8), # SM_80
 ]
 
 def matmul(writer, C, A, B, M, N, K, kx, threads, dtype, sparse, ctx, shmptr, shmsize):

tensorforge/backend/instructions/memory/load.py

@@ -2,11 +2,12 @@
 import math
 from tensorforge.common.matrix.tensor import Tensor
 from . import AbstractShrMemWrite, MemoryInstruction
-from tensorforge.backend.symbol import SymbolType, Symbol, DataView, LeadIndex
+from tensorforge.backend.symbol import Symbol, SymbolType, DataView, LeadIndex, write_loops, LeadLoop, Loop
 from tensorforge.common.exceptions import InternalError
 from tensorforge.backend.writer import Writer
 from tensorforge.common.matrix.boundingbox import BoundingBox
 from tensorforge.common.context import Context
+from tensorforge.backend.data_types import RegMemObject
 from typing import Union, List
 
 # to find a number coprime to the number of shared memory banks
@@ -278,24 +279,24 @@ def __init__(self,
                num_threads: int):
     super(GlbToRegLoader, self).__init__(context)
 
-    if src.stype != SymbolType.Register:
-      raise InternalError('store: operand `src` is not in reg mem')
+    if dest.stype != SymbolType.Register:
+      raise InternalError('store: operand `dest` is not in reg mem')
 
-    if not isinstance(src.obj, RegMemObject):
-      raise InternalError(f'store: operand `src` is registers, instead: {type(src.obj)}')
+    if not isinstance(dest.obj, RegMemObject):
+      raise InternalError(f'store: operand `dest` is registers, instead: {type(dest.obj)}')
 
-    if dest.stype != SymbolType.Global:
-      raise InternalError('store: operand `dest` is not in global memory.')
+    if src.stype != SymbolType.Global:
+      raise InternalError('store: operand `src` is not in global memory.')
 
-    if not isinstance(dest.obj, Tensor):
-      raise InternalError('store: operand `dest` is not a matrix')
+    if not isinstance(src.obj, Tensor):
+      raise InternalError('store: operand `src` is not a matrix')
 
     src.add_user(self)
     dest.add_user(self)
 
-    dest.data_view = DataView(shape=dest.obj.shape,
+    dest.data_view = DataView(shape=src.obj.shape,
                               permute=None,
-                              bbox=dest.obj.get_bbox())
+                              bbox=src.obj.get_bbox())
 
     # if dest.data_view.get_dim_size(0) > src.data_view.get_dim_size(0):
     #   raise InternalError('store: `src` and `dest` do not match in size aling dim `0`')
@@ -305,13 +306,12 @@ def __init__(self,
     self._num_threads: int = num_threads
     self._is_ready: bool = True
 
-  def gen_code(self, writer: Writer) -> None:
+  def gen_code_inner(self, writer: Writer) -> None:
     writer.new_line()
     dest_view = self._dest.data_view
 
     allow_nontemporal = len(self._src.get_user_list()) == 1
 
-    writer(f'// {self}')
     src_bbox = self._src.data_view.get_bbox()
 
     loops = []
@@ -326,4 +326,4 @@ def inner(indices):
     write_loops(self._context, writer, loops, inner)
 
   def __str__(self) -> str:
-    return f'{self._dest.name} = store{{g>r}}({self._src.name});'
+    return f'{self._dest.name} = load{{g>r}}({self._src.name});'

tensorforge/backend/symbol.py

@@ -145,7 +145,9 @@ def write_nonlead(self):
     return f'{self._nonlead}'
 
   def write(self, context: Context):
-    if self._block > 1:
+    if context.get_vm().get_lexic().simd_mode:
+      return f'({self._nonlead} * {self._block})'
+    elif self._block > 1:
       return f'(({context.get_vm().get_lexic().thread_idx_x} / {self._stride}) % {self._block}) + {self._nonlead} * {self._block}'
     elif self._block == 1:
       return f'{self._nonlead}'
@@ -443,11 +445,14 @@ def encode_values(self, pos, runIdx, writer, context: Context, variable, index:
     return wrote
 
   def load_linear(self, writer, context: Context, variable, index):
-    if self.stype == SymbolType.Register:
-      access = f'{self.name}[{index // 32}]' # TODO
+    if context.get_vm().get_lexic().simd_mode:
+      writer(f'{context.get_vm().get_lexic().simd(self.get_fptype(), 16)} {variable}({index});')
     else:
-      access = f'{self.name}[{index} + threadIdx.x]'
-    writer(f'{self.get_fptype()} {variable} = {access};')
+      if self.stype == SymbolType.Register:
+        access = f'{self.name}[{index // 32}]' # TODO
+      else:
+        access = f'{self.name}[{index} + threadIdx.x]'
+      writer(f'{self.get_fptype()} {variable} = {access};')
 
   def load(self, writer, context: Context, variable, index: List[Union[str, int, Immediate, Variable, LeadIndex]], nontemp):
     if self.stype == SymbolType.Data or (not self.obj.is_dense() and not isinstance(self.obj.spp, BoundingBoxSPP)):
@@ -473,7 +478,9 @@ def load(self, writer, context: Context, variable, index: List[Union[str, int, I
       if self.stype == SymbolType.Register or self.stype == SymbolType.Scratch:
         assert len(self.lead_dims) == 1
         idx = index[self.lead_dims[0]]
-        if not idx.is_thread_dependent():
+        if isinstance(idx, (float, int, np.int32)) or not idx.is_thread_dependent():
+          if isinstance(idx, (float, int, np.int32)):
+            idx = Immediate(idx, Datatype.I32)
           # doesn't work
           if isinstance(idx, Variable):
             writevar = idx.write_nonlead()
@@ -490,7 +497,9 @@ def load(self, writer, context: Context, variable, index: List[Union[str, int, I
           access = pre_access
       else:
         access = pre_access
-      if self.stype == SymbolType.Global:
+      if context.get_vm().get_lexic().simd_mode:
+        writer(f'{context.get_vm().get_lexic().simd(self.get_fptype(), 16)} {variable}({access});')
+      elif self.stype == SymbolType.Global:
         writer(f'{self.get_fptype()} {variable};')
         writer(context.get_vm().get_lexic().glb_load(variable, access, nontemp))
       else:
@@ -502,19 +511,26 @@ def store(self, writer, context, variable, index: List[Union[str, int, Immediate
 
     access = self.access(context, index)
 
-    if self.stype == SymbolType.Global:
-      assign = context.get_vm().get_lexic().glb_store(access, variable, nontemp)
+    if context.get_vm().get_lexic().simd_mode:
+      if self.stype == SymbolType.Global:
+        writer(f'{variable}.copy_to({access});')
+      else:
+        writer(f'{variable} = {access};')
     else:
-      assign = f'{access} = {variable};'
-    if self.stype == SymbolType.Register or self.stype == SymbolType.Scratch:
-      assert len(self.lead_dims) == 1
-      if isinstance(index[self.lead_dims[0]], LeadIndex):
-        writer(assign)
+      if self.stype == SymbolType.Global:
+        assign = context.get_vm().get_lexic().glb_store(access, variable, nontemp)
       else:
-        with writer.If(f'{context.get_vm().get_lexic().thread_idx_x} == {index[self.lead_dims[0]]}'):
+        assign = f'{access} = {variable};'
+
+      if self.stype == SymbolType.Register or self.stype == SymbolType.Scratch:
+        assert len(self.lead_dims) == 1
+        if isinstance(index[self.lead_dims[0]], LeadIndex):
           writer(assign)
-    else:
-      writer(assign)
+        else:
+          with writer.If(f'{context.get_vm().get_lexic().thread_idx_x} == {index[self.lead_dims[0]]}'):
+            writer(assign)
+      else:
+        writer(assign)
 
   def add_user(self, user):
     self._users.append(user)

tensorforge/common/vm/lexic/lexic.py

@@ -18,6 +18,7 @@ def __init__(self, underlying_hardware):
     self.block_idx_x = None
     self.stream_type = None
     self.restrict_kw = None
+    self.simd_mode = False
 
   @abstractmethod
   def multifile(self):

tensorforge/common/vm/lexic/sycl_lexic.py

@@ -18,6 +18,8 @@ def __init__(self, backend, underlying_hardware):
     self.stream_type = "sycl::queue"
     self.restrict_kw = "__restrict__"
 
+    self.simd_mode = self._underlying_hardware == 'intel' and self._backend == 'oneapi'
+
   def multifile(self):
     return False
 
@@ -45,7 +47,10 @@ def kernel_definition(self, file, kernel_bounds, base_name, params, precision=No
       localmem = None
 
     if self._underlying_hardware == 'intel' and self._backend == 'oneapi':
-      add_items = '[[intel::reqd_sub_group_size(16)]] [[intel::kernel_args_restrict]]'
+      if self.simd_mode:
+        add_items = '[[intel::sycl_explicit_simd]] [[intel::kernel_args_restrict]]'
+      else:
+        add_items = '[[intel::reqd_sub_group_size(16)]] [[intel::kernel_args_restrict]]'
     else:
       add_items = ''
 
@@ -62,7 +67,10 @@ def sync_block(self):
     return "item.barrier()"
 
   def sync_simd(self):
-    return "item.barrier()" # TODO make better
+    if self.simd_mode:
+      return ""
+    else:
+      return "item.barrier()" # TODO make better
 
   def sync_grid(self):
     raise NotImplementedError() # TODO
@@ -74,8 +82,11 @@ def get_sub_group_id(self, sub_group_size):
   def active_sub_group_mask(self):
     return f'item.get_sub_group()'
 
-  def broadcast(self, variable, lane, block=None, subblock=None):
-    return f'group_broadcast(-1, {variable}, {lane})'
+  def broadcast(self, variable, lane, block=None, subblock=1):
+    if self.simd_mode:
+      return f'{variable}.select<{block}, {subblock}>({lane})'
+    else:
+      return f'group_broadcast(-1, {variable}, {lane})'
 
   def kernel_range_object(self, name, values):
     return f"sycl::range<3> {name} ({values})"
@@ -96,6 +107,9 @@ def get_headers(self):
   def get_fptype(self, fptype, length=1):
     return f'sycl::vec<{fptype}, {length}>'
 
+  def get_simd(self, fptype, size):
+    return f'tensorforge::intel_esimd::simd<{fptype}, {size}>'
+
   def get_operation(self, op: Operation, fptype, value1, value2):
     if op == Operation.COPY:
       return value1

tensorforge/include/tensorforge_device/isycl.h

@@ -0,0 +1,14 @@
+#pragma once
+
+#include <sycl/ext/intel/esimd.hpp>
+#include <sycl/ext/intel/experimental/esimd/tfloat32.hpp>
+#include <sycl/sycl.hpp>
+
+#include "base.h"
+
+namespace tensorforge {
+namespace intel_esimd = sycl::ext::intel::esimd;
+namespace intel_xmx = iesimd::xmx;
+
+using TF32 = sycl::ext::intel::experimental::esimd::tfloat32;
+} // namespace tensorforge