[PATCH] D21937: DAGCombiner: Fold away vector extract of insert with the same index

[Justin Bogner via llvm-commits]

Matt Arsenault <Matthew.Arsenault at amd.com> writes:
> arsenm created this revision.
> arsenm added a reviewer: bogner.
> arsenm added a subscriber: llvm-commits.
> Herald added a reviewer: tstellarAMD.
>
> This only really matters when the index is non-constant since the
> constant case already gets taken care of by other combines.

LGTM.

> http://reviews.llvm.org/D21937
>
> Files:
>   lib/CodeGen/SelectionDAG/DAGCombiner.cpp
>   test/CodeGen/AMDGPU/vector-extract-insert.ll
>
> Index: test/CodeGen/AMDGPU/vector-extract-insert.ll
> ===================================================================
> --- /dev/null
> +++ test/CodeGen/AMDGPU/vector-extract-insert.ll
> @@ -0,0 +1,66 @@
> +; RUN: llc -march=amdgcn -verify-machineinstrs < %s | FileCheck -check-prefix=GCN %s
> +
> +; Test that when extracting the same unknown vector index from an
> +; insertelement the dynamic indexing is folded away.
> +
> +declare i32 @llvm.amdgcn.workitem.id.x() #0
> +
> +; No dynamic indexing required
> +; GCN-LABEL: {{^}}extract_insert_same_dynelt_v4i32:
> +; GCN: s_load_dword [[VAL:s[0-9]+]], s{{\[[0-9]+:[0-9]+\]}}, 0xd{{$}}
> +; GCN-NOT buffer_load_dword
> +; GCN-NOT: [[VAL]]
> +; GCN: v_mov_b32_e32 [[VVAL:v[0-9]+]], [[VAL]]
> +; GCN-NOT: [[VVAL]]
> +; GCN: buffer_store_dword [[VVAL]]
> +define void @extract_insert_same_dynelt_v4i32(i32 addrspace(1)* %out, <4 x i32> addrspace(1)* %in, i32 %val, i32 %idx) #1 {
> +  %id = call i32 @llvm.amdgcn.workitem.id.x()
> +  %id.ext = sext i32 %id to i64
> +  %gep.in = getelementptr inbounds <4 x i32>, <4 x i32> addrspace(1)* %in, i64 %id.ext
> +  %gep.out = getelementptr inbounds i32, i32 addrspace(1)* %out, i64 %id.ext
> +  %vec = load <4 x i32>, <4 x i32> addrspace(1)* %gep.in
> +  %insert = insertelement <4 x i32> %vec, i32 %val, i32 %idx
> +  %extract = extractelement <4 x i32> %insert, i32 %idx
> +  store i32 %extract, i32 addrspace(1)* %gep.out
> +  ret void
> +}
> +
> +; GCN-LABEL: {{^}}extract_insert_different_dynelt_v4i32:
> +; GCN: buffer_load_dwordx4
> +; GCN: v_movreld_b32
> +; GCN: v_movrels_b32
> +; GCN: buffer_store_dword v
> +define void @extract_insert_different_dynelt_v4i32(i32 addrspace(1)* %out, <4 x i32> addrspace(1)* %in, i32 %val, i32 %idx0, i32 %idx1) #1 {
> +  %id = call i32 @llvm.amdgcn.workitem.id.x()
> +  %id.ext = sext i32 %id to i64
> +  %gep.in = getelementptr inbounds <4 x i32>, <4 x i32> addrspace(1)* %in, i64 %id.ext
> +  %gep.out = getelementptr inbounds i32, i32 addrspace(1)* %out, i64 %id.ext
> +  %vec = load <4 x i32>, <4 x i32> addrspace(1)* %gep.in
> +  %insert = insertelement <4 x i32> %vec, i32 %val, i32 %idx0
> +  %extract = extractelement <4 x i32> %insert, i32 %idx1
> +  store i32 %extract, i32 addrspace(1)* %gep.out
> +  ret void
> +}
> +
> +; GCN-LABEL: {{^}}extract_insert_same_elt2_v4i32:
> +; GCN: s_load_dword [[VAL:s[0-9]+]], s{{\[[0-9]+:[0-9]+\]}}, 0xd{{$}}
> +; GCN-NOT buffer_load_dword
> +; GCN-NOT: [[VAL]]
> +; GCN: v_mov_b32_e32 [[VVAL:v[0-9]+]], [[VAL]]
> +; GCN-NOT: [[VVAL]]
> +; GCN: buffer_store_dword [[VVAL]]
> +define void @extract_insert_same_elt2_v4i32(i32 addrspace(1)* %out, <4 x i32> addrspace(1)* %in, i32 %val, i32 %idx) #1 {
> +  %id = call i32 @llvm.amdgcn.workitem.id.x()
> +  %id.ext = sext i32 %id to i64
> +  %gep.in = getelementptr inbounds <4 x i32>, <4 x i32> addrspace(1)* %in, i64 %id.ext
> +  %gep.out = getelementptr inbounds i32, i32 addrspace(1)* %out, i64 %id.ext
> +  %vec = load <4 x i32>, <4 x i32> addrspace(1)* %gep.in
> +  %insert = insertelement <4 x i32> %vec, i32 %val, i32 %idx
> +  %extract = extractelement <4 x i32> %insert, i32 %idx
> +  store i32 %extract, i32 addrspace(1)* %gep.out
> +  ret void
> +}
> +
> +
> +attributes #0 = { nounwind readnone }
> +attributes #1 = { nounwind }
> \ No newline at end of file
> Index: lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> ===================================================================
> --- lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> +++ lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> @@ -12392,6 +12392,14 @@
>        return DAG.getNode(ISD::TRUNCATE, SDLoc(N), NVT, BCSrc);
>    }
>  
> +  // extract_vector_elt (insert_vector_elt vec, val, idx), idx) -> val
> +  //
> +  // This only really matters if the index is non-constant since other combines
> +  // on the constant elements already work.
> +  if (InVec.getOpcode() == ISD::INSERT_VECTOR_ELT &&
> +      EltNo == InVec.getOperand(2))
> +    return InVec.getOperand(1);
> +
>    // Transform: (EXTRACT_VECTOR_ELT( VECTOR_SHUFFLE )) -> EXTRACT_VECTOR_ELT.
>    // We only perform this optimization before the op legalization phase because
>    // we may introduce new vector instructions which are not backed by TD