[fix] fix the bug where 'tvm_callback_hip_postproc' missing target (#3)

* [Dev][AMD]add tensor_ * [fix]fix the bug of 'tvm_callback_hip_postproc' lacking target * Revert "[Dev][AMD]add tensor_" This reverts commit e04eecd.
TileLang · Oct 29, 2024 · 42b4550 · 42b4550
1 parent 27078af
commit 42b4550
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/target/opt/build_rocm_nort.cc b/src/target/opt/build_rocm_nort.cc
@@ -107,7 +107,7 @@ runtime::Module BuildHIP(IRModule mod, Target target) {
   std::string code = cg.Finish();
 
   if (const auto* f = Registry::Get("tvm_callback_hip_postproc")) {
-    code = (*f)(code).operator std::string();
+    code = (*f)(code, target).operator std::string();
   }
   std::string fmt = "ptx";
   std::string ptx;

diff --git a/src/target/opt/build_rocm_on.cc b/src/target/opt/build_rocm_on.cc
@@ -160,14 +160,14 @@ runtime::Module BuildHIP(IRModule mod, Target target) {
   std::string code = cg.Finish();
 
   if (const auto* f = Registry::Get("tvm_callback_hip_postproc")) {
-    code = (*f)(code).operator std::string();
+    code = (*f)(code, target).operator std::string();
   }
   std::string fmt = "ptx";
   std::string ptx;
   const auto* f_enter = Registry::Get("target.TargetEnterScope");
   (*f_enter)(target);
   if (const auto* f = Registry::Get("tvm_callback_hip_compile")) {
-    ptx = (*f)(code).operator std::string();
+    ptx = (*f)(code, target).operator std::string();
     // Dirty matching to check PTX vs hsaco.
     // TODO(leiwang1999) more reliable checks
     if (ptx[0] != '/') fmt = "hsaco";