gemini-code-assist[bot] commented on code in PR #20024:
URL: https://github.com/apache/tvm/pull/20024#discussion_r3603257990
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1333,17 +1333,22 @@ def _impl_v13(cls, bb, inputs, attr, params):
output = _np.take(data.data.numpy(), indices.data.numpy(),
axis=axis)
return relax.const(output, output.dtype)
- # If input is a shape expression, take a value from that shape and
return it as a constant.
+ # If input is a shape expression, take a value from that shape. A 0-D
+ # scalar constant index resolves to one dimension that we return as a
+ # PrimValue to keep shape-specialized handling in downstream
+ # shape-construction patterns. Any other index materializes the shape
as
+ # an int64 tensor and gathers from it at runtime, reusing the
+ # negative-index normalization below. ONNX Gather defines the output
rank
+ # as q + r - 1 (q = rank of indices); since the shape is rank 1, a
+ # non-scalar index such as (1,) must keep its rank, so only a true
+ # 0-D index collapses to a scalar.
if isinstance(data, relax.ShapeExpr):
- assert isinstance(indices, relax.Constant), (
- "Only constant indices supported for shape gather."
- )
- np_index = indices.data.numpy()
- if len(np_index.shape) == 1:
- np_index = np_index[0]
- np_index = int(np_index)
- shape_val = data[np_index]
- return relax.prim_value(shape_val)
+ if isinstance(indices, relax.Constant) and
indices.data.numpy().ndim == 0:
+ np_index = int(indices.data.numpy().item())
+ shape_val = data[np_index]
+ return relax.prim_value(shape_val)
Review Comment:
To improve readability and avoid calling `indices.data.numpy()` twice, you
could store its result in a variable before checking its dimension and using
its value.
```suggestion
if isinstance(indices, relax.Constant):
indices_np = indices.data.numpy()
if indices_np.ndim == 0:
np_index = int(indices_np.item())
shape_val = data[np_index]
return relax.prim_value(shape_val)
```
##########
src/target/llvm/codegen_llvm.cc:
##########
@@ -376,9 +377,38 @@ std::unique_ptr<llvm::Module> CodeGenLLVM::Finish() {
<< "Failed to link modules";
}
link_modules_.clear();
+
+ // Verify the freshly generated (pre-optimization) module. This is TVM's own
+ // codegen output; if it fails to verify that is a genuine TVM bug, so keep
+ // throwing here.
this->Verify();
+
+ // LLVM's optimizer has been observed to miscompile perfectly valid IR into a
+ // structurally-invalid module for certain shapes -- e.g. avg_pool2d with
+ // C == 4 and W == 3 sends LLVM's vectorizer down a buggy path that emits a
+ // shufflevector whose definition does not dominate its use (TVM issue
#20015;
+ // reproducible with stock `opt -passes='default<O2>'`, an upstream LLVM
bug).
+ // To avoid hard-crashing the whole build on such backend-optimizer bugs,
keep
+ // a verified copy of the pre-optimization module and fall back to it if
+ // optimization produces an invalid module. The result is correct (just
+ // unoptimized for the affected module) rather than a crash, and we never
emit
+ // the broken module.
+ std::unique_ptr<llvm::Module> pre_opt_module = llvm::CloneModule(*module_);
Review Comment:
`llvm::CloneModule` could potentially fail and return a `nullptr`. While
this might be unlikely, it's safer to add a check for the returned pointer. If
`pre_opt_module` is null and the subsequent optimization verification fails,
`module_` would be assigned a null `unique_ptr`, which would likely cause a
crash later on.
```c
std::unique_ptr<llvm::Module> pre_opt_module = llvm::CloneModule(*module_);
ICHECK(pre_opt_module) << "Failed to clone LLVM module before
optimization";
```
--
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]
