Skip to content

Commit

Permalink
reset MatMulTileConstraint to check if testmatmul run
Browse files Browse the repository at this point in the history
  • Loading branch information
@runwangdl
runwangdl committed Feb 12, 2025
1 parent 5907d6d commit 131f8d7
Showing 1 changed file with 193 additions and 29 deletions.
222 changes: 193 additions & 29 deletions Deeploy/Targets/PULPOpen/TileConstraints/MatMulTileConstraint.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,170 @@
# # ----------------------------------------------------------------------
# #
# # File: MatMulTileConstraint.py
# #
# # Last edited: 04.07.2023
# #
# # Copyright (C) 2023, ETH Zurich and University of Bologna.
# #
# # Author:
# # - Victor Jung, jungvi@iis.ee.ethz.ch, ETH Zurich
# # - Moritz Scherer, scheremo@iis.ee.ethz.ch, ETH Zurich
# #
# # ----------------------------------------------------------------------
# # SPDX-License-Identifier: Apache-2.0
# #
# # Licensed under the Apache License, Version 2.0 (the License); you may
# # not use this file except in compliance with the License.
# # You may obtain a copy of the License at
# #
# #
# # Unless required by applicable law or agreed to in writing, software
# # distributed under the License is distributed on an AS IS BASIS, WITHOUT
# # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# # See the License for the specific language governing permissions and
# # limitations under the License.

# from typing import Dict, List, Tuple

# from Deeploy.AbstractDataTypes import PointerClass
# from Deeploy.CommonExtensions.DataTypes import int8_t
# from Deeploy.DeeployTypes import NetworkContext, OperatorRepresentation
# from Deeploy.TilingExtension.MemoryConstraints import NodeMemoryConstraint
# from Deeploy.TilingExtension.TileConstraint import TileConstraint
# from Deeploy.TilingExtension.TilerModel import TilerModel
# from Deeploy.TilingExtension.TilingCodegen import AbsoluteHyperRectangle, HyperRectangle, TilingSchedule, \
# VariableReplacementScheme


# class MatMulTileConstraint(TileConstraint):

# @staticmethod
# def addGeometricalConstraint(tilerModel: TilerModel, parseDict: Dict, ctxt: NetworkContext) -> TilerModel:

# # Get to-be-tiled tensor's buffers
# bufferA = ctxt.lookup(name = parseDict['A'])
# bufferB = ctxt.lookup(name = parseDict['B'])
# outputBuffer = ctxt.lookup(name = parseDict['data_out'])

# # Add I/O dimensions to the model as variables
# for _buffer in [bufferA, bufferB, outputBuffer]:
# tilerModel.addTensorDimToModel(ctxt, _buffer.name)

# dimOffsetA = len(bufferA.shape) - 2
# dimOffsetB = len(bufferB.shape) - 2
# dimOffsetOut = len(outputBuffer.shape) - 2

# AFirstDimVar = tilerModel.getTensorDimVar(tensorName = bufferA.name, dimIdx = dimOffsetA)
# ASecondDimVar = tilerModel.getTensorDimVar(tensorName = bufferA.name,
# dimIdx = dimOffsetA + 1 )
# BFirstDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name, dimIdx = dimOffsetB)
# BSecondDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name, dimIdx = dimOffsetB + 1)
# outputFirstDimVar = tilerModel.getTensorDimVar(tensorName = outputBuffer.name, dimIdx = dimOffsetOut)
# outputSecondDimVar = tilerModel.getTensorDimVar(tensorName = outputBuffer.name, dimIdx = dimOffsetOut + 1)

# # Map output dims to inputs dims
# tilerModel.addConstraint(outputFirstDimVar == AFirstDimVar)
# tilerModel.addConstraint(outputSecondDimVar == BSecondDimVar)

# # Add GEMM Geometrical constraints
# tilerModel.addConstraint(ASecondDimVar == BFirstDimVar)

# return tilerModel

# @staticmethod
# def addPolicyConstraint(tilerModel: TilerModel, parseDict: Dict, ctxt: NetworkContext) -> TilerModel:

# bufferA = ctxt.lookup(name = parseDict['A'])
# bufferB = ctxt.lookup(name = parseDict['B'])

# dimOffsetA = len(bufferA.shape) - 2
# dimOffsetB = len(bufferB.shape) - 2

# ASecondDimVar = tilerModel.getTensorDimVar(tensorName = bufferA.name,
# dimIdx = dimOffsetA + 1 )
# BFirstDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name, dimIdx = dimOffsetB)
# BSecondDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name, dimIdx = dimOffsetB + 1)

# # VIC: We don't want to deal with intermediate results between kernel calls
# tilerModel.addConstraint(ASecondDimVar == parseDict['N'])
# tilerModel.addConstraint(BFirstDimVar == parseDict['N'])

# # VIC: For now we tile only one of the inputs
# tilerModel.addConstraint(BSecondDimVar == parseDict['O'])

# return tilerModel

# @classmethod
# def serializeTilingSolution(
# cls, tilingSolution: NodeMemoryConstraint, absoluteOutputCubes: List[AbsoluteHyperRectangle],
# targetMemLevel: str, ctxt: NetworkContext,
# operatorRepresentation: OperatorRepresentation) -> Tuple[VariableReplacementScheme, TilingSchedule]:
# outputCubes = [cube.rectangle for cube in absoluteOutputCubes]

# addrNames = ['A', 'B', 'data_out']
# inputBaseOffsets, outputBaseOffsets = cls.extractBaseAddr(tilingSolution, targetMemLevel,
# operatorRepresentation, addrNames)

# varA = operatorRepresentation['A']

# NSize = ctxt.lookup(varA).shape[-1]
# NOffset = 0

# inputACubes = []
# inputBCubes = []

# replacements = {"M": [], "O": [], "batch": []}

# # Every output is constructed by a pair of inputs. Reconstruct this pair.
# for cube in outputCubes:
# 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
# else:
# (BatchOffset, BOffset, MOffset, OOffset) = cube.offset
# (BatchSize, BSize, MSize, OSize) = cube.dims


# replacements["M"].append(MSize)
# replacements["O"].append(OSize)
# replacements["batch"].append(BSize)

# ACube = HyperRectangle((BatchOffset, BOffset, MOffset, NOffset), (BatchSize, BSize, MSize, NSize))
# BCube = HyperRectangle((BatchOffset, BOffset, NOffset, OOffset), (BatchSize, BSize, NSize, OSize))

# inputACubes.append(ACube)
# inputBCubes.append(BCube)

# inputLoadSchedule = []
# outputLoadSchedule = []

# replacements["N"] = [NSize] * len(outputCubes)

# replacementTypes = {
# "M": PointerClass(int8_t),
# "N": PointerClass(int8_t),
# "O": PointerClass(int8_t),
# "batch": PointerClass(int8_t)
# }

# for a, b in zip(inputACubes, inputBCubes):
# inputLoadSchedule.append({"A": a, "B": b})

# for out in outputCubes:
# outputLoadSchedule.append({"data_out": out})

# schedule = TilingSchedule(inputBaseOffsets, outputBaseOffsets, inputLoadSchedule, outputLoadSchedule)

# return VariableReplacementScheme(replacements, replacementTypes), schedule

# ----------------------------------------------------------------------
#
# File: MatMulTileConstraint.py
Expand Down Expand Up @@ -50,19 +217,28 @@ def addGeometricalConstraint(tilerModel: TilerModel, parseDict: Dict, ctxt: Netw
for _buffer in [bufferA, bufferB, outputBuffer]:
tilerModel.addTensorDimToModel(ctxt, _buffer.name)

dimOffsetA = len(bufferA.shape) - 2
dimOffsetB = len(bufferB.shape) - 2
dimOffsetOut = len(outputBuffer.shape) - 2
tensorsShapeLen = len(bufferA.shape)

AFirstDimVar = tilerModel.getTensorDimVar(tensorName = bufferA.name, dimIdx = dimOffsetA)
AFirstDimVar = tilerModel.getTensorDimVar(tensorName = bufferA.name,
dimIdx = (tensorsShapeLen - 2) + parseDict['transA'])
ASecondDimVar = tilerModel.getTensorDimVar(tensorName = bufferA.name,
dimIdx = dimOffsetA + 1 )
BFirstDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name, dimIdx = dimOffsetB)
BSecondDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name, dimIdx = dimOffsetB + 1)
outputFirstDimVar = tilerModel.getTensorDimVar(tensorName = outputBuffer.name, dimIdx = dimOffsetOut)
outputSecondDimVar = tilerModel.getTensorDimVar(tensorName = outputBuffer.name, dimIdx = dimOffsetOut + 1)
dimIdx = (tensorsShapeLen - 1) - parseDict['transA'])
BFirstDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name,
dimIdx = (tensorsShapeLen - 2) + parseDict['transB'])
BSecondDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name,
dimIdx = (tensorsShapeLen - 1) - parseDict['transB'])
outputFirstDimVar = tilerModel.getTensorDimVar(tensorName = outputBuffer.name, dimIdx = (tensorsShapeLen - 2))
outputSecondDimVar = tilerModel.getTensorDimVar(tensorName = outputBuffer.name, dimIdx = (tensorsShapeLen - 1))

# Map output dims to inputs dims
for idx in range(tensorsShapeLen - 2):
tilerModel.addConstraint(
tilerModel.getTensorDimVar(tensorName = outputBuffer.name, dimIdx = idx) == tilerModel.getTensorDimVar(
tensorName = bufferA.name, dimIdx = idx))
tilerModel.addConstraint(
tilerModel.getTensorDimVar(tensorName = outputBuffer.name, dimIdx = idx) == tilerModel.getTensorDimVar(
tensorName = bufferB.name, dimIdx = idx))

tilerModel.addConstraint(outputFirstDimVar == AFirstDimVar)
tilerModel.addConstraint(outputSecondDimVar == BSecondDimVar)

Expand All @@ -77,13 +253,14 @@ def addPolicyConstraint(tilerModel: TilerModel, parseDict: Dict, ctxt: NetworkCo
bufferA = ctxt.lookup(name = parseDict['A'])
bufferB = ctxt.lookup(name = parseDict['B'])

dimOffsetA = len(bufferA.shape) - 2
dimOffsetB = len(bufferB.shape) - 2
tensorsShapeLen = len(bufferA.shape)

ASecondDimVar = tilerModel.getTensorDimVar(tensorName = bufferA.name,
dimIdx = dimOffsetA + 1 )
BFirstDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name, dimIdx = dimOffsetB)
BSecondDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name, dimIdx = dimOffsetB + 1)
dimIdx = (tensorsShapeLen - 1) - parseDict['transA'])
BFirstDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name,
dimIdx = (tensorsShapeLen - 2) + parseDict['transB'])
BSecondDimVar = tilerModel.getTensorDimVar(tensorName = bufferB.name,
dimIdx = (tensorsShapeLen - 1) - parseDict['transB'])

# VIC: We don't want to deal with intermediate results between kernel calls
tilerModel.addConstraint(ASecondDimVar == parseDict['N'])
Expand Down Expand Up @@ -117,21 +294,8 @@ def serializeTilingSolution(

# Every output is constructed by a pair of inputs. Reconstruct this pair.
for cube in outputCubes:
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
else:
(BatchOffset, BOffset, MOffset, OOffset) = cube.offset
(BatchSize, BSize, MSize, OSize) = cube.dims

(BatchOffset, BOffset, MOffset, OOffset) = cube.offset
(BatchSize, BSize, MSize, OSize) = cube.dims

replacements["M"].append(MSize)
replacements["O"].append(OSize)
Expand Down

0 comments on commit 131f8d7

Please sign in to comment.