Senior Software Engineer, Tensor Processing Units Compiler

Full-time - Mid Level

Sunnyvale, CA

10,001+ employees

Web Search Portals, Libraries, Archives, and Other Information Services

About the position

The Senior Software Engineer for Tensor Processing Units (TPUs) Compiler at Google will be responsible for developing and optimizing compilers that enable TPUs to accelerate machine learning and scientific computing workloads. This role involves working on the MLIR/LLVM based TPU compiler, supporting new workloads, optimizing for new models, and collaborating with hardware designers to co-design future processors. The position requires versatility and leadership qualities, as engineers will tackle various challenges across the full-stack in a fast-paced environment.

Responsibilities

• Contribute to a compiler for a novel processor designed to accelerate machine learning workloads.
• Compile high-performance implementations of operations at a distributed scale.
• Work closely with users of TPUs to improve performance/efficiency and hardware designers to co-design future processors.
• Investigate high-level representations to effectively program large-scale, distributed, and heterogeneous systems.

Requirements

• Bachelor's degree or equivalent practical experience.
• 5 years of experience testing, maintaining, or launching software products, and 1 year of experience with software design and architecture.
• 2 years of experience working with CUDA C++ application development and 1 year of experience with Native Code, Just-In-Time (JIT), Cross, Source-to-Source or any other type of compilers.
• 2 years of experience with data structures or algorithms, with experience with machine learning algorithms and tools (e.g., TensorFlow), artificial intelligence, deep learning, or natural language processing.

Nice-to-haves

• Master's degree or PhD in Computer Science or related technical fields.
• Experience with performance, large-scale systems data analysis, visualization tools, or debugging.
• Experience with debugging correctness and performance issues at all levels of the stack.
• Experience with optimizations in mid-level and low-level architecture.
• Experience with hardware/software co-design.
• Experience in GPU integrating low-level CUDA work into higher-level frameworks (e.g., TF, JAX, PyTorch).

Benefits

• Health insurance
• Dental insurance
• Vision insurance
• 401(k) plan
• Paid holidays
• Paid time off
• Employee stock purchase plan
• Tuition reimbursement
• Professional development opportunities
• Flexible scheduling
• Wellness programs