已合并
[AscendNPU IR]: fix: Fix PlanMemory overflow because of 0-size temp buffer when memory equals size of memory scope #1367
[AscendNPU IR]: fix: Fix PlanMemory overflow because of 0-size temp buffer when memory equals size of memory scope #1367
已合并
hujiajun创建于 7月5日
2 个文件变更+60-1
@@ -1758,7 +1758,7 @@ bool MemPlan::IsEnoughForBuffersNoReuse(StorageEntry *rootStorageEntry,
1758 auto iter =1758 auto iter =
1759 bufferScope2RequiredSize.find(rootStorageEntry->bufInfo->bufferScope);1759 bufferScope2RequiredSize.find(rootStorageEntry->bufInfo->bufferScope);
1760 assert(iter != bufferScope2RequiredSize.end());1760 assert(iter != bufferScope2RequiredSize.end());
1761- if (iter->second < restBufferSize) {1761+ if (iter->second <= restBufferSize) {
1762 PlanBuffersWithoutReuse(rootStorageEntry, alignUnit);1762 PlanBuffersWithoutReuse(rootStorageEntry, alignUnit);
1763 return true;1763 return true;
1764 }1764 }
@@ -2117,6 +2117,10 @@ LogicalResult MemPlan::SpecAlloc(MemBoundList &outline, PlanRecHis &his,
2117 return success();2117 return success();
2118 }2118 }
2119 assert(e && "StorageEntry should not be null");2119 assert(e && "StorageEntry should not be null");
2120+ if (e->alignedConstBits == 0) {
2121+ e->bitsOffset = 0;
2122+ return success();
2123+ }
2120 for (MemBoundListConstIter start = outline.begin(); start != outline.end();2124 for (MemBoundListConstIter start = outline.begin(); start != outline.end();
2121 ++start) {2125 ++start) {
2122 uint64_t size = 0;2126 uint64_t size = 0;
@@ -32,6 +32,61 @@ module {
32 }32 }
33}33}
34 34 
35+// -----
36+module {
37+ func.func @test_mem_noreuse_max(%src : memref<16384xi64, #hivm.address_space<gm>>,
38+ %dst : memref<16384xi32, #hivm.address_space<gm>>) {
39+ // CHECK-NOT: memref.alloc()
40+ // CHECK: {{.*}} = hivm.hir.pointer_cast(%[[CONST0:.*]])
41+ %alloc = memref.alloc() : memref<16384xi64, #hivm.address_space<ub>>
42+ // CHECK: {{.*}} = hivm.hir.pointer_cast(%[[CONST1:.*]])
43+ %alloc_0 = memref.alloc() : memref<16384xi32, #hivm.address_space<ub>>
44+ // CHECK: {{.*}} = hivm.hir.pointer_cast(%[[CONST2:.*]])
45+ %alloc_1 = memref.alloc() : memref<0xi64, #hivm.address_space<ub>>
46+ hivm.hir.load ins(%src : memref<16384xi64, #hivm.address_space<gm>>)
47+ outs(%alloc : memref<16384xi64, #hivm.address_space<ub>>)
48+ hivm.hir.vcast ins(%alloc : memref<16384xi64, #hivm.address_space<ub>>)
49+ outs(%alloc_0 : memref<16384xi32, #hivm.address_space<ub>>)
50+ temp_buffer (%alloc_1 : memref<0xi64, #hivm.address_space<ub>>) round_mode = <truncwithoverflow>
51+ hivm.hir.store ins(%alloc_0 : memref<16384xi32,#hivm.address_space<ub>>)
52+ outs(%dst: memref<16384xi32,#hivm.address_space<gm>>)
53+ return
54+ }
55+}
56+ 
57+// -----
58+module {
59+ func.func @test_mem_specalloc_max(%src1 : memref<16384xi64, #hivm.address_space<gm>>,
60+ %src2 : memref<16384xi32, #hivm.address_space<gm>>,
61+ %dst : memref<16384xi32, #hivm.address_space<gm>>) {
62+ // CHECK-NOT: memref.alloc()
63+ // CHECK: %[[CONST2:.*]] = arith.constant 131072 : i64
64+ // CHECK: %[[CONST0:.*]] = arith.constant 0 : i64
65+ // CHECK: {{.*}} = hivm.hir.pointer_cast(%[[CONST1:.*]])
66+ %alloc = memref.alloc() : memref<16384xi64, #hivm.address_space<ub>>
67+ // CHECK: {{.*}} = hivm.hir.pointer_cast(%[[CONST2]])
68+ %alloc_0 = memref.alloc() : memref<16384xi32, #hivm.address_space<ub>>
69+ // CHECK: {{.*}} = hivm.hir.pointer_cast(%[[CONST0]])
70+ %alloc_1 = memref.alloc() : memref<0xi64, #hivm.address_space<ub>>
71+ hivm.hir.load ins(%src1 : memref<16384xi64, #hivm.address_space<gm>>)
72+ outs(%alloc : memref<16384xi64, #hivm.address_space<ub>>)
73+ hivm.hir.vcast ins(%alloc : memref<16384xi64, #hivm.address_space<ub>>)
74+ outs(%alloc_0 : memref<16384xi32, #hivm.address_space<ub>>)
75+ temp_buffer (%alloc_1 : memref<0xi64, #hivm.address_space<ub>>) round_mode = <truncwithoverflow>
76+ // CHECK: {{.*}} = hivm.hir.pointer_cast(%[[CONST1]])
77+ %alloc_2 = memref.alloc() : memref<16384xi32, #hivm.address_space<ub>>
78+ hivm.hir.load ins(%src2 : memref<16384xi32, #hivm.address_space<gm>>)
79+ outs(%alloc_2 : memref<16384xi32, #hivm.address_space<ub>>)
80+ // CHECK: {{.*}} = hivm.hir.pointer_cast(%[[CONST2]])
81+ %alloc_3 = memref.alloc() : memref<16384xi32, #hivm.address_space<ub>>
82+ hivm.hir.vadd ins(%alloc_0, %alloc_2 : memref<16384xi32, #hivm.address_space<ub>>, memref<16384xi32, #hivm.address_space<ub>>)
83+ outs(%alloc_3 : memref<16384xi32, #hivm.address_space<ub>>)
84+ hivm.hir.store ins(%alloc_3 : memref<16384xi32,#hivm.address_space<ub>>)
85+ outs(%dst: memref<16384xi32,#hivm.address_space<gm>>)
86+ return
87+ }
88+}
89+ 
35// -----90// -----
36module {91module {
37 // CHECK-LABEL: func.func @test_infer_mem_allocate_loop_conflict92 // CHECK-LABEL: func.func @test_infer_mem_allocate_loop_conflict