MLIR dialect design method for static memory planning

Abstract

The invention discloses an MMIR dialect design method for static memory planning. The compilation link is a compilation link which is located after graph calculation fusion and before back-end code generation in an MLIR graph compiler. Compared with a one-shot-buffer architecture of a memref dialect in an existing open source framework, the method disclosed by the invention has the innovation points that the buffering of a static single assignment structure is kept; alias analysis of auxiliary operation does not need to be inserted; deducing the life cycle of the optimized topological sorting; the invention discloses a memory planning solver supporting multi-level storage. Experiments show that on ResNet, Yolov5 and other typical models, compared with a traditional scheme, the method has the advantages that SRAM peak memory occupation can be reduced by 10-15%, compared with a memref scheme, the number of MLIR intermediate representation instructions is reduced by 35%, and a lightweight compiling scheme is provided for efficient memory management on end side AI hardware.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of artificial intelligence model compilation data processing, in particular to a MLIR dialect design method for static memory planning. BACKGROUND

[0002] With the rapid development of artificial intelligence technology, deep neural network models are increasingly widely used in image recognition, natural language processing, autonomous driving and other fields. However, the exponential growth of model parameter size and computational complexity poses a serious challenge to the memory management of underlying heterogeneous storage architectures. Modern heterogeneous computing platforms usually adopt a multi-level storage system, including global storage (DDR), high-bandwidth video memory (HBM) and on-chip storage (SRAM), each level of storage has significant differences in bandwidth, delay and capacity. Under this background, AI compilers need to build cross-level memory planning strategies based on data access characteristics, and distribute data to the optimal storage level by predicting memory usage patterns, thereby effectively alleviating memory bandwidth bottlenecks and storage fragmentation problems.

[0003] Static memory planning, as a core optimization technology of AI compilers, uses fixed input dimension information to perform tensor space deduction before the execution of the computation graph, and determines the input-output address mapping relationship of each operator in advance. This technology has double advantages: on the one hand, it can effectively utilize high-bandwidth storage space, and on the other hand, it can reduce peak memory occupancy through memory reuse strategies. Its key value lies in providing memory optimization decision basis for the compiler through accurate memory trajectory analysis.

[0004] MLIR (Multi-Level Intermediate Representation) is an open-source domain-specific compilation framework developed by Google. Its layered architecture design based on multi-level intermediate representation realizes the decoupling of front-end model analysis, middle-end optimization and compilation, and back-end hardware adaptation. In view of the heterogeneous memory access characteristics of deep learning operators, the MLIR standard dialect memref provides memory reference layout abstraction and multi-level storage access interface for tensor storage. However, in practical applications, it is found that the existing MLIR memory planning scheme has the following technical limitations:

[0005] First, regarding the compilation process, existing solutions heavily rely on the memref dialect and one-shot-buffering conversion mechanism. This dependency requires all upstream dialect operations to implement the BufferizableOpInterface interface, significantly increasing the system's debugging complexity and usage threshold. More seriously, the buffering process converts the tensor type to a memref pointer type, breaking the static single assignment (SSA) form of the intermediate representation. This forces a mem2reg conversion to reconstruct the SSA properties before backend instruction translation, which not only adds additional compilation overhead but also introduces potential correctness risks.

[0006] Secondly, in terms of implementation mechanism, existing solutions require the insertion of numerous auxiliary memory-effect operations (MemoryEffectOp), such as AllocOp and SubviewOp. These operations become redundant after memory allocation, not participating in the final instruction generation and binary translation, yet significantly increasing the complexity and processing overhead of the IR file. Especially when processing large-scale neural network models, this redundancy can cause the IR file to swell dramatically, severely impacting compilation efficiency.

[0007] Third, existing frameworks have significant shortcomings in terms of customizability. The publicly available development documentation examples for the one-shot-buffering framework are limited, and its software framework and interface design are complex, making it difficult for developers to quickly understand and customize memory allocation algorithms for specific application scenarios. This lack of customizability severely limits the compiler's adaptability across different hardware platforms and application scenarios.

[0008] Memory allocation, as the final step in the graph compiler's middleware architecture, directly adapts to the backend hardware memory architecture and has a high degree of customization requirements. Existing technologies have limitations in maintaining the SSA structure, eliminating redundant operations, and supporting algorithm customization, which severely restrict the performance of AI compilers on heterogeneous computing platforms. Therefore, there is an urgent need to build a novel custom intermediate representation and dialect conversion mechanism that can achieve an efficient memory allocation scheme while maintaining the SSA structure, and simultaneously provide flexible algorithm access capabilities to meet the diverse needs of different hardware platforms and application scenarios.

[0009] In summary, the technical deficiencies in static memory planning of existing MLIR standard dialects have become a key bottleneck restricting further performance improvements of AI compilers. How to design a static memory planning scheme that maintains the SSA structure while supporting flexible customization has become an important technical problem that urgently needs to be solved by those skilled in the art. Summary of the Invention

[0010] To address the aforementioned technical problems, this invention provides an MLIR dialect design method oriented towards static memory planning, as shown in the appendix. Figure 1 As shown, the compilation stage in the mlir-opt compiler, which is located after graph computation fusion and before backend code generation, specifically includes the following main steps:

[0011] Step S01: Buffering conversion; extend operation attributes and value type attributes, and retain the static single assignment form of the intermediate representation during the conversion process.

[0012] Step S02: Alias ​​analysis; using the sparse data flow analysis framework of MLIR, gradually construct the alias equivalence set of values; determine whether the value definition operation requires explicit buffer allocation.

[0013] Step S03: Lifecycle analysis; Using the MLIR lifecycle analysis framework, the lifecycle interval representation of the buffer is optimized by topological sorting, and then the integer intervals of the alias set are merged.

[0014] Step S04: Memory planning solution; Configure the amount of hardware memory resources, add the explicit buffer to the allocation list, given the space size and lifespan range, run the memory reuse strategy and memory pool management algorithm to obtain the memory address and memory level of the buffer;

[0015] Step S05: Tensor address allocation; Copy the explicitly allocated buffer address and memory level to the alias operand, and read the offset in the buffer of the alias operand, and sum them to obtain the address of the tensor.

[0016] In step S01, the buffering transformation, the extended operation attributes and value type attributes are intended to replace auxiliary memory-effect operation instructions (MemoryEffectOp) such as allocOp and subviewOp in the memref dialect. In contrast, the memref dialect in the MLIR open-source framework requires implementing BufferizableOpInterface for each computational operation in the preceding graph computation fusion dialect before applying the dialect rewriting mode to replace the tensor value type with a memref pointer class in situ. Because memref's built-in aliasing determines whether a buffer object needs to be newly created or directly copied from an existing buffer, memref operands may be copied. Therefore, the IR file of the memref dialect does not have a static single assignment (SSA) structure. The open-source framework reconstructs the required static single assignment structure using mem2reg transformation in the LLVM compiler.

[0017] In one embodiment of the present invention, the extended operation attributes are obtained by parsing the semantics of the operation type in the IR file, determining the buffer behavior of the return value operands of computational operations, and adding two new Boolean attributes for operations, such as is-inplace and need-alloc.

[0018] In one embodiment of the present invention, the need-alloc attribute is only initialized during buffering transformation, updated after alias analysis, and read during memory planning to determine whether to allocate a new buffer address. Here, need-alloc being true indicates that a new memory buffer needs to be allocated for the return value, and need-alloc being false indicates that the existing memory buffer is directly copied.

[0019] In one embodiment of the present invention, the is-inplace attribute is assigned during buffering. During alias analysis, the is-inplace attribute is read to infer the alias set of operands, and the explicitly allocated buffer address is broadcast to the alias operands through the alias set. Here, is-inplace being true indicates that the return value and the input operand share the same buffer, i.e., memory blocks with the same lifetime and contiguous addresses; the offset addresses of different operands may be different. is-inplace being false indicates that the memory blocks of the return value and the input operand are different.

[0020] In one embodiment of the present invention, the extended value attribute is obtained by parsing the semantics of computation operations, dimension splitting operations (Slice, Split, etc.), dimension concatenation operations (deSlice, Concat, etc.), and dimension transformation operations (Transpose, Reshape, etc.) in the IR file, and calculating the layout information of the return value operands in memory according to the characteristics of the backend hardware. Two new long integer attributes, such as bufferOffset and dataSize, are added. These two attributes are one of the key variables for memory allocation of the computation graph.

[0021] In one embodiment of the present invention, the layout information refers to the ability to convert high-dimensional coordinates into one-dimensional linear indices in memory based on affine parameters. Whether the layout is compact determines the size of the tensor in the memory block. Affine parameters include, for example, the number of channels in the thread bundle, line span, and surface span.

[0022] In one embodiment of the present invention, the dataSize attribute is calculated based on the tensor dimension of the operator after graph computation fusion and the affine parameters of the hardware layout. The bufferOffset attribute is the buffer offset address determined based on the operation semantics and the dataSize attribute of other tensors.

[0023] In step S02, the alias analysis, the sparse data flow analysis framework using MLIR refers to sparsely establishing lattice points on a static single-assignment structure, with each lattice point having a corresponding program point and analysis state.

[0024] The program point is the position where a grid point is attached in the program.

[0025] The AnalysisState is a data stream that evolves with each analysis iteration, and a program point is the only identifiable object attached to the AnalysisState.

[0026] The sparse data flow analysis treats each SSA value as a program point and iterates forward through all operations. During the iteration, the grid list of operand SSA values ​​and the grid list of return operand SSA values ​​are passed as parameters to the state transition function. The analysis function selects whether to update the analysis state of the return operand SSA values ​​based on a user-defined choice. The update of the grid analysis state triggers a callback function, which causes other grid points affected by the current SSA value grid point to re-execute the state transition function and modify the operation attributes associated with the current SSA value grid point. This process iterates until the analysis state of the grid points reaches a fixed point.

[0027] In one embodiment of the present invention, the stepwise construction of the alias equivalence set of values ​​in alias analysis refers to using the input and output SSA values ​​of computational operations as program points, and the alias equivalence set of values ​​as analysis states. The program point position attached to the grid point is the SSA value, and the analysis state corresponding to the grid point is the alias set of the SSA value. The alias set consists of other SSA values ​​that share the same buffer as the current SSA value. Updating the analysis state refers to updating the alias set of the grid point, including appending and deleting SSA values ​​from the set.

[0028] In one embodiment of the present invention, determining whether the value definition operation in alias analysis requires explicit buffer allocation means determining whether the operand and return value share a memory buffer based on the in-place attribute of the operation, and ensuring that there is one and only one SSA value in the alias set that requires explicit memory allocation.

[0029] In step S03, lifecycle analysis, the application of topological sorting refers to the application of the depth-first search (DFS) algorithm to re-mark the execution order of operations.

[0030] In one embodiment of the present invention, the range representation of the optimized buffer refers to, after topological sorting, expanding the integer attribute op-index of the operation, calling the lifecycle analysis interface of MLIR to obtain the active operation set of SSA value, and determining the integer range of lifecycle (LiveRange) based on the minimum op-index and maximum op-index of the active operation set, denoted as [begin, end], and expanding the attribute of the SSA value.

[0031] The active operation set includes live-in and live-out operations for SSA values.

[0032] The live-in operation refers to an operation whose return operand is the current SSA value, which is usually the definer.

[0033] The live-out operation refers to an operation whose input operand is the current SSA value, typically the user.

[0034] The integer range of the merged alias set refers to the fact that the current SSA value alias set is not empty, the LiveRange of the alias SSA value is unioned, and the LiveRange attribute of itself and the alias SSA value is updated.

[0035] In step S04, the memory planning solution, the number of configured hardware memory resources refers to the preset types and sizes of backend hardware memory resources, with different types of memory resources, such as SRAM and DDR, managed by memory pools respectively.

[0036] In one embodiment of the present invention, adding the explicit buffer to the allocation list means appending the SSA value of the operand when the operation attribute need-alloc is true to the allocation list.

[0037] In one embodiment of the present invention, the known space size and lifecycle range refer to allocating memory for the SSA values ​​in the list to be allocated in the order of insertion, and obtaining the dataSize and liveRange attributes of the SSA values.

[0038] In one embodiment of the present invention, the runtime memory reuse strategy and memory pool management algorithm refer to the following: the memory reuse strategy employs the classic greedy eviction algorithm, reading the `liveRange` variable and performing memory reclamation; the memory pool management algorithm employs the classic buddyAlloc buddy allocation algorithm, reading the `dataSize` variable and performing allocation and reclamation. The memory reuse strategy has a high API call level, while the memory pool management algorithm has a low API call level. After the memory reuse strategy determines the memory block to be reclaimed, it calls the memory pool API to perform a memory release operation (allocate), and the memory allocation request calls the memory pool API to perform a memory allocation operation (deallocate). During offline model compilation, memory allocation and release are algorithm-level operations, not physical-level operations.

[0039] In one embodiment of the present invention, obtaining the memory address and memory level of the buffer means that each request for a buffer carries a priority selection. Taking SRAM (L1) and DDR (L2) two-level storage hierarchy as an example, when L2 priority is configured, the memory allocation request is distributed to the buddy allocation algorithm of the DDR memory pool and the allocated address is returned; when L1 priority is configured, the memory allocation request is distributed to the buddy allocation algorithm of the SRAM memory pool and the allocated address is returned; when the L1 memory pool space is insufficient, the request will be forwarded to the alternative L2 memory pool for application and the allocated address and final memory level will be returned.

[0040] This invention also provides an MLIR dialect system oriented towards static memory planning, serving the MLIR dialect design method, including:

[0041] The Bufferize Pass, connected to the graph computation fusion compiler, implements the buffering, aliasing, and lifecycle analysis steps. The input is configured as the IR file output by the graph computation fusion compiler. Static analysis is run, extending the attributes of computational operations to include op-index, is-inplace, and need-alloc; and extending the attributes of the input and output operand tensors of computational operations to include bufferOffset, dataSize, and liveRange. The output is configured as the IR file with extended operations and value attributes after static analysis.

[0042] The MemoryAlloc Pass, connected to the buffered analysis pass, implements the memory planning solution steps and tensor address allocation steps. The input is configured as an IR file containing update operations and value attributes. Memory planning is run, and the attributes of the input and output operand tensors for computational operations are expanded to include addrOffset and memType. The output is configured as an IR file containing expanded value attributes after memory planning.

[0043] In one embodiment of the present invention, it further includes:

[0044] The mlir-opt-like compiler is connected to the graph computation fusion compiler, buffered analysis pass, and static memory planning transformation pass, respectively, and is configured to coordinate the execution order and data transfer of each pass;

[0045] The memory resource management module is connected to the static memory planning conversion pass, configured as the memory resource parameters of the system hardware, manages memory pool objects at different levels, manages active buffer objects in the memory pool, and handles buffer allocation and release requests.

[0046] The present invention also provides a computer-readable storage medium having a computer program stored thereon, wherein the program, when executed by a processor, implements any of the MLIR dialect design methods described herein.

[0047] Compared to the one-shot-buffering architecture of memref dialects in existing open-source frameworks, this invention offers the following advantages: it reduces the number of instructions in MLIR intermediate representation files, the lightweight software module design shortens the debugging cycle for developers, and ensures the correctness of model execution results. Experiments show that on typical models such as ResNet and Yolov5, this method can reduce peak SRAM memory usage by 10-15% and the number of MLIR intermediate representation instructions by 35% compared to traditional solutions, providing a lightweight compilation solution for efficient memory management on edge AI hardware. Attached Figure Description

[0048] To make the content of this invention easier to understand, the invention will be further described in detail below with reference to specific embodiments and accompanying drawings.

[0049] Figure 1 This is a schematic diagram of the static memory planning compiler dialect of the present invention;

[0050] Figure 2 This is the input intermediate representation of the IR file content in this embodiment of the invention;

[0051] Figure 3 This is an input intermediate representation of the IR control flow graph in an embodiment of the present invention;

[0052] Figure 4 This is the output intermediate representation of the IR file content in this embodiment of the invention;

[0053] Figure 5 This is the design class diagram of the buffering S01, aliasing analysis S02, and lifecycle analysis S03 of the present invention;

[0054] Figure 6This is the design class diagram of memory planning solution S04 and tensor address allocation S05 of the present invention;

[0055] Figure 7 This is an example diagram illustrating the calculation of alias buffer and explicit buffer addresses in an embodiment of the present invention. Detailed Implementation

[0056] The technical solution of the present invention will be described in detail below with reference to the accompanying drawings. This embodiment takes the deployment of a neural network model on a dedicated AI chip as an application scenario to demonstrate the complete static memory planning process.

[0057] The compiled object processed in this invention is an intermediate representation (IR) file generated after front-end optimizations such as graph computation fusion. This IR file organizes the computation graph structure using a static single assignment (SSA) format, where each operator corresponds to one or more computational operations. The input and output tensors of the operator correspond to the source operand (OperandValue) and destination operand (ResultValue) of the operation, respectively. Each destination operand is defined only once throughout the entire program, ensuring the correctness of the SSA format.

[0058] The MLIR dialect designed in this invention aims to complete the memory address allocation of all activation value tensors by performing static analysis on the aforementioned IR, ensuring that the allocation results meet the address validity requirements under hardware resource constraints. The static analysis specifically includes: performing compile-time analysis on the computation graph to determine the total number of buffers required; identifying tensors that do not require independent buffer allocation; determining whether multiple tensors share the same buffer; representing the lifetime of each buffer using integer ranges to analyze their temporal overlap; and extracting the memory usage attribute information of each buffer. The "address validity" includes, but is not limited to: avoiding read / write conflicts at the same memory address by multiple instructions at the same time, preventing memory access out of bounds, and ensuring the stable availability of addresses throughout the buffer's lifetime.

[0059] To achieve the above objectives, this invention proposes and implements two core compilation passes, which are registered in the pass pipeline of the mlir-opt compiler. These passes are the bufferization pass and the memory allocation pass, which work together to complete the entire process from static analysis to memory layout planning.

[0060] Appendix Figure 2The diagram provides an example of the IR content of a computational subgraph in a YOLOv5s model. To highlight key information, non-core fields such as tensor dimensions and operator attributes have been omitted, focusing instead on the data dependencies between its SSA variables.

[0061] As attached Figure 3 As shown, this embodiment further illustrates the appendix. Figure 2 The control-flow graph (CFG) of the computational subgraph IR file is shown below. In this CFG, nodes strictly correspond to basic blocks or single operations in the computational graph, representing the execution units of the computational logic; directed edges explicitly represent the flow relationship of tensor SSA (Static Single Assignment Value) values ​​between operations, i.e., data dependency paths. Specifically, special visual annotations are provided in the attached figures below to illustrate the results of the aliasing analysis:

[0062] The bolded nodes (as shown in the figure, marked as "Swish" and "Concat") indicate that the operation introduces an alias dependency between SSA values, meaning that its output tensor shares the same physical memory region as the input tensor (for example, the Concat operation does not change the data content but only adjusts the view, causing the input / output SSA values ​​to point to the same buffer).

[0063] Dashed edges (such as the edge connecting "Conv" and "Swish") indicate that the corresponding SSA value uses the AliasBuffer mechanism, and its memory space is implemented by referencing the existing buffer, without needing to make an explicit allocation request to the memory manager (i.e. need_alloc = false).

[0064] A bold solid line (such as the "Concat" output line) indicates that the SSA value requires an explicit buffer, which will trigger the memory allocation algorithm to request independent memory space for it (i.e., need_alloc = true).

[0065] As attached Figure 4 The image shows the output IR segment after compilation and optimization using the method of this invention. It illustrates the final memory layout information obtained by static memory planning for each tensor, including key attribute values ​​such as MemoryType and OffsetInMemory. The original metadata has been appropriately simplified to clearly express the memory allocation results.

[0066] Memory address attributes: MemoryType (e.g., SRAM or DDR) for each tensor, and the hexadecimal address of each tensor (the result of summing buffer offset addresses);

[0067] Evidence of memory reuse: By comparing the beginSlot and endSlot attributes of the left endpoint of a tensor's interval, it can be verified whether the lifecycles of tensors overlap and whether the conditions for address reuse are met.

[0068] Address validity guarantee: All allocation results satisfy the conflict-free constraint—no two tensors have overlapping lifetimes in their explicit buffers, and all offset addresses are within the hardware memory resource boundaries.

[0069] The output IR, as the final product of static memory planning, has been automatically verified by the MLIR verifier, proving the feasibility of the method of this invention to achieve address allocation with zero runtime overhead in resource-constrained scenarios.

[0070] In a specific embodiment of the present invention, Bufferize-Pass is first executed to construct a semantic model of the tensor buffer and derive its lifecycle information. Its main function is to complete the semantic modeling, alias analysis, and lifecycle interval derivation of the tensor buffer, providing necessary metadata support for subsequent memory allocation.

[0071] As attached Figure 5 The class diagram structure of the rewritePattern module, AliasAnalysis module, and LivenessAnaylsis module in this embodiment is shown, illustrating the calling relationships during the static analysis process;

[0072] Furthermore, such as Figure 1 The static memory plan shown is located in the MLIR compiler and includes the following steps:

[0073] Step S01: Buffering and conversion stage, marking memory access attributes.

[0074] These attributes will serve as the basis for subsequent alias analysis and buffer type determination. All computational operations, except for constant operations, participate in this attribute initialization process. Specifically, a boolean attribute `need_alloc` is set for each operation, indicating whether the result tensor of the operation needs to be explicitly allocated independent memory space; and a boolean attribute `is_inplace` is set, indicating whether the operation reuses the memory buffer of the input tensor in-place.

[0075] When this pass is implemented based on the MLIR framework, in the initial stage, a custom `setMemoryEffectOpAttr` interface is called to mark the memory attributes of each non-ConstOp operation in the computation graph. Specifically, an `is_inplace` boolean attribute is added to the operation (set to true when the operation semantics support in-place computation, such as activation function fusion operations, in-place splitting, and in-place concatenation operations in convolutional layers). For operations whose source operands and destination operands do not share a memory buffer, `is_inplace` is false. The `in_place` attribute of an operation is initialized once and is not updated again; the `need_alloc` attribute of an operation is initialized to true and awaits updating.

[0076] Step S02: Custom Alias ​​Analysis Execution Phase.

[0077] Construct and run a custom alias analyzer, AliasAnalysis, which inherits from the MLIR sparse forward data flow analysis framework SparseForwardDataFlowAnalysis, to infer alias relationships between SSA values ​​and update the need_alloc attribute value.

[0078] Define the AliasLattice type as a program grid point, which encapsulates the AnalysisState, representing the alias set state of a certain SSA value on the control flow path.

[0079] Define the AliasInfo data structure, which is the concrete implementation of AnalysisState. It is used to store the alias set associated with the current grid point, that is, a set of SSA values ​​regarded as a shared memory buffer, and supports the merging operation of two alias sets through the join() interface.

[0080] Each AliasLattice instance can obtain the corresponding AliasInfo analysis state implementation through getValue(), and obtain the SSA value (i.e., analysis object) it is anchored to through the getAnchor() method.

[0081] The conservative flag Pessimistic is introduced during the analysis: in the initial state, Pessimistic is true for all grid points, indicating that no aliasing relationship has been found; when a potential alias is detected, it is set to false, allowing state propagation.

[0082] Next, the `visitOperation(Operation*, OperaandLatticeList, ResultLatticeList*)` method is overridden to process each operation during the forward traversal of the IR: It extracts the `is_inplace` attribute of the current operation; if `is_inplace` is true, it considers that there is a subview relationship between the source operand (Operand) and its destination operand (Result), sharing the same contiguous memory region; it creates a new `AliasInfo` instance and adds all source operands to this alias collection; it calls the `propagateIfChanged()` method to attempt to update the state of the corresponding grid point of the destination operand, triggering the `join()` operation to merge the existing alias collection; and it calls the `addDependency()` method to add the grid point corresponding to the defining operation of the source operand to the dependency queue, thereby triggering the state update of the relevant anchor points in subsequent backpropagation, ensuring the convergence of the alias collection throughout the entire use-def chain.

[0083] Next, the onUpdate(AliasLattice*) method is overridden to implement the callback logic after the grid state changes: the alias set of the destination operand corresponding to the current anchor point is traversed; based on the partial order relationship, it is determined that in the alias set, there is exactly one SSA value definition operation that satisfies need_alloc = true, and the need_alloc = false for all other SSA values ​​in the same set; this mechanism ensures that only one representative explicit buffer is allocated memory in each alias set, and the address of the other alias tensors is reused through offset mapping.

[0084] Subsequently, the custom AliasAnalysis alias analyzer is run. In visitOperation, when is_inplace=true is detected for an operation, an AliasInfo instance is created to add all source operand SSA values ​​(such as %input) to the current alias set, and propagateIfChanged is called to trigger an AliasLattice state update for the destination operand (such as %output). At this time, the AliasInfo::join operation is executed to merge the new set with the existing alias set (for example, merging {%swish_input,%concat_input} and {%concat_input,%concat_output} into {%swish_input,%concat_input,%concat_output}). At the same time, the definition operation anchor of the source operand is added to the dependency queue through addDependency(getDefiningOp(operand)), so that the data flow analysis framework can trigger the onUpdate callback of the upstream operation in reverse when the state changes. In the onUpdate implementation, a constraint is imposed on the alias set of the destination operand: all SSA values ​​in the set (such as {%a,%b,%c}) are traversed, and according to the OpIndex order of their defined operations, only the need_alloc attribute of the SSA value with the smallest OpIndex (such as %a) is set to true, while the remaining values ​​(%b,%c) are set to false. This mechanism ensures that each alias group has one and only one explicit buffer that needs to be allocated memory.

[0085] Step S03: Lifecycle analysis phase, operations are sorted and operation indexes are generated, active intervals are derived, and new memory layout and trajectory attribute fields are attached to each tensor.

[0086] The operation sequence in the computation graph is sorted to establish a unified time-dimensional reference frame. Optionally, a depth-first search (DFS) order or the default IR traversal order (such as intra-block order) can be used. An integer type OpIndex attribute is appended to each operation, representing its relative position in the entire execution flow, for subsequent active interval computation.

[0087] The built-in LivenessAnalysis module of MLIR is invoked, and combined with the generated OpIndex, the liveRange of each tensor is calculated. It is a closed-open integer interval of the form [beginSlot, endSlot], which represents the range of operations of the tensor from the first definition to the last use.

[0088] Furthermore, based on the AliasSet obtained from the aforementioned alias analysis, the overall lifecycle range of the associated buffer is updated synchronously: if multiple tensors belong to the same alias group, the lifecycle of their shared buffer is the union of the liveRange of all member tensors. For example, if %a and %b belong to the same AliasSet and %a's liveRange = [3,8] and %b's liveRange = [6,10], then the lifecycle of their shared buffer is the union of [3,10].

[0089] Attach new memory layout and trajectory attribute fields to each tensor, including:

[0090] MemoryType: An enumeration type indicating whether the tensor is ultimately allocated to SRAM or DDR;

[0091] OffsetInMemory: The absolute byte offset of this tensor in the target memory pool;

[0092] OffsetInBuffer: The offset relative to the starting address of its explicit buffer;

[0093] Strides and Dims: Describe the multidimensional stride and dimension information of a tensor, used for address calculation;

[0094] dataSize: Describes the size of the tensor in memory;

[0095] BeginSlot and EndSlot: Record the start and end operation indexes of their active range.

[0096] The above attributes will be read or updated during subsequent memory allocation phases.

[0097] Based on the above analysis, Memory-Alloc-Pass performs static memory allocation. Using known buffer lifetime and size information, Memory-Alloc-Pass rationally maps intermediate tensor address spaces to multi-level storage systems (such as L1-SRAM and L2-DDR), and optimizes the memory reuse efficiency of the entire computation graph.

[0098] As attached Figure 5 The class diagram structure of the StaticMemoryPlanner module in this embodiment is shown, illustrating the calling relationships between MemoryPool, ResourceManager, BuddyAllocator, and GreedyEvict.

[0099] Step S04: Memory planning and solution stage.

[0100] All tensors that need to allocate memory are divided into two categories: the list of TaskI / O tensors, which contains task input and output tensors, and is default allocated to the DDR memory pool according to the system policy; the list of intermediate result tensors contains all intermediate calculation result tensors except ConstOp. They are sorted in ascending order of the beginSlot of their active intervals to perform dynamic allocation decisions in chronological order.

[0101] For tensors with alias relationships (represented by AliasSet), the following rules are executed: converge the need_alloc attributes of the defOp operations (defOp) of all Values in the AliasSet; there is exactly one defOp with need_alloc = true, and the tensor generated by it is regarded as the explicit buffer (ExplicitBuffer) of this alias group; the remaining aliasBuffers do not allocate memory separately, but share the address space of the explicitBuffer through offset calculation.

[0102] The designed static memory planning solver StaticMemoryPlanner contains three sub-modules.

[0103] The ResourceManger hardware memory resource parameter module configures the SRAM memory pool capacity (e.g., 64KB) and the DDR memory pool capacity (e.g., 1GB); maintains the set of aliveBuffers of currently active buffers to track buffers whose life cycle has not ended.

[0104] The GreedyEviction greedy eviction policy module traverses the list of intermediate result tensors in the order of beginSlot; for each tensor to be allocated, first check whether its dependent input tensors have exited the active interval (i.e., endSlot < currentSlot), if so, call checkRelease() to recycle the memory space they occupy; execute the resolvePlacement(size_t space, MemoryType preferredType) function, and first try to allocate a specified size of space in the SRAM memory pool; if the SRAM is not idle enough, degrade to DDR allocation; return the type of memory pool (SRAM / DDR) successfully allocated and the starting offset address offset.

[0105] The BuddyAllocator module is equipped with an independent BuddyAllocator instance for each memory pool to manage a fixed-size contiguous memory space. It uses a complete binary tree structure (BitBinaryTree) to maintain the allocation state, with nodes containing fxfData (full allocation flag), splitData (split status), and allocData (actual allocation information). Free blocks are organized in linked lists on different levels of buddyRules. During allocation, borrowFromBuddy(size) is called to find the most suitable partner block for splitting. During release, reclaimFromPartners(offset) is called to determine if adjacent blocks are free; if they can be merged, they are merged upwards to form a larger free block, reducing fragmentation.

[0106] This Pass first creates two types of tensor lists: the task input / output tensor list (defaultly allocated to the DDR memory pool, with its MemoryType forced to DDR); the intermediate result tensor list (excluding ConstOp, sorted in ascending order of beginSlot, for example, the tensor with beginSlot = 4 is ranked before the tensor with beginSlot = 7). For the tensor groups in the alias set AliasSet (such as {%0, %1, %2}), the need_alloc attribute is used to constrain the memory allocation request to be initiated only for the tensor with the largest dataSize (such as %2). Subsequently, the StaticMemoryPlanner solver is initialized, passing in the above two types of tensors to be allocated, and two-level memory pool resources are configured: the ResourceManger module configures the SRAM memory pool capacity and the DDR memory pool capacity, and maintains the aliveBuffer set (stores the SSA values of the current active buffers); when the GreedyEviction policy module processes the intermediate result tensor list in order, for each tensor, it first executes checkRelease(currentSlot): traverses the aliveBuffer set, marks the buffers with endSlot ≤ currentSlot as recyclable (for example, when currentSlot = 9, the buffer with endSlot = 8 is recycled), and calls BuddyAllocator::deallocate to release its memory space; then it executes resolvePlacement(space, preferredType = SRAM): if preferredType is SRAM, it calls BuddyAllocator::allocate(space) to attempt to allocate in the SRAM pool; when the free space in SRAM is insufficient (for example, the remaining space < space), it downgrades preferredType to DDR and retries the allocation.The BuddyAllocator uses a buddy system algorithm to manage the memory pool: during initialization, a complete binary tree structure is constructed, where fxfData is a bit array marking fully allocated nodes (1 indicates allocated), splitData is a bit array recording the split status (1 indicates split), and allocData stores the actual allocated block information; during allocation, borrowFromBuddy(size) is executed: calculate the minimum k value that satisfies 2^k≥size, select a block from the free linked list of buddyRule at level k, and if the linked list is empty, recursively split higher-level blocks (for example, when 2048 bytes need to be allocated, select a 2048-byte block from level k=11); during release, reclaimFromPartners(offset) is executed: check if the buddy block (offsetXORblock_size) is free, if free, merge it into a larger block and move it up to a higher-level linked list (for example, merge two 1024-byte blocks to form a 2048-byte block).

[0107] Step S05: Tensor address allocation.

[0108] Tensor address allocation is achieved through the dialect rewriting function updateSubValue(): iterates through all aliasBuffers in the AliasSet where the explicit buffer is located; broadcasts MemoryType and OffsetInMemory from explicitBuffer to all aliasBuffers; and calculates the final effective offset address finalOffset = OffsetInMemory + splitOffset by combining the splitOffset of each aliasBuffer.

[0109] As attached Figure 7 As shown, the three input and output tensors of concat are related as alias buffers and explicit buffers, respectively, sharing the same buffer address space. By properly setting the OffsetInBuffer and dataSize of the input tensors, even if the three convolution activation functions are executed simultaneously, no conflict will occur in the actual memory address space.

[0110] Calling the getMeminfo() interface can print key memory information for each tensor, including: memory type (memtype), memory offset (offset), number of bytes required (size), and live range (liveRange). The output can be used for debugging, performance evaluation, or generating hardware deployment configuration files.

[0111] After completing the address allocation for the explicit buffer (e.g., %buf0), the updateSubValue process is executed: the MemoryType and OffsetInMemory of %buf0 are broadcast to all associated tensors in the alias set (e.g., %buf1, %buf2), and the final offset address is calculated based on the splitOffset of each tensor (e.g., finalOffset of %buf1 = OffsetInMemory + 128); if the tensor has a subview relationship (e.g., tensors generated by SubviewOp), the dimension offset calculation is additionally added (e.g., finalOffset += (dim0*stride0 + dim1*stride1)). Finally, all operations in the task are traversed (excluding ConstOp), and the getMeminfo interface is called to output the complete memory information of each tensor: including memtype (e.g., SRAM), offset (hexadecimal representation, e.g., 0x1A00), size (number of bytes, e.g., 4096), and liveRange (e.g., [4,15]). This information is used to generate the hardware memory mapping configuration file and verify the validity of the allocation results.

[0112] This embodiment, by building a custom dialect conversion within a mlir-opt-like compiler, reduces peak SRAM memory usage by 10-15% compared to traditional non-mlir compilation schemes, and reduces the number of MLIR intermediate representation instructions by 35% compared to memref dialects, providing a lightweight compilation solution for efficient memory management on edge AI hardware.

[0113] Obviously, the above embodiments are merely illustrative examples for clear explanation and are not intended to limit the implementation. Those skilled in the art will recognize that other variations or modifications can be made based on the above description. It is neither necessary nor possible to exhaustively list all possible implementations here. However, obvious variations or modifications derived therefrom are still within the scope of protection of this invention.

Claims

1. A method for designing MLIR dialects oriented towards static memory planning, characterized in that, Includes the following steps: Step S01: Buffering conversion; extend operation attributes and value type attributes, and retain the static single assignment form of the intermediate representation during the conversion process; Step S02: Alias ​​Analysis; Using the sparse data flow analysis framework of MLIR, gradually construct the alias equivalence set of values; determine whether the value definition operation requires explicit buffer allocation; Step S03: Lifecycle analysis; Using the MLIR lifecycle analysis framework, the lifecycle interval representation of the buffer is optimized by topological sorting, and then the integer intervals of the alias set are merged. Step S04: Solve the memory planning problem; Configure the amount of hardware memory resources, add the explicit buffer to the allocation list, know the space size and lifespan range, run the memory reuse strategy and memory pool management algorithm to obtain the memory address and memory level of the buffer; Step S05: Tensor address allocation; Copy the explicitly allocated buffer address and memory level to the alias operand, and read the offset in the buffer of the alias operand, and sum them to obtain the address of the tensor.

2. The MLIR dialect design method according to claim 1, characterized in that: The memory access attributes are marked; the extended operation attributes are added by parsing the semantics of the operation types in the IR file, and two new memory access attributes for operations are added; the extended value attributes are added by parsing the semantics of the operations in the IR file, and key variables for memory allocation are added.

3. The MLIR dialect design method according to claim 1, characterized in that: The alias analysis is implemented using the MLIR sparse forward data flow analysis framework. The static single assignments of the input and output of computational operations are taken as program points, and the alias equivalence set of values ​​is taken as the analysis state. Based on the memory access attributes of the operation, it is determined whether the operands and return values ​​share the memory buffer. The alias equivalence set of values ​​is constructed step by step, and it is ensured that there is one and only one static single assignment in the alias set that requires explicit memory allocation.

4. The MLIR dialect design method according to claim 1, characterized in that: In the lifecycle analysis, a depth-first search algorithm is used to re-mark the execution order of operations, and the lifecycle intervals of static single assignments in the alias set are merged to update the lifecycle attributes of itself and the static single assignments of aliases.

5. The MLIR dialect design method according to claim 1, characterized in that: The memory planning solution employs the classic greedy eviction algorithm for memory reuse, reading lifecycle variables and performing reclamation; the memory pool management algorithm employs the classic buddy allocation algorithm, reading memory space size variables and performing allocation and reclamation.

6. An MLIR dialect system oriented towards static memory planning, serving the method of any one of claims 1 to 5, characterized in that, include: The buffering pass, connected to the graph computation fusion compiler, implements the buffering step, alias analysis step, and lifecycle analysis step. The memory allocation pass, connected to the buffered analysis pass, implements the memory planning solution step and the tensor address allocation step.

7. The MLIR dialect design system according to claim 6, characterized in that, Also includes: The mlir-opt-like compiler is connected to the graph computation fusion compiler, buffered analysis pass, and static memory planning transformation pass, respectively, and is configured to coordinate the execution order and data transfer of each pass; The memory resource management module is connected to the static memory planning conversion module, configured as the memory resource parameters of the system hardware, manages memory pool objects at different levels, manages active buffer objects in the memory pool, and handles buffer allocation and release requests.

8. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the program is executed by the processor, it implements the method as described in any one of claims 1 to 5.

Images