Preserve purity in RunCodegen

apache · Apr 3, 2023 · 1e34266 · 1e34266
1 parent d55292f
commit 1e34266
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/src/relax/transform/run_codegen.cc b/src/relax/transform/run_codegen.cc
@@ -102,7 +102,13 @@ class CodeGenRunner : ExprMutator {
           func = (*RemoveFuncAttrFunc)(func, tvm::attr::kGlobalSymbol);
           func = (*RemoveFuncAttrFunc)(func, attr::kCodegen);
           builder_->UpdateFunction(gvar, func);
-          return create_call_dps_packed(new_func, func->ret_struct_info);
+          // preserve the purity: if the func was originally pure, wrap call_pure
+          bool purity = GetStructInfoAs<FuncStructInfoNode>(gvar)->purity;
+          auto ret = create_call_dps_packed(new_func, func->ret_struct_info);
+          if (purity) {
+            return WrapCallPure(ret);
+          }
+          return ret;
         }
       }
     }

diff --git a/tests/python/relax/test_codegen_dnnl.py b/tests/python/relax/test_codegen_dnnl.py
@@ -72,8 +72,6 @@ def test_dnnl_offload():
 
     seq = tvm.transform.Sequential(
         [
-            relax.transform.ToNonDataflow(),
-            relax.transform.RemovePurityChecking(),
             relax.transform.FuseOpsByPattern([("dnnl.conv2d_relu", pat)]),
             relax.transform.MergeCompositeFunctions(),
             relax.transform.RunCodegen(),