AMDGPU: Account for dmask when computing image mem size

[arsenm]

Summary

Only the number of elements in the dmask will really be accessed.

Tests will be included in future GlobalISel patch.

[rtaylor]

Adding David Stuttard as I think he worked on dmask originally.

[nhaehnle]

The memory size of image loads is a bit of bizarre concept, but I assume this helps with the GlobalISel stuff, so sure.

LGTM, one stylistic nitpick.

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

950–953

As a stylistic note, I think it's nicer to have braces around both sides of an if-else, if one side requires it.

[arsenm]

570a578e46c8290a7f28b1602dadea750e147928

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

if (Size == 16 && Subtarget->has16BitInsts()) {
       return NumIntermediates;
     }
   }
 
   return TargetLowering::getVectorTypeBreakdownForCallingConv(
     Context, CC, VT, IntermediateVT, NumIntermediates, RegisterVT);
 }
 
+static EVT memVTFromImageData(Type *Ty, unsigned DMaskLanes) {
+  assert(DMaskLanes != 0);
+
+  if (auto *VT = dyn_cast<VectorType>(Ty)) {
+    unsigned NumElts = std::min(DMaskLanes,
+                                static_cast<unsigned>(VT->getNumElements()));
+    return EVT::getVectorVT(Ty->getContext(),
+                            EVT::getEVT(VT->getElementType()),
+                            NumElts);
+  }
+
+  return EVT::getEVT(Ty);
+}
+
 // Peek through TFE struct returns to only use the data size.
-static EVT memVTFromImageReturn(StructType &Ty) {
+static EVT memVTFromImageReturn(StructType &Ty, unsigned DMaskLanes) {
   // Only limited forms of aggregate type currently expected.
   if (Ty.getNumContainedTypes() != 2 ||
       !Ty.getContainedType(1)->isIntegerTy(32))
     return EVT();
-  return EVT::getEVT(Ty.getContainedType(0));
+  return memVTFromImageData(Ty.getContainedType(0), DMaskLanes);
 }
 
 bool SITargetLowering::getTgtMemIntrinsic(IntrinsicInfo &Info,
                                           const CallInst &CI,
                                           MachineFunction &MF,
                                           unsigned IntrID) const {
   if (const AMDGPU::RsrcIntrinsic *RsrcIntr =
           AMDGPU::lookupRsrcIntrinsic(IntrID)) {
if (const AMDGPU::RsrcIntrinsic *RsrcIntr =
     } else {
       Info.ptrVal = MFI->getBufferPSV(
         *MF.getSubtarget<GCNSubtarget>().getInstrInfo(),
         CI.getArgOperand(RsrcIntr->RsrcArg));
     }
 
     Info.flags = MachineMemOperand::MODereferenceable;
     if (Attr.hasFnAttribute(Attribute::ReadOnly)) {
-      Info.opc = ISD::INTRINSIC_W_CHAIN;
-      // TODO: Account for dmask reducing loaded size.
+      unsigned DMaskLanes = 4;
+
+      if (RsrcIntr->IsImage) {
+        const AMDGPU::ImageDimIntrinsicInfo *Intr
+          = AMDGPU::getImageDimIntrinsicInfo(IntrID);
+        const AMDGPU::MIMGBaseOpcodeInfo *BaseOpcode =
+          AMDGPU::getMIMGBaseOpcodeInfo(Intr->BaseOpcode);
+
+        if (!BaseOpcode->Gather4) {
+          // If this isn't a gather, we may have excess loaded elements in the
+          // IR type. Check the dmask for the real number of elements loaded.
+
+          unsigned DMask
+            = cast<ConstantInt>(CI.getArgOperand(0))->getZExtValue();
+          DMaskLanes = DMask == 0 ? 1 : countPopulation(DMask);
+        }
+
         if (auto *ST = dyn_cast<StructType>(CI.getType())) {
           // Some intrinsics return an aggregate type - special case to work out
           // the correct memVT
-        Info.memVT = memVTFromImageReturn(*ST);
+          Info.memVT = memVTFromImageReturn(*ST, DMaskLanes);
         } else
-        Info.memVT = MVT::getVT(CI.getType(), true);
+          Info.memVT = memVTFromImageData(CI.getType(), DMaskLanes);
+      } else
+        Info.memVT = EVT::getEVT(CI.getType());
 
+      // FIXME: What does alignment mean for an image?
+      Info.opc = ISD::INTRINSIC_W_CHAIN;
       Info.flags |= MachineMemOperand::MOLoad;
     } else if (Attr.hasFnAttribute(Attribute::WriteOnly)) {
       Info.opc = ISD::INTRINSIC_VOID;
-      Info.memVT = MVT::getVT(CI.getArgOperand(0)->getType());
+
+      Type *DataTy = CI.getArgOperand(0)->getType();
+      if (RsrcIntr->IsImage) {
+        unsigned DMask = cast<ConstantInt>(CI.getArgOperand(1))->getZExtValue();
+        unsigned DMaskLanes = DMask == 0 ? 1 : countPopulation(DMask);
+        Info.memVT = memVTFromImageData(DataTy, DMaskLanes);
+      } else
+        Info.memVT = EVT::getEVT(DataTy);
+
       Info.flags |= MachineMemOperand::MOStore;
     } else {
       // Atomic
       Info.opc = ISD::INTRINSIC_W_CHAIN;
       Info.memVT = MVT::getVT(CI.getType());
       Info.flags = MachineMemOperand::MOLoad |
                    MachineMemOperand::MOStore |
                    MachineMemOperand::MODereferenceable;