From 5d34d7f1a0cae8367066ce2b55afe42b94a7466f Mon Sep 17 00:00:00 2001
From: Amara Emerson <amara@apple.com>
Date: Fri, 18 Sep 2020 15:13:12 -0700
Subject: [PATCH] [GlobalISel] Add lowering support for G_ABS and use for
 AArch64.

Differential Revision: https://reviews.llvm.org/D87952
---
 .../CodeGen/GlobalISel/LegalizerHelper.cpp    | 16 +++++++++
 .../AArch64/GISel/AArch64LegalizerInfo.cpp    |  3 ++
 .../AArch64/GlobalISel/legalize-abs.mir       | 34 +++++++++++++++++++
 .../GlobalISel/legalizer-info-validation.mir  |  4 +--
 4 files changed, 55 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/GlobalISel/legalize-abs.mir

diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index e8ddfc8e083e..f4c8acc9e82a 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -3090,6 +3090,22 @@ LegalizerHelper::lower(MachineInstr &MI, unsigned TypeIdx, LLT LowerHintTy) {
   case G_SSHLSAT:
   case G_USHLSAT:
     return lowerShlSat(MI);
+  case G_ABS: {
+    // Expand %res = G_ABS %a into:
+    // %v1 = G_ASHR %a, scalar_size-1
+    // %v2 = G_ADD %a, %v1
+    // %res = G_XOR %v2, %v1
+    LLT DstTy = MRI.getType(MI.getOperand(0).getReg());
+    Register OpReg = MI.getOperand(1).getReg();
+    auto ShiftAmt =
+        MIRBuilder.buildConstant(DstTy, DstTy.getScalarSizeInBits() - 1);
+    auto Shift =
+        MIRBuilder.buildAShr(DstTy, OpReg, ShiftAmt);
+    auto Add = MIRBuilder.buildAdd(DstTy, OpReg, Shift);
+    MIRBuilder.buildXor(MI.getOperand(0).getReg(), Add, Shift);
+    MI.eraseFromParent();
+    return Legalized;
+  }
   }
 }
 
diff --git a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
index 99a012308606..a0af333db528 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
@@ -657,6 +657,9 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
 
   getActionDefinitionsBuilder({G_MEMCPY, G_MEMMOVE, G_MEMSET}).libcall();
 
+  getActionDefinitionsBuilder(G_ABS).lowerIf(
+      [=](const LegalityQuery &Query) { return Query.Types[0].isScalar(); });
+
   computeTables();
   verify(*ST.getInstrInfo());
 }
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/legalize-abs.mir b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-abs.mir
new file mode 100644
index 000000000000..25bfbc1fe17b
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-abs.mir
@@ -0,0 +1,34 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=aarch64 -run-pass=legalizer -global-isel-abort=1 %s -o - | FileCheck %s
+---
+name:            abs_s32
+liveins:
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: abs_s32
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 31
+    ; CHECK: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[COPY]], [[C]](s64)
+    ; CHECK: [[ADD:%[0-9]+]]:_(s32) = G_ADD [[COPY]], [[ASHR]]
+    ; CHECK: [[XOR:%[0-9]+]]:_(s32) = G_XOR [[ADD]], [[ASHR]]
+    ; CHECK: $w0 = COPY [[XOR]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s32) = G_ABS %0(s32)
+    $w0 = COPY %1(s32)
+...
+---
+name:            abs_s64
+liveins:
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: abs_s64
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 63
+    ; CHECK: [[ASHR:%[0-9]+]]:_(s64) = G_ASHR [[COPY]], [[C]](s64)
+    ; CHECK: [[ADD:%[0-9]+]]:_(s64) = G_ADD [[COPY]], [[ASHR]]
+    ; CHECK: [[XOR:%[0-9]+]]:_(s64) = G_XOR [[ADD]], [[ASHR]]
+    ; CHECK: $x0 = COPY [[XOR]](s64)
+    %0:_(s64) = COPY $x0
+    %1:_(s64) = G_ABS %0(s64)
+    $x0 = COPY %1(s64)
+...
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir b/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir
index 058a8fcb0890..c6c9b9e25cae 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir
@@ -497,8 +497,8 @@
 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: G_ABS (opcode {{[0-9]+}}): 1 type index, 0 imm indices
-# DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
-# DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
+# DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
+# DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: G_BR (opcode {{[0-9]+}}): 0 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined