Skip to content

Fix unimplemented functions mapping and CUDA TF32 precision setup#35

Merged
janpfeifer merged 4 commits into
mainfrom
fix-notimplemented
Jun 1, 2026
Merged

Fix unimplemented functions mapping and CUDA TF32 precision setup#35
janpfeifer merged 4 commits into
mainfrom
fix-notimplemented

Conversation

@janpfeifer

Copy link
Copy Markdown
Contributor

This PR introduces two main enhancements/fixes:

  1. Unimplemented Functions Integration:

    • Embeds notimplemented.Function into xla.Function.
    • Adds a newFunction helper to set up standard XLA error handling, raising an informative message ("{op} not implemented for XLA backend") when an unimplemented operation is called on the function.
    • Refactors b.Main(), b.NewFunction(), and f.Closure() to use newFunction.
  2. CUDA TF32 Precision Handling:

    • Updates DotGeneral precision settings for CUDA:
      • Explicitly requests DotGeneralPrecisionDefault when DotGeneralUseTF32 is enabled and accumulation dtype is F32 (retaining TF32 performance).
      • Otherwise, falls back to DotGeneralPrecisionHighest for maximum accuracy.
    • Sets DotGeneralUseTF32 to enable by default on CUDA backends (isPluginType(pluginName, "cuda")).

@janpfeifer
janpfeifer merged commit 8d8e857 into main Jun 1, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant