JIT: fold wasm R2R address constants into load memarg offsets

src/coreclr/jit/codegenwasm.cpp

@@ -2791,18 +2791,29 @@ void CodeGen::genCodeForIndir(GenTreeIndir* tree)
     assert(tree->OperIs(GT_IND));
 
     var_types type = tree->TypeGet();
+    GenTree*  addr = tree->Addr();
 
-    genConsumeAddress(tree->Addr());
+    genConsumeAddress(addr);
 
     if ((tree->gtFlags & GTF_IND_NONFAULTING) == 0)
     {
-        regNumber addrReg = GetMultiUseOperandReg(tree->Addr());
+        regNumber addrReg = GetMultiUseOperandReg(addr);
         genEmitNullCheck(addrReg);
     }
 
     // TODO-WASM: Memory barriers
 
-    if (type == TYP_SIMD12)
+    if (addr->isContained())
+    {
+        // A contained address constant folds into the memarg offset, so emit just the image base here.
+        //
+        assert(addr->IsIconHandle() && !tree->TypeIs(TYP_SIMD12));
+        assert(addr->AsIntConCommon()->ImmedValNeedsReloc(m_compiler));
+        GetEmitter()->emitImageBase();
+        GetEmitter()->emitIns_MemargAddress(ins_Load(type), emitActualTypeSize(type),
+                                            (void*)addr->AsIntConCommon()->IntegralValue());
+    }
+    else if (type == TYP_SIMD12)
     {
         genLoadIndTypeSimd12(tree);
     }

src/coreclr/jit/emitwasm.cpp

@@ -182,18 +182,46 @@ bool emitter::emitInsIsStore(instruction ins)
     return false;
 }
 
+//------------------------------------------------------------------------
+// emitImageBase: Emit the module base (imageBase global) onto the stack.
+//
+void emitter::emitImageBase()
+{
+    emitIns_I(INS_global_get, EA_HANDLE_CNS_RELOC,
+              (cnsval_ssize_t)(size_t)m_compiler->eeGetWasmWellKnownGlobals()->imageBase);
+}
+
 //------------------------------------------------------------------------
 // emitAddressConstant: Emit a memory address constant, like an indirection cell.
 // This will automatically make use of relocations and the module base (imageBase).
 void emitter::emitAddressConstant(void* address)
 {
     // Load our module base from the image base global, then load our address constant, then sum them.
-    emitIns_I(INS_global_get, EA_HANDLE_CNS_RELOC,
-              (cnsval_ssize_t)(size_t)m_compiler->eeGetWasmWellKnownGlobals()->imageBase);
+    emitImageBase();
     emitIns_I(INS_i32_const_address, EA_SET_FLG(EA_PTRSIZE, EA_CNS_RELOC_FLG), (cnsval_ssize_t)address);
     emitIns(INS_i32_add);
 }
 
+//------------------------------------------------------------------------
+// emitIns_MemargAddress: Emit a load or store whose memarg offset is a relocated address constant.
+//
+// Arguments:
+//   ins     - the load or store instruction
+//   attr    - emit attributes
+//   address - the address constant, relocated relative to the module base
+//
+// Notes:
+//   The module base must already be on the stack; see emitImageBase. This folds the addition that
+//   emitAddressConstant would otherwise emit into the memarg offset. The memarg offset is always a
+//   relocation, never a raw address, so the caller must only reach here when relocating.
+//
+void emitter::emitIns_MemargAddress(instruction ins, emitAttr attr, void* address)
+{
+    assert(emitInsFormat(ins) == IF_MEMARG);
+    assert(m_compiler->opts.compReloc);
+    emitIns_I(ins, EA_SET_FLG(attr, EA_CNS_RELOC_FLG), (cnsval_ssize_t)address);
+}
+
 void emitter::emitFuncletAddressConstant(cnsval_ssize_t funcletId)
 {
     // Load our table base, then load our funclet pointer offset, then sum them.
@@ -1003,12 +1031,13 @@ size_t emitter::emitOutputInstr(insGroup* ig, instrDesc* id, BYTE** dp)
         case IF_MEMARG:
         {
             dst += emitOutputOpcode(dst, ins);
-            uint64_t align  = emitGetAlignHintLog2(id);
-            uint64_t offset = emitGetInsSC(id);
+            uint64_t align = emitGetAlignHintLog2(id);
             assert(align <= UINT32_MAX); // spec says memarg alignment is u32
             assert(align < 64);          // spec says align > 2^6 produces a memidx for multiple memories.
             dst += emitOutputULEB128(dst, align);
-            dst += emitOutputULEB128(dst, offset);
+            // TODO-WASM: as with IF_MEMADDR, this reloc is specific to R2R and assumes the address we want is an
+            // offset from __image_base.
+            dst += emitOutputConstant(dst, id, UNSIGNED, CorInfoReloc::WASM_MEMORY_ADDR_REL_LEB);
             break;
         }
         case IF_LOCAL_DECL:
@@ -1332,7 +1361,14 @@ void emitter::emitDispIns(
         {
             unsigned       log2align = emitGetAlignHintLog2(id);
             cnsval_ssize_t offset    = emitGetInsSC(id);
-            printf(" %u %llu", log2align, (uint64_t)offset);
+            if (id->idIsCnsReloc())
+            {
+                printf(" %u reloc 0x%llx", log2align, (uint64_t)offset);
+            }
+            else
+            {
+                printf(" %u %llu", log2align, (uint64_t)offset);
+            }
             dispLclVarInfoIfAny();
         }
         break;

src/coreclr/jit/emitwasm.h

@@ -37,8 +37,10 @@ void emitIns_V128Imm(instruction ins, const uint8_t bytes[16]);
 void emitIns_Lane(instruction ins, uint8_t laneIdx);
 void emitIns_MemargLane(instruction ins, emitAttr attr, cnsval_ssize_t offset, uint8_t laneIdx);
 
+void emitImageBase();
 void emitAddressConstant(void* address);
 void emitFuncletAddressConstant(cnsval_ssize_t funcletId);
+void emitIns_MemargAddress(instruction ins, emitAttr attr, void* address);
 
 static unsigned SizeOfULEB128(uint64_t value);
 static unsigned SizeOfSLEB128(int64_t value);

src/coreclr/jit/lowerwasm.cpp

@@ -472,6 +472,22 @@ void Lowering::ContainCheckIndir(GenTreeIndir* indirNode)
 
     // TODO-WASM-CQ: contain suitable LEAs here. Take note of the fact that for this to be correct we must prove the
     // LEA doesn't overflow. It will involve creating a new frontend node to represent "nuw" (offset) addition.
+
+    // Contain a relocatable address constant so it folds into the load's memarg offset.
+    //
+    // Codegen for such a constant is "global.get $imageBase; i32.const <reloc>; i32.add"; containing it lets
+    // genCodeForIndir emit just the image base and put the relocated address in the memarg instead. Only loads
+    // can do this: the parent emits the base in place of the address, and a store cannot because its value
+    // operand has already been pushed. A multiply-used address is excluded because codegen re-materializes it
+    // from a register that would no longer hold the full address.
+    //
+    GenTree* addr = indirNode->Addr();
+    if (indirNode->OperIs(GT_IND) && !indirNode->TypeIs(TYP_SIMD12) && addr->IsIconHandle() &&
+        addr->AsIntConCommon()->ImmedValNeedsReloc(m_compiler) &&
+        ((addr->gtLIRFlags & LIR::Flags::MultiplyUsed) == LIR::Flags::None))
+    {
+        MakeSrcContained(indirNode, addr);
+    }
 }
 
 //------------------------------------------------------------------------