Minor fix for the debug interface of using PTX directly (pytorch#1917)

jjsjann123 · Aug 17, 2022 · ee8ef33 · ee8ef33
1 parent 6e8f953
commit ee8ef33
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/torch/csrc/jit/codegen/cuda/executor.cpp b/torch/csrc/jit/codegen/cuda/executor.cpp
@@ -1098,7 +1098,8 @@ std::vector<at::Tensor> FusionExecutor::runFusion(
 void FusionExecutor::compileRtc(
     const std::string& code,
     const std::string& name,
-    bool structured) {
+    bool structured,
+    CompileOptions options) {
   FUSER_PERF_SCOPE("ExecutorRunFusion::compileRtc");
   std::string scode;
   if (!structured) {
@@ -1107,7 +1108,7 @@ void FusionExecutor::compileRtc(
     scode = code;
   }
   fusion_id_ = 1;
-  options_ = CompileOptions();
+  options_ = options;
 
   std::tie(compiled_kernel_, last_compiler_log_) =
       executor_utils::nvrtcCompile(scode, name, fusion_id_);

diff --git a/torch/csrc/jit/codegen/cuda/executor.h b/torch/csrc/jit/codegen/cuda/executor.h
@@ -141,7 +141,8 @@ class TORCH_CUDA_CU_API FusionExecutor : public NonCopyable {
   void compileRtc(
       const std::string& code,
       const std::string& name,
-      bool structured = false);
+      bool structured = false,
+      CompileOptions options = CompileOptions());
 
   //! Internal tests only. Runs the compiled CUDA kernel from compileRtc.
   void runRtc(