gemini-code-assist[bot] commented on code in PR #18776:
URL: https://github.com/apache/tvm/pull/18776#discussion_r2804797749
##########
python/tvm/s_tir/analysis/__init__.py:
##########
@@ -118,3 +119,96 @@ def find_anchor_sblock(mod: IRModule) -> Optional[SBlock]:
The anchor block if found, None otherwise.
"""
return _ffi_api.find_anchor_sblock(mod) # type: ignore # pylint:
disable=no-member
+
+
+def verify_gpu_code(func: PrimFunc, constraints: Dict[str, int]) -> None:
+ """Verify if module contains illegal host side direct memory access.
+
+ Parameters
+ ----------
+ func: tvm.tir.PrimFunc
+ The module to be verified.
+
+ constraints : Dict[str, int]
+ The attribute constraints.
+
+ Returns
+ -------
+ result : bool
+ The result of verification.
+ """
+ return _ffi_api.verify_gpu_code(func, constraints) # type: ignore
+
+
+def calculate_allocated_bytes(
+ func_or_mod: Union[PrimFunc, IRModule],
+) -> Union[Dict[str, int], Dict[str, Dict[str, int]]]:
Review Comment:
The return type hint `Union[Dict[str, int], Dict[str, Dict[str, int]]]` is
incorrect. The underlying C++ implementation always returns a `Map<String,
Map<String, Integer>>`, which translates to `Dict[str, Dict[str, int]]` in
Python for both `PrimFunc` and `IRModule` inputs. For a `PrimFunc`, the result
is wrapped in a dictionary with the key "main". The type hint should be
simplified to `Dict[str, Dict[str, int]]` to accurately reflect the return type.
```suggestion
def calculate_allocated_bytes(
func_or_mod: Union[PrimFunc, IRModule],
) -> Dict[str, Dict[str, int]]:
```
##########
python/tvm/s_tir/analysis/__init__.py:
##########
@@ -118,3 +119,96 @@ def find_anchor_sblock(mod: IRModule) -> Optional[SBlock]:
The anchor block if found, None otherwise.
"""
return _ffi_api.find_anchor_sblock(mod) # type: ignore # pylint:
disable=no-member
+
+
+def verify_gpu_code(func: PrimFunc, constraints: Dict[str, int]) -> None:
Review Comment:
The return type hint for `verify_gpu_code` is `None`, but the docstring and
the underlying C++ implementation indicate it should return a `bool`. Please
correct the type hint to match the actual return type.
```suggestion
def verify_gpu_code(func: PrimFunc, constraints: Dict[str, int]) -> bool:
```
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
