Fix more long-kernel bugs

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._preload_registers = False #self._context.get_vm().get_hw_descr().vendor in ['amd']
     self._deferred_stores = {}
     self._temporaries = {}
 
@@ -112,9 +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:
+          if self._preload_registers and self._ops[i].symbol.obj.is_dense() and not (self._ops[i].symbol in self._loaders_cache.keys()):
+            # 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._mem_regions[i]] = load_op
+            self._loaders_cache[self._ops[i].symbol] = load_op
             self._instructions.append(load_op)
           else:
             # Note: operand will reside in glb. mem for gemm operation
@@ -156,7 +157,7 @@ def _make_loader_and_symbol_reg(self, operand, is_transpose) -> Tuple[Symbol, Gl
     threads = self._num_threads
     lead_dim = [0] # [t for t in self._descr.target[0] if t >= 0]
 
-    for d, dim in enumerate(operand.bbox.sizes()):
+    for d, dim in enumerate(operand.data_view._bbox.sizes()):
       if d not in lead_dim or threads == 0:
         regsize *= dim
       else:
@@ -251,7 +252,8 @@ def _make_store(self):
         #                                        dest=dest_symbol,
         #                                        shr_mem=self._shr_mem,
         #                                        num_threads=self._num_threads))
-        pass # see note below
+        # see note below (but update to the new temp regs)
+        self._deferred_stores[dest_symbol.name] = (self._temp_regs, dest_symbol)
       elif dest_symbol.stype == SymbolType.Global:
         if self._use_registers_always:
           self._deferred_stores[dest_symbol.name] = (self._temp_regs, dest_symbol)

tensorforge/backend/instructions/compute/multilinear.py

@@ -10,6 +10,8 @@
 from tensorforge.common.basic_types import Datatype
 from tensorforge.backend.writer import Writer
 
+from tensorforge.common.matrix.tensor import Tensor
+
 from .primitives import nvidia as nvidia
 from .primitives import amd as amd
 
@@ -80,8 +82,11 @@ def _analyze(self):
             opdim = [''] * op.bbox.rank()
             for j in range(op.bbox.rank()):
                 # TODO: check adding the data_view box here again
-                lower = op.bbox.lower()[j] #+ op.symbol.data_view._bbox.lower()[j]
-                upper = op.bbox.upper()[j] #+ op.symbol.data_view._bbox.lower()[j]
+                lower = op.bbox.lower()[j] + op.symbol.data_view._bbox.lower()[j]
+                upper = op.bbox.upper()[j] + op.symbol.data_view._bbox.lower()[j]
+                if op.symbol.obj and isinstance(op.symbol.obj, Tensor):
+                    lower -= op.symbol.obj.get_bbox().lower()[j]
+                    upper -= op.symbol.obj.get_bbox().lower()[j]
                 #if self._target[i][j] != 0:
                 #    lower -= op.offset[j]
                 #    upper -= op.offset[j]

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

@@ -486,13 +486,16 @@ def threadrange(start, size):
                                         with writer.AnonymousScope():
                                             writer('__syncwarp();')
                                             with threadrange(ii, atom.m):
-                                                for kkk in range(0, atom.k):
-                                                    writer(f'{shmptr}[{aoffs} + (threadIdx.x - {ii}) % {atom.m} + {kkk * atom.m}] = {Areg}_{kkk};')
+                                                # for kkk in range(0, atom.k):
+                                                #     writer(f'{shmptr}[{aoffs} + (threadIdx.x - {ii}) % {atom.m} + {kkk * atom.m}] = {Areg}_{kkk};')
+                                                for kkk in range(0, atom.k, ktile):
+                                                    writer(f'*(float4*)&{shmptr}[{aoffs} + ((threadIdx.x - {ii}) % {atom.m}) * {ktile} + {kkk * atom.m}] = make_float4({Areg}_{kkk}, {Areg}_{kkk+1}, {Areg}_{kkk+2}, {Areg}_{kkk+3});')
                                             writer('__syncwarp();')
 
                                             for kk in range(0, kregs):
                                                 for iii in range(0, mregs):
-                                                    writer(f'{atom.d.ctype()} {Areg2}_{iii + kk * mregs} = {shmptr}[{aoffs} + (threadIdx.x / {ktile}) + (threadIdx.x % {ktile} + {kk * ktile}) * {atom.m} + {iii * mtile}];')
+                                                    #writer(f'{atom.d.ctype()} {Areg2}_{iii + kk * mregs} = {shmptr}[{aoffs} + (threadIdx.x / {ktile}) + (threadIdx.x % {ktile} + {kk * ktile}) * {atom.m} + {iii * mtile}];')
+                                                    writer(f'{atom.d.ctype()} {Areg2}_{iii + kk * mregs} = {shmptr}[{aoffs} + threadIdx.x + {(iii + kk * mregs) * 32}];')
 
                                             atom.generate(writer, ctx, [f'{Areg2}_{i}' for i in range (aregs)], [f'{Breg2}_{i}' for i in range (bregs)], [f'{Creg}[{i}][{ii // atom.m}]' for i in range (cregs)])
 

tensorforge/backend/instructions/memory/store.py

@@ -3,7 +3,7 @@
 from tensorforge.common.matrix.tensor import Tensor
 from tensorforge.common.matrix.boundingbox import BoundingBox
 from tensorforge.backend.data_types import RegMemObject
-from tensorforge.backend.symbol import Symbol, SymbolType, DataView, LeadIndex, write_loops, LeadLoop, Loop
+from tensorforge.backend.symbol import Symbol, SymbolType, DataView, LeadIndex, write_loops, LeadLoop, Loop, Immediate
 from tensorforge.common.exceptions import InternalError
 from tensorforge.backend.writer import Writer
 from . import AbstractShrMemWrite, MemoryInstruction
@@ -175,15 +175,25 @@ def gen_code(self, writer: Writer) -> None:
 
     writer(f'// {self}')
     src_bbox = self._src.data_view.get_bbox()
+    dest_bbox = self._dest.data_view.get_bbox()
     with writer.Scope():
+      manual = [False]
       loops = []
       loops += [LeadLoop('i0', src_bbox.lower()[0], src_bbox.upper()[0], self._num_threads, 1)]
       for i in range(1, src_bbox.rank()):
-        loops += [Loop(f'i{i}', src_bbox.lower()[i], src_bbox.upper()[i], 1)]
+        unroll = (src_bbox.lower()[i], src_bbox.upper()[i]) != (dest_bbox.lower()[i], dest_bbox.upper()[i])
+        lower = min(src_bbox.lower()[i], dest_bbox.lower()[i])
+        upper = max(src_bbox.upper()[i], dest_bbox.upper()[i])
+        loops += [Loop(f'i{i}', lower, upper, 1, unroll)]
+        manual += [unroll]
 
       def inner(indices):
-        self._src.load(writer, self._context, 'value', indices, False)
-        self._dest.store(writer, self._context, 'value', indices, allow_nontemporal)
+        needsLoad = all(not isinstance(index, Immediate) or (src_bbox.lower()[i] <= index._value and src_bbox.upper()[i] > index._value) for i,index in enumerate(indices))
+        if needsLoad:
+          self._src.load(writer, self._context, 'value', indices, False)
+          self._dest.store(writer, self._context, 'value', indices, allow_nontemporal)
+        else:
+          self._dest.store(writer, self._context, '0', indices, allow_nontemporal)
 
       write_loops(self._context, writer, loops, inner)
 

tensorforge/generators/generator.py

@@ -360,9 +360,13 @@ def _emit_ir(self, descr_list):
                           self._num_threads)
     # builder.build_prologue()
 
+    def get_symbol_view(op):
+      symbol = self._scopes.get_symbol(op.tensor)
+      return SymbolView(symbol, op.bbox, op.offset)
+
     for gemm_descr in descr_list:
       if isinstance(gemm_descr, MultilinearDescr):
-        builder.build(ops=[SymbolView(self._scopes.get_symbol(op.tensor), op.bbox, op.offset) for op in gemm_descr.ops],
+        builder.build(ops=[get_symbol_view(op) for op in gemm_descr.ops],
                         dest_obj=gemm_descr.dest,
                         descr=gemm_descr)
         self._section.ir.extend(builder.get_instructions())