已合并
[AscendNPU IR]: fix: Fix PlanMemory overflow because of 0-size temp buffer when memory equals size of memory scope #1367
hujiajun创建于 7月5日
[AscendNPU IR]: fix: Fix PlanMemory overflow because of 0-size temp buffer when memory equals size of memory scope #1367
已合并
共 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 | // ----- |
| 36 | module { | 91 | module { |
| 37 | // CHECK-LABEL: func.func @test_infer_mem_allocate_loop_conflict | 92 | // CHECK-LABEL: func.func @test_infer_mem_allocate_loop_conflict |