How to obtain information of llvm.func's arguments

Beginners
#1

Here is the llvm dialect code:

module {
       llvm.func @mat_test(%arg0:i32)
       {
           ......
           llvm.return
       }
   }

I want to get the memory space type of %arg0, and save it as a metadata info, so the corresponding llvm ir code will be:

; ModuleID = 'LLVMDialectModule'
  source_filename = "LLVMDialectModule"
  
  declare i8* @malloc(i64)
  
  declare void @free(i8*)

 ; Function Attrs: noinline alignstack(4)
 define void @mat_test(i32 %0) !dbg !3 !arg_addr_space !6 {
     ......
 }
.......
 !6 = !{i32 0} // here "0" represent the address space info of arg0
......

But, with the ‘mlir-translate’ tool , the obtained llvm ir code is shown below, which don’t have the ‘arg_addr_space’ metadata:

; ModuleID = 'LLVMDialectModule'
  source_filename = "LLVMDialectModule"
  
  declare i8* @malloc(i64)
  
  declare void @free(i8*)

 ; Function Attrs: noinline alignstack(4)
 define void @mat_test(i32 %0) !dbg !3 !arg_addr_space !6 {
     ......
 }
......

At present, I add the following ‘tranlateAttrtoMetadata’ function to ModuleTranslation.cpp ,whici will be called by the ‘checkSupportedModuleOps’ function, to translate the passthrough attributes into the llvm ir metadata.

static LogicalResult
tranlateAttrtoMetadata(LLVMFuncOp function, Optional<ArrayAttr> attributes,
                             llvm::Function *llvmFunc) {

  if (function.isExternal()){
    return success();
  }

  // MDNode for the kernel argument address space qualifiers.
  SmallVector<llvm::Metadata *, 8> addressQuals;
  Location loc = function.getLoc();
  if (!attributes){
    llvmFunc->setMetadata("kernel_arg_addr_space",
                  llvm::MDNode::get(llvmFunc->getContext(), addressQuals));
    return success();
  }

  for (Attribute attr : *attributes) {
    if (auto stringAttr = attr.dyn_cast<StringAttr>()) {
      continue;
    }

    auto arrayAttr = attr.dyn_cast<ArrayAttr>();
    if (!arrayAttr || arrayAttr.size() != 2)
      return emitError(loc)
             << "expected 'passthrough' to contain string or array attributes";

    auto keyAttr = arrayAttr[0].dyn_cast<StringAttr>();
    auto valueAttr = arrayAttr[1].dyn_cast<StringAttr>();
    if (!keyAttr || !valueAttr)
      return emitError(loc)
             << "expected arrays within 'passthrough' to contain two strings";

    StringRef key = keyAttr.getValue();
    StringRef value = valueAttr.getValue();
    if (key.str().find("arg_addr_space") != std::string::npos) {
      if (value.empty())
        return emitError(loc) << "LLVM attribute '" << key << "' expects a value";

      int result;
      if (!value.getAsInteger(0, result)) {
        addressQuals.push_back(
            llvm::ConstantAsMetadata::get((new llvm::IRBuilder<>(llvmFunc->getContext()))->getInt32(result)));
      } else {
        return emitError(loc) << "The value of kernel_arg_addr_space must be an Integer object";
      }
    }
  }

  llvmFunc->setMetadata("arg_addr_space",
                  llvm::MDNode::get(llvmFunc->getContext(), addressQuals));
  return success();
}

Based on the above method ,I have to add the argment address space info into the llvm dialect code, like this:

module {
  llvm.func @mat_test(%arg0:i32) attributes {passthrough = [ ["arg_addr_space", "4"]]  }
  {
          ......
          llvm.return
   }
 }

So, Is there a better way to get the argment address space info and pass it to llvm ir?