Surprisingly bad codegen fixed by `extern "C"`

code review
1

Only because I read The Rust Calling Convention We Deserve · mcyoung a while back I had a hunch that something really wrong is going on here:

pub type Reg = u8;

// `repr` does not matter here, but added for completeness
#[repr(C)]
pub struct Rs1Rs2Operands {
    pub rs1: Reg,
    pub rs2: Reg,
}

#[repr(u16, align(4))]
pub enum ContractInstruction {
    CAddi4spn {
        rs1: Reg,
        rs2: Reg,
        rd: Reg,
        nzuimm: u16,
    },
    CLw {
        rs1: Reg,
        rs2: Reg,
        rd: Reg,
        uimm: u8,
    },
}

impl ContractInstruction {
    #[unsafe(no_mangle)]
    pub fn get_rs1_rs2_operands(self) -> Rs1Rs2Operands {
        match self {
            Self::CAddi4spn { rs1, rs2, .. } => Rs1Rs2Operands { rs1, rs2 },
            Self::CLw { rs1, rs2, .. } => Rs1Rs2Operands { rs1, rs2 },
        }
    }

    #[unsafe(no_mangle)]
    pub extern "C" fn get_rs1_rs2_operands_c(self) -> Rs1Rs2Operands {
        match self {
            Self::CAddi4spn { rs1, rs2, .. } => Rs1Rs2Operands { rs1, rs2 },
            Self::CLw { rs1, rs2, .. } => Rs1Rs2Operands { rs1, rs2 },
        }
    }
}

I would have expected the return value to occupy a single register (after all it is just two u8s) and the two methods to generate exactly the same code. Surprisingly (and not in a good way) the code is not the same:

get_rs1_rs2_operands:
        mov     rax, rdi
        mov     edx, eax
        shr     edx, 24
        shr     eax, 16
        ret

get_rs1_rs2_operands_c:
        mov     rax, rdi
        shr     eax, 16
        ret

https://rust.godbolt.org/z/qbWTsb6v1

In a tight loop this is, obviously, very bad for performance to have twice as many instructions for no reason.

I'm sure there were discussions about this in the past, but did really nothing happen about it yet? I'm really not looking forward to slapping extern "C" everywhere just to get a more reasonable assembly.

2

In a tight loop this will be inlined and completely irrelevant.

Rust used to pack things into integers more often; it made things worse overall because the packing and unpacking code ended up causing more problems than it helped.

3

There is very little that is obvious when it comes to code generation. More instructions can be better.

4

Interesting. Whether this is inlined depends on a few factors and is not guaranteed, but the fact that the return value with just 2 bytes occupies two registers is concerning, especially since there is a chance they will be packed back into a single register later on (and depending on circumstances it might be highly desirable).

This is just a simplified case for demonstration purposes that I found while looking for something else.

Very true, I should not have said that.

5

It would help to hear more about why it's "concerning". Do you have a demonstration of it being a problem? What has you more concerned about the two ISA registers being used than about the shifting and masking instructions that would be needed to split a single u16 back into multiple u8s?

(Thanks to register renaming what it means to "occupy" a register is pretty complicated anyway.)

6

In the code I'm dealing with more generally GPRs are a very scarce resource. Relieving even one results in big performance improvements on x86-64. Force-inlining helps a lot, at least until LLVM stops generating fast code after a certain threshold of the function size (or at least it appears that way).

A somewhat reliable indicator so far is trying to compile with PGO. When PGO results in significantly worse performance, that often means non-PGO builds are struggling too.

In such cases return value from functions occupying as few GPRs as possible until the moment when the split is actually needed is very desirable. I may extract a pair of fields like that before a large match and those fields may only be used in a relatively few branches.

It is hard for me to reason about register spilling in such cases though. For me it is mostly a lot of trial and error until I find a specific permutation LLVM is happy with and finally get a performance improvement.

7

If it's critical that values are packed in this way, why don't you make your instruction type a single integer and extract fields when you need them?

9

I could, I could also write the whole thing in assembly. But I don't think I have to elaborate on ergonomic implications of such decisions.

Not sure why I have to mention it, but such response is unhelpful at best.

My context is deterministic cross-platform interpreter for untrusted code in blockchain environment. Writing a secure JIT is actually pretty hard. I don't need "a primitive JIT" that barely works with trusted code only. Not to mention the complexity of supporting scalable vector extensions in such JIT.

It is not helpful to be so dismissive.

10

Writing "the whole thing in assembly" doesn't follow from what I suggested.

11

It is a matter of picking abstractions. Having a simple struct with two fields is the most natural approach for me. The fields are not u8s in the real code BTW, they are #[repr(u8)] enums. I could make an opaque struct with u16 inside and have getters/setters, extract individual u8s and transmute them to an enum, but that is a lot of not very pleasant boilerplate.

12

Could you elaborate on this a bit? To be clear, this isn't an argument that you're wrong, necessarily, but this looks to me like like you're at least wrong about the reasons.

So far as I can tell from my own simple JITs and the craziness of "real" JITs, a basic JIT is actually much easier to make safe than an advanced one, and I can't see why it would necessarily "barely work" unless you're being particularly lazy: something like dynasm makes it pretty simple (including vector ops) to spit out snippets for each op in order and patch up control flow as you go, and even spilling every single intermediate to the stack to avoid needing fiddly register assignment can beat the pants off of most interpreters (which is the mortal enemy of branch prediction). At that point, it's dangerously close to trivial to JIT, certainly more so than optimizing interpreters at the point where you're trying to optimize register pressure.

The being cross platform thing is definitely a bigger concern, but generally people just mean Windows, Mac and Linux across x64 and ARM, which isn't terribly much additional work to support either.

To cover the bases: it generally is better to go shopping especially for security, but that's just as true for interpreters, and I'm sure you have some reason you're not already using something obvious like WASM. Equally, it's entirely possible there's something you're necessarily doing that's notably tricky when JITting, eg being able to swap out running functions.

13

The interpreter is modular and composable, base ISA and every extension is implemented separately and then composed together into a combined flat enum, combined flat execution function*, combined instruction decoding function, etc. You can then pick and choose standard and custom extensions you want and see how that goes.

It took quite a long time to implement vector instructions even with the help of LLMs and even then I only really got it to the decent state once vector certification tests were added to riscv/riscv-arch-test so I can actually test my implementation against it and fix remaining bugs.

Integer vector extension for RISC-V is hundreds of instructions with multiple variants of each with complex requirements. And it is ~10k lines (with tests) just for instruction decoding and another ~70k for execution logic. And this is for a relatively straightforward interpreter, now imagine JIT with a lot of platform-specific code.

I'd like to have JIT one day, but no way it is easy or quick to do. If you believe otherwise, I might want you to collaborate on the project together to have both interpreter and JIT using the same overall architecture, if you're interested :wink:

With interpreter it is much easier to ensure it is safe and deterministic, but I am considering using Cranelift one day for JIT too.

14

This is, obviously, wrong. It depends on how the results are used in further processing, even in a tight loop.

Instruction count is a bad proxy for performance and on a modern out-of-order CPU what matters is the dependency chain and the execution ports, not how many instructions you counted.

The problem is something else: a compiler that's optimal in all use cases is impossible in the general case. Optimal instruction scheduling and register allocation are NP-hard, and "fully optimal" is undecidable.

Maybe the compiler could be modified in such a way that your use case would compile to an optimal result. But your (single) use case is not the only one. There are literally millions, if not billions, of other source code files out there that the compiler should handle optimally.

So there are always trade-offs, and the goal of making a compiler optimal for the average case will probably miss some specific use cases.

15

This kind of sounds like you're taking arbitrary raw RISC-V as your untrusted input and are trying to sandbox it? Yeah, that's possible (eg. NaCl) but stupidly hard relative to a purpose-built bytecode, as I'm sure you're aware.

That said, in theory once you've done all that work to get it verified and are mapping to equivalent code for an interpreter, a simple JIT should just be for-each instruction, emit roughly whatever your compiler emits for the interpreter arm for that instruction (including bounds checking :face_savoring_food:). Certainly a lot of boring drudge work, but not complex and shouldn't introduce anything new to your security modeling (that is, if they could fool the JIT, they could probably fool your interpreter)

16

A primitive JIT would easily work with anything untrusted. It's easier to make it than to make fast interpreter. It's when you try to add optimizations that you have trouble.

Then you can transform the code to some bytecode that may then be easily interpreted. Even that would be faster than trying to convince LLVM to generate acceptable code for RISC-V compressed instructions parsing.

If you don't plan to go with efficient implementation that matches vtype to instruction and generates one, single, implementation that doesn't check vtype for every instruction, then making JIT for vector instructions is not too hard. It's when you want to avoid looking on vtype on each instruction it becomes hard.

You are applying pretty significant amount of force to push square peg into a round hole. One may tell how to apply even more force or show the square hole that you may use. The choice is, ultimately, yours, but you are asking to change the compiler to work with your, pretty unique, needs—that's a much larger change than what I am proposing.

Not “wrong”, not “pointless”, mind you. Just much larger.

Asking someone to do that while completely dismissing the possibility of changing your own code doesn't sound honest, to me.

RISC-V vectors are a bit of an odd beast there. To ensure that vector instructions wouldn't consume all the opcode spate RVV designers stuffed half of the opcode into a separate CPU register called vtype. Handling that in an interpreter is trivial: you just interpret instructions that sets vtype and then uses vtype separately — and that's it.

With JIT there are obvious temptation to, somehow, locate and merge two instructions: vsetv… (there are three ones, if I remember correctly) and actual vector instructions. And that immediately leads to problems with hostile code that jumps in the gap between these two instructions.

But if you simply leave that issue unresolved and just make JIT that faithfully reads vtype again and again—it would still be easier to make it secure and performant than endless fights with compiler about efficient register allocations in the interpreter.

17

If your JIT failed to track vtype correctly then you would presumably also fail to track it in the interpreter, right? It's not like stateful execution is unique to RISC-V here, pretty much all ISAs have something like that.

Do you mean you have more temptation to add tracking and emit the equivalent target op, than the equivalent interpreter match arm would if you're directly interpreting the RISC-V bytes? I guess that's true but that would also be the kind of thing optimizing the interpreter could get fooled by too, you can't do pretty much any significant optimizations even without JIT without tracking state like that (I'm thinking emitting intermediate bytecode to avoid needing to look up state, "unspill" stack variables, that sort of thing)