Start implementing float4 loads

Author: davschneller

tensorforge/backend/instructions/memory/load.py

@@ -184,25 +184,23 @@ def inner(indices):
         writer(f'__syncwarp();')
         writer(f'{self._pipeline}.producer_commit();')
 
-    #if False:
-    #  writer('cooperative_groups::wait(cooperative_groups::this_thread_block());')
-
   def _write_datatransfer(self, writer, src_offset, dst_offset, index, length, nontemporal, linscale=None):
-    if not self._use_cuda_memcpy or linscale is not None or True:
-      pos = 0
-      for vecsize in [1]:
-        if src_offset % vecsize == 0:
-          num_hops = ((length - pos * self._num_threads) // (self._num_threads * vecsize)) * vecsize
-          self._write_hop(writer, src_offset, dst_offset, index, pos, pos + num_hops, vecsize, nontemporal, linscale)
-          pos += num_hops
-      rest = length % self._num_threads
-      if rest > 0:
-        with writer.If(f'{self._linear_idx()} < {rest}'):
-          self._write_hop(writer, src_offset, dst_offset, index, pos, pos+1, 1, nontemporal, linscale)
+    pos = 0
+
+    if self._use_cuda_memcpy:
+      granularities = [1]
     else:
-      dest_access_index = self._dest.access_address(self._context, index)
-      src_access_index = self._src.access_address(self._context, index)
-      writer(f'cuda::memcpy_async(cooperative_groups::this_thread_block(), &{self._dest.name}[{dst_offset} + {dest_access_index}], &{self._src.name}[{src_offset} + {src_access_index}], cuda::aligned_size_t<{self._dest.get_fptype().size()}>({length * self._dest.get_fptype().size()}), {self._pipeline});')
+      granularities = [4, 2, 1]
+
+    for vecsize in granularities:
+      if src_offset % vecsize == 0:
+        num_hops = ((length - pos * self._num_threads) // (self._num_threads * vecsize)) * vecsize
+        self._write_hop(writer, src_offset, dst_offset, index, pos, pos + num_hops, vecsize, nontemporal, linscale)
+        pos += num_hops
+    rest = length % self._num_threads
+    if rest > 0:
+      with writer.If(f'{self._linear_idx()} < {rest}'):
+        self._write_hop(writer, src_offset, dst_offset, index, pos, pos+1, 1, nontemporal, linscale)
 
   def _write_hop(self, writer, src_offset, dst_offset, index, start, end, increment, nontemporal, linscale):
     if end > start:
@@ -325,9 +323,66 @@ def gen_code_inner(self, writer: Writer) -> None:
       for dim in src_bbox.sizes():
         total_size *= dim
 
-      for i in range(0, total_size, self._num_threads):
-        self._src.load_linear(writer, self._context, f'v{i}', i)
-        self._dest.store_linear(writer, self._context, f'v{i}', i)
+      start = 0
+      for g in [1]: #[4, 2, 1]:
+        granularity = self._num_threads * g
+        for i in range(start, total_size, granularity):
+          self._src.load_linear(writer, self._context, f'v{i}', i, g)
+          self._dest.store_linear(writer, self._context, f'v{i}', i, g)
+
+        start = (total_size // granularity) * granularity
+
+    elif self._context.get_vm().get_hw_descr().vendor in ['amd']:
+
+      # float4 load
+
+      # for now: use  0 1 2 3, transpose4x4
+
+      # TODO: sort into 4x4x4 blocks
+
+      lead_size = src_bbox.size(0)
+      lead_count = (lead_size + self._num_threads - 1) // self._num_threads
+
+      total_count = lead_count
+      for dim in src_bbox.sizes()[1:]:
+        total_count *= dim
+
+      start = 0
+
+      prec = 'float'
+
+      for g in [4, 2, 1]: # [4, 3, 2, 1]
+        # 4x4
+        # writer(f'const auto f{g}idx = (threadIdx.x % {g}) * {self._num_threads} + (threadIdx.x / {g}) * {g};')
+
+        writer(f'const auto f{g}idx = ((threadIdx.x / {16 // g}) % {g}) * {self._num_threads} + (threadIdx.x % {16 // g}) * {g} + (threadIdx.x / 16) * 16;')
+
+        total_count_g = (total_count // g) * g
+        for i in range(start, total_count_g, g):
+          sidx = i // lead_count
+          ridx = i % lead_count
+          index = sidx * lead_size + ridx * self._num_threads
+          writer(f'const auto v{i} = *(tensorforge::VectorT<{prec}, {g}>*)&{self._src.name}[{index} + f{g}idx];')
+
+          args2 = ', '.join(f'v{i}[{k}]' for k in range(g))
+
+          for k in range(g):
+            writer(f'{prec} v{i}w{k} = 0;')
+
+          args1 = ', '.join(f'v{i}w{k}' for k in range(g))
+
+          if g == 4:
+            writer(f'tensorforge::transpose16x4({args1}, {args2});')
+          if g == 2:
+            writer(f'tensorforge::transpose16x2({args1}, {args2});')
+          if g == 1:
+            writer(f'{args1} = {args2};')
+
+          # TODO: generalize
+          for k in range(g):
+            writer(f'{self._dest.name}[{i + k}] = v{i}w{k};')
+
+        start = total_count_g
 
     else:
       loops = []

tensorforge/backend/instructions/memory/store.py

@@ -195,7 +195,10 @@ def inner(indices):
         else:
           self._dest.store(writer, self._context, '0', indices, allow_nontemporal)
 
-      write_loops(self._context, writer, loops, inner)
+      if not any(manual) and self._context.get_vm().get_hw_descr().vendor in ['amd'] and False:
+        pass
+      else:
+        write_loops(self._context, writer, loops, inner)
 
   def __str__(self) -> str:
     return f'{self._dest.name} = store{{r>g}}({self._src.name});'

tensorforge/backend/symbol.py

@@ -444,17 +444,21 @@ def encode_values(self, pos, runIdx, writer, context: Context, variable, index:
               wrote |= self.encode_values(pos + 1, runIdx, writer, context, variable, index, nontemp, leadidx)
     return wrote
 
-  def load_linear(self, writer, context: Context, variable, index):
+  def load_linear(self, writer, context: Context, variable, index, vec = 1):
     if context.get_vm().get_lexic().simd_mode:
       writer(f'{context.get_vm().get_lexic().simd(self.get_fptype(), self.num_threads)} {variable}({index});')
     else:
       if self.stype == SymbolType.Register:
         access = f'{self.name}[{index // self.num_threads}]'
       else:
-        access = f'{self.name}[{index} + threadIdx.x]'
-      writer(f'{self.get_fptype()} {variable} = {access};')
+        access = f'{self.name}[{index} + threadIdx.x * {vec}]'
 
-  def store_linear(self, writer, context: Context, variable, index):
+      if vec == 1:
+        writer(f'{self.get_fptype()} {variable} = {access};')
+      else:
+        writer(f'tensorforge::VectorT<{self.get_fptype()}, {vec}> {variable} = *(tensorforge::VectorT<{self.get_fptype()}, {vec}>*)&{access};')
+
+  def store_linear(self, writer, context: Context, variable, index, vec = 1):
     if context.get_vm().get_lexic().simd_mode:
       pass
       # TODO:
@@ -463,8 +467,13 @@ def store_linear(self, writer, context: Context, variable, index):
       if self.stype == SymbolType.Register:
         access = f'{self.name}[{index // self.num_threads}]'
       else:
-        access = f'{self.name}[{index} + threadIdx.x]'
-      writer(f'{access} = {variable};')
+        access = f'{self.name}[{index} + threadIdx.x * {vec}]'
+
+      if vec == 1:
+        writer(f'{access} = {variable};')
+      else:
+        convert = f'*(tensorforge::VectorT<{self.get_fptype()}, {vec}>*)&'
+        writer(f'{convert}{access} = {convert}{variable};')
 
   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)):

tensorforge/common/vm/lexic/hip_lexic.py

@@ -69,7 +69,7 @@ def get_headers(self):
     return ["hip/hip_runtime.h", "tensorforge_device/hip.h"]
 
   def get_fptype(self, fptype, length=1):
-    return f'HIP_vector_type<{fptype}, {length}>'
+    return f'tensorforge::VectorT<{fptype}, {length}>'
 
   def glb_store(self, lhs, rhs, nontemporal=False):
     if nontemporal and self._underlying_hardware == 'amd':

tensorforge/generators/generator.py

@@ -191,6 +191,8 @@ def generate_inner():
 
           if self._persistent_threading:
             # TODO: OMP target
+            # TODO: maybe iterate over adjacent elements? (for indirect pointers)
+
             offset = []
             idx = i - 1
             for ssection in reversed(self._sections[:i]):
@@ -299,12 +301,6 @@ def _emit_global_ir(self):
     self._section.shr_mem_obj = shmbuilder.get_resultant_obj()
     self._section.global_ir.extend(shmbuilder.get_instructions())
 
-    builder = GetElementPtrBuilder(self._context, self._scopes)
-    for symbol in self._scopes.get_global_scope().values():
-      if symbol.obj.addressing == Addressing.SCALAR or (symbol.obj.addressing == Addressing.NONE and symbol.stype == SymbolType.Data):
-        builder.build(symbol)
-        self._section.global_ir.extend(builder.get_instructions())
-
     # load globals to shared memory (if requested)
     if self._preload_globals:
       load_ir = []
@@ -336,6 +332,17 @@ def _emit_global_ir(self):
         self._scopes.remove_scope()
         self._preload_globals = False
 
+    builder = GetElementPtrBuilder(self._context, self._scopes)
+    for symbol in self._scopes.get_global_scope().values():
+      if symbol.obj.addressing == Addressing.SCALAR or (symbol.obj.addressing == Addressing.NONE and (symbol.stype == SymbolType.Data or not self._preload_globals)):
+        builder.build(symbol)
+        self._section.global_ir.extend(builder.get_instructions())
+
+    # pipelines
+    for symbol in self._scopes.get_global_scope().values():
+      if symbol.obj.addressing in [Addressing.STRIDED, Addressing.PTR_BASED]:
+        pass
+
     if not self._preload_globals:
       if last_barrier:
         self._section.global_ir.append(SyncGrid(self._context))
@@ -349,8 +356,8 @@ def _emit_ir(self, descr_list):
     builder = GetElementPtrBuilder(self._context, self._scopes)
     self._scopes.add_scope()
     for symbol in self._scopes.get_global_scope().values():
-      firstptr = symbol.obj.addressing == Addressing.SCALAR or (symbol.obj.addressing == Addressing.NONE and symbol.stype == SymbolType.Data)
-      if not firstptr and not (self._preload_globals and symbol.obj.addressing == Addressing.NONE):
+      firstptr = symbol.obj.addressing == Addressing.SCALAR or symbol.obj.addressing == Addressing.NONE
+      if not firstptr:
         builder.build(symbol)
         self._section.ir.extend(builder.get_instructions())
 

tensorforge/include/tensorforge_device/hip.h

@@ -721,6 +721,24 @@ transpose16x16b32(T &w1, T &w2, T &w3, T &w4, T &w5, T &w6, T &w7, T &w8, T &w9,
   w16 = dppUpdate<0x128, 0b1111, 0b0011, true>(u8, u16);
 }
 
+template <typename T>
+__device__ __forceinline__ void transpose16x2(T &w1, T &w2, T v1, T v2) {
+  w1 = dppUpdate<0x128, 0b1111, 0b1100, true>(v2, v1);
+  w2 = dppUpdate<0x128, 0b1111, 0b0011, true>(v1, v2);
+}
+
+template <typename T>
+__device__ __forceinline__ void transpose16x4(T &w1, T &w2, T &w3, T &w4, T v1,
+                                              T v2, T v3, T v4) {
+  const T u1 = dppUpdate<0x124, 0b1111, 0b1010, true>(v2, v1);
+  const T u2 = dppUpdate<0x12c, 0b1111, 0b0101, true>(v1, v2);
+  const T u3 = dppUpdate<0x124, 0b1111, 0b1010, true>(v4, v3);
+  const T u4 = dppUpdate<0x12c, 0b1111, 0b0101, true>(v3, v4);
+
+  transpose16x2(w1, w3, u1, u3);
+  transpose16x2(w2, w4, u2, u4);
+}
+
 #define CM4STR(p1, p2, p3, p4, c, a, b)                                        \
   "v_cndmask_b32_dpp " c ", " a ", " b CMVCC                                   \
   " quad_perm:[" STR(p1) "," STR(p2) "," STR(p3) "," STR(                      \