search

webonnx / wonnx

A WebGPU-accelerated ONNX inference run-time written 100% in Rust, ready for native and the web
Other
1.65k stars 60 forks source link

feat: support Unsqueeze and Identity ops #31

Closed pixelspark closed 2 years ago

pixelspark commented 2 years ago

These ops appear to only change shape (which we don't care about as we require the dimensions to be inferred up front anyway), not data. Squeeze was already there so I added Unsqueeze and Identity as well.

We should probably not bother copying data between buffers but have a step in the compiler that simply aliases buffers, i.e. in input_a -> Identity -> SomeOtherOp, the SomeOtherOp should just take input_a as input rather than the output from identity (which we then won't generate).

pixelspark commented 2 years ago

This PR now also includes logic to skip identity ops altogether (when they are followed by some other op) by stringing together the input of the identity op and the input of the next op. The copy shader is still executed when an identity op is last, or is the only op in the graph (tested by the additional test). Fixes #33

haixuanTao commented 2 years ago

Would it be possible to just skip the node?. The thing is by passing on the next node. There is a chance that we miss on some sequence optimisiation.

pixelspark commented 2 years ago

Would it be possible to just skip the node?. The thing is by passing on the next node. There is a chance that we miss on some sequence optimisiation.

Technically the current implementation already skips the node, only it does so inside the sequencer. Are you suggesting to instead add a preprocessing step before we feed the model into the sequencer, which removes these nodes (by editing the ModelProto basically)? Fine with me (just a little bit more fiddly I guess), I can have a go at this later hopefully.

haixuanTao commented 2 years ago

Ok, i'm fine with the current implementation for now.

Yeah, adding a new step is going to be a lot of dev. Let's keep it simple for now.

Shall i merge?

pixelspark commented 2 years ago

Ok, i'm fine with the current implementation for now.

Yeah, adding a new step is going to be a lot of dev. Let's keep it simple for now.

Shall i merge?

Agreed, let's merge this, it's complicated enough already 👍