wasm-tools icon indicating copy to clipboard operation
wasm-tools copied to clipboard

Impossible to implement `VisitOperator::simd_visitor` generically

Open nagisa opened this issue 8 months ago • 5 comments

I have a snippet of code similar to this in my code base, after I made changes to upgrade to a more recent version of wasm-tools:

pub(crate) struct JoinVisitor<L, R>(pub(crate) L, pub(crate) R);

macro_rules! join_visit {
    ($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*))*) => {
        $(fn $visit(&mut self $($(,$arg: $argty)*)?) -> Self::Output {
            (self.0.$visit($($($arg.clone()),*)?), self.1.$visit($($($arg),*)?))
        })*
    }
}

impl<'a, L, R> wasmparser::VisitOperator<'a> for JoinVisitor<L, R>
where
    L: wasmparser::VisitOperator<'a>,
    R: wasmparser::VisitOperator<'a>,
{
    type Output = (L::Output, R::Output);
    wasmparser::for_each_visit_operator!(join_visit);
}

impl<'a, L, R> wasmparser::VisitSimdOperator<'a> for JoinVisitor<L, R>
where
    L: wasmparser::VisitSimdOperator<'a>,
    R: wasmparser::VisitSimdOperator<'a>,
{
    wasmparser::for_each_visit_simd_operator!(join_visit);
}

But currently this seems to still fail to visit SIMD opcodes if I'm calling JoinVisitor::visit_operator. From the looks of it that's because there is a new function I need to implement VisitOperator::simd_visitor.

However I cannot implement this method in a generic-preserving way. If I do the naive thing

impl<'a, L, R> VisitOperator<'a> for JoinVisitor<L, R>
where
    L: VisitOperator<'a> + VisitSimdOperator<'a>,
    R: VisitOperator<'a> + VisitSimdOperator<'a>,
{
    fn simd_visitor(&mut self) -> Option<...> { Some(self) }
}

then I will have this code build, but the JoinVisitor: VisitOperator will now require an implementation of VisitSimdOperator. As far as I can tell there is no way to construct this such a way that simd_visitor returns Some only if L + R: VisitSimdOperator. An obvious

    fn simd_visitor(
        &mut self,
    ) -> Option<&mut dyn wasmparser::VisitSimdOperator<'a, Output = Self::Output>>
    where
        L: wasmparser::VisitSimdOperator<'a>,
        R: wasmparser::VisitSimdOperator<'a>,
    {
        Some(self)
    }

won't work because this method has stricter requirements than the trait being impld, nor will this:

    fn simd_visitor(
        &mut self,
    ) -> Option<&mut dyn wasmparser::VisitSimdOperator<'a, Output = Self::Output>> {
        Some(JoinVisitor(self.0.simd_visitor()?, self.1.simd_visitor()?))
    }

due to the signature requiring &mut dyn VisitSimdOperator to be returned.

Besides this particular snag, simd_visitor is also a plain footgun in that implementing a VisitSimdOperator is not all that's required on each of the implementors to support visiting SIMD operators, and the failure mode is a runtime error.


In the world where we know we'll only have VisitOperator and VisitSimdOperator this could be addressed by adding visit_operator to VisitSimdOperator as well (which would delegate to its supertrait's visit_operator if it isn't a SIMD opcode,) but that won't work out nicely if future delivers another Visit*Operator that's unrelated to VisitSimdOperator.

The only other way to address this that I can see is to add a new associated type to VisitOperator along the lines of type SimdVisitor that simd_visitor would return. Defaulting to Infallible or something.

nagisa avatar May 07 '25 10:05 nagisa

I attempted prototying adding an associated SimdVisitor type to VisitOperator but that readily wants to make VisitOperator not-object-safe :/ I'm out of ideas. I'll just make JoinVisitor be overzealous with generic requirements for the time being.

nagisa avatar May 07 '25 15:05 nagisa

I'm personally all for improving the trait design here. Object-safety isn't really a hard requirement if the simd_visitor method returns the associated type, so that seems like an ok tradeoff. I'm not sure how that would solve your use case though of L or R conditionally implementing VisitSimdOperator?

Otherwise though the split you mention also seems reasonable to me where you visit through the simd operator trait, no the general operator trait.

alexcrichton avatar May 07 '25 20:05 alexcrichton

Otherwise though the split you mention also seems reasonable to me where you visit through the simd operator trait, no the general operator trait.

I attempted prototyping this as well. I've a branch that passes tests but I don't think I want to push it to land as the downsides of this approach well outweigh the benefits (see the commit message for some thoughts there.)

I'm not sure how that would solve your use case though of L or R conditionally implementing VisitSimdOperator?

I'd set SimdVisitor = JoinVisitor<L::SimdVisitor, R::SimdVisitor> which happens to implement VisitSimdOperator if both L & R do so and use one of the implementations from OP above:

    fn simd_visitor(
        self,
    ) -> Option<Self::SimdVisitor> {
        Some(JoinVisitor(self.0.simd_visitor()?, self.1.simd_visitor()?))
    }

But as I said before this needing to return something by value means Self has to be taken by value too, and that has unfortunate implications. I don't think we can get rid of object safety for this. I'm pretty sure there are at least a some places that use dyn VisitOperator -- including a few in my own code that would be pretty difficult to deobjectify.


I looked at the original PR that introduced the split and in my eyes the improvements to compile time in particular aren't that significant percentage-wise. Code size-wise, though, yes. I wonder what the outcome would be if we merged the Visit*Operator traits back together, but left everything else gated… That would simplify the users' life a bunch while likely keeping most of the compile time/artifact size benefits…

nagisa avatar May 08 '25 11:05 nagisa

The main reason for the split traits was for cargo-feature compatibility. Code built with a simd-less wasmparser still needs to work when the feature is enabled. I'm not opposed to having everything in one trait myself, but that's the needle to thread at least.

Personally I'm not overly sold on the simd-less version of a wasmparser if the maintenance is pretty high or has downsides like this. I don't think it's necessarily unreasonable to just unconditionally include things and tackle binary size via other vectors.

cc @Robbepop on this regardless though as you're likely interested

alexcrichton avatar May 08 '25 15:05 alexcrichton

In Wasmi we achieved big wins in compile times from having optional SIMD support in wasmparser. Though it has to be said that Wasmi itself has optional SIMD and thus can take full advantage of it. This results in ~35-40% compile time reduction which is quite significant.

One alternative solution that could work similarly with other benefits is to implement trait-based WasmFeatures, thus allowing to make wasmparser aware of always enabled or disabled WasmFeatures at compile time. It might be worth a try to see where we land once this is implemented.

Also, maybe once Rust's parallel frontend has been landed compile time weight from having to implement all of wasmparser's VisitOperator trait could be lifted.

Another design idea is to provide default method implementations for wasmparser::VisitOperator to for example default implement all SIMD related operators to forward to a single new visit_simd_operator(op: SimdOp) method (and same could be applied elsewhere). The big downside is that there is no longer a compile-time check that a user implements all of the necessary trait methods.

Robbepop avatar May 08 '25 20:05 Robbepop