[codegen] Allow using the pinned register as the heap base via a setting;

cranelift/codegen/meta/src/shared/settings.rs

@@ -95,6 +95,21 @@ pub fn define() -> SettingGroup {
         false,
     );
 
+    settings.add_bool(
+        "use_pinned_reg_as_heap_base",
+        r#"Use the pinned register as the heap base.
+
+        Enabling this requires the enable_pinned_reg setting to be set to true. It enables a custom
+        legalization of the `heap_addr` instruction so it will use the pinned register as the heap
+        base, instead of fetching it from a global value.
+
+        Warning! Enabling this means that the pinned register *must* be maintained to contain the
+        heap base address at all times, during the lifetime of a function. Using the pinned
+        register for other purposes when this is set is very likely to cause crashes.
+        "#,
+        false,
+    );
+
     settings.add_bool("enable_simd", "Enable the use of SIMD instructions.", false);
 
     settings.add_bool(

cranelift/codegen/src/legalizer/heap.rs

@@ -14,7 +14,7 @@ pub fn expand_heap_addr(
     inst: ir::Inst,
     func: &mut ir::Function,
     cfg: &mut ControlFlowGraph,
-    _isa: &dyn TargetIsa,
+    isa: &dyn TargetIsa,
 ) {
     // Unpack the instruction.
     let (heap, offset, access_size) = match func.dfg[inst] {
@@ -32,16 +32,24 @@ pub fn expand_heap_addr(
 
     match func.heaps[heap].style {
         ir::HeapStyle::Dynamic { bound_gv } => {
-            dynamic_addr(inst, heap, offset, access_size, bound_gv, func)
-        }
-        ir::HeapStyle::Static { bound } => {
-            static_addr(inst, heap, offset, access_size, bound.into(), func, cfg)
+            dynamic_addr(isa, inst, heap, offset, access_size, bound_gv, func)
         }
+        ir::HeapStyle::Static { bound } => static_addr(
+            isa,
+            inst,
+            heap,
+            offset,
+            access_size,
+            bound.into(),
+            func,
+            cfg,
+        ),
     }
 }
 
 /// Expand a `heap_addr` for a dynamic heap.
 fn dynamic_addr(
+    isa: &dyn TargetIsa,
     inst: ir::Inst,
     heap: ir::Heap,
     offset: ir::Value,
@@ -82,11 +90,12 @@ fn dynamic_addr(
     }
     pos.ins().trapnz(oob, ir::TrapCode::HeapOutOfBounds);
 
-    compute_addr(inst, heap, addr_ty, offset, offset_ty, pos.func);
+    compute_addr(isa, inst, heap, addr_ty, offset, offset_ty, pos.func);
 }
 
 /// Expand a `heap_addr` for a static heap.
 fn static_addr(
+    isa: &dyn TargetIsa,
     inst: ir::Inst,
     heap: ir::Heap,
     offset: ir::Value,
@@ -134,11 +143,12 @@ fn static_addr(
         pos.ins().trapnz(oob, ir::TrapCode::HeapOutOfBounds);
     }
 
-    compute_addr(inst, heap, addr_ty, offset, offset_ty, pos.func);
+    compute_addr(isa, inst, heap, addr_ty, offset, offset_ty, pos.func);
 }
 
 /// Emit code for the base address computation of a `heap_addr` instruction.
 fn compute_addr(
+    isa: &dyn TargetIsa,
     inst: ir::Inst,
     heap: ir::Heap,
     addr_ty: ir::Type,
@@ -165,7 +175,12 @@ fn compute_addr(
     }
 
     // Add the heap base address base
+    let base = if isa.flags().enable_pinned_reg() && isa.flags().use_pinned_reg_as_heap_base() {
+        pos.ins().get_pinned_reg(isa.pointer_type())
+    } else {
         let base_gv = pos.func.heaps[heap].base;
-    let base = pos.ins().global_value(addr_ty, base_gv);
+        pos.ins().global_value(addr_ty, base_gv)
+    };
+
     pos.func.dfg.replace(inst).iadd(base, offset);
 }

cranelift/codegen/src/settings.rs

@@ -389,6 +389,7 @@ mod tests {
              enable_float = true\n\
              enable_nan_canonicalization = false\n\
              enable_pinned_reg = false\n\
+             use_pinned_reg_as_heap_base = false\n\
              enable_simd = false\n\
              enable_atomics = true\n\
              enable_safepoints = false\n\