[llvm] [AArch64][CodeGen] Fold tbx(splat(0), table, idxs) to tbl(table, idxs) (PR #214146)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: Durgesh Nandan Mohanty (dnmohanty)

<details>
<summary>Changes</summary>

When the destination/background vector of a TBX instruction is a splat of zero, the operation is equivalent to a TBL instruction. TBL implicitly zeroes out any elements where the index is out of bounds, matching the behavior of TBX with a zero background vector. 

This patch adds a DAG combine to optimize this case, reducing instruction latency and register pressure.

Fixes #<!-- -->214077

---
Full diff: https://github.com/llvm/llvm-project/pull/214146.diff


3 Files Affected:

- (modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+20) 
- (added) llvm/test/CodeGen/AArch64/test_tbx.ll (+11) 
- (added) test_tbx.ll (+11) 


``````````diff
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 51be0e66b19b0..f41c0cae79043 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -25035,6 +25035,26 @@ static SDValue performIntrinsicCombine(SDNode *N,
   switch (IID) {
   default:
     break;
+    case Intrinsic::aarch64_neon_tbx1:
+  case Intrinsic::aarch64_neon_tbx2:
+  case Intrinsic::aarch64_neon_tbx3:
+  case Intrinsic::aarch64_neon_tbx4: {
+    if (ISD::isBuildVectorAllZeros(N->getOperand(1).getNode())) {
+      unsigned TblIID = 0;
+      if (IID == Intrinsic::aarch64_neon_tbx1) TblIID = Intrinsic::aarch64_neon_tbl1;
+      else if (IID == Intrinsic::aarch64_neon_tbx2) TblIID = Intrinsic::aarch64_neon_tbl2;
+      else if (IID == Intrinsic::aarch64_neon_tbx3) TblIID = Intrinsic::aarch64_neon_tbl3;
+      else if (IID == Intrinsic::aarch64_neon_tbx4) TblIID = Intrinsic::aarch64_neon_tbl4;
+
+      SmallVector<SDValue, 4> Ops;
+      Ops.push_back(DAG.getTargetConstant(TblIID, SDLoc(N), MVT::i32));
+      for (unsigned i = 2; i < N->getNumOperands(); ++i)
+        Ops.push_back(N->getOperand(i));
+
+      return DAG.getNode(ISD::INTRINSIC_WO_CHAIN, SDLoc(N), N->getValueType(0), Ops);
+    }
+    break;
+  }
   case Intrinsic::aarch64_neon_vcvtfxs2fp:
   case Intrinsic::aarch64_neon_vcvtfxu2fp:
     return tryCombineFixedPointConvert(N, DCI, DAG);
diff --git a/llvm/test/CodeGen/AArch64/test_tbx.ll b/llvm/test/CodeGen/AArch64/test_tbx.ll
new file mode 100644
index 0000000000000..d7752e0d07494
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/test_tbx.ll
@@ -0,0 +1,11 @@
+; RUN: llc -mtriple=aarch64 -mattr=+neon < %s | FileCheck %s
+
+define <16 x i8> @test_zero_splat(<16 x i8> %tbl, <16 x i8> %idx) {
+; CHECK-LABEL: test_zero_splat:
+; CHECK: tbl
+; CHECK: ret
+  %res = call <16 x i8> @llvm.aarch64.neon.tbx1.v16i8(<16 x i8> zeroinitializer, <16 x i8> %tbl, <16 x i8> %idx)
+  ret <16 x i8> %res
+}
+
+declare <16 x i8> @llvm.aarch64.neon.tbx1.v16i8(<16 x i8>, <16 x i8>, <16 x i8>)
\ No newline at end of file
diff --git a/test_tbx.ll b/test_tbx.ll
new file mode 100644
index 0000000000000..d7752e0d07494
--- /dev/null
+++ b/test_tbx.ll
@@ -0,0 +1,11 @@
+; RUN: llc -mtriple=aarch64 -mattr=+neon < %s | FileCheck %s
+
+define <16 x i8> @test_zero_splat(<16 x i8> %tbl, <16 x i8> %idx) {
+; CHECK-LABEL: test_zero_splat:
+; CHECK: tbl
+; CHECK: ret
+  %res = call <16 x i8> @llvm.aarch64.neon.tbx1.v16i8(<16 x i8> zeroinitializer, <16 x i8> %tbl, <16 x i8> %idx)
+  ret <16 x i8> %res
+}
+
+declare <16 x i8> @llvm.aarch64.neon.tbx1.v16i8(<16 x i8>, <16 x i8>, <16 x i8>)
\ No newline at end of file

``````````

</details>


https://github.com/llvm/llvm-project/pull/214146