[ROCM] View llvm ir and gcn asm with module.get_source(...) (#590)

* view llvm ir and gcn asm with module.get_source(...)

* fix lint

src/codegen/llvm/codegen_amdgpu.cc

@@ -196,6 +196,13 @@ runtime::Module BuildAMDGPU(Array<LoweredFunc> funcs, std::string target) {
   pass.run(*mObj);
   std::string obj(dataObj.begin(), dataObj.end());
 
+  llvm::legacy::PassManager passAsm;
+  CHECK(tm->addPassesToEmitFile(passAsm, destAsm,
+                                llvm::TargetMachine::CGFT_AssemblyFile) == 0)
+      << "Cannot emit target CGFT_AssemblyFile";
+  passAsm.run(*mAsm);
+  std::string assembly(dataAsm.begin(), dataAsm.end());
+
   const auto* f = tvm::runtime::Registry::Get("tvm_callback_rocm_link");
   CHECK(f != nullptr) << "Require tvm_callback_rocm_link to exist, do import tvm.contrib.rocm";
 
@@ -206,7 +213,7 @@ runtime::Module BuildAMDGPU(Array<LoweredFunc> funcs, std::string target) {
   std::string hsaco = (*f)(arr);
   std::string ll(data_ll.begin(), data_ll.end());
 
-  return ROCMModuleCreate(hsaco, "hsaco", ExtractFuncInfo(funcs), ll);
+  return ROCMModuleCreate(hsaco, "hsaco", ExtractFuncInfo(funcs), ll, assembly);
 }
 
 TVM_REGISTER_API("codegen.build_rocm")

src/runtime/rocm/rocm_module.cc

@@ -30,8 +30,9 @@ class ROCMModuleNode : public runtime::ModuleNode {
   explicit ROCMModuleNode(std::string data,
                           std::string fmt,
                           std::unordered_map<std::string, FunctionInfo> fmap,
-                          std::string hip_source)
-      : data_(data), fmt_(fmt), fmap_(fmap), hip_source_(hip_source) {
+                          std::string hip_source,
+                          std::string assembly)
+    : data_(data), fmt_(fmt), fmap_(fmap), hip_source_(hip_source), assembly_(assembly) {
     std::fill(module_.begin(), module_.end(), nullptr);
   }
   // destructor
@@ -61,7 +62,8 @@ class ROCMModuleNode : public runtime::ModuleNode {
 
   std::string GetSource(const std::string& format) final {
     if (format == fmt_) { return data_; }
-    if (fmt_ == "hsaco") { return data_; }
+    if (format == "llvm") { return hip_source_; }
+    if (format == "asm") { return assembly_; }
     return "";
   }
 
@@ -116,6 +118,8 @@ class ROCMModuleNode : public runtime::ModuleNode {
   std::unordered_map<std::string, FunctionInfo> fmap_;
   // The hip source.
   std::string hip_source_;
+  // The gcn asm.
+  std::string assembly_;
   // the internal modules per GPU, to be lazily initialized.
   std::array<hipModule_t, kMaxNumGPUs> module_;
   // internal mutex when updating the module
@@ -202,9 +206,10 @@ Module ROCMModuleCreate(
     std::string data,
     std::string fmt,
     std::unordered_map<std::string, FunctionInfo> fmap,
-    std::string hip_source) {
+    std::string hip_source,
+    std::string assembly) {
   std::shared_ptr<ROCMModuleNode> n =
-      std::make_shared<ROCMModuleNode>(data, fmt, fmap, hip_source);
+    std::make_shared<ROCMModuleNode>(data, fmt, fmap, hip_source, assembly);
   return Module(n);
 }
 
@@ -216,7 +221,7 @@ Module ROCMModuleLoadBinary(void* strm) {
   stream->Read(&fmt);
   stream->Read(&fmap);
   stream->Read(&data);
-  return ROCMModuleCreate(data, fmt, fmap, std::string());
+  return ROCMModuleCreate(data, fmt, fmap, std::string(), std::string());
 }
 
 

src/runtime/rocm/rocm_module.h

@@ -31,7 +31,8 @@ Module ROCMModuleCreate(
     std::string data,
     std::string fmt,
     std::unordered_map<std::string, FunctionInfo> fmap,
-    std::string rocm_source);
+    std::string rocm_source,
+    std::string assembly);
 }  // namespace runtime
 }  // namespace tvm
 #endif  // TVM_RUNTIME_ROCM_ROCM_MODULE_H_