Improve register allocation (part 1)

src/Tinyrossa-POWER/TRPPC64CodeGenerator.class.st

@@ -50,35 +50,30 @@ TRPPC64CodeGenerator >> loadConstant32: value into: reg [
 	].
 ]
 
-{ #category : #registers }
-TRPPC64CodeGenerator >> virtualRegistersModifiedBy: instruction do: block [
-	"Evaluate block for each virtual register modified by
-	 given instruction."
-
-	instruction isPseudoInstruction "such as label" ifTrue: [ 
-		^ self
-	]. 
-
-	instruction externalBindings keysAndValuesDo: [ :name :value |
-		name = 'rt' ifTrue: [ 
-			(value isBitVector and: [ value isSymbolic and: [ value isConstant ] ]) ifTrue: [ 
-				| vReg |
+{ #category : #utilities }
+TRPPC64CodeGenerator >> registerCopyFrom: srcReg to: dstReg [
+	generate mr: dstReg, srcReg
+]
 
-				vReg := virtualRegisters at: value sym ifAbsent: nil.
-				block value: vReg.
+{ #category : #'registers-private' }
+TRPPC64CodeGenerator >> virtualRegistersAssignedByProcessorInstruction: instruction do: block [ 
+	self assert: instruction isProcessorInstruction.
+	instruction externalBindings 
+		keysAndValuesDo: [:name :value | 
+			name = 'rt' ifTrue: [
+				(value isBitVector and: [ value isSymbolic and: [ value isConstant ] ]) ifTrue: [
+					| vReg |
+
+					vReg := virtualRegisters at: value sym ifAbsent: nil.
+					block value: vReg.
+				].
 			].
 		].
-	].
 ]
 
-{ #category : #registers }
-TRPPC64CodeGenerator >> virtualRegistersReadBy: instruction do: block [
-	"Evaluate block for each virtual register read by
-	 given instruction."
-
-	instruction isPseudoInstruction "such as label" ifTrue: [ 
-		^ self
-	]. 
+{ #category : #'registers-private' }
+TRPPC64CodeGenerator >> virtualRegistersReadByProcessorInstruction: instruction do: block [
+	self assert: instruction isProcessorInstruction.
 
 	instruction externalBindings keysAndValuesDo: [ :name :value |
 		(#('ra' 'rb' 'rc' 'rs') includes: name) ifTrue: [ 

src/Tinyrossa-POWER/TRPPC64PSABILinkage.class.st

@@ -55,10 +55,6 @@ TRPPC64PSABILinkage >> generateCall: node [
 
 { #category : #'code generation' }
 TRPPC64PSABILinkage >> generateEpilogue: valReg [
-	"Move value to ABI return register..."
-	generate
-		mr: gr4, valReg.
-
 	"Restore preserved registers"
 	self preservedRegisters do: [:reg |
 		(self allocatedRegisters includes: reg) ifTrue: [ 
@@ -120,6 +116,23 @@ TRPPC64PSABILinkage >> generatePrologue [
 	].
 ]
 
+{ #category : #'code generation' }
+TRPPC64PSABILinkage >> generateReturn: node [
+	| srcReg leave deps |
+
+	self assert: codegen compilation functionType == node child1 type.
+
+	srcReg := codegen evaluator evaluate: node child1.
+
+	deps := TRRegisterDependencies new.
+	deps pre addDependency: srcReg on: gr4.
+
+	leave := generate leave: srcReg.   
+	leave dependencies: deps.
+
+	^ nil
+]
+
 { #category : #mapping }
 TRPPC64PSABILinkage >> mapParameters: parameterTypes [ 
 	"Map parameters to argument registers.

src/Tinyrossa-RISCV/TRRV64GCodeGenerator.class.st

@@ -2,6 +2,7 @@ Class {
 	#name : #TRRV64GCodeGenerator,
 	#superclass : #TRCodeGenerator,
 	#pools : [
+		'TRDataTypes',
 		'TRIntLimits',
 		'TRRV64GISALimits',
 		'TRRV64GRegisters'
@@ -114,35 +115,56 @@ TRRV64GCodeGenerator >> loadConstant64: value into: reg [
 	].
 ]
 
-{ #category : #registers }
-TRRV64GCodeGenerator >> virtualRegistersModifiedBy: instruction do: block [
-	"Evaluate block for each virtual register modified by
-	 given instruction."
+{ #category : #utilities }
+TRRV64GCodeGenerator >> registerCopyFrom: srcReg to: dstReg [
+	"Copy value of source register into destination register"
 
-	instruction isPseudoInstruction "such as label" ifTrue: [ 
-		^ self
-	]. 
+	generate addi: dstReg, srcReg, 0
+]
 
-	instruction externalBindings keysAndValuesDo: [ :name :value |
-		name = 'rd' ifTrue: [ 
-			(value isBitVector and: [ value isSymbolic and: [ value isConstant ] ]) ifTrue: [ 
-				| vReg |
+{ #category : #utilities }
+TRRV64GCodeGenerator >> registerLoad: reg from: sym [
+	| offset |
 
-				vReg := virtualRegisters at: value sym ifAbsent: nil.
-				block value: vReg.
+	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.
+	instruction externalBindings 
+		keysAndValuesDo: [:name :value | 
+			name = 'rd' ifTrue: [
+				(value isBitVector and: [ value isSymbolic and: [ value isConstant ] ]) ifTrue: [
+					| vReg |
+
+					vReg := virtualRegisters at: value sym ifAbsent: nil.
+					block value: vReg.
+				].
 			].
 		].
-	].
 ]
 
-{ #category : #registers }
-TRRV64GCodeGenerator >> virtualRegistersReadBy: instruction do: block [
-	"Evaluate block for each virtual register read by
-	 given instruction."
-
-	instruction isPseudoInstruction "such as label" ifTrue: [ 
-		^ self
-	]. 
+{ #category : #'registers-private' }
+TRRV64GCodeGenerator >> virtualRegistersReadByProcessorInstruction: instruction do: block [
+	self assert: instruction isProcessorInstruction.
 
 	instruction externalBindings keysAndValuesDo: [ :name :value |
 		(#('rs1' 'rs2' 'rs3') includes: name) ifTrue: [ 

src/Tinyrossa-RISCV/TRRV64GPSABILinkage.class.st

@@ -21,8 +21,19 @@ Class {
 
 { #category : #accessing }
 TRRV64GPSABILinkage >> allocatableRegisters [
-	"^ { t0 . t1 . t2 . t3 . t4 . t5 . t6 }"
-	^ self preservedRegisters
+	(codegen linkRegisterKilled not or:[codegen compilation config stressRA])  ifTrue:[
+		"For leaf methods (or when we want to stress RA), we prefer volatile registers
+		 over preserved registers. This might save us a need to spill / reload (preserved) 
+		 registers in prologue / epilogue for small methods."
+
+		^ self volatileRegisters , self preservedRegisters
+	] ifFalse:[
+		"For non-leaf methods we prefer preserved registers over volatile registers. 
+		 This might save us a need to spill / reload (volatile) registers at call 
+		 instructions for small functions."
+
+		^ self preservedRegisters , self volatileRegisters
+	]
 ]
 
 { #category : #accessing }
@@ -35,14 +46,18 @@ TRRV64GPSABILinkage >> allocatedPreservedRegisters [
 
 { #category : #'code generation' }
 TRRV64GPSABILinkage >> generateCall: node [
-	"Also handles indirect calls"
-	| indirect firstParameterIndex parameterVregs parameterTypes parameterRregs retVreg|
+	"Also handles indirect calls."
+	| indirect deps firstParameterIndex parameterVregs parameterTypes parameterRregs call retVreg |
 
 	indirect := node opcode isIndirect.
+	deps := TRRegisterDependencies new.
+
+	"Step 1: collect parameters and map parameters to parameter registers
+	 (and stack slots, thought that's not supported yet)"
+
 	"If this is an indirect call, the first child is the address of the callee."
 	firstParameterIndex := indirect ifTrue: [ 2 ] ifFalse: [ 1 ].
 
-
 	parameterVregs := Array new: node children size - firstParameterIndex + 1.
 	parameterTypes := Array new: node children size - firstParameterIndex + 1.
 
@@ -55,30 +70,47 @@ TRRV64GPSABILinkage >> generateCall: node [
 		parameterTypes at: i - firstParameterIndex + 1 
 					  put: parameter type.
 	].
-
 	parameterRregs := self mapParameters: parameterTypes.
-	parameterVregs with: parameterRregs do: [:valReg :paramReg | 
-		generate addi: paramReg , valReg, 0
+
+	"Step 2: map parameters into parameter registers and
+	 thrash all other unused parameter registers and volatile
+	 registers."
+	parameterVregs with: parameterRregs do: [:vReg :rReg |
+		deps pre addDependency: vReg on: rReg. 
+		rReg ~~ a0 ifTrue: [
+			deps post addDependency: vReg on: rReg.
+		]
+	].
+	self parameterRegisters do: [:rReg |  
+		(parameterRregs includes: rReg) ifFalse: [  
+			deps pre addTrashed: rReg. 
+			deps post addTrashed: rReg.
+		].
+	].
+	self volatileRegisters do: [:rReg | 
+		deps pre addTrashed: rReg.
+		deps post addTrashed: rReg.
 	].
 
+	"Step 3: generate all instruction"
 	indirect ifTrue: [ 
 		| addrReg |
 
 		addrReg := codegen evaluator evaluate: node child1.
-		generate jalr: ra, addrReg, 0.
+		call := generate jalr: ra, addrReg, 0.
 	] ifFalse: [
 		"If the call a recursive call..."
 		node symbol == codegen compilation functionSymbol ifTrue: [ 
 			"...then generate 'jal ra, <function>'..."
-			generate jal: ra, node symbol.
+			call := generate jal: ra, node symbol.
 		] ifFalse: [
 			"...otherwise..."
 			codegen compilation isAOT ifTrue: [ 
 				"
 				In AOT mode we generate call and let the (runtime) linker 
 				to properly relocate it.
 				"
-				generate call: node symbol
+				call := generate call: node symbol
 			] ifFalse: [ 
 				"
 				In JIT mode we load address directly into 'ra' 
@@ -93,17 +125,24 @@ TRRV64GPSABILinkage >> generateCall: node [
 				self assert: node symbol address notNil description: 'No address set for function symbol'.
 
 				codegen loadConstant64: node symbol address into: ra.
-				generate jalr: ra, ra, 0.
+				call := generate jalr: ra, ra, 0.
 			].
 		].
 	].
+	call dependencies: deps.
 
+
+	"Note that link register has been overwritten"
+	codegen linkRegisterKilled: true.
+
+	"Step 4: map return value into return register (if any) and finish"
 	node symbol type == Void ifTrue:[
 		retVreg := nil.
 	] ifFalse:[
 		retVreg := codegen allocateRegister.
-		generate addi: retVreg , a0, 0.
+		deps post addDependency: retVreg on: a0.
 	].
+
 	^ retVreg
 ]
 
@@ -116,11 +155,10 @@ TRRV64GPSABILinkage >> generateCallIndirect: node [
 TRRV64GPSABILinkage >> generateEpilogue: valReg [
 	| preserved offset |
 
-	"Move value to ABI return register...
-	 ...and reload link register."
-	generate
-		addi: a0, valReg, 0;
-		ld: ra, (sp + 0).      
+	"Reload link register if needed"
+	codegen linkRegisterKilled ifTrue:[
+		generate ld: ra, (sp + 0).      
+	].
 
 	"Restore preserved registers"
 	offset := framePreservedOffset.
@@ -145,7 +183,9 @@ TRRV64GPSABILinkage >> generatePrologue [
 	generate addi: sp, sp, frameSize negated.
 
 	"Save link register"
-	generate sd: ra, (sp + 0).
+	codegen linkRegisterKilled ifTrue:[
+		generate sd: ra, (sp + 0).
+	].
 
 	"Save parameters"
 	parameters := codegen compilation symbolManager lookupSymbolsByType: TRParameterSymbol.
@@ -164,6 +204,23 @@ TRRV64GPSABILinkage >> generatePrologue [
 	].
 ]
 
+{ #category : #'code generation' }
+TRRV64GPSABILinkage >> generateReturn: node [
+	| srcReg leave deps |
+
+	self assert: codegen compilation functionType == node child1 type.
+
+	srcReg := codegen evaluator evaluate: node child1.
+
+	deps := TRRegisterDependencies new.
+	deps pre addDependency: srcReg on: a0.
+
+	leave := generate leave: srcReg.   
+	leave dependencies: deps.
+
+	^ nil
+]
+
 { #category : #initialization }
 TRRV64GPSABILinkage >> initializeWithCodeGenerator: aTRCodeGenerator [ 
 	super initializeWithCodeGenerator: aTRCodeGenerator.
@@ -260,3 +317,8 @@ TRRV64GPSABILinkage >> parameterRegisters [
 TRRV64GPSABILinkage >> preservedRegisters [
 	^ { s0 . s1 . s2 . s3 . s4 . s5 . s6 . s7 . s8 . s9 . s10 . s11 }
 ]
+
+{ #category : #accessing }
+TRRV64GPSABILinkage >> volatileRegisters [
+	^ { t0 . t1 . t2 . t3 . t4 . t5 . t6 }
+]

src/Tinyrossa/AcInstruction.extension.st

@@ -1,5 +1,15 @@
 Extension { #name : #AcInstruction }
 
+{ #category : #'*Tinyrossa' }
+AcInstruction >> dependencies [
+	^ self annotationAt: TRRegisterDependencies
+]
+
+{ #category : #'*Tinyrossa' }
+AcInstruction >> dependencies: aTRRegisterDependencies [
+	self annotationAddOrReplace: aTRRegisterDependencies
+]
+
 { #category : #'*Tinyrossa' }
 AcInstruction >> isLeaveInstruction [
 	^ false

src/Tinyrossa/TRCodeEvaluator.class.st

@@ -156,13 +156,7 @@ TRCodeEvaluator >> evaluate_lreturn: node [
 
 { #category : #evaluation }
 TRCodeEvaluator >> evaluate_return: node [
-	| retReg |
-
-	self assert: codegen compilation functionType == node child1 type.
-
-	retReg := self evaluate: node child1.
-	generate leave: retReg.
-	^ nil
+	^ codegen linkage generateReturn: node.
 ]
 
 { #category : #evaluation }