allocs.go raw

   1  package transform
   2  
   3  // This file implements an escape analysis pass. It looks for calls to
   4  // runtime.alloc and replaces these calls with a stack allocation if the
   5  // allocated value does not escape. It uses the LLVM nocapture flag for
   6  // interprocedural escape analysis.
   7  
   8  import (
   9  	"fmt"
  10  	"go/token"
  11  	"regexp"
  12  
  13  	"tinygo.org/x/go-llvm"
  14  )
  15  
  16  // OptimizeAllocs tries to replace heap allocations with stack allocations
  17  // whenever possible. It relies on the LLVM 'nocapture' flag for interprocedural
  18  // escape analysis, and within a function looks whether an allocation can escape
  19  // to the heap.
  20  // If printAllocs is non-nil, it indicates the regexp of functions for which a
  21  // heap allocation explanation should be printed (why the object can't be stack
  22  // allocated).
  23  func OptimizeAllocs(mod llvm.Module, printAllocs *regexp.Regexp, maxStackAlloc uint64, logger func(token.Position, string)) {
  24  	allocator := mod.NamedFunction("runtime.alloc")
  25  	if allocator.IsNil() {
  26  		// nothing to optimize
  27  		return
  28  	}
  29  
  30  	targetData := llvm.NewTargetData(mod.DataLayout())
  31  	defer targetData.Dispose()
  32  	ctx := mod.Context()
  33  	builder := ctx.NewBuilder()
  34  	defer builder.Dispose()
  35  
  36  	// Determine the maximum alignment on this platform.
  37  	complex128Type := ctx.StructType([]llvm.Type{ctx.DoubleType(), ctx.DoubleType()}, false)
  38  	maxAlign := int64(targetData.ABITypeAlignment(complex128Type))
  39  
  40  	for _, heapalloc := range getUses(allocator) {
  41  		logAllocs := printAllocs != nil && printAllocs.MatchString(heapalloc.InstructionParent().Parent().Name())
  42  		if heapalloc.Operand(0).IsAConstantInt().IsNil() {
  43  			// Do not allocate variable length arrays on the stack.
  44  			if logAllocs {
  45  				logAlloc(logger, heapalloc, "size is not constant")
  46  			}
  47  			continue
  48  		}
  49  
  50  		size := heapalloc.Operand(0).ZExtValue()
  51  		if size > maxStackAlloc {
  52  			// The maximum size for a stack allocation.
  53  			if logAllocs {
  54  				logAlloc(logger, heapalloc, fmt.Sprintf("object size %d exceeds maximum stack allocation size %d", size, maxStackAlloc))
  55  			}
  56  			continue
  57  		}
  58  
  59  		if size == 0 {
  60  			// If the size is 0, the pointer is allowed to alias other
  61  			// zero-sized pointers. Use the pointer to the global that would
  62  			// also be returned by runtime.alloc.
  63  			zeroSizedAlloc := mod.NamedGlobal("runtime.zeroSizedAlloc")
  64  			if !zeroSizedAlloc.IsNil() {
  65  				heapalloc.ReplaceAllUsesWith(zeroSizedAlloc)
  66  				heapalloc.EraseFromParentAsInstruction()
  67  			}
  68  			continue
  69  		}
  70  
  71  		// In general the pattern is:
  72  		//     %0 = call i8* @runtime.alloc(i32 %size, i8* null)
  73  		//     %1 = bitcast i8* %0 to type*
  74  		//     (use %1 only)
  75  		// But the bitcast might sometimes be dropped when allocating an *i8.
  76  		// The 'bitcast' variable below is thus usually a bitcast of the
  77  		// heapalloc but not always.
  78  		bitcast := heapalloc // instruction that creates the value
  79  		if uses := getUses(heapalloc); len(uses) == 1 && !uses[0].IsABitCastInst().IsNil() {
  80  			// getting only bitcast use
  81  			bitcast = uses[0]
  82  		}
  83  
  84  		if at := valueEscapesAt(bitcast); !at.IsNil() {
  85  			if logAllocs {
  86  				atPos := getPosition(at)
  87  				msg := "escapes at unknown line"
  88  				if atPos.Line != 0 {
  89  					msg = fmt.Sprintf("escapes at line %d", atPos.Line)
  90  				}
  91  				logAlloc(logger, heapalloc, msg)
  92  			}
  93  			continue
  94  		}
  95  		// The pointer value does not escape.
  96  
  97  		// Determine the appropriate alignment of the alloca.
  98  		attr := heapalloc.GetCallSiteEnumAttribute(0, llvm.AttributeKindID("align"))
  99  		alignment := int(maxAlign)
 100  		if !attr.IsNil() {
 101  			// 'align' return value attribute is set, so use it.
 102  			// This is basically always the case, but to be sure we'll default
 103  			// to maxAlign if it isn't.
 104  			alignment = int(attr.GetEnumValue())
 105  		}
 106  
 107  		// Insert alloca in the entry block. Do it here so that mem2reg can
 108  		// promote it to a SSA value.
 109  		fn := bitcast.InstructionParent().Parent()
 110  		builder.SetInsertPointBefore(fn.EntryBasicBlock().FirstInstruction())
 111  		allocaType := llvm.ArrayType(mod.Context().Int8Type(), int(size))
 112  		alloca := builder.CreateAlloca(allocaType, "stackalloc")
 113  		alloca.SetAlignment(alignment)
 114  
 115  		// Zero the allocation inside the block where the value was originally allocated.
 116  		zero := llvm.ConstNull(alloca.AllocatedType())
 117  		builder.SetInsertPointBefore(bitcast)
 118  		store := builder.CreateStore(zero, alloca)
 119  		store.SetAlignment(alignment)
 120  
 121  		// Replace heap alloc bitcast with stack alloc bitcast.
 122  		bitcast.ReplaceAllUsesWith(alloca)
 123  		if heapalloc != bitcast {
 124  			bitcast.EraseFromParentAsInstruction()
 125  		}
 126  		heapalloc.EraseFromParentAsInstruction()
 127  	}
 128  }
 129  
 130  // valueEscapesAt returns the instruction where the given value may escape and a
 131  // nil llvm.Value if it definitely doesn't. The value must be an instruction.
 132  func valueEscapesAt(value llvm.Value) llvm.Value {
 133  	uses := getUses(value)
 134  	for _, use := range uses {
 135  		if use.IsAInstruction().IsNil() {
 136  			panic("expected instruction use")
 137  		}
 138  		switch use.InstructionOpcode() {
 139  		case llvm.GetElementPtr:
 140  			if at := valueEscapesAt(use); !at.IsNil() {
 141  				return at
 142  			}
 143  		case llvm.BitCast:
 144  			// A bitcast escapes if the casted-to value escapes.
 145  			if at := valueEscapesAt(use); !at.IsNil() {
 146  				return at
 147  			}
 148  		case llvm.Load:
 149  			// Load does not escape.
 150  		case llvm.Store:
 151  			// Store only escapes when the value is stored to, not when the
 152  			// value is stored into another value.
 153  			if use.Operand(0) == value {
 154  				return use
 155  			}
 156  		case llvm.Call:
 157  			if !hasFlag(use, value, "nocapture") {
 158  				return use
 159  			}
 160  		case llvm.ICmp:
 161  			// Comparing pointers don't let the pointer escape.
 162  			// This is often a compiler-inserted nil check.
 163  		default:
 164  			// Unknown instruction, might escape.
 165  			return use
 166  		}
 167  	}
 168  
 169  	// Checked all uses, and none let the pointer value escape.
 170  	return llvm.Value{}
 171  }
 172  
 173  // logAlloc prints a message to stderr explaining why the given object had to be
 174  // allocated on the heap.
 175  func logAlloc(logger func(token.Position, string), allocCall llvm.Value, reason string) {
 176  	logger(getPosition(allocCall), "object allocated on the heap: "+reason)
 177  }
 178