[llvm] [RISCV] Hoist immediate addresses from loads/stores (PR #83644)

[Visoiu Mistrih Francis via llvm-commits]

https://github.com/francisvm updated https://github.com/llvm/llvm-project/pull/83644

>From 344b54552fef9f9f9bb5cecc2cf0f7cee6259e36 Mon Sep 17 00:00:00 2001
From: Francis Visoiu Mistrih <francisvm at apple.com>
Date: Fri, 1 Mar 2024 18:05:56 -0800
Subject: [PATCH] [RISCV] Hoist immediate addresses from loads/stores

In case of loads/stores from an immediate address, avoid rematerializing
the constant for every block and allow consthoist to hoist it to the
entry block.
---
 .../Target/RISCV/RISCVTargetTransformInfo.cpp |  8 +++
 .../ConstantHoisting/RISCV/immediates.ll      | 58 +++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp b/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
index 504970c0a9ac74..ecd373649e2c79 100644
--- a/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
@@ -165,6 +165,14 @@ InstructionCost RISCVTTIImpl::getIntImmCostInst(unsigned Opcode, unsigned Idx,
     // split up large offsets in GEP into better parts than ConstantHoisting
     // can.
     return TTI::TCC_Free;
+  case Instruction::Store:
+    // If the address is a constant, use the materialization cost.
+    if (Idx == 1)
+      return getIntImmCost(Imm, Ty, CostKind);
+    return TTI::TCC_Free;
+  case Instruction::Load:
+    // If the address is a constant, use the materialization cost.
+    return getIntImmCost(Imm, Ty, CostKind);
   case Instruction::And:
     // zext.h
     if (Imm == UINT64_C(0xffff) && ST->hasStdExtZbb())
diff --git a/llvm/test/Transforms/ConstantHoisting/RISCV/immediates.ll b/llvm/test/Transforms/ConstantHoisting/RISCV/immediates.ll
index 131ef673a61f82..0a477d5b648475 100644
--- a/llvm/test/Transforms/ConstantHoisting/RISCV/immediates.ll
+++ b/llvm/test/Transforms/ConstantHoisting/RISCV/immediates.ll
@@ -150,3 +150,61 @@ define i64 @test16(i64 %a) nounwind {
   ret i64 %2
 }
 
+; Check that we hoist the absolute address of the stores to the entry block.
+define void @test17(ptr %s, i32 %size) nounwind {
+; CHECK-LABEL: test17
+; CHECK: %const = bitcast i32 -1073741792 to i32
+; CHECK: %0 = inttoptr i32 %const to ptr
+; CHECK: store i32 20, ptr %0
+; CHECK: %1 = inttoptr i32 %const to ptr
+; CHECK: store i32 10, ptr %1
+entry:
+  %cond = icmp eq i32 %size, 0
+  br i1 %cond, label %if.true, label %exit
+if.true:
+  store i32 20, ptr inttoptr (i32 -1073741792 to ptr)
+  br label %exit
+exit:
+  store i32 10, ptr inttoptr (i32 -1073741792 to ptr)
+  ret void
+}
+
+; Check that we hoist the absolute address of the loads to the entry block.
+define i32 @test18(ptr %s, i32 %size) nounwind {
+; CHECK-LABEL: test18
+; CHECK: %const = bitcast i32 -1073741792 to i32
+; CHECK: %0 = inttoptr i32 %const to ptr
+; CHECK: %1 = load i32, ptr %0
+; CHECK: %2 = inttoptr i32 %const to ptr
+; CHECK: %3 = load i32, ptr %2
+entry:
+  %cond = icmp eq i32 %size, 0
+  br i1 %cond, label %if.true, label %if.false
+if.true:
+  %0 = load i32, ptr inttoptr (i32 -1073741792 to ptr)
+  br label %return
+if.false:
+  %1 = load i32, ptr inttoptr (i32 -1073741792 to ptr)
+  br label %return
+return:
+  %val = phi i32 [%0, %if.true], [%1, %if.false]
+  ret i32 %val
+}
+
+
+; For addresses between [0, 2048), we can use ld/sd xN, address(zero), so don't
+; hoist.
+define void @test19(ptr %s, i32 %size) nounwind {
+; CHECK-LABEL: test19
+; CHECK: store i32 20, ptr inttoptr (i32 2044 to ptr)
+; CHECK: store i32 10, ptr inttoptr (i32 2044 to ptr)
+entry:
+  %cond = icmp eq i32 %size, 0
+  br i1 %cond, label %if.true, label %exit
+if.true:
+  store i32 20, ptr inttoptr (i32 2044 to ptr)
+  br label %exit
+exit:
+  store i32 10, ptr inttoptr (i32 2044 to ptr)
+  ret void
+}