# Licensed to the Apache Software Foundation (ASF) under one# or more contributor license agreements. See the NOTICE file# distributed with this work for additional information# regarding copyright ownership. The ASF licenses this file# to you 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## http://www.apache.org/licenses/LICENSE-2.0## 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."""Configurable VTA Hareware Environment scope."""# pylint: disable=invalid-name, exec-usedfrom__future__importabsolute_importas_absimportosimportjsonimportcopyimporttvmfromtvmimporttefrom.importintrindefget_vta_hw_path():"""Get the VTA HW path."""curr_path=os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))vta_hw_default=os.path.abspath(os.path.join(curr_path,"../../../3rdparty/vta-hw"))VTA_HW_PATH=os.getenv("VTA_HW_PATH",vta_hw_default)returnos.path.abspath(VTA_HW_PATH)defpkg_config(cfg):"""Returns PkgConfig pkg config object."""pkg_config_py=os.path.join(get_vta_hw_path(),"config/pkg_config.py")libpkg={"__file__":pkg_config_py}exec(compile(open(pkg_config_py,"rb").read(),pkg_config_py,"exec"),libpkg,libpkg)PkgConfig=libpkg["PkgConfig"]returnPkgConfig(cfg)classDevContext(object):"""Internal development context This contains all the non-user facing compiler internal context that is hold by the Environment. Parameters ---------- env : Environment The environment hosting the DevContext Note ---- This class is introduced so we have a clear separation of developer related, and user facing attributes. """# Memory id for DMAMEM_ID_UOP=0MEM_ID_WGT=1MEM_ID_INP=2MEM_ID_ACC=3MEM_ID_OUT=4MEM_ID_ACC_8BIT=5# VTA ALU OpcodesALU_OPCODE_MIN=0ALU_OPCODE_MAX=1ALU_OPCODE_ADD=2ALU_OPCODE_SHR=3ALU_OPCODE_MUL=4# Task queue id (pipeline stage)QID_LOAD_INP=1QID_LOAD_WGT=1QID_LOAD_OUT=2QID_STORE_OUT=3QID_COMPUTE=2def__init__(self,env):self.vta_axis=te.thread_axis("vta")self.vta_push_uop=tvm.tir.StringImm("VTAPushGEMMOp")ctx=tvm.tir.call_intrin("handle","tir.vta.command_handle")self.command_handle=tvm.tir.Call("handle","tir.tvm_thread_context",[ctx])self.DEBUG_NO_SYNC=Falseenv._dev_ctx=selfself.gemm=intrin.gemm(env,env.mock_mode)defget_task_qid(self,qid):"""Get transformed queue index."""return1ifself.DEBUG_NO_SYNCelseqid
[文档]classEnvironment(object):"""Hardware configuration object. This object contains all the information needed for compiling to a specific VTA backend. Parameters ---------- cfg : dict of str to value. The configuration parameters. Example -------- .. code-block:: python # the following code reconfigures the environment # temporarily to attributes specified in new_cfg.json new_cfg = json.load(json.load(open("new_cfg.json"))) with vta.Environment(new_cfg): # env works on the new environment env = vta.get_env() """current=None# constantsMAX_XFER=1<<22# debug flagsDEBUG_DUMP_INSN=1<<1DEBUG_DUMP_UOP=1<<2DEBUG_SKIP_READ_BARRIER=1<<3DEBUG_SKIP_WRITE_BARRIER=1<<4# memory scopesinp_scope="local.inp_buffer"wgt_scope="local.wgt_buffer"acc_scope="local.acc_buffer"# initialization functiondef__init__(self,cfg):# Produce the derived parameters and update dictself.pkg=pkg_config(cfg)self.__dict__.update(self.pkg.cfg_dict)# data type widthself.INP_WIDTH=1<<self.LOG_INP_WIDTHself.WGT_WIDTH=1<<self.LOG_WGT_WIDTHself.ACC_WIDTH=1<<self.LOG_ACC_WIDTHself.OUT_WIDTH=1<<self.LOG_OUT_WIDTH# tensor intrinsic shapeself.BATCH=1<<self.LOG_BATCHself.BLOCK_IN=1<<self.LOG_BLOCK_INself.BLOCK_OUT=1<<self.LOG_BLOCK_OUT# buffer sizeself.UOP_BUFF_SIZE=1<<self.LOG_UOP_BUFF_SIZEself.INP_BUFF_SIZE=1<<self.LOG_INP_BUFF_SIZEself.WGT_BUFF_SIZE=1<<self.LOG_WGT_BUFF_SIZEself.ACC_BUFF_SIZE=1<<self.LOG_ACC_BUFF_SIZEself.OUT_BUFF_SIZE=1<<self.LOG_OUT_BUFF_SIZE# bytes per bufferself.INP_ELEM_BITS=self.BATCH*self.BLOCK_IN*self.INP_WIDTHself.WGT_ELEM_BITS=self.BLOCK_OUT*self.BLOCK_IN*self.WGT_WIDTHself.ACC_ELEM_BITS=self.BATCH*self.BLOCK_OUT*self.ACC_WIDTHself.OUT_ELEM_BITS=self.BATCH*self.BLOCK_OUT*self.OUT_WIDTHself.INP_ELEM_BYTES=self.INP_ELEM_BITS//8self.WGT_ELEM_BYTES=self.WGT_ELEM_BITS//8self.ACC_ELEM_BYTES=self.ACC_ELEM_BITS//8self.OUT_ELEM_BYTES=self.OUT_ELEM_BITS//8# dtypesself.acc_dtype="int%d"%self.ACC_WIDTHself.inp_dtype="int%d"%self.INP_WIDTHself.wgt_dtype="int%d"%self.WGT_WIDTHself.out_dtype="int%d"%self.OUT_WIDTH# bistream nameself.BITSTREAM=self.pkg.bitstream# model stringself.MODEL=self.TARGET+"_"+self.BITSTREAM# lazy cached membersself.mock_mode=Falseself._mock_env=Noneself._dev_ctx=Noneself._last_env=Nonedef__enter__(self):self._last_env=Environment.currentEnvironment.current=selfreturnselfdef__exit__(self,ptype,value,trace):Environment.current=self._last_env@propertydefcfg_dict(self):returnself.pkg.cfg_dict@propertydefdev(self):"""Developer context"""ifself._dev_ctxisNone:self._dev_ctx=DevContext(self)returnself._dev_ctx@propertydefmock(self):"""A mock version of the Environment The ALU, dma_copy and intrinsics will be mocked to be nop. """ifself.mock_mode:returnselfifself._mock_envisNone:self._mock_env=copy.copy(self)self._mock_env._dev_ctx=Noneself._mock_env.mock_mode=Truereturnself._mock_env@propertydefdma_copy(self):"""DMA copy pragma"""return"dma_copy"ifnotself.mock_modeelse"skip_dma_copy"@propertydefalu(self):"""ALU pragma"""return"alu"ifnotself.mock_modeelse"skip_alu"@propertydefgemm(self):"""GEMM intrinsic"""returnself.dev.gemm@propertydeftarget(self):returntvm.target.vta(model=self.MODEL)@propertydeftarget_host(self):"""The target host"""ifself.TARGETin["pynq","de10nano"]:return"llvm -mtriple=armv7-none-linux-gnueabihf"ifself.TARGET=="ultra96":return"llvm -mtriple=aarch64-linux-gnu"ifself.TARGETin["sim","tsim","intelfocl"]:return"llvm"raiseValueError("Unknown target %s"%self.TARGET)@propertydeftarget_vta_cpu(self):returntvm.target.arm_cpu(model=self.TARGET)
[文档]defget_env():"""Get the current VTA Environment. Returns ------- env : Environment The current environment. """returnEnvironment.current
def_init_env():"""Initialize the default global env"""config_path=os.path.join(get_vta_hw_path(),"config/vta_config.json")ifnotos.path.exists(config_path):raiseRuntimeError("Cannot find config in %s"%str(config_path))cfg=json.load(open(config_path))returnEnvironment(cfg)Environment.current=_init_env()