tkonolige commented on code in PR #13003:
URL: https://github.com/apache/tvm/pull/13003#discussion_r995200316
##########
python/tvm/utils/roofline/cuda.py:
##########
@@ -161,6 +165,51 @@ def peak_flops_tensorcore_tir(
return n * 16 * 16 * 16 * 2 * sms * 8 / times.min
[email protected]_peak_flops.register("cuda")
+def estimate_peak_flops(
+ func: PrimFunc, # pylint: disable=unused-argument
+ features: Dict[str, np.ndarray],
+ target: Target,
+ dev: Device,
+ remote: Optional[RPCSession],
+) -> Tuple[float, float, str]:
+ """Estimate the peak FLOP/s of a cuda device.
+
+ Parameters
+ ----------
+ func : PrimFunc
+ Function to estimate peak flops for. Used to check if a specific kind
+ intrinsic or dtype could be used with this function.
+ features : Dict[str, np.ndarry]
+ Features extracted from `func`. Used to check if a specific kind
+ intrinsic or dtype could be used with this function.
+ target : Target
+ Target to run on. This should be as specific to the actual hardware as
+ possible.
+ dev : Device
+ Device to run on.
+ remote : Optional[RPCSession]
+ Remote session used to upload artifacts for runtime evaluation. Must be
+ the same session used to create `dev`.
+
+ Returns
+ -------
+ flops : float
+ Estimated number of flops used by `func`.
+ peak_flops : float
+ Approximate sustained FLOP/s of this target/device combo. Addition and
+ multiplications are each counted as separate FLOPs.
+ name : str
+ Dtype/intrinsic used by `func` to achieve peak flops.
+ """
+ assert nvcc.have_tensorcore(
+ dev.compute_version
+ ), "CUDA roofline only works with devices that have tensorcores"
+ flops = np.sum(features["float_addsub"] + features["float_mul"] +
features["float_mad"])
Review Comment:
Good points, added.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
