Add support for mismatched array sizes

Author: Icohedron

llvm/lib/Target/DirectX/DXILLegalizePass.cpp

@@ -268,36 +268,41 @@ static void emitMemcpyExpansion(IRBuilder<> &Builder, Value *Dst, Value *Src,
     return nullptr;
   };
 
-  ArrayType *ArrTy = GetArrTyFromVal(Dst);
-  assert(ArrTy && "Expected Dst of memcpy to be a Pointer to an Array Type");
+  ArrayType *DstArrTy = GetArrTyFromVal(Dst);
+  assert(DstArrTy && "Expected Dst of memcpy to be a Pointer to an Array Type");
   if (auto *DstGlobalVar = dyn_cast<GlobalVariable>(Dst))
     assert(!DstGlobalVar->isConstant() &&
            "The Dst of memcpy must not be a constant Global Variable");
-
   [[maybe_unused]] ArrayType *SrcArrTy = GetArrTyFromVal(Src);
   assert(SrcArrTy && "Expected Src of memcpy to be a Pointer to an Array Type");
 
+  Type *DstElemTy = DstArrTy->getElementType();
+  uint64_t DstElemByteSize = DL.getTypeStoreSize(DstElemTy);
+  assert(DstElemByteSize > 0 && "Dst element type store size must be set");
+  Type *SrcElemTy = SrcArrTy->getElementType();
+  [[maybe_unused]] uint64_t SrcElemByteSize = DL.getTypeStoreSize(SrcElemTy);
+  assert(SrcElemByteSize > 0 && "Src element type store size must be set");
+
   // This assumption simplifies implementation and covers currently-known
   // use-cases for DXIL. It may be relaxed in the future if required.
-  assert(ArrTy == SrcArrTy &&
-         "Array Types of Src and Dst in memcpy must match");
-
-  Type *ElemTy = ArrTy->getElementType();
-  uint64_t ElemSize = DL.getTypeStoreSize(ElemTy);
-  assert(ElemSize > 0 && "Size must be set");
-
-  [[maybe_unused]] uint64_t Size = ArrTy->getArrayNumElements();
-  assert(ElemSize * Size >= ByteLength &&
-         "Array size must be at least as large as the memcpy length");
-
-  uint64_t NumElemsToCopy = ByteLength / ElemSize;
-  assert(ByteLength % ElemSize == 0 &&
+  assert(DstElemTy == SrcElemTy &&
+         "The element types of Src and Dst arrays must match");
+
+  [[maybe_unused]] uint64_t DstArrNumElems = DstArrTy->getArrayNumElements();
+  assert(DstElemByteSize * DstArrNumElems >= ByteLength &&
+         "Dst array size must be at least as large as the memcpy length");
+  [[maybe_unused]] uint64_t SrcArrNumElems = SrcArrTy->getArrayNumElements();
+  assert(SrcElemByteSize * SrcArrNumElems >= ByteLength &&
+         "Src array size must be at least as large as the memcpy length");
+
+  uint64_t NumElemsToCopy = ByteLength / DstElemByteSize;
+  assert(ByteLength % DstElemByteSize == 0 &&
          "memcpy length must be divisible by array element type");
   for (uint64_t I = 0; I < NumElemsToCopy; ++I) {
     Value *Offset = ConstantInt::get(Type::getInt32Ty(Ctx), I);
-    Value *SrcPtr = Builder.CreateInBoundsGEP(ElemTy, Src, Offset, "gep");
-    Value *SrcVal = Builder.CreateLoad(ElemTy, SrcPtr);
-    Value *DstPtr = Builder.CreateInBoundsGEP(ElemTy, Dst, Offset, "gep");
+    Value *SrcPtr = Builder.CreateInBoundsGEP(SrcElemTy, Src, Offset, "gep");
+    Value *SrcVal = Builder.CreateLoad(SrcElemTy, SrcPtr);
+    Value *DstPtr = Builder.CreateInBoundsGEP(DstElemTy, Dst, Offset, "gep");
     Builder.CreateStore(SrcVal, DstPtr);
   }
 }

llvm/test/CodeGen/DirectX/legalize-memcpy.ll

@@ -43,6 +43,27 @@ define void @replace_3int_memcpy_test() #0 {
   ret void
 }
 
+define void @replace_mismatched_size_int_memcpy_test() #0 {
+; CHECK-LABEL: define void @replace_mismatched_size_int_memcpy_test(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = alloca [2 x i32], align 4
+; CHECK-NEXT:    [[TMP2:%.*]] = alloca [3 x i32], align 4
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, ptr [[TMP1]], i32 0
+; CHECK-NEXT:    [[TMP3:%.*]] = load i32, ptr [[GEP]], align 4
+; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds i32, ptr [[TMP2]], i32 0
+; CHECK-NEXT:    store i32 [[TMP3]], ptr [[GEP1]], align 4
+; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i32, ptr [[TMP1]], i32 1
+; CHECK-NEXT:    [[TMP4:%.*]] = load i32, ptr [[GEP2]], align 4
+; CHECK-NEXT:    [[GEP3:%.*]] = getelementptr inbounds i32, ptr [[TMP2]], i32 1
+; CHECK-NEXT:    store i32 [[TMP4]], ptr [[GEP3]], align 4
+; CHECK-NEXT:    ret void
+;
+  %1 = alloca [2 x i32], align 4
+  %2 = alloca [3 x i32], align 4
+  call void @llvm.memcpy.p0.p0.i32(ptr nonnull align 4 dereferenceable(12) %2, ptr align 4 dereferenceable(8) %1, i32 8, i1 false)
+  ret void
+}
+
 define void @replace_int16_memcpy_test() #0 {
 ; CHECK-LABEL: define void @replace_int16_memcpy_test(
 ; CHECK-SAME: ) #[[ATTR0]] {