blob: 58bbdc76eae9cf36745141d81b1b02f3ee4c75bd (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
|
// slang-ir-autodiff-region.cpp
#include "slang-ir-autodiff-region.h"
namespace Slang
{
RefPtr<IndexedRegionMap> buildIndexedRegionMap(IRGlobalValueWithCode* func)
{
RefPtr<IndexedRegionMap> regionMap = new IndexedRegionMap;
List<IRBlock*> workList;
regionMap->mapBlock(func->getFirstBlock(), nullptr);
workList.add(func->getFirstBlock());
while (workList.getCount() > 0)
{
auto currentBlock = workList.getLast();
workList.removeLast();
auto terminator = currentBlock->getTerminator();
auto currentRegion = regionMap->getRegion(currentBlock);
switch (terminator->getOp())
{
case kIROp_Loop:
{
auto loopRegion = regionMap->newRegion(as<IRLoop>(terminator), currentRegion);
auto condBlock = as<IRLoop>(terminator)->getTargetBlock();
regionMap->mapBlock(condBlock, loopRegion);
workList.add(condBlock);
auto ifElse = as<IRIfElse>(condBlock->getTerminator());
SLANG_RELEASE_ASSERT(ifElse);
// TODO: this is one of the places we'll need to change if we support loops that
// loop on either the true or false side. For now, we assume the loop is on the
// true side only.
//
regionMap->mapBlock(ifElse->getFalseBlock(), currentRegion);
workList.add(ifElse->getFalseBlock());
}
}
for (auto successor : currentBlock->getSuccessors())
{
// If already mapped, skip.
if (regionMap->hasMapping(successor))
continue;
regionMap->mapBlock(successor, currentRegion);
workList.add(successor);
}
}
return regionMap;
}
}; // namespace Slang
|
