spirv/nir: Add more stage asserts