[Intel] Add getAddressSpace to TargetInfo

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

third_party/intel/lib/TritonIntelGPUToLLVM/TargetInfo.cpp

@@ -321,6 +321,16 @@ Value TargetInfo::getStackPointer(RewriterBase &rewriter,
   return funcOp.getArgument(funcOp.getNumArguments() - 1);
 }
 
+int TargetInfo::getAddressSpace(Attribute addressSpace) const {
+  int spaceId = 0;
+  if (isa<triton::gpu::SharedMemorySpaceAttr>(addressSpace)) {
+    spaceId = 3;
+  } else {
+    llvm::report_fatal_error("Only support SharedMemorySpace for now");
+  }
+  return spaceId;
+}
+
 Value TargetInfo::getGlobalStringStart(Location loc, RewriterBase &rewriter,
                                        StringRef name, StringRef value,
                                        unsigned addressSpace) const {

third_party/intel/lib/TritonIntelGPUToLLVM/TargetInfo.h

@@ -71,6 +71,8 @@ class TargetInfo : public mlir::triton::TargetInfoBase {
   Value getStackPointer(RewriterBase &rewriter,
                         FunctionOpInterface funcOp) const override;
 
+  int getAddressSpace(Attribute addressSpace) const override;
+
   Value getGlobalStringStart(Location loc, RewriterBase &rewriter,
                              StringRef name, StringRef value,
                              unsigned addressSpace) const;