Company Description
Apex Compute is at the forefront of revolutionizing AI compute. Our mission is to redefine AI compute architecture by designing efficient, innovative, and cutting-edge hardware and software solutions. We are a passionate team of engineers pushing the boundaries of what's possible in AI hardware, and we offer company stock options to ensure our team shares in our long-term success.
Role Description
We are seeking a highly motivated Compiler Engineer to join our team on a full-time basis. This role is based in Los Altos, CA, with flexibility for a hybrid work arrangement. As a Compiler Engineer, you will play a pivotal role in developing the software stack for our next-generation AI compute hardware. Our custom architecture implements many high-level operations directly in hardware, and this role focuses on efficiently mapping modern machine learning workloads onto these capabilities. You will work closely with hardware and systems engineers to design compiler infrastructure that lowers machine learning workloads into MLIR dialects or similar intermediate representations, enabling efficient execution on our architecture.
Responsibilities
Compiler Development
- Design and implement compiler infrastructure and optimization passes that map machine learning workloads onto architecture-specific primitives.
- Develop lowering pipelines using MLIR, LLVM, or similar compiler frameworks.
Framework & System Integration
- Integrate the compiler stack with machine learning frameworks such as PyTorch or similar graph-based ML frameworks.
- Collaborate with hardware and software teams to enable efficient execution of ML models on the accelerator.
Operator Mapping & Optimization
- Map high-level ML operations such as matrix multiplication, normalization, attention primitives, and tensor operations to hardware-supported primitives.
- Develop strategies for operator fusion, tiling, scheduling, and memory movement optimization.
Performance Analysis & Optimization
- Analyze compute graphs and optimize execution by improving operator scheduling, memory locality, and data movement.
- Work closely with hardware architects to co-design compiler optimizations aligned with the accelerator architecture.
Documentation & Collaboration
- Develop documentation, testing infrastructure, and validation tools for compiler components.
- Participate in cross-functional team discussions to improve compiler performance and system architecture.
Qualifications
Educational Background
- Bachelor’s or Master’s degree in Computer Science, Computer Engineering, or a related field.
Compiler Expertise
- Proven experience with compiler frameworks such as MLIR, LLVM, TVM, XLA, or similar systems.
- Strong understanding of intermediate representations, compiler passes, and optimization pipelines.
Technical Skills
- Strong programming skills in C++ and/or Python.
- Experience working with tensor computation graphs and machine learning workloads.
Analytical & Problem-Solving
- Strong analytical skills with the ability to troubleshoot and optimize complex software systems.
- Excellent communication skills and ability to collaborate effectively in a fast-paced engineering environment.
Desirable Experience
- Experience building compilers for AI accelerators, GPUs, or specialized hardware.
- Familiarity with transformer models and attention kernels.
- Experience with hardware-software co-design or performance-critical systems.
Why Join Us?
You’ll have the opportunity to work on revolutionary AI hardware alongside a talented, passionate, and ambitious team. This role provides a chance to grow your technical skills, gain valuable hands-on experience, and contribute to groundbreaking innovations in AI compute hardware. Additionally, as part of our team, you’ll be eligible for company stock options, allowing you to share in our long-term success.