[AMDGPU] prohibit >4 bytes private memory access.

[vpykhtin]

Summary

Fixes mlopen 5x5 case.

This is a preview, not for submitting.

Failed tests:

LLVM :: CodeGen/AMDGPU/indirect-private-64.ll
LLVM :: CodeGen/AMDGPU/insert_vector_elt.ll
LLVM :: CodeGen/AMDGPU/private-element-size.ll

[kanarayan]

If this fixes correctness issue, commit please. Optimization can come after, if not along.

[vpykhtin]

In D19371#408064, @kanarayan wrote:

If this fixes correctness issue, commit please. Optimization can come after, if not along.

I think what Matt suggests is more reasonable here. I would wait for his patch.

include/llvm/Target/TargetLowering.h

virtual bool allowsMisalignedMemoryAccesses(EVT,
                                               bool * /*Fast*/ = nullptr) const {
     return false;
   }
 
   /// Return true if the target supports a memory access of this type for the
   /// given address space and alignment. If the access is allowed, the optional
   /// final parameter returns if the access is also fast (as defined by the
   /// target).
-  bool allowsMemoryAccess(LLVMContext &Context, const DataLayout &DL, EVT VT,
+  virtual bool allowsMemoryAccess(LLVMContext &Context, const DataLayout &DL, EVT VT,
                                   unsigned AddrSpace = 0, unsigned Alignment = 1,
                                   bool *Fast = nullptr) const;
 
   /// Returns the target specific optimal type for load and store operations as
   /// a result of memset, memcpy, and memmove lowering.
   ///
   /// If DstAlign is zero that means it's safe to destination alignment can
   /// satisfy any constraint. Similarly if SrcAlign is zero it means there isn't
   /// a need to check it against alignment requirement, probably because the
   /// source does not need to be loaded. If 'IsMemset' is true, that means it's

lib/Target/AMDGPU/SIISelLowering.h

MachineSDNode *buildRSRC(SelectionDAG &DAG,
                            SDValue Ptr,
                            uint32_t RsrcDword1,
                            uint64_t RsrcDword2And3) const;
   std::pair<unsigned, const TargetRegisterClass *>
   getRegForInlineAsmConstraint(const TargetRegisterInfo *TRI,
                                StringRef Constraint, MVT VT) const override;
   ConstraintType getConstraintType(StringRef Constraint) const override;
   SDValue copyToM0(SelectionDAG &DAG, SDValue Chain, SDLoc DL, SDValue V) const;
+
+  bool allowsMemoryAccess(LLVMContext &Context, const DataLayout &DL, EVT VT,
+    unsigned AddrSpace = 0, unsigned Alignment = 1,
+    bool *Fast = nullptr) const override;
 };
 
 } // End namespace llvm
 
 #endif

lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::LowerATOMIC_CMP_SWAP(SDValue Op, SelectionDAG &DAG) const {
   return DAG.getMemIntrinsicNode(AMDGPUISD::ATOMIC_CMP_SWAP, DL,
                                  VTList, Ops, VT, AtomicNode->getMemOperand());
 }
 
 //===----------------------------------------------------------------------===//
 // Custom DAG optimizations
 //===----------------------------------------------------------------------===//
 
+bool SITargetLowering::allowsMemoryAccess(LLVMContext &Context,
+                                          const DataLayout &DL, EVT VT,
+                                          unsigned AddrSpace,
+                                          unsigned Alignment,
+                                          bool *Fast) const {
+
+  if (AddrSpace == AMDGPUAS::PRIVATE_ADDRESS && VT.getStoreSize() > 4)
+    return false;
+
+  return AMDGPUTargetLowering::allowsMemoryAccess(Context, DL, VT, AddrSpace,
+                                                  Alignment, Fast);
+}
+
 SDValue SITargetLowering::performUCharToFloatCombine(SDNode *N,
                                                      DAGCombinerInfo &DCI) const {
   EVT VT = N->getValueType(0);
   EVT ScalarVT = VT.getScalarType();
   if (ScalarVT != MVT::f32)
     return SDValue();
 
   SelectionDAG &DAG = DCI.DAG;