[GitHub] [tvm] mbaret commented on a diff in pull request #10871: [CUDNN] Add cuDNN as a Relay partitioning target (BYOC)

[GitBox <gi...@apache.org>]

mbaret commented on code in PR #10871:
URL: https://github.com/apache/tvm/pull/10871#discussion_r842877342


##########
python/tvm/relay/op/contrib/cudnn.py:
##########
@@ -0,0 +1,129 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-argument
+"""cuDNN Relay integration."""
+from typing import Callable, List, Tuple, Dict, Optional
+
+import tvm
+import tvm.ir
+from tvm import relay
+from tvm import te
+from tvm.relay import transform
+from tvm.contrib import cudnn
+
+from ...dataflow_pattern import is_op, wildcard
+from .register import register_pattern_table
+
+
+def partition_for_cudnn(
+    mod: tvm.IRModule, params: Optional[Dict[str, tvm.runtime.NDArray]] = None
+) -> tvm.IRModule:
+    """Partition the graph to offload for cuDNN.
+
+    Parameters
+    ----------
+    mod : tvm.IRModule
+        The module to partition.
+    params : Optional[Dict[str, tvm.runtime.NDArray]]
+        Constant input parameters.
+
+    Returns
+    -------
+    tvm.IRModule
+        The partitioned module.
+    """
+
+    seq = tvm.transform.Sequential(
+        [
+            transform.InferType(),
+            transform.MergeComposite(pattern_table()),
+            transform.AnnotateTarget("cudnn"),
+            transform.PartitionGraph(),
+            transform.InferType(),
+        ]
+    )
+    return seq(mod)
+
+
+@register_pattern_table("cudnn")
+def pattern_table() -> List[Tuple[str, relay.Pattern, Callable[[relay.Call], bool]]]:
+    """Get the cuDNN pattern table."""
+
+    def softmax_pattern() -> relay.Pattern:
+        """Create pattern for softmax."""
+        return is_op("nn.softmax")(wildcard())
+
+    def check_softmax(matched: relay.Call) -> bool:
+        """Check if softmax is supported by cuDNN."""
+        if matched.args[0].checked_type.dtype not in ["float64", "float32", "float16"]:
+            return False
+
+        return True
+
+    return [
+        ("cudnn.softmax", softmax_pattern(), check_softmax),
+    ]
+
+
+_LowerFunc = Callable[[relay.Call, List[te.Tensor]], te.Tensor]
+_LOWER_MAP: Dict[str, _LowerFunc] = {}
+
+
+def _lower_composite(comp_name: str) -> Callable[[_LowerFunc], _LowerFunc]:
+    """Register a lowering function for a given composite function name."""
+
+    def _register(f: _LowerFunc) -> _LowerFunc:
+        _LOWER_MAP[comp_name] = f
+        return f
+
+    return _register
+
+
+@tvm._ffi.register_func("relay.ext.cudnn")
+def relay_to_runtime(partition: relay.Function) -> tvm.runtime.Module:
+    """Compile cuDNN Relay functions to a runtime module."""
+    assert isinstance(partition, relay.Function)
+    assert isinstance(partition.body, relay.Call)
+    assert isinstance(partition.body.op, relay.Function)
+
+    global_name = str(partition.attrs.global_symbol)
+    target = tvm.target.cuda()

Review Comment:
   I had a go using this based on your suggestion, but it seems the Target only ends up in the context if you use the 'with' way of specifying targets (not just directly passing them to relay.build). In an ideal world, this somehow gets plumbed directly through the BYOC/partitioning mechanism - but I think that exceeds the immediate scope of this PR. Perhaps we could leave as a future improvement for now?



-- 
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