Optimizations on arm64 for Unsafe.put*

Generate better code if some of the inputs are constant.

Bug: 362820038
Test: 858-checker-unsafe

Change-Id: I263b48e0c60aa6c57cbb5ac2fdcf016e6cf83895

5 files changed, 81 insertions, 10 deletions

diff --git a/compiler/optimizing/intrinsics_arm64.cc b/compiler/optimizing/intrinsics_arm64.cc
index 903713eadc..80ec1b629a 100644
--- a/compiler/optimizing/intrinsics_arm64.cc
+++ b/compiler/optimizing/intrinsics_arm64.cc
@@ -57,6 +57,7 @@ using helpers::CPURegisterFrom;
 using helpers::DRegisterFrom;
 using helpers::HeapOperand;
 using helpers::LocationFrom;
+using helpers::Int64FromLocation;
 using helpers::InputCPURegisterAt;
 using helpers::InputCPURegisterOrZeroRegAt;
 using helpers::OperandFrom;
@@ -875,10 +876,21 @@ void IntrinsicCodeGeneratorARM64::VisitJdkUnsafeGetByte(HInvoke* invoke) {
 static void CreateUnsafePutLocations(ArenaAllocator* allocator, HInvoke* invoke) {
   LocationSummary* locations =
       new (allocator) LocationSummary(invoke, LocationSummary::kNoCall, kIntrinsified);
-  locations->SetInAt(0, Location::NoLocation());        // Unused receiver.
+  static constexpr int kOffsetIndex = 2;
+  static constexpr int kValueIndex = 3;
+  // Unused receiver.
+  locations->SetInAt(0, Location::NoLocation());
+  // The object.
   locations->SetInAt(1, Location::RequiresRegister());
-  locations->SetInAt(2, Location::RequiresRegister());
-  locations->SetInAt(3, Location::RequiresRegister());
+  // The offset.
+  locations->SetInAt(
+      kOffsetIndex, Location::RegisterOrConstant(invoke->InputAt(kOffsetIndex)));
+  // The value.
+  if (IsZeroBitPattern(invoke->InputAt(kValueIndex))) {
+    locations->SetInAt(kValueIndex, Location::ConstantLocation(invoke->InputAt(kValueIndex)));
+  } else {
+    locations->SetInAt(kValueIndex, Location::RequiresRegister());
+  }
 }
 
 void IntrinsicLocationsBuilderARM64::VisitUnsafePut(HInvoke* invoke) {
@@ -960,18 +972,27 @@ static void GenUnsafePut(HInvoke* invoke,
   LocationSummary* locations = invoke->GetLocations();
   MacroAssembler* masm = codegen->GetVIXLAssembler();
 
+  static constexpr int kOffsetIndex = 2;
+  static constexpr int kValueIndex = 3;
   Register base = WRegisterFrom(locations->InAt(1));    // Object pointer.
-  Register offset = XRegisterFrom(locations->InAt(2));  // Long offset.
-  Register value = RegisterFrom(locations->InAt(3), type);
-  Register source = value;
-  MemOperand mem_op(base.X(), offset);
+  Location offset = locations->InAt(kOffsetIndex);       // Long offset.
+  CPURegister value = InputCPURegisterOrZeroRegAt(invoke, kValueIndex);
+  CPURegister source = value;
+  MemOperand mem_op;
+  if (offset.IsConstant()) {
+    mem_op = MemOperand(base.X(), Int64FromLocation(offset));
+  } else {
+    mem_op = MemOperand(base.X(), XRegisterFrom(offset));
+  }
 
   {
     // We use a block to end the scratch scope before the write barrier, thus
     // freeing the temporary registers so they can be used in `MarkGCCard`.
     UseScratchRegisterScope temps(masm);
 
-    if (kPoisonHeapReferences && type == DataType::Type::kReference) {
+    if (kPoisonHeapReferences &&
+        type == DataType::Type::kReference &&
+        !IsZeroBitPattern(invoke->InputAt(kValueIndex))) {
       DCHECK(value.IsW());
       Register temp = temps.AcquireW();
       __ Mov(temp.W(), value.W());
@@ -986,9 +1007,9 @@ static void GenUnsafePut(HInvoke* invoke,
     }
   }
 
-  if (type == DataType::Type::kReference) {
+  if (type == DataType::Type::kReference && !IsZeroBitPattern(invoke->InputAt(kValueIndex))) {
     bool value_can_be_null = true;  // TODO: Worth finding out this information?
-    codegen->MaybeMarkGCCard(base, value, value_can_be_null);
+    codegen->MaybeMarkGCCard(base, Register(source), value_can_be_null);
   }
 }
 
diff --git a/test/858-checker-unsafe/expected-stderr.txt b/test/858-checker-unsafe/expected-stderr.txt
new file mode 100644
index 0000000000..e69de29bb2
--- /dev/null
+++ b/test/858-checker-unsafe/expected-stderr.txt
diff --git a/test/858-checker-unsafe/expected-stdout.txt b/test/858-checker-unsafe/expected-stdout.txt
new file mode 100644
index 0000000000..e69de29bb2
--- /dev/null
+++ b/test/858-checker-unsafe/expected-stdout.txt
diff --git a/test/858-checker-unsafe/info.txt b/test/858-checker-unsafe/info.txt
new file mode 100644
index 0000000000..2c91897412
--- /dev/null
+++ b/test/858-checker-unsafe/info.txt
@@ -0,0 +1 @@
+Test optimizations we generate for Unsafe.put.
diff --git a/test/858-checker-unsafe/src/Main.java b/test/858-checker-unsafe/src/Main.java
new file mode 100644
index 0000000000..01cde39f77
--- /dev/null
+++ b/test/858-checker-unsafe/src/Main.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import java.lang.reflect.Field;
+import sun.misc.Unsafe;
+
+public class Main {
+  private static Unsafe getUnsafe() throws NoSuchFieldException, IllegalAccessException {
+    Class<?> unsafeClass = Unsafe.class;
+    Field f = unsafeClass.getDeclaredField("theUnsafe");
+    f.setAccessible(true);
+    return (Unsafe) f.get(null);
+  }
+
+  private static Unsafe unsafe;
+
+  public static void main(String[] args) throws NoSuchFieldException, IllegalAccessException {
+    unsafe = getUnsafe();
+    testPutZero();
+    testPutFixedOffset();
+  }
+
+  /// CHECK-START-ARM64: void Main.testPutZero() disassembly (after)
+  /// CHECK:                  str wzr, [x{{[0-9]+}}, x{{[0-9]+}}]
+  private static void testPutZero() {
+    int[] object = new int[42];
+    unsafe.putInt(object, unsafe.arrayBaseOffset(int[].class), 0);
+  }
+
+  /// CHECK-START-ARM64: void Main.testPutFixedOffset() disassembly (after)
+  /// CHECK:                  stur w{{[0-9]+}}, [x{{[0-9]+}}, #38]
+  private static void testPutFixedOffset() {
+    int[] object = new int[42];
+    unsafe.putInt(object, 38, 12);
+  }
+}