onnxruntime::Stream needs an improvement

[cloudhan]

Currently to use newly introduced onnxruntime::Stream, the experience is pretty awkward.

The main problem is that you need to manually static_cast the native stream handle before calling any function accepts native stream.
A aggragate wrapper with conversion function implemented will solve this problem at large margin as shown below:

https://godbolt.org/z/hMqax8vxv

__global__ void kernel_take_native_stream() {}

namespace ort {
struct Stream {
    explicit Stream(cudaStream_t native_stream) : stream{native_stream} {}
    void* GetHandle() { return stream; }

    cudaStream_t stream;
};
}

struct BetterStream {
    ort::Stream* s;
    operator cudaStream_t() { return s ? nullptr : static_cast<cudaStream_t>(s->GetHandle()); }
};

void func_take_native_stream(cudaStream_t s) {}
void func_take_ort_stream(ort::Stream* s) {}
void func_take_better_stream(BetterStream s) {}


int main() {
    cudaStream_t native_stream = nullptr;
    ort::Stream ort_stream_singleton{native_stream};
    auto ort_stream = &ort_stream_singleton;
    BetterStream better_stream{ort_stream};

    // before the ort Stream PR
    func_take_native_stream(native_stream);  

    // status quo
    func_take_ort_stream(ort_stream);  // good
    func_take_native_stream(static_cast<cudaStream_t>(ort_stream->GetHandle()));  // awkward
    kernel_take_native_stream<<<1,1,0,static_cast<cudaStream_t>(ort_stream->GetHandle())>>>();  // awkward

    // proposed
    func_take_better_stream(better_stream);  // good
    func_take_native_stream(better_stream);  // good
    kernel_take_native_stream<<<1,1,0,better_stream>>>();  // good
}

With this, I think to move to ort stream at whole repo will be as simple as a sed replace in the function signature.

@jslhcl @souptc

[jslhcl]

Could you write down the full example that involves rocm stream? Are you going to add

operator hipStream_t() { ... }

in BetterStream as well? Why not do the same trick in the Stream class?

[cloudhan]

Yeah, we shall do the same with hipstream. Because you use Stream as pointer, where conversion function does not work. You need to call ort->operator cudaStream_t() explicitly and it defeat our purpose.