[DirectX] Scalarizer is producing GEP chains that fail validation

[Icohedron]

The validator doesn't like GEP chains where an array GEP is followed by a scalar GEP.

If you compile the following LLVM IR to DXIL using llc -filetype=obj -mtriple=dxil-pc-shadermodel6.3-library

define i32 @gep_chain() #0 {
    %1 = alloca [4 x i32], align 4
    %2 = getelementptr inbounds nuw [4 x i32], ptr %1, i32 0, i32 2
    %3 = getelementptr inbounds nuw i32, ptr %2, i32 1
    ; Effectively, %3 is equivalent to getelementptr inbounds nuw [4 x i32], ptr %1, i32 0, i32 3
    %4 = load i32, ptr %3
    ret i32 %4
}

attributes #0 = { convergent norecurse nounwind "hlsl.export"}

and then validate the resulting dxil using dxv, you get the following validation error (#140417):

Function: gep_chain: error: Access to out-of-bounds memory is disallowed.
note: at '%3 = getelementptr inbounds i32, i32* %2, i32 1' in block '#0' of function 'gep_chain'.
Validation failed.

These GEP chains can appear when compiling HLSL shaders such as this one: https://godbolt.org/z/ExeGKb7MG

[llvmbot]

@llvm/issue-subscribers-bug

Author: Deric C. (Icohedron)

The validator doesn't like GEP chains where an array GEP is followed by a scalar GEP.

If you compile the following LLVM IR to DXIL using llc -filetype=obj -mtriple=dxil-pc-shadermodel6.3-library

define i32 @<!-- -->gep_chain() #<!-- -->0 {

    %1 = alloca [4 x i32], align 4
    %2 = getelementptr inbounds nuw [4 x i32], ptr %1, i32 0, i32 2
    %3 = getelementptr inbounds nuw i32, ptr %2, i32 1
    %4 = load i32, ptr %3
    ret i32 %4
}

attributes #<!-- -->0 = { convergent norecurse nounwind "hlsl.export"}

and then validate the resulting dxil using dxv, you get the following validation error (#140417):

Function: gep_chain: error: Access to out-of-bounds memory is disallowed.
note: at '%3 = getelementptr inbounds i32, i32* %2, i32 1' in block '#<!-- -->0' of function 'gep_chain'.
Validation failed.

These GEP chains can appear when compiling HLSL shaders such as this one: https://godbolt.org/z/ExeGKb7MG
The scalarizer pass is responsible for emitting the scalar GEP in this case, which itself occurs after the dxil-flatten-arrays pass that had attempted to collapse GEP chains. Therefore, we need a later pass, such as the dxil-legalize pass, or a new pass to collapse the GEP chains again.

I propose that to solve this cleanly, we should shift the responsibility of collapsing GEP chains from the dxil-flatten-arrays pass over to the dxil-legalize pass.

Another reason to move GEP-chain-collapsing over to the dxil-legalize pass is to consolidate i8 GEP transformations as well. Currently i8 GEPs can appear during dxil-data-scalarization and dxil-flatten-arrays (see #145780), which occurs before the dxil-legalize pass that is responsible for i8 GEP legalization. By moving the GEP-chain-collapsing logic from the dxil-flatten-arrays pass over to the dxil-legalize pass, the dxil-legalize pass can simultaneously collapse GEP chains and legalize any i8 and scalar GEPs in those GEP chains.

The dxil-flatten-arrays pass' visitGetElementPtrInst can be simplified down to the following:

/// This visitor simply converts GEPs on multidimensional arrays into GEPs on
/// flattened arrays. A later "GEP chain collapser" pass will be used to combine
/// GEP chains into single flat GEPs, including cases of i8 GEPs, and scalar
/// GEPs introduced by the scalarizer.
bool DXILFlattenArraysVisitor::visitGetElementPtrInst(GetElementPtrInst &GEP) {
  if (!isMultiDimensionalArray(GEP.getSourceElementType())) {
    return false;
  }

  ArrayType *ArrType = cast<ArrayType>(GEP.getSourceElementType());
  IRBuilder<> Builder(&GEP);
  auto [TotalElements, BaseType] = getElementCountAndType(ArrType);
  ArrayType *FlattenedArrayType = ArrayType::get(BaseType, TotalElements);
  Value *PtrOperand = GEP.getPointerOperand();

  Value *FlattenedIndex;

  const DataLayout &DL = GEP.getDataLayout();
  unsigned BitWidth = DL.getIndexTypeSizeInBits(GEP.getType());
  APInt ConstantOffset(BitWidth, 0);
  SmallMapVector<Value *, APInt, 4> VariableOffsets;
  bool Success = GEP.collectOffset(DL, BitWidth, VariableOffsets, ConstantOffset);
  (void)Success;
  assert(Success && "Failed to collect offsets for GEP");

  // GEP.collectOffset returns the offset in bytes. So we need to divide its
  // offsets by the size in bytes of the BaseType
  unsigned BaseTypeSizeInBytes = BaseType->getPrimitiveSizeInBits() / 8;

  Value *ZeroIndex = Builder.getInt({BitWidth, 0});
  FlattenedIndex = Builder.getInt(ConstantOffset.udiv(BaseTypeSizeInBytes));
  for (auto [VarIndex, Multiplier] : VariableOffsets){
    Value *ConstIntMul = Builder.getInt(Multiplier.udiv(BaseTypeSizeInBytes));
    Value *MulVarIndex = Builder.CreateMul(VarIndex, ConstIntMul);
    FlattenedIndex = Builder.CreateAdd(FlattenedIndex, MulVarIndex);
  }

  Value *NewGEP = Builder.CreateGEP(
      FlattenedArrayType, PtrOperand, {ZeroIndex, FlattenedIndex},
      GEP.getName(), GEP.getNoWrapFlags());

  GEP.replaceAllUsesWith(NewGEP);
  GEP.eraseFromParent();
  return true;
}

This will flatten all GEPs, but not collapse any GEP chains. Again, the collapsing of GEP chains could be handled by the dxil-legalize pass (or a new pass).

[Icohedron]

As per the design discussion today, we will try to solve the issue by moving the scalarizer and legalize pass before the flatten arrays pass. The memcpy and memset legalization in the legalize pass can be moved to the flatten arrays pass as it is somewhat related.

[Icohedron]

If #146173 gets merged, then we may only need to move the scalarizer pass before the dxil-flatten-arrays pass because #146173's implementation of GEP chain collapsing / GEP flattening naturally handles GEPs regardless of the source element types given to them.

llvm-project/llvm/test/CodeGen/DirectX/flatten-array.ll

Lines 258 to 275 in dee77e9

	; This test demonstrates that the collapsing of GEP chains occurs regardless of
	; the source element type given to the GEP. As long as the root pointer being
	; indexed to is an aggregate data structure, the GEP will be flattened.
	define void @gep_scalar_flatten() {
	; CHECK-LABEL: gep_scalar_flatten
	; CHECK-NEXT: [[ALLOCA:%.*]] = alloca [24 x i32]
	; CHECK-NEXT: getelementptr inbounds nuw [24 x i32], ptr [[ALLOCA]], i32 0, i32 17
	; CHECK-NEXT: getelementptr inbounds nuw [24 x i32], ptr [[ALLOCA]], i32 0, i32 17
	; CHECK-NEXT: getelementptr inbounds nuw [24 x i32], ptr [[ALLOCA]], i32 0, i32 23
	; CHECK-NEXT: ret void
	%a = alloca [2 x [3 x [4 x i32]]], align 4
	%i8root = getelementptr inbounds nuw i8, [2 x [3 x [4 x i32]]]* %a, i32 68 ; %a[1][1][1]
	%i32root = getelementptr inbounds nuw i32, [2 x [3 x [4 x i32]]]* %a, i32 17 ; %a[1][1][1]
	%c0 = getelementptr inbounds nuw [2 x [3 x [4 x i32]]], [2 x [3 x [4 x i32]]]* %a, i32 0, i32 1 ; %a[1]
	%c1 = getelementptr inbounds nuw i32, [3 x [4 x i32]]* %c0, i32 8 ; %a[1][2]
	%c2 = getelementptr inbounds nuw i8, [4 x i32]* %c1, i32 12 ; %a[1][2][3]
	ret void
	}

We could still move the dxil-legalize pass to an earlier stage if we want to along with the scalarizer, but it won't be required anymore to legalize the i8 GEPs for the dxil-flatten-arrays pass.