RLSRA: spill / reload trashed live registers

This commit add support for spilling / reloading trashed registers that
contain live values. Mostly, this is the case of volatile (and argument)
registers being (potentially) thrashed by calls.

This is only supported by reverse linear scan allocator.

Implementation note: for now, spill slot is allocated as unnamed automatic
on a stack for each (spilled) interval. For actually spilling, it used
`TRCodeGenerator >> registerStore:to:` and `#registerLoad:from:`. This may
not work in all cases with all linkages - some linkages may prescribe
where to spill registers (PPC64 ELF v2 ABI?, Microsoft x64?).

src/Tinyrossa-RISCV/TRRV64GCodeGenerator.class.st

@@ -2,6 +2,7 @@ Class {
  #name : #TRRV64GCodeGenerator,
  #superclass : #TRCodeGenerator,
  #pools : [
+ 'TRDataTypes',
  'TRIntLimits',
  'TRRV64GISALimits',
  'TRRV64GRegisters'
@@ -121,6 +122,30 @@ TRRV64GCodeGenerator >> registerCopyFrom: srcReg to: dstReg [
  generate addi: dstReg, srcReg, 0
 ]
 
+{ #category : #utilities }
+TRRV64GCodeGenerator >> registerLoad: reg from: sym [
+ | offset |
+
+ self assert: reg isTRRegister.
+ self assert: sym isTRAutomaticSymbol.
+ self assert: sym type == Address.
+
+ offset := AcDSLSymbol value: sym name.
+ generate ld: reg, (sp + offset).
+]
+
+{ #category : #utilities }
+TRRV64GCodeGenerator >> registerStore: reg to: sym [
+ | offset |
+
+ self assert: reg isTRRegister.
+ self assert: sym isTRAutomaticSymbol.
+ self assert: sym type == Address.
+
+ offset := AcDSLSymbol value: sym name.
+ generate sd: reg, (sp + offset).
+]
+
 { #category : #'registers-private' }
 TRRV64GCodeGenerator >> virtualRegistersAssignedByProcessorInstruction: instruction do: block [ 
  self assert: instruction isProcessorInstruction.

src/Tinyrossa/TRCodeGenerator.class.st

@@ -305,6 +305,26 @@ TRCodeGenerator >> registerCopyFrom: srcReg to: dstReg [
  self subclassResponsibility
 ]
 
+{ #category : #utilities }
+TRCodeGenerator >> registerLoad: reg from: sym [
+ "Load value of given symbol `sym` into a register `reg`.
+
+ Currently, symbol must be an automatic and of type Address.
+ This method is used to implement register reloads.
+ "
+ self subclassResponsibility.
+]
+
+{ #category : #utilities }
+TRCodeGenerator >> registerStore: reg to: sym [
+ "Store value of given register `reg` into a symbol `sym`.
+
+ Currently, symbol must be an automatic and of type Address.
+ This method is used to implement register spills.
+ " 
+ ^ self subclassResponsibility
+]
+
 { #category : #accessing }
 TRCodeGenerator >> virtualRegisters [
  ^ virtualRegisters

src/Tinyrossa/TRRegisterLiveInterval.class.st

@@ -11,7 +11,9 @@ Class {
  #instVars : [
  'register',
  'start',
- 'stop'
+ 'stop',
+ 'spilled',
+ 'spillSlot'
  ],
  #category : #'Tinyrossa-Codegen-Register Allocation'
 }
@@ -33,6 +35,7 @@ TRRegisterLiveInterval >> initializeWithRegister: aTRVirtualRegister [
  register := aTRVirtualRegister.
  start := SmallInteger maxVal.
  stop := 0.
+ spilled := false.
 ]
 
 { #category : #'printing & storing' }
@@ -54,6 +57,35 @@ TRRegisterLiveInterval >> register [
  ^ register
 ]
 
+{ #category : #accessing }
+TRRegisterLiveInterval >> spillSlot [
+ ^ spillSlot
+]
+
+{ #category : #accessing }
+TRRegisterLiveInterval >> spillSlot: aTRAutomaticSymbol [
+ self assert: aTRAutomaticSymbol isTRAutomaticSymbol.
+ self assert: spillSlot isNil.
+
+ spillSlot := aTRAutomaticSymbol.
+]
+
+{ #category : #accessing }
+TRRegisterLiveInterval >> spilled [
+ ^ spilled
+]
+
+{ #category : #accessing }
+TRRegisterLiveInterval >> spilled: aBoolean [
+ self assert: aBoolean = spilled not.
+ self assert: spillSlot isTRAutomaticSymbol.
+
+ spilled := aBoolean.
+ spilled ifTrue: [ 
+ spillSlot incUseCount.
+ ].
+]
+
 { #category : #accessing }
 TRRegisterLiveInterval >> start [
  ^ start

src/Tinyrossa/TRReverseLinearScanRegisterAllocator.class.st

@@ -24,33 +24,33 @@ Class {
  'allocatableRegisters',
  'availableRegisters'
  ],
+ #pools : [
+ 'TRDataTypes',
+ 'TRRegisterKinds'
+ ],
  #category : #'Tinyrossa-Codegen-Register Allocation'
 }
 
 { #category : #allocation }
 TRReverseLinearScanRegisterAllocator >> allocateRegister: interval [
  "Allocate register for given `interval`."
 
+ | assigned |
+
  self assert: interval register allocation isNil.
- self assert: availableRegisters notEmpty.
 
- allocatableRegisters do: [:rReg | 
- (availableRegisters includes: rReg) ifTrue: [ 
- interval register allocation: rReg.
- availableRegisters remove: rReg.
- live add: interval.
- ^ self.
- ].
- ].
- self error: 'Should not happen!'.
+ assigned := self pickRegister: interval.
+ assigned isNil ifTrue: [ self error: 'No available register!' ].
+
+ live add: interval.
 ]
 
 { #category : #allocation }
 TRReverseLinearScanRegisterAllocator >> allocateRegisters [
  instructions := codegen instructions.
  allocatableRegisters := codegen linkage allocatableRegisters.
  codegen compilation config stressRA ifTrue: [ 
- allocatableRegisters := allocatableRegisters copyFrom: 1 to: 2.
+ "allocatableRegisters := allocatableRegisters copyFrom: 1 to: 2."
  ].
  availableRegisters := allocatableRegisters asSet. 
  live := SortedCollection sortBlock: [ :a :b | a start < b start ].
@@ -96,19 +96,40 @@ TRReverseLinearScanRegisterAllocator >> allocateRegistersAt: insnIndex [
  When reading this code, keep in mind that we progress in reverse order, 
  from last to first instruction!
  "
- | insn deps |
+ | insn deps liveAcross thrashed |
 
  insn := instructions at: insnIndex.
  deps := insn dependencies.
 
- "Satisfy post-dependencies, i.e., move values from fixed (real)
- registers to desired (virtual) registers"
+ "Satisfy post-dependencies, i.e., 
+ (i) move values from fixed (real) registers to desired 
+ (virtual) registers and...
+ (ii) ...reload all trashed registers live across 
+ this instruction
+ "
  deps notEmptyOrNil ifTrue: [
+ "Compute 'live-across' intervals, that is intervals that are
+ assigned before this instruction and used after this instruction."
+ liveAcross := Set new: live size.
+ live do: [:i | (i start < insnIndex and: [ i stop > insnIndex ]) ifTrue:[liveAcross add:i] ].
+
+ thrashed := OrderedCollection new: deps post size.
  codegen cursor: insnIndex.
  deps post do: [:dep |
  dep isDependency ifTrue:[
  self insertMoveFrom: dep rreg to: dep vreg.
  ].
+ dep isTrash ifTrue: [
+ liveAcross do:[:i |
+ (i register allocation == dep rreg) ifTrue: [
+ "Live-across register is trashed, we have to spill and reload. 
+ So reload here and note the it has to be spilled before this 
+ instruction executes (see handling of pre-dependencies below)"
+ self insertReload: i.
+ thrashed add: i.
+ ].
+ ]
+ ].
  ].
  ].
 
@@ -123,15 +144,29 @@ TRReverseLinearScanRegisterAllocator >> allocateRegistersAt: insnIndex [
  self allocateRegister: intervals removeLast.
  ].
 
- "Satisfy pre-dependencies, i.e., move values from (virtual) registers
- to desired (real) registers. They're placed *before* the instruction
- being processed, hence the `insnIndex - 1`"
+ "Satisfy pre-dependencies, i.e., 
+ (i) move values from (virtual) registers to desired 
+ (real) registers and...
+ (ii) spill all thrashed live registers.
+
+ Moves and spills must be placed placed *before* 
+ the instruction being processed, hence the `insnIndex - 1`"
  deps notEmptyOrNil ifTrue: [
  codegen cursor: insnIndex - 1.
- deps pre do: [:dep |
+ deps pre reverseDo: [:dep |
  dep isDependency ifTrue:[
  self insertMoveFrom: dep vreg to: dep rreg.
  ].
+ dep isTrash ifTrue: [
+ thrashed do:[:i |
+ (i register allocation == dep rreg) ifTrue: [
+ "Live register is trashed and has to be spilled. 
+ See handling of post-dependencies above where `spilled` set
+ is populated."
+ self insertSpill: i.
+ ].
+ ]
+ ].
  ].
  ].
 ]
@@ -143,6 +178,97 @@ TRReverseLinearScanRegisterAllocator >> expireRegistersAt: insnIndex [
  | expired |
 
  expired := live removeLast.
- availableRegisters add: expired register allocation
+ self freeRegister: expired.
  ].
 ]
+
+{ #category : #utilities }
+TRReverseLinearScanRegisterAllocator >> freeRegister: interval [
+ "Free register assigned to given interval, i.e.,
+ put it back to list of available registers."
+
+ self assert: interval register allocation notNil.
+
+ availableRegisters add: interval register allocation
+]
+
+{ #category : #utilities }
+TRReverseLinearScanRegisterAllocator >> insertReload: interval [
+ | slot |
+
+ self assert: interval spilled not.
+ self assert: interval register kind == GPR description: 'FIXME: FPRs not yet supported'.
+ self assert: interval register allocation notNil.
+
+ slot := interval spillSlot.
+ slot isNil ifTrue: [ 
+ slot := codegen compilation symbolManager defineAutomatic: nil type: Address.
+ interval spillSlot: slot.
+ ].
+ codegen registerLoad: interval register from: slot.
+ interval spilled: true.
+]
+
+{ #category : #utilities }
+TRReverseLinearScanRegisterAllocator >> insertSpill: interval [
+ | slot |
+
+ self assert: interval spilled.
+ self assert: interval spillSlot isTRAutomaticSymbol.
+ self assert: interval register kind == GPR description: 'FIXME: FPRs not yet supported'.
+
+ slot := interval spillSlot.
+ codegen registerStore: interval register to: slot.
+ interval spilled: false.
+]
+
+{ #category : #utilities }
+TRReverseLinearScanRegisterAllocator >> pickRegister: interval [
+ "Pick (choose) and assign the best real register for given live interval.
+ Return the chosen register.
+
+ Internal list of currently available registers is updated accordingly.
+
+ If there's no available register at this point, return `nil`. Caller
+ is responsible for handling this case and schedule a spill / reload.
+ "
+ self assert: interval register allocation isNil.
+
+ availableRegisters isEmpty ifTrue: [ ^ nil ].
+
+ allocatableRegisters do: [:rReg | 
+ (availableRegisters includes: rReg) ifTrue: [ 
+ interval register allocation: rReg.
+ self takeRegister: interval.
+ ^ rReg
+ ].
+ ]. 
+ self assert: false description: 'Should never be reached'.
+]
+
+{ #category : #utilities }
+TRReverseLinearScanRegisterAllocator >> pickSpill: interval [
+ "Pick (choose) and return the best spill FIXME: TBW"
+
+ | insn candidates |
+
+ insn := instructions at: interval stop.
+
+ candidates := live reject: [:each | each spilled ].
+ codegen virtualRegistersReadBy: insn do: [:vReg |
+ live do: [:e | e register == vReg ifTrue:[ candidates remove: e ifAbsent:nil] ]
+ ].
+ candidates isEmpty ifTrue: [ 
+ ^ nil.
+ ].
+ ^ candidates first.
+]
+
+{ #category : #utilities }
+TRReverseLinearScanRegisterAllocator >> takeRegister: interval [
+ "Mark register assigned to given interval as used."
+
+ self assert: interval register allocation notNil.
+
+ availableRegisters remove: interval register allocation
+]