Compiler Engineer - Deep Learning

Key Skills

Aartificial intelligence, C Programming, C++, CUDA kernels, CUDA/GPU programming, Deep Learning, JAX framework, OpenCL, PyTorch, TensorFlow

Job Description

We are looking for Deep Learning Compiler Engineers. NVIDIA is hiring software engineers for its Deep Learning Compiler team. Academic and commercial groups around the world are using GPUs to power a revolution in deep learning, enabling breakthroughs in problems from image classification to speech recognition to natural language processing and artificial intelligence. Join the team which is building software which will be used by the entire deep learning community.

As a member of the Deep Learning Compiler Team, you will be responsible for developing compiler optimization algorithms for deep learning networks. You will be driving inference and training performance of JAX framework and XLA and OpenXLA compilers on NVIDIA GPUs at scale. You’ll collaborate with our partners in deep learning framework teams and our hardware architecture teams to accelerate the next generation of deep learning software.

What you'll be doing:

• Crafting and implementing compiler optimization techniques for deep learning network graphs

• Designing novel graph partitioning and tensor sharding techniques for distributed training and inference

• Performance tuning and analysis

• Code-generation for NVIDIA GPU backends using open-source compilers such as MLIR, LLVM and OpenAI Triton.

• Defining APIs in JAX and related libraries and other general software engineering work

What we need to see:

• Bachelors, Masters or Ph.D. in Computer Science, Computer Engineering, related field (or equivalent experience)

• 2+ years of relevant work or research experience in performance analysis and compiler optimizations.

• Ability to work independently, define project goals and scope, and lead your own development effort adopting clean software engineering and testing practices.

• Excellent C/C++ programming and software design skills, including debugging, performance analysis, and test design.

• Strong foundation in CPU and/or GPU architecture. Knowledge of high-performance computing and distributed programming. CUDA or OpenCL programming experience is desired but not required.

• Experience with the following technologies is a huge plus: XLA, TVM, MLIR, LLVM, OpenAI Triton, deep learning models and algorithms, and deep learning framework design.

• Strong interpersonal skills are required along with the ability to work in a dynamic product-oriented team. A history of mentoring junior engineers and interns is a bonus.

Ways to stand out from the crowd:

• Worked on a deep learning framework such as JAX, Pytorch or Tensorflow.

• Experience with CUDA or with GPUs

• Proficient with open-source compilers such as LLVM and MLIR.

With highly competitive salaries and a comprehensive benefits package, NVIDIA is widely considered to be one of the technology industry's most desirable employers. We have some of the most brilliant and hardworking people in the world working with us and our product lines are growing fast in some of the hottest state of the art fields such as Virtual Reality, Artificial Intelligence, Deep Learning and Autonomous Vehicles.

The base salary range is 116,000 USD - 224,250 USD. Your base salary will be determined based on your location, experience, and the pay of employees in similar positions.

You will also be eligible for equity and benefits. NVIDIA accepts applications on an ongoing basis.

NVIDIA is committed to fostering a diverse work environment and proud to be an equal opportunity employer. As we highly value diversity in our current and future employees, we do not discriminate (including in our hiring and promotion practices) on the basis of race, religion, color, national origin, gender, gender expression, sexual orientation, age, marital status, veteran status, disability status or any other characteristic protected by law.