Fix more bugs; make work

Author: davschneller

tensorforge/backend/instructions/compute/multilinear.py

@@ -10,7 +10,7 @@
 from tensorforge.common.basic_types import Datatype
 from tensorforge.backend.writer import Writer
 
-from .primitives import nvidia as nv
+from .primitives import nvidia as nvidia
 from .primitives import amd as amd
 
 class MultilinearInstruction(ComputeInstruction):
@@ -66,10 +66,6 @@ def __init__(self,
 
         self._analyze()
 
-    def _choose_lead_dim(self):
-        self._shm_volume = 0
-        pass
-
     def _analyze(self):
         targetrank = 0
         for i, op in enumerate(self._ops):
@@ -207,7 +203,7 @@ def nonlead_writer(varlist):
         write_loops(self._context, writer, loopstack, nonlead_writer)
 
     def _nonleading_dim_test(self, writer: Writer):
-        can_use = self._context.get_vm().get_hw_descr().vendor == 'amd'
+        can_use = self._context.get_vm().get_hw_descr().vendor in ['amd']
         can_use &= len(self._ops) == 2
 
         if can_use:
@@ -222,6 +218,7 @@ def _nonleading_dim_test(self, writer: Writer):
                 N *= mx - mi
 
             M *= -(-(self._ns[0][1] - self._ns[0][0]) // self._num_threads)
+            Mx = (self._ns[0][1] - self._ns[0][0])
 
             def unwindJ(j):
                 idx = [None]
@@ -236,12 +233,15 @@ def unwindI(i):
                 idx = [LeadIndex(i % size + self._ns[0][0] // self._num_threads, self._num_threads, 1)]
                 return idx
 
+            # TODO: remove
+            kx = self._ks[0][0]
+
             def unwindK(k, full):
                 size = self._ks[0][1] - self._ks[0][0]
                 if full:
                     idx = [k % size + self._ks[0][0]]
                 else:
-                    sizeL = -(-size // self._num_threads)
+                    sizeL = -(-(size + kx) // self._num_threads)
                     idx = [LeadIndex(k % sizeL + self._ks[0][0] // self._num_threads, self._num_threads, 1)]
                 k //= size
                 for mi, mx in self._ks[1:]:
@@ -250,9 +250,6 @@ def unwindK(k, full):
                     k //= size
                 return idx
 
-            # TODO: remove
-            kx = self._ks[0][0]
-
             def unwindOp(i, j, k, opid, full):
                 iidx = unwindI(i)
                 jidx = unwindJ(j)
@@ -301,7 +298,7 @@ def A(writer, var, i, k):
             if self._context.get_vm().get_hw_descr().vendor == 'amd':
                 amd.matmul(writer, C, A, B, M, N, K, kx, self._num_threads, self._dest.datatype, sparse, self._context)
             elif self._context.get_vm().get_hw_descr().vendor == 'nvidia':
-                return nvidia.matmul(writer, C, A, B, M, N, K, kx, self._num_threads, self._dest.datatype, sparse, self._context, 'TODO', 0)
+                return nvidia.matmul(writer, C, A, B, Mx, N, K, kx, self._num_threads, self._dest.datatype, sparse, self._context, 'TODO', 0)
             return True
         return False
 

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

@@ -447,7 +447,7 @@ def threadrange(start, size):
                     B(writer, f'{Breg}_{k//threads}_{jj}', j + jj, k // threads)
                 for jj in range(min(atom.n, N - j), atom.n):
                     writer(f'{atom.d.ctype()} {Breg}_{k//threads}_{jj}{"{}"};')
-            for ix in range(0, M):
+            for i in range(0, M, threads):
                 with writer.AnonymousScope():
                     writer(f'{atom.d.ctype()} {Creg}[{cregs}][{threads // atom.m}]{"{}"};')
                     for k in range(0, K, threads):
@@ -471,40 +471,40 @@ def threadrange(start, size):
                                             writer(f'{atom.d.ctype()} {Breg2}_{kkk + jj * kregs} = {shmptr}[{boffs} + (threadIdx.x % {ktile}) + (threadIdx.x / {ktile} + {jj * ntile}) * {atom.k} + {kkk * ktile}];')
 
                                     for kkk in range(0, min(atom.k, K - k - kk)):
-                                        A(writer, f'{Areg}_{kkk}', ix, k + kk + kkk)
+                                        A(writer, f'{Areg}_{kkk}', i // threads, k + kk + kkk)
                                     for kkk in range(min(atom.k, K - k - kk), atom.k):
                                         writer(f'{atom.d.ctype()} {Areg}_{kkk}{"{}"};')
 
-                                    for iix in range(0, threads, atom.m):
+                                    for ii in range(0, min(threads, M - i), atom.m):
                                         with writer.AnonymousScope():
                                             writer('__syncwarp();')
-                                            with threadrange(iix, atom.m):
+                                            with threadrange(ii, atom.m):
                                                 for kkk in range(0, atom.k):
-                                                    writer(f'{shmptr}[{aoffs} + (threadIdx.x - {iix}) % {atom.m} + {kkk * atom.m}] = {Areg}_{kkk};')
+                                                    writer(f'{shmptr}[{aoffs} + (threadIdx.x - {ii}) % {atom.m} + {kkk * atom.m}] = {Areg}_{kkk};')
                                             writer('__syncwarp();')
 
                                             for kk in range(0, kregs):
-                                                for ii in range(0, mregs):
-                                                    writer(f'{atom.d.ctype()} {Areg2}_{ii + kk * mregs} = {shmptr}[{aoffs} + (threadIdx.x / {ktile}) + (threadIdx.x % {ktile} + {kk * ktile}) * {atom.m} + {ii * mtile}];')
+                                                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}];')
 
-                                            atom.generate(writer, ctx, [f'{Areg2}_{i}' for i in range (aregs)], [f'{Breg2}_{i}' for i in range (bregs)], [f'{Creg}[{i}][{iix // atom.m}]' for i in range (cregs)])
+                                            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)])
 
                     for jj in range(0, atom.n):
                         writer(f'{atom.d.ctype()} {Creg}_{jj}{"{}"};')
 
-                    for i in range(0, threads, atom.m):
+                    for ii in range(0, threads, atom.m):
                         with writer.AnonymousScope():
                             for jj in range(0, nregs * 2):
-                                for ii in range(0, mregs):
-                                    writer(f'{shmptr}[{coffs} + threadIdx.x * 2 + {ii} + {jj * 64}] = {Creg}[{ii + mregs * jj}][{i // atom.m}];')
+                                for iii in range(0, mregs):
+                                    writer(f'{shmptr}[{coffs} + threadIdx.x * 2 + {iii} + {jj * 64}] = {Creg}[{iii + mregs * jj}][{ii // atom.m}];')
 
                             writer('__syncwarp();')
-                            with threadrange(i, atom.m):
+                            with threadrange(ii, atom.m):
                                 for jj in range(0, atom.n):
                                     writer(f'{Creg}_{jj} = {shmptr}[{coffs} + (threadIdx.x % {atom.m}) * {atom.n} + {jj}];')
                             writer('__syncwarp();')
 
                     for jj in range(0, min(atom.n, N - j)):
-                        C(writer, f'{Creg}_{jj}', ix, j + jj)
+                        C(writer, f'{Creg}_{jj}', i // threads, j + jj)
 
     return True