You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2022/06/17 21:01:36 UTC

[GitHub] [tvm] mbs-octoml opened a new pull request, #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

mbs-octoml opened a new pull request, #11770:
URL: https://github.com/apache/tvm/pull/11770

   I tried to do to the TensorRT integration what #11631 did to the CUTLASS integration, viz:
    - Make sure all compilation options are passed in Target instances. This helps Collage.
    - Use a custom pass invoked via RelayToTIRTargetHooks instead of the relay.ext.$toolchain mechanism.
      This helps use decouple external codegen from lowering.
   
   This PR collects the prep for that change:
    - TensorRT uses the JSONSerializer visitor to encode each partition function. Previously, when the
      visitor encountered a Constant it simply generated and recorded a name for the constant. Then,
      completely separately, and via a callback in TECompiler, the function is visited again in the
      same order and with the same name generation convention by a ConstantUpdater to actually collect the
      bindings, which are then encoded into a ConstLoaderModule to be made available at runtime.
   
      However if all TensorRT compilation is to be done by a stand-alone pass there's no TECompiler callback
      hackery available. So I've added a "const_name_to_ndarray" attribute to the IRModule of type
      Map<String, runtime::NDArray> so that named constants can be accumulated throughout compilation by
      any pass which needs to do so. Then the Graph, AOT and VM executors are all updated to merge those
      constants into the final runtime artifact
   
      (Compare with "Constants", the equivalent attribute for extracting TIR AllocateConsts.)
   
    - The TensorRT tests use the create_executor interface but it wasn't quite ready for the
      new more general form of passing list-of-targets.
   
    - I want TensorRT compilation to work out of the box without the need for any special targets if
      all the default options should apply. Go back and make the CUTLASS integration I did follow the
      same convention.
   
    - TensorRT actually needs to 'undo' partitionings in some situations. Add an InlineCompilerFunctions
      pass to make that robust. In particular, it must undo both the 'partitioning' (ie separating out
      the "Compiler" function) and any 'compositing' (ie separating out small sub-graphs as
      "Composite" functions).
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on a diff in pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on code in PR #11770:
URL: https://github.com/apache/tvm/pull/11770#discussion_r907641605


##########
src/relay/backend/contrib/cutlass/codegen.cc:
##########
@@ -43,6 +43,16 @@ namespace cutlass {
 
 namespace {
 
+/*! \brief Return the "cutlass" Target instance to use to guide compilation. */
+Target GetCutlassTarget() {
+  Target target = Target::Current(/*allow_not_defined=*/true);
+  if (!target.defined() || target->kind->name != "cutlass") {
+    // Use the default CUTLASS compilation options.
+    target = Target("cutlass");

Review Comment:
   I'll beef up the comment. This is so folks who just want to use the default CUTLASS compilation options (which are all very sensible) don't need to futz with multi targets and can just do the usual target="cuda" or whatever shorthand they prefer.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] areusch merged pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
areusch merged PR #11770:
URL: https://github.com/apache/tvm/pull/11770


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on a diff in pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on code in PR #11770:
URL: https://github.com/apache/tvm/pull/11770#discussion_r907628048


##########
python/tvm/relay/build_module.py:
##########
@@ -607,16 +607,16 @@ class AotExecutor(_interpreter.Executor):
     device : :py:class:`Device`
         The runtime device to run the code on.
 
-    target : :py:class:`Target`
-        The target option to build the function.
+    raw_targets : Array[tvm.target.Target]
+        The available targets.
     """
 
-    def __init__(self, mod, device, target):
+    def __init__(self, mod, device, raw_targets):
         assert mod is not None
         self.mod = mod
         self.device = device
-        self.target = target
-        assert target.attrs.get("executor", "graph") == "aot"
+        self.raw_targets = raw_targets
+        assert raw_targets[0].attrs.get("executor", "graph") == "aot"

Review Comment:
   See following comment.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on a diff in pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on code in PR #11770:
URL: https://github.com/apache/tvm/pull/11770#discussion_r907629807


##########
src/relay/backend/aot_executor_codegen.cc:
##########
@@ -1328,17 +1333,11 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {
   runtime::NDArray get_param_by_name(String key) {
     auto it = this->output_.params.find(key);
     CHECK(it != this->output_.params.end()) << "no such parameter " << key;
-    return (*it).second.second;
+    return (*it).second;

Review Comment:
   The params (constants) map was capturing both the ndarray and the internal id, but the id was never used. So I removed the id. This means the constants map is always string -> ndarray so the "const_name_to_ndarray" attribute can be used consistently irrespective of executor. It also means there's no need for the extra projection.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on PR #11770:
URL: https://github.com/apache/tvm/pull/11770#issuecomment-1159465775

   tests/python/frontend/pytorch/test_forward.py::test_argsort appears to be a flake.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on a diff in pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on code in PR #11770:
URL: https://github.com/apache/tvm/pull/11770#discussion_r907636570


##########
python/tvm/relay/build_module.py:
##########
@@ -707,24 +708,24 @@ def create_executor(kind="debug", mod=None, device=None, target="llvm", params=N
     -------
     executor : :py:class:`~tvm.relay.backend.interpreter.Executor`
     """
+    raw_targets = Target.canon_multi_target(target)

Review Comment:
   As a rule it would be good to have an 'anything goes' API which supports all the legacy ways of conveying target/targets, an an internal 'canonical' API which assumes the most general form of Array<Target>. However I don't see how to enforce such a distinction given the chaotic layering. So as the next best thing I've made sure to canonicalize at all the right points, but canonicalization is idempotent.
   
   Meanwhile, I've reverted the changes to require raw_targets in canonical form in favor of the generic 'any multi-target like object' used elsewhere. I've then moved the assert into create_executor since the targets are already required to be in canonical form to manufacture the default device. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on a diff in pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on code in PR #11770:
URL: https://github.com/apache/tvm/pull/11770#discussion_r907641605


##########
src/relay/backend/contrib/cutlass/codegen.cc:
##########
@@ -43,6 +43,16 @@ namespace cutlass {
 
 namespace {
 
+/*! \brief Return the "cutlass" Target instance to use to guide compilation. */
+Target GetCutlassTarget() {
+  Target target = Target::Current(/*allow_not_defined=*/true);
+  if (!target.defined() || target->kind->name != "cutlass") {
+    // Use the default CUTLASS compilation options.
+    target = Target("cutlass");

Review Comment:
   I'll beef up the comment. This is so folks who just want to use the default CUTLASS compilation options (which are all very sensible) don't need to futz with multi targets



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on PR #11770:
URL: https://github.com/apache/tvm/pull/11770#issuecomment-1170499789

   PTAL @areusch 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on a diff in pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on code in PR #11770:
URL: https://github.com/apache/tvm/pull/11770#discussion_r907640562


##########
src/relay/backend/contrib/codegen_json/codegen_json.h:
##########
@@ -245,11 +258,14 @@ class JSONSerializer : public MemoizedExprTranslator<std::vector<JSONGraphNodeEn
     return memo_[GetRef<Expr>(vn)];
   }
 
-  std::vector<JSONGraphNodeEntry> VisitExpr_(const ConstantNode* cn) {
-    std::string name = symbol_ + "_const_" + std::to_string(params_.size());
-    params_.push_back(name);
-    auto node = std::make_shared<JSONGraphNode>(name, "const" /* op_type_ */);
-    return AddNode(node, GetRef<Expr>(cn));
+  std::vector<JSONGraphNodeEntry> VisitExpr_(const ConstantNode* constant_node) {
+    std::string name = symbol_ + "_const_" + std::to_string(const_names_.size());
+    VLOG(1) << "Will require parameter '" << name << "' to be available at runtime";

Review Comment:
   Done -- this was me trying to reverse engineer the consts protocol. It is very unfortunate we've bound everything together as 'params'.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on PR #11770:
URL: https://github.com/apache/tvm/pull/11770#issuecomment-1170354716

   It turns out switch the 'ccompiler' demo external codegen to use target hooks is enough to exercise all this, all be it the PR has grown quite a bit. Let's see what CI makes of it since probably broken something else.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] areusch commented on a diff in pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
areusch commented on code in PR #11770:
URL: https://github.com/apache/tvm/pull/11770#discussion_r902764117


##########
include/tvm/ir/module.h:
##########
@@ -507,10 +509,32 @@ constexpr const char* kRuntime = "runtime";
 constexpr const char* kWorkspaceMemoryPools = "workspace_memory_pools";
 
 /*
- * \brief Module attribute for tir constants
+ * \brief All the runtime::NDArrays extracted from PrimFunc tir::AllocateConst nodes. The
+ * node will record the index into this array. See also kConstNameToNDArray below, which is
+ * the analog for Realy Functions.
+ *
+ * Type: Array<runtime::NDArray>
  */
 constexpr const char* kConstantsArray = "Constants";

Review Comment:
   if you have time, want to fix the naming here? (e.g. either kConstants and "constants" or kConstantsArray and "constants_array")



##########
python/tvm/relay/build_module.py:
##########
@@ -707,24 +708,24 @@ def create_executor(kind="debug", mod=None, device=None, target="llvm", params=N
     -------
     executor : :py:class:`~tvm.relay.backend.interpreter.Executor`
     """
+    raw_targets = Target.canon_multi_target(target)

Review Comment:
   so we require canon_multi_target, but we then immediate re-canonicalize. should we either assert it's canonicalized rather than force-converting it, or update docstring? perhaps here and in the executors, you really want a CanonicalizedTarget type?



##########
src/relay/backend/contrib/codegen_json/codegen_json.h:
##########
@@ -245,11 +258,14 @@ class JSONSerializer : public MemoizedExprTranslator<std::vector<JSONGraphNodeEn
     return memo_[GetRef<Expr>(vn)];
   }
 
-  std::vector<JSONGraphNodeEntry> VisitExpr_(const ConstantNode* cn) {
-    std::string name = symbol_ + "_const_" + std::to_string(params_.size());
-    params_.push_back(name);
-    auto node = std::make_shared<JSONGraphNode>(name, "const" /* op_type_ */);
-    return AddNode(node, GetRef<Expr>(cn));
+  std::vector<JSONGraphNodeEntry> VisitExpr_(const ConstantNode* constant_node) {
+    std::string name = symbol_ + "_const_" + std::to_string(const_names_.size());
+    VLOG(1) << "Will require parameter '" << name << "' to be available at runtime";

Review Comment:
   could we say more about "available"? does that mean the user needs to supply it to set_input()?



##########
src/relay/backend/contrib/cutlass/codegen.cc:
##########
@@ -43,6 +43,16 @@ namespace cutlass {
 
 namespace {
 
+/*! \brief Return the "cutlass" Target instance to use to guide compilation. */
+Target GetCutlassTarget() {
+  Target target = Target::Current(/*allow_not_defined=*/true);
+  if (!target.defined() || target->kind->name != "cutlass") {
+    // Use the default CUTLASS compilation options.
+    target = Target("cutlass");

Review Comment:
   how come target->kind->name != "cutlass" makes this happen? seems like it should icheck-fail?



##########
python/tvm/relay/build_module.py:
##########
@@ -607,16 +607,16 @@ class AotExecutor(_interpreter.Executor):
     device : :py:class:`Device`
         The runtime device to run the code on.
 
-    target : :py:class:`Target`
-        The target option to build the function.
+    raw_targets : Array[tvm.target.Target]
+        The available targets.
     """
 
-    def __init__(self, mod, device, target):
+    def __init__(self, mod, device, raw_targets):
         assert mod is not None
         self.mod = mod
         self.device = device
-        self.target = target
-        assert target.attrs.get("executor", "graph") == "aot"
+        self.raw_targets = raw_targets
+        assert raw_targets[0].attrs.get("executor", "graph") == "aot"

Review Comment:
   shall we add a corresponding assert for graph?
   
   what's the expected form of this array of Targets (seems the host should be the first one)? how should we know what to supply here (and for GraphExecutor)?



##########
src/relay/backend/aot_executor_codegen.cc:
##########
@@ -1328,17 +1333,11 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {
   runtime::NDArray get_param_by_name(String key) {
     auto it = this->output_.params.find(key);
     CHECK(it != this->output_.params.end()) << "no such parameter " << key;
-    return (*it).second.second;
+    return (*it).second;

Review Comment:
   huh?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on a diff in pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on code in PR #11770:
URL: https://github.com/apache/tvm/pull/11770#discussion_r907618811


##########
include/tvm/ir/module.h:
##########
@@ -507,10 +509,32 @@ constexpr const char* kRuntime = "runtime";
 constexpr const char* kWorkspaceMemoryPools = "workspace_memory_pools";
 
 /*
- * \brief Module attribute for tir constants
+ * \brief All the runtime::NDArrays extracted from PrimFunc tir::AllocateConst nodes. The
+ * node will record the index into this array. See also kConstNameToNDArray below, which is
+ * the analog for Realy Functions.
+ *
+ * Type: Array<runtime::NDArray>
  */
 constexpr const char* kConstantsArray = "Constants";

Review Comment:
   Done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] mbs-octoml commented on pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
mbs-octoml commented on PR #11770:
URL: https://github.com/apache/tvm/pull/11770#issuecomment-1167723426

   Thanks @masahi  & @areusch, PTAL.
   
   Re value_or: done, thanks, that's much nicer.
   Re split PRs, done. This was already a split of a split so sorry for the grab bag.
   Re end-to-end test: I think this should be done in test_target_hooks.py, but to really exercise everything requires a new fake JSON-based runtime module, a fake custom codegen pass and a new fake external codegen target. I started down that path but it will be quite a lift. Unfortunately however we can't rely on the tensorrt unit test for this (the first user of the new const_name_to_constants attribute) since the functionality won't be testable without running, which is not enabled in CI. WDYT?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [tvm] areusch commented on pull request #11770: [BYOC] Handle constants in IRModule-at-a-time external codegen

Posted by GitBox <gi...@apache.org>.
areusch commented on PR #11770:
URL: https://github.com/apache/tvm/pull/11770#issuecomment-1168987437

   i hate to be a stickler, but i feel like _something_ should test this in CI. would such a fake module become useful in other compiler testing?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org