tkonolige commented on code in PR #13003:
URL: https://github.com/apache/tvm/pull/13003#discussion_r996099556
##########
python/tvm/utils/roofline/cuda.py:
##########
@@ -234,3 +259,63 @@ def estimate_peak_bandwidth(
b = nd.empty((blocks, 4, warp_size), dtype="float32", device=dev)
times = f.time_evaluator(f.entry_name, dev, repeat=10, number=1)(a, b)
return a.numpy().size * 4 / times.min # 4 bytes per float32
+
+
+@registry.estimate_peak_bandwidth.register("cuda")
+def estimate_peak_bandwidth(
+ func: PrimFunc, # pylint: disable=unused-argument
+ features: Dict[str, np.ndarray],
+ target: Target,
+ dev: Device,
+ remote: Optional[RPCSession] = None,
+) -> Tuple[float, float, str]:
+ """Estimate peak memory bandwidth of a target/device combo.
+
+ Peak bandwidth is estimated by running a small experiment on the underlying
+ hardware. The peak bandwidth measurement assumes that vector instructions
+ are being used to load the data.
+
+ Parameters
+ ----------
+ func : PrimFunc
+ Function to estimate peak bandwidth for. Used to check if a specific
+ kind of memory could be used with this function.
+ features : Dict[str, np.ndarry]
+ Features extracted from `func`. Used to check if a specific kind of
+ memory could be used with this function.
+ target : Target
+ Target to use for measurement. This target should be as specific to the
+ underlying hardware as possible.
+ dev : Device
+ Device to measure peak bandwidth on.
+ remote : Optional[RPCSession]
+ Remote session used to upload artifacts for runtime evaluation. Must be
+ the same session used to create `dev`.
+
+ Returns
+ -------
+ loaded_bytes : float
+ Estimated bytes loaded by `func`.
+ peak_bandwidth : float
+ Peak memory bandwidth in bytes/seconds.
+ name : str
+ Name of the memory being used.
+ """
+ loaded_bytes = 0.0
+ # assume no more than 100 buffers
+ for i in range(100):
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-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
