[Bug] Slice indexing in ONNX

[soodoshll]

Please refer to pytorch/pytorch#24251

Basically, ONNX uses extremely large numbers to represent slicing until the end of certain dimensions, which will be prohibited the defensive conditions in

hidet/python/hidet/graph/ops/definitions/transform.py

Line 481 in 80a35d6

if not (-n <= i <= n and -n <= j <= n):

[yaoyaoding]

The indexing of hidet tensor would follow the specification of ArrayAPI standard. Thus, we need to do deal with the difference between onnx sematics and ArrayAPI standard when importing onnx model.

Could you please provide some examples to trigger this error? If not, we could leave it to the future when some actual model triggers this error.

[soodoshll]

Yes, a very simple snippet (torch->onnx->hidet) can trigger this error:

import torch
import hidet
import onnx

class Foo(torch.nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x):
        return x[2:]

device = 'cuda'

model = Foo()
model.to(device)

x = torch.ones([100], dtype=torch.int32, device=device)
z = model(x)

torch.onnx.export(model, (x,), 'tmp.onnx', input_names = ['x'],
                  output_names = ['z'])
model = onnx.load('tmp.onnx')

hidet.torch.dynamo_config.search_space(1)

x = hidet.from_torch(x)
symbol_data = [hidet.symbol_like(x)]
hidet_onnx_module = hidet.graph.frontend.from_onnx(model)
symbol_output = hidet_onnx_module(*symbol_data)
graph: hidet.FlowGraph = hidet.trace_from(symbol_output, inputs=symbol_data)
with hidet.graph.PassContext() as ctx:
    graph_opt: hidet.FlowGraph = hidet.graph.optimize(graph)
cuda_graph = graph_opt.cuda_graph()
outputs = cuda_graph.run([x])
print(outputs[0])

[yaoyaoding]

Thanks @soodoshll, working on it.

[yaoyaoding]

Thanks @soodoshll, this error should be fixed in #106.