Start optimizing the memory loading

Author: davschneller

tensorforge/backend/instructions/builders/multilinear_builder.py

@@ -42,7 +42,7 @@ def __init__(self,
     self._dest_regs = None
 
     self._use_registers_always = self._context.get_vm().get_hw_descr().vendor in ['amd']
-    self._preload_registers = False #self._context.get_vm().get_hw_descr().vendor in ['amd']
+    self._preload_registers = self._context.get_vm().get_hw_descr().vendor in ['amd']
     self._deferred_stores = {}
     self._temporaries = {}
 
@@ -112,10 +112,10 @@ def _make_load_op(self, i):
           self._loaders_cache[self._mem_regions[i]] = load_op
           self._instructions.append(load_op)
         else:
-          if self._preload_registers and self._ops[i].symbol.obj.is_dense() and not (self._ops[i].symbol in self._loaders_cache.keys()):
+          if self._preload_registers and self._ops[i].symbol.obj.is_dense():
             # only register-preload dense matrices for now
             self._mem_regions[i], load_op = self._make_loader_and_symbol_reg(self._ops[i].symbol, is_transpose=self._descr.permute[i])
-            self._loaders_cache[self._ops[i].symbol] = load_op
+            self._deferred_stores[self._ops[i].symbol.name] = self._mem_regions[i].symbol, self._mem_regions[i].symbol
             self._instructions.append(load_op)
           else:
             # Note: operand will reside in glb. mem for gemm operation
@@ -204,16 +204,19 @@ def _alloc_register_array(self):
 
     # TODO: shrink to enumerate(self._dest_obj.bbox.sizes())
     if self._add:
-      sizes = self._get_target_symbol().data_view._bbox.sizes()
+      bbox = self._get_target_symbol().data_view._bbox
     else:
-      sizes = self._dest_obj.bbox.sizes()
+      bbox = self._dest_obj.bbox
 
-    for d, dim in enumerate(sizes):
+    for d in range(bbox.rank()):
+      dim = bbox.size(d)
       if d not in lead_dim or threads == 0:
         regsize *= dim
       else:
-        regsize *= (dim + threads - 1) // threads
-        threads //= dim
+        r_start = bbox.lower()[d] // threads
+        r_end = (bbox.upper()[d] + threads - 1) // threads
+        regsize *= r_end - r_start
+        threads //= dim # TODO?
     name = self._name_registers()
     regmem = RegMemObject(name, regsize)
     registers = Symbol(name=name, stype=SymbolType.Register, obj=regmem)

tensorforge/backend/instructions/memory/load.py

@@ -16,7 +16,10 @@ def _find_next_coprime(number, conumber):
     if math.gcd(i, conumber) == 1:
       return i
 
-class GlbToShrLoader(AbstractShrMemWrite):
+class LoadInstruction:
+  pass
+
+class GlbToShrLoader(AbstractShrMemWrite, LoadInstruction):
   def __init__(self, **kwargs):
     super(GlbToShrLoader, self).__init__(kwargs['context'])
     self._dest = kwargs['dest']
@@ -51,7 +54,7 @@ def __init__(self, **kwargs):
     self._shr_mem.add_user(self)
     self._is_ready: bool = False
 
-    self._use_cuda_memcpy = False #self._context.get_vm().get_hw_descr().vendor == 'nvidia'
+    self._use_cuda_memcpy = self._context.get_vm().get_hw_descr().vendor == 'nvidia'
 
     if self._permute is None:
       self._permute = [i for i in range(len(self._src.obj.shape))]
@@ -178,6 +181,7 @@ def inner(indices):
         loop.__exit__(None, None, None)
 
       if self._use_cuda_memcpy:
+        writer(f'__syncwarp();')
         writer(f'{self._pipeline}.producer_commit();')
 
     #if False:
@@ -271,7 +275,7 @@ def get_headers(self) -> List[str]:
   def __str__(self):
     return f'{self._dest.name} = load{{g>s}}({self._src.name}[{", ".join(str(p) for p in self._permute)}])'
 
-class GlbToRegLoader(MemoryInstruction):
+class GlbToRegLoader(MemoryInstruction, LoadInstruction):
   def __init__(self,
                context: Context,
                src: Symbol,
@@ -327,3 +331,18 @@ def inner(indices):
 
   def __str__(self) -> str:
     return f'{self._dest.name} = load{{g>r}}({self._src.name});'
+
+class LoadWait(MemoryInstruction, LoadInstruction):
+  def __init__(self, instr):
+    super(LoadWait, self).__init__(instr._context)
+    self._instr = instr
+    self._is_ready = True
+
+  def gen_code_inner(self, writer: Writer) -> None:
+    if isinstance(self._instr, GlbToShrLoader):
+      if self._instr._use_cuda_memcpy:
+        writer(f'{self._instr._pipeline}.consumer_wait();')
+        writer(f'{self._instr._pipeline}.consumer_release();')
+
+  def __str__(self) -> str:
+    return f'wait({self._instr});'

tensorforge/backend/opt/memmove.py

@@ -0,0 +1,33 @@
+from typing import List
+from .abstract import AbstractTransformer, Context, AbstractInstruction
+from tensorforge.backend.instructions.compute import ComputeInstruction
+from tensorforge.backend.instructions.memory import AbstractShrMemWrite, MemoryInstruction
+from tensorforge.backend.instructions.memory.load import LoadInstruction, LoadWait
+from tensorforge.backend.instructions.ptr_manip import GetElementPtr
+from tensorforge.backend.symbol import SymbolType
+
+class MoveLoads(AbstractTransformer):
+  def __init__(self,
+               context: Context,
+               instructions: List[AbstractInstruction]):
+    super(MoveLoads, self).__init__(context, instructions)
+
+  def apply(self) -> None:
+    instrsOut = []
+    stored = []
+    for instr in reversed(self._instrs):
+        if not isinstance(instr, ComputeInstruction):
+            while len(stored) > 0:
+                delayed = stored.pop()
+                instrsOut += [delayed]
+        if isinstance(instr, LoadInstruction):
+            instrsOut += [LoadWait(instr)]
+            while len(stored) > 0:
+                delayed = stored.pop()
+                instrsOut += [delayed]
+            stored.append(instr)
+        else:
+            instrsOut += [instr]
+    instrsOut += stored[::-1]
+
+    self._instrs = instrsOut[::-1]

tensorforge/backend/opt/optimizer.py

@@ -8,7 +8,7 @@
 from .shr_mem_analyzer import ShrMemOpt
 from .sync_block import SyncThreadsOpt
 from .remove_redundancy import RemoveRedundancyOpt
-
+from .memmove import MoveLoads
 
 class OptimizationStage:
   def __init__(self,
@@ -24,6 +24,10 @@ def __init__(self,
     self._num_threads = num_threads
 
   def optimize(self):
+    opt = MoveLoads(self._context, self._instrs)
+    opt.apply()
+    self._instrs = opt.get_instructions()
+
     opt = LivenessAnalysis(self._context, self._instrs)
     opt.apply()
     live_map: Dict[int, Set[Symbol]] = opt.get_live_map()

tensorforge/generators/generator.py

@@ -181,6 +181,8 @@ def _generate_kernel(self):
 
           def generate_inner():
             with writer.If(f'{self._get_flag_guard(writer, i)}'):
+              if self._context.get_vm().get_hw_descr().vendor == 'nvidia':
+                writer(f'cuda::pipeline<cuda::thread_scope_thread> pipeline = cuda::make_pipeline();')
               for instruction in section.ir:
                 if instruction.is_ready():
                   instruction.gen_code(writer)
@@ -430,7 +432,7 @@ def _populate_global_scope(self):
     for matrix in self._matrix_list:
       if matrix not in self._tmp_list:
         # temporary. For now, take only the selector matrices
-        if matrix.has_values() and len(matrix.get_values()) < 16:
+        if matrix.has_values() and len(matrix.get_values()) < 16 and False:
           stype = SymbolType.Data
         elif matrix.addressing == Addressing.SCALAR:
           stype = SymbolType.Scalar