gemini-code-assist[bot] commented on code in PR #19649:
URL: https://github.com/apache/tvm/pull/19649#discussion_r3329796610
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -2162,6 +2163,44 @@ def _convert_stablehlo_sort(self, op):
relax.op.sort(data, axis=int(opts.Dimension()),
descending=descending)
)
+ def _convert_stablehlo_custom_call(self, op):
+ """Convert supported annotation-only STABLEHLO_CUSTOM_CALL targets."""
+ from tflite.StablehloCustomCallOptions import
StablehloCustomCallOptions
+
+ input_tensors = self.get_input_tensors(op)
+ output_tensors = self.get_output_tensors(op)
+ opts = self._get_stablehlo_options(op, StablehloCustomCallOptions)
+ call_target_name = self._decode_tflite_string(opts.CallTargetName())
Review Comment:
If `opts` is `None` (which can happen if the operator's builtin options are
not set in the TFLite model), calling `opts.CallTargetName()` will raise an
`AttributeError`. It is safer to guard against `None` and raise a clean
`OpNotImplemented` error.
```python
opts = self._get_stablehlo_options(op, StablehloCustomCallOptions)
if opts is None:
raise tvm.error.OpNotImplemented(
"STABLEHLO_CUSTOM_CALL without options is not supported"
)
call_target_name = self._decode_tflite_string(opts.CallTargetName())
```
##########
tests/python/relax/test_frontend_tflite.py:
##########
@@ -6812,6 +6867,62 @@ def test_stablehlo_scatter_update_window_unsupported():
from_tflite(tflite_model)
+def test_stablehlo_custom_call_sharding():
+ """TFLite StableHLO CUSTOM_CALL Sharding annotation lowers to identity."""
+ mod = _load_model_from_buffer(_build_stablehlo_custom_call_model())
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 2), dtype="float32")) -> R.Tensor((2, 2),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((2, 2), dtype="float32") = x
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_stablehlo_custom_call_unsupported_target():
+ """TFLite StableHLO CUSTOM_CALL rejects unknown external call targets."""
+ buf = _build_stablehlo_custom_call_model(call_target_name="custom_backend")
+ if hasattr(tflite.Model, "Model"):
+ tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0)
+ else:
+ tflite_model = tflite.Model.GetRootAsModel(buf, 0)
+
+ with pytest.raises(
+ tvm.error.OpNotImplemented,
+ match="STABLEHLO_CUSTOM_CALL target custom_backend is not supported",
+ ):
+ from_tflite(tflite_model)
+
+
+def test_stablehlo_custom_call_sharding_side_effect_unsupported():
+ """TFLite StableHLO CUSTOM_CALL rejects side-effecting Sharding calls."""
+ buf = _build_stablehlo_custom_call_model(has_side_effect=True)
+ if hasattr(tflite.Model, "Model"):
+ tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0)
+ else:
+ tflite_model = tflite.Model.GetRootAsModel(buf, 0)
+
+ with pytest.raises(tvm.error.OpNotImplemented, match="side effects"):
+ from_tflite(tflite_model)
Review Comment:
We can simplify the test by using the existing helper function
`_load_model_from_buffer` instead of manually extracting the root model with
`hasattr` checks. This reduces boilerplate and improves readability.
```python
with pytest.raises(tvm.error.OpNotImplemented, match="side effects"):
_load_model_from_buffer(buf)
```
##########
tests/python/relax/test_frontend_tflite.py:
##########
@@ -6812,6 +6867,62 @@ def test_stablehlo_scatter_update_window_unsupported():
from_tflite(tflite_model)
+def test_stablehlo_custom_call_sharding():
+ """TFLite StableHLO CUSTOM_CALL Sharding annotation lowers to identity."""
+ mod = _load_model_from_buffer(_build_stablehlo_custom_call_model())
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 2), dtype="float32")) -> R.Tensor((2, 2),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((2, 2), dtype="float32") = x
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_stablehlo_custom_call_unsupported_target():
+ """TFLite StableHLO CUSTOM_CALL rejects unknown external call targets."""
+ buf = _build_stablehlo_custom_call_model(call_target_name="custom_backend")
+ if hasattr(tflite.Model, "Model"):
+ tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0)
+ else:
+ tflite_model = tflite.Model.GetRootAsModel(buf, 0)
+
+ with pytest.raises(
+ tvm.error.OpNotImplemented,
+ match="STABLEHLO_CUSTOM_CALL target custom_backend is not supported",
+ ):
+ from_tflite(tflite_model)
+
+
+def test_stablehlo_custom_call_sharding_side_effect_unsupported():
+ """TFLite StableHLO CUSTOM_CALL rejects side-effecting Sharding calls."""
+ buf = _build_stablehlo_custom_call_model(has_side_effect=True)
+ if hasattr(tflite.Model, "Model"):
+ tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0)
+ else:
+ tflite_model = tflite.Model.GetRootAsModel(buf, 0)
+
+ with pytest.raises(tvm.error.OpNotImplemented, match="side effects"):
+ from_tflite(tflite_model)
+
+
+def test_stablehlo_custom_call_sharding_metadata_mismatch_unsupported():
+ """TFLite StableHLO CUSTOM_CALL rejects Sharding calls that change tensor
metadata."""
+ buf =
_build_stablehlo_custom_call_model(output_tensor_type=_tfl_tensor_type.INT32)
+ if hasattr(tflite.Model, "Model"):
+ tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0)
+ else:
+ tflite_model = tflite.Model.GetRootAsModel(buf, 0)
+
+ with pytest.raises(tvm.error.OpNotImplemented, match="Sharding tensor
metadata mismatch"):
+ from_tflite(tflite_model)
Review Comment:
We can simplify the test by using the existing helper function
`_load_model_from_buffer` instead of manually extracting the root model with
`hasattr` checks. This reduces boilerplate and improves readability.
```python
with pytest.raises(tvm.error.OpNotImplemented, match="Sharding tensor
metadata mismatch"):
_load_model_from_buffer(buf)
```
##########
tests/python/relax/test_frontend_tflite.py:
##########
@@ -6812,6 +6867,62 @@ def test_stablehlo_scatter_update_window_unsupported():
from_tflite(tflite_model)
+def test_stablehlo_custom_call_sharding():
+ """TFLite StableHLO CUSTOM_CALL Sharding annotation lowers to identity."""
+ mod = _load_model_from_buffer(_build_stablehlo_custom_call_model())
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 2), dtype="float32")) -> R.Tensor((2, 2),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((2, 2), dtype="float32") = x
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_stablehlo_custom_call_unsupported_target():
+ """TFLite StableHLO CUSTOM_CALL rejects unknown external call targets."""
+ buf = _build_stablehlo_custom_call_model(call_target_name="custom_backend")
+ if hasattr(tflite.Model, "Model"):
+ tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0)
+ else:
+ tflite_model = tflite.Model.GetRootAsModel(buf, 0)
+
+ with pytest.raises(
+ tvm.error.OpNotImplemented,
+ match="STABLEHLO_CUSTOM_CALL target custom_backend is not supported",
+ ):
+ from_tflite(tflite_model)
Review Comment:
We can simplify the test by using the existing helper function
`_load_model_from_buffer` instead of manually extracting the root model with
`hasattr` checks. This reduces boilerplate and improves readability.
```suggestion
with pytest.raises(
tvm.error.OpNotImplemented,
match="STABLEHLO_CUSTOM_CALL target custom_backend is not supported",
):
_load_model_from_buffer(buf)
```
--
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]
