Fix batch replacement in pulp gemm

Author: lukamac

Deeploy/Targets/PULPOpen/TileConstraints/GEMMTileConstraint.py

@@ -24,6 +24,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 from typing import Dict, List, Tuple
 
 from Deeploy.AbstractDataTypes import PointerClass
@@ -135,25 +136,21 @@ def serializeTilingSolution(
 
         # Every output is constructed by a pair of inputs. Reconstruct this pair.
         for cube in outputCubes:
+            MOffset, OOffset = cube.offset[-2:]
+            MSize, OSize = cube.dims[-2:]
 
-            BSize = 1
-            BOffset = 0
-            BatchSize = 1
-            BatchOffset = 0
-
-            if len(cube.offset) == 2:
-                (MOffset, OOffset) = cube.offset
-                (MSize, OSize) = cube.dims
-            elif len(cube.offset) == 3:
-                (BatchOffset, MOffset, OOffset) = cube.offset
-                (BatchSize, MSize, OSize) = cube.dims
+            if len(cube.offset) > 2:
+                BatchSize = math.prod(cube.dims[:-2])
+
+                # Check that we don't tile upper dimensions
+                if len(cube.offset) > 3:
+                    assert all(off == 0 for off in cube.offset[:-3])
             else:
-                (BatchOffset, BOffset, MOffset, OOffset) = cube.offset
-                (BatchSize, BSize, MSize, OSize) = cube.dims
+                BatchSize = 1
 
             replacements["M"].append(MSize)
             replacements["O"].append(OSize)
-            replacements["batch"].append(BSize)
+            replacements["batch"].append(BatchSize)
 
             if transA == 0:
                 AMatrixOffsets = (MOffset, NOffset)
@@ -162,44 +159,26 @@ def serializeTilingSolution(
                 AMatrixOffsets = (NOffset, MOffset)
                 AMatrixShape = (NSize, MSize)
 
+            if len(buffA.shape) > 2:
+                batchDimCount = len(buffA.shape) - 2
+                AMatrixOffsets = tuple(cube.offset[:-2][-batchDimCount:]) + AMatrixOffsets
+                AMatrixShape = tuple(cube.dims[:-2][-batchDimCount:]) + AMatrixShape
+
+            ACube = HyperRectangle(AMatrixOffsets, AMatrixShape)
+
             if transB == 0:
                 BMatrixOffsets = (NOffset, OOffset)
                 BMatrixShape = (NSize, OSize)
             else:
                 BMatrixOffsets = (OOffset, NOffset)
                 BMatrixShape = (OSize, NSize)
 
-            if len(buffA.shape) == 2:
-                ACube = HyperRectangle(AMatrixOffsets, AMatrixShape)
-            elif len(buffA.shape) == 3:
-                ACube = HyperRectangle((BatchOffset,) + AMatrixOffsets, (BatchSize,) + AMatrixShape)
-            else:
-                ACube = HyperRectangle(
-                    (
-                        BatchOffset,
-                        BOffset,
-                    ) + AMatrixOffsets,
-                    (
-                        BatchSize,
-                        BSize,
-                    ) + AMatrixShape,
-                )
-
-            if len(buffB.shape) == 2:
-                BCube = HyperRectangle(BMatrixOffsets, BMatrixShape)
-            elif len(buffB.shape) == 3:
-                BCube = HyperRectangle((BatchOffset,) + BMatrixOffsets, (BatchSize,) + BMatrixShape)
-            else:
-                BCube = HyperRectangle(
-                    (
-                        BatchOffset,
-                        BOffset,
-                    ) + BMatrixOffsets,
-                    (
-                        BatchSize,
-                        BSize,
-                    ) + BMatrixShape,
-                )
+            if len(buffB.shape) > 2:
+                batchDimCount = len(buffB.shape) - 2
+                BMatrixOffsets = tuple(cube.offset[:-2][-batchDimCount:]) + BMatrixOffsets
+                BMatrixShape = tuple(cube.dims[:-2][-batchDimCount:]) + BMatrixShape
+
+            BCube = HyperRectangle(BMatrixOffsets, BMatrixShape)
 
             RequantCube = HyperRectangle((OOffset,), (OSize,))
 
@@ -367,25 +346,21 @@ def serializeTilingSolution(
 
         # Every output is constructed by a pair of inputs. Reconstruct this pair.
         for cube in outputCubes:
+            MOffset, OOffset = cube.offset[-2:]
+            MSize, OSize = cube.dims[-2:]
+
+            if len(cube.offset) > 2:
+                BatchSize = math.prod(cube.dims[:-2])
 
-            BSize = 1
-            BOffset = 0
-            BatchSize = 1
-            BatchOffset = 0
-
-            if len(cube.offset) == 2:
-                (MOffset, OOffset) = cube.offset
-                (MSize, OSize) = cube.dims
-            elif len(cube.offset) == 3:
-                (BatchOffset, MOffset, OOffset) = cube.offset
-                (BatchSize, MSize, OSize) = cube.dims
+                # Check that we don't tile upper dimensions
+                if len(cube.offset) > 3:
+                    assert all(off == 0 for off in cube.offset[:-3])
             else:
-                (BatchOffset, BOffset, MOffset, OOffset) = cube.offset
-                (BatchSize, BSize, MSize, OSize) = cube.dims
+                BatchSize = 1
 
             replacements["M"].append(MSize)
             replacements["O"].append(OSize)
-            replacements["batch"].append(BSize)
+            replacements["batch"].append(BatchSize)
 
             if transA == 0:
                 AMatrixOffsets = (MOffset, NOffset)
@@ -394,54 +369,36 @@ def serializeTilingSolution(
                 AMatrixOffsets = (NOffset, MOffset)
                 AMatrixShape = (NSize, MSize)
 
+            if len(buffA.shape) > 2:
+                batchDimCount = len(buffA.shape) - 2
+                AMatrixOffsets = tuple(cube.offset[:-2][-batchDimCount:]) + AMatrixOffsets
+                AMatrixShape = tuple(cube.dims[:-2][-batchDimCount:]) + AMatrixShape
+
+            ACube = HyperRectangle(AMatrixOffsets, AMatrixShape)
+
             if transB == 0:
                 BMatrixOffsets = (NOffset, OOffset)
                 BMatrixShape = (NSize, OSize)
             else:
                 BMatrixOffsets = (OOffset, NOffset)
                 BMatrixShape = (OSize, NSize)
 
-            if len(buffA.shape) == 2:
-                ACube = HyperRectangle(AMatrixOffsets, AMatrixShape)
-            elif len(buffA.shape) == 3:
-                ACube = HyperRectangle((BatchOffset,) + AMatrixOffsets, (BatchSize,) + AMatrixShape)
-            else:
-                ACube = HyperRectangle(
-                    (
-                        BatchOffset,
-                        BOffset,
-                    ) + AMatrixOffsets,
-                    (
-                        BatchSize,
-                        BSize,
-                    ) + AMatrixShape,
-                )
-
-            if len(buffB.shape) == 2:
-                BCube = HyperRectangle(BMatrixOffsets, BMatrixShape)
-            elif len(buffB.shape) == 3:
-                BCube = HyperRectangle((BatchOffset,) + BMatrixOffsets, (BatchSize,) + BMatrixShape)
-            else:
-                BCube = HyperRectangle(
-                    (
-                        BatchOffset,
-                        BOffset,
-                    ) + BMatrixOffsets,
-                    (
-                        BatchSize,
-                        BSize,
-                    ) + BMatrixShape,
-                )
+            if len(buffB.shape) > 2:
+                batchDimCount = len(buffB.shape) - 2
+                BMatrixOffsets = tuple(cube.offset[:-2][-batchDimCount:]) + BMatrixOffsets
+                BMatrixShape = tuple(cube.dims[:-2][-batchDimCount:]) + BMatrixShape
+
+            BCube = HyperRectangle(BMatrixOffsets, BMatrixShape)
 
             CMatrixOffsets = (MOffset, OOffset)
             CMatrixShape = (MSize, OSize)
 
-            if len(buffC.shape) == 2:
-                CCube = HyperRectangle(CMatrixOffsets, CMatrixShape)
-            elif len(buffC.shape) == 3:
-                CCube = HyperRectangle((BatchOffset,) + CMatrixOffsets, (BatchSize,) + CMatrixShape)
-            else:
-                CCube = HyperRectangle((BatchOffset, BOffset) + CMatrixOffsets, (BatchSize, BSize) + CMatrixShape)
+            if len(buffC.shape) > 2:
+                batchDimCount = len(buffC.shape) - 2
+                CMatrixOffsets = tuple(cube.offset[:-2][-batchDimCount:]) + CMatrixOffsets
+                CMatrixShape = tuple(cube.dims[:-2][-batchDimCount:]) + CMatrixShape
+
+            CCube = HyperRectangle(CMatrixOffsets, CMatrixShape)
 
             inputACubes.append(ACube)
             inputBCubes.append(BCube)