[MPS] Fast math env var (pytorch#129007)

Allow users to decide whether they want to have fast math enabled via env var
Pull Request resolved: pytorch#129007
Approved by: https://github.com/malfet
ghstack dependencies: pytorch#129006, pytorch#129008

aten/src/ATen/native/mps/OperationUtils.mm

@@ -655,13 +655,14 @@ void executeMPSAllocatorCallback(void* ptr, EventType event) override {}
 }
 
 id<MTLLibrary> MetalShaderLibrary::compileLibrary(const std::string& src) {
+  static const char* fast_math = std::getenv("PYTORCH_MPS_FAST_MATH");
   NSError* error = nil;
   MTLCompileOptions* options = compile_options;
   if (!options) {
     options = [[MTLCompileOptions new] autorelease];
     [options setLanguageVersion:is_macos_13_or_newer(MacOSVersion::MACOS_VER_14_0_PLUS) ? MTLLanguageVersion3_1
                                                                                         : MTLLanguageVersion2_3];
-    [options setFastMathEnabled:NO];
+    [options setFastMathEnabled:(!fast_math || std::stoi(fast_math) == 0) ? NO : YES];
   }
 
   const auto str = [NSString stringWithCString:src.c_str() encoding:NSASCIIStringEncoding];

docs/source/mps_environment_variables.rst

@@ -16,6 +16,8 @@ MPS Environment Variables
     - High watermark ratio for MPS allocator. By default, it is set to 1.7.
   * - ``PYTORCH_MPS_LOW_WATERMARK_RATIO``
     - Low watermark ratio for MPS allocator. By default, it is set to 1.4 if the memory is unified and set to 1.0 if the memory is discrete.
+  * - ``PYTORCH_MPS_FAST_MATH``
+    - If set to ``1``, enable fast math for MPS metal kernels. See section 1.6.3 in https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf for precision implications.
   * - ``PYTORCH_MPS_PREFER_METAL``
     - If set to ``1``, force using metal kernels instead of using MPS Graph APIs. For now this is only used for matmul op.
   * - ``PYTORCH_ENABLE_MPS_FALLBACK``