More bugfixes; kernel chaining

Author: davschneller

tensorforge/backend/instructions/builders/multilinear_builder.py

@@ -3,7 +3,7 @@
 from tensorforge.backend.scopes import Scopes
 from tensorforge.backend.symbol import Symbol, SymbolType, SymbolView
 from tensorforge.backend.instructions.allocate import RegisterAlloc
-from tensorforge.backend.instructions.memory.load import GlbToShrLoader
+from tensorforge.backend.instructions.memory.load import GlbToShrLoader, GlbToRegLoader
 from tensorforge.backend.instructions.clear_registers import ClearRegisters
 from tensorforge.backend.instructions.memory.store import StoreRegToGlb, StoreRegToShr, StoreRegToReg
 from tensorforge.backend.instructions.sync_block import SyncThreads
@@ -19,8 +19,6 @@
 
 
 class MultilinearBuilder(AbstractBuilder):
-  GemmClass = None
-
   def __init__(self,
                context: Context,
                scopes: Scopes,
@@ -44,6 +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._deferred_stores = {}
     self._temporaries = {}
 
@@ -71,7 +70,7 @@ def build(self, ops: List[Symbol], dest_obj: Tensor, descr: MultilinearDescr):
   # TODO: check if we always can allow a direct global memory load
   def _make_load_op(self, i):
 
-    prefer_broadcast = self._context.get_vm().get_hw_descr().vendor == 'amd'
+    prefer_broadcast = self._context.get_vm().get_hw_descr().vendor in ['amd']
 
     has_lead_dim = 0 in self._descr.target[i]
     transpose = self._descr.permute[i] != [j for j in range(len(self._descr.target[i]))]
@@ -113,8 +112,13 @@ def _make_load_op(self, i):
           self._loaders_cache[self._mem_regions[i]] = load_op
           self._instructions.append(load_op)
         else:
-          # Note: operand will reside in glb. mem for gemm operation
-          self._mem_regions[i] = self._ops[i]
+          if self._preload_registers:
+            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._instructions.append(load_op)
+          else:
+            # Note: operand will reside in glb. mem for gemm operation
+            self._mem_regions[i] = self._ops[i]
 
       elif self._ops[i].symbol.stype == SymbolType.SharedMem or self._ops[i].symbol.stype == SymbolType.Register:
         if self._ops[i].symbol in self._loaders_cache.keys():
@@ -147,6 +151,32 @@ def _make_load_op(self, i):
       else:
         raise InternalError(f'gemm-builder: op{i} ({self._ops[i].symbol.name}) must be either in shr or glb mem, given: {self._ops[i].symbol.stype}')
 
+  def _make_loader_and_symbol_reg(self, operand, is_transpose) -> Tuple[Symbol, GlbToRegLoader]:
+    regsize = 1
+    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()):
+      if d not in lead_dim or threads == 0:
+        regsize *= dim
+      else:
+        regsize *= (dim + threads - 1) // threads
+        threads //= dim
+    name = self._name_registers()
+    regmem = RegMemObject(name, regsize)
+    registers = Symbol(name=name, stype=SymbolType.Register, obj=regmem)
+    registers.num_threads = self._num_threads
+    registers.datatype = self._context.fp_type
+    self._scopes.add_symbol(registers)
+    registerAlloc = RegisterAlloc(self._context, registers, regsize, 0.0)
+    self._instructions.append(registerAlloc)
+
+    load_op = GlbToRegLoader(context=self._context,
+                                     dest=registers,
+                                     src=operand,
+                                     num_threads=self._num_threads)
+    return SymbolView(registers), load_op
+
   def _make_loader_and_symbol(self, operand, is_transpose) -> Tuple[Symbol, GlbToShrLoader]:
     shr_mem_region = Symbol(name=self._name_shr_reg(),
                             stype=SymbolType.SharedMem,
@@ -170,7 +200,14 @@ def _alloc_register_array(self):
     regsize = 1
     threads = self._num_threads
     lead_dim = [0] # [t for t in self._descr.target[0] if t >= 0]
-    for d, dim in enumerate(self._dest_obj.bbox.sizes()):
+
+    # TODO: shrink to enumerate(self._dest_obj.bbox.sizes())
+    if self._add:
+      sizes = self._get_target_symbol().data_view._bbox.sizes()
+    else:
+      sizes = self._dest_obj.bbox.sizes()
+
+    for d, dim in enumerate(sizes):
       if d not in lead_dim or threads == 0:
         regsize *= dim
       else:

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

@@ -356,16 +356,16 @@ def reduction(writer: Writer, source, target, operation, blocks):
             var = tempvar
 
 def cdna1(ctx):
-    arch = ctx.get_vm().get_hw_descr().name
+    arch = ctx.get_vm().get_hw_descr().model
     return arch in ('gfx908', 'gfx90a', 'gfx942', 'gfx950')
 
 def cdna2(ctx):
-    arch = ctx.get_vm().get_hw_descr().name
+    arch = ctx.get_vm().get_hw_descr().model
     return arch in ('gfx90a', 'gfx942', 'gfx950')
 
 def amdarch(ctx):
-    archstr = ctx.get_vm().get_hw_descr().name
-    return int(arch[3:], base=16)
+    archstr = ctx.get_vm().get_hw_descr().model
+    return int(archstr[3:], base=16)
 
 def mfma_emu_int8(writer: Writer, C, B, A, c, a, b):
     # cf. the Ozaki II paper
@@ -416,10 +416,10 @@ def mfma_emu_bf16_f32(writer: Writer, C, B, A, c, a, b):
     writer(f'const bfloat16x4 {B3} = bfloat16x4({Br} - {B2});')
     writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4bf16(get_native_vector({A1}), get_native_vector({B1}), {C}, {c}, {a}, {b});')
     writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4bf16(get_native_vector({A1}), get_native_vector({B2}), {C}, {c}, {a}, {b});')
-    writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4bf16(get_native_vector({A1}), get_native_vector({B3}), {C}, {c}, {a}, {b});')
     writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4bf16(get_native_vector({A2}), get_native_vector({B1}), {C}, {c}, {a}, {b});')
-    writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4bf16(get_native_vector({A2}), get_native_vector({B2}), {C}, {c}, {a}, {b});')
+    writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4bf16(get_native_vector({A1}), get_native_vector({B3}), {C}, {c}, {a}, {b});')
     writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4bf16(get_native_vector({A3}), get_native_vector({B1}), {C}, {c}, {a}, {b});')
+    writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4bf16(get_native_vector({A2}), get_native_vector({B2}), {C}, {c}, {a}, {b});')
 
 def matmul32(writer: Writer, C, B, A, M, N, K, kx, threads):
     with writer.AnonymousScope():
@@ -544,8 +544,39 @@ def hfma(writer: Writer, C, A, B, repeat, datatype, threads, ctx):
                 if b is not None:
                     func(writer, c, a, b, j)
 
+def wmma3atom(threads):
+    assert threads == 32
+
+    N = 16
+    M = 16
+    K = 16
+
+    for i in range(N):
+        writer(f'const auto {a}_{i} = tensorforge::broadcast<32, 16, 0>({A}_{i});')
+    for j in range(N):
+        writer(f'const auto {b}_{j} = tensorforge::broadcast<32, 16, 0>({B}_{j});')
+
+    writer(f'tensorforge::transpose16x16({",".join(f"{b}_{i}" for i in range(N))});')
+
+    writer(f'VectorT<short, 16> {a}_p1;')
+    writer(f'VectorT<short, 16> {a}_p2;')
+    writer(f'VectorT<short, 16> {a}_p3;')
+    writer(f'VectorT<short, 16> {b}_p1;')
+    writer(f'VectorT<short, 16> {b}_p2;')
+    writer(f'VectorT<short, 16> {b}_p3;')
+
+    writer(f'VectorT<float, 8> {c}{"{}"};')
+    writer(f'{c} = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32({a}_p1, {b}_p1, {c});')
+    writer(f'{c} = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32({a}_p2, {b}_p1, {c});')
+    writer(f'{c} = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32({a}_p1, {b}_p2, {c});')
+    writer(f'{c} = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32({a}_p3, {b}_p1, {c});')
+    writer(f'{c} = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32({a}_p1, {b}_p3, {c});')
+    writer(f'{c} = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32({a}_p2, {b}_p2, {c});')
+
+    # TODO: gfx1200, f'__builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12'
+
 def matmul(writer, C, A, B, M, N, K, kx, threads, dtype, sparse, ctx):
-    if cdna1(ctx) and not sparse and dtype == Datatype.F32:
+    if amdarch(ctx) >= 0x908 and amdarch(ctx) < 0x1000 and not sparse and dtype == Datatype.F32:
         matmul32(writer, C, A, B, M, N, K, kx, threads)
     else:
         ab = []

tensorforge/backend/symbol.py

@@ -444,7 +444,7 @@ def encode_values(self, pos, runIdx, writer, context: Context, variable, index:
 
   def load_linear(self, writer, context: Context, variable, index):
     if self.stype == SymbolType.Register:
-      access = f'{self.name}[{index}]'
+      access = f'{self.name}[{index // 32}]' # TODO
     else:
       access = f'{self.name}[{index} + threadIdx.x]'
     writer(f'{self.get_fptype()} {variable} = {access};')

tensorforge/frontend/yateto.py

@@ -7,7 +7,7 @@
 from tensorforge.common.matrix.spp import FullSPP, BoundingBoxSPP, ListSPP
 from tensorforge.common.matrix.boundingbox import BoundingBox as BBox
 from tensorforge.generators.generator import Generator as TensorForgeGenerator
-from tensorforge.generators.descriptions import MultilinearDescr, ElementwiseDescr, GridBarrierDescr, GridFenceDescr
+from tensorforge.generators.descriptions import MultilinearDescr, ElementwiseDescr, GridBarrierDescr, GridFenceDescr, RegionDescription
 
 from tensorforge.ir.data.variable import TensorView, TensorAlloc
 from tensorforge.ir.data.variable import TensorData
@@ -27,6 +27,9 @@ def __init__(self, arch):
     self._ir_list = []
     self._tensor_list = {}
 
+    # TODO: maybe remove again
+    self._prefix = ""
+
   def add_operation(self, dest, ops, target, permute, add):
     self._cache_matrices(dest, ops, target, permute)
     can_be_aligned = self._can_be_aligned(dest, ops, target, permute)
@@ -96,9 +99,9 @@ def get_tensor(self, op, can_be_aligned, dims):
     if isinstance(op, (float, int)):
       return SubTensor(tensor = Tensor([], Addressing.SCALAR, data = [op]))
     elif self.is_scalar(op):
-      return SubTensor(self._cache[op.name()])
+      return SubTensor(self._cache[f'{self._prefix}{op.name()}'])
     else:
-      tensor = self._cache[op.name]
+      tensor = self._cache[f'{self._prefix}{op.name}']
       currentPreShape = BBox([s for s, _ in op.eqspp.nnzbounds()], [e+1 for _, e in op.eqspp.nnzbounds()])
 
       tml = op.memoryLayout
@@ -136,10 +139,10 @@ def assigner(pretensor):
       if self.is_scalar(pretensor):
         self.make_tensor(pretensor, False, None)
         indicesIndexed[pretensor.name()] = []
-        subTensor = SubTensor(self._cache[pretensor.name()], BBox([], []))
+        subTensor = SubTensor(self._cache[f'{self._prefix}{pretensor.name()}'], BBox([], []))
       else:
         bbox = BBox([s for s, _ in pretensor.eqspp().nnzbounds()], [e+1 for _, e in pretensor.eqspp().nnzbounds()])
-        subTensor = SubTensor(self._cache[pretensor.name()], bbox)
+        subTensor = SubTensor(self._cache[f'{self._prefix}{pretensor.name()}'], bbox)
       return subTensor, indicesIndexed[pretensor.name()]
 
     for statement in statements:
@@ -203,13 +206,17 @@ def make_tensor(self, op, can_be_aligned, dims):
       entry = self._get_tensorforge_matrix(op)
       entry_name = op.name
 
+    entry_name = f'{self._prefix}{entry_name}'
+
     if not (entry_name in self._cache and entry.is_same(self._cache[entry_name])):
       self._cache[entry_name] = entry
 
   def tensor_ref(self, d):
     name = d['name']
     eqspp = d['spp']
 
+    name = f'{self._prefix}{name}'
+
     assert(name in self._cache)
 
     return SubTensor(self._cache[name], self._cache[name].bbox)
@@ -226,6 +233,8 @@ def tensor_ref_new(self, d):
 
   def add_tensor(self, d):
     name = d['name']
+    name = f'{self._prefix}{name}'
+
     datatype = Datatype.ytt2enum(d['datatype'])
 
     datatype_new = BaseDatatype.ytt2enum(d['datatype'])
@@ -276,16 +285,17 @@ def _cache_matrices(self, dest, ops, target, permute):
 
     if dest.is_temporary: # (dest is never a scalar---for the time being)
       self.make_tensor(dest, can_be_aligned, [i for i in range(len(dest.indices))])
-      self._tmp_matrices[dest.name] = self._cache[dest.name]
+      self._tmp_matrices[f'{self._prefix}{dest.name}'] = self._cache[f'{self._prefix}{dest.name}']
     else:
       self.make_tensor(dest, can_be_aligned, [i for i in range(len(dest.indices))])
 
 
 
   def _add_scalar(self, scalar):
-    tensor = Tensor([], Addressing.SCALAR, alias=scalar.name(), datatype=self._datatype(scalar.datatype))
-    self._tmp_matrices[scalar.name()] = tensor # SubTensor(tensor, tensor.bbox)
-    return self._tmp_matrices[scalar.name()]
+    name = f'{self._prefix}{scalar.name()}'
+    tensor = Tensor([], Addressing.SCALAR, alias=name, datatype=self._datatype(scalar.datatype))
+    self._tmp_matrices[name] = tensor # SubTensor(tensor, tensor.bbox)
+    return self._tmp_matrices[name]
 
   def deduce_addresing(self, term):
     if term.is_compute_constant:
@@ -323,7 +333,7 @@ def _get_tensorforge_matrix(self, tensor):
     return yi.gen_matrix(shape,
                                bboxrange,
                                addressing=addr_mode,
-                               name=tensor.name,
+                               name=f'{self._prefix}{tensor.name}',
                                is_tmp=tensor.is_temporary,
                                permute=None,
                                pattern=pattern,
@@ -345,28 +355,35 @@ def _gen_call_site(self, generator):
       if matrix.is_tmp or matrix.addressing == Addressing.NONE:
         offset_name_map[name] = '0'
       else:
-        offset_name_map[name] = f'extraOffset_{name}'
+        parts = name.split('.')
+        assert len(parts) <= 2
+        varname = f'extraOffset_{parts[-1]}'
+        if len(parts) == 2:
+          offset_name_map[name] = f'{parts[0]}.{varname}'
+        else:
+          offset_name_map[name] = varname
 
     return generator.generate_call_site(mat_name_map,
-                                        offset_name_map,
-                                        'numElements',
-                                        'flags',
-                                        'streamPtr')
+                                        offset_name_map)
 
   def _append_operation(self, op):
     if isinstance(op, (float, int)):
       return Tensor([], Addressing.SCALAR, data = op)
     elif self.is_scalar(op):
-      return self._cache[op.name()]
+      return self._cache[f'{self._prefix}{op.name()}']
     else:
-      return self._cache[op.name]
+      return self._cache[f'{self._prefix}{op.name}']
 
   def switch_region(self, barrier):
     if barrier:
       self._descr_list += [GridBarrierDescr()]
     else:
       self._descr_list += [GridFenceDescr()]
 
+  def set_region_name(self, name):
+    self._prefix = f"{name}."
+    self._descr_list += [RegionDescription(name)]
+
 class TensorForgeWriter:
   def __init__(self, tensorforge_generator, headers):
     self._headers = list(headers) + list(tensorforge_generator.get_helper_headers())
@@ -410,6 +427,9 @@ def region_switch(self, barrier):
     self.generator.switch_region(barrier)
     return 0
 
+  def set_region_name(self, name):
+    self.generator.set_region_name(name)
+
   def add_operation(self, description):
     return self.generator.add_operation_new(description)
 

tensorforge/generators/descriptions.py

@@ -189,3 +189,16 @@ def __str__(self):
 
   def trueBarrier(self):
     return True
+
+class RegionDescription(OperationDescription):
+  def __init__(self, name):
+    self.name = name
+
+  def matrix_list(self):
+    return []
+
+  def get_num_threads(self, ctx):
+    return 32, 32
+
+  def __str__(self):
+    return f'region "{self.name}"'