[AMD] Update shared memory size for cdna4 (#5964)

New architecture has increased shared memory size.

Author: SJW

third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp

@@ -64,7 +64,10 @@ llvm::AMDGPU::GPUKind TargetInfo::getGPUKind() const {
   return llvm::AMDGPU::parseArchAMDGCN(arch);
 }
 
-int TargetInfo::getSharedMemorySize() const { return 64 * 1024; }
+int TargetInfo::getSharedMemorySize() const {
+  int kbytes = getISAFamily() == ISAFamily::CDNA4 ? 160 : 64;
+  return kbytes * 1024;
+}
 
 bool TargetInfo::supportMaximumMinimum() const { return false; }
 
@@ -102,7 +105,8 @@ bool TargetInfo::canUseStMatrix(RankedTensorType tensorTy,
 }
 
 bool TargetInfo::canUseLDSTransLoad(int bitwidth) const {
-  return arch == "gfx950" && llvm::is_contained({16, 8, 4, 6}, bitwidth);
+  return getISAFamily() == ISAFamily::CDNA4 &&
+         llvm::is_contained({16, 8, 4, 6}, bitwidth);
 }
 
 void TargetInfo::storeMatrixShared(RewriterBase &rewriter, Location loc,