2023-02-26 nightly release (a88bfc6)

pytorch · ezyang · Sep 11, 2022 · Sep 12, 2022 · Sep 13, 2022 · Sep 14, 2022
commit 3313e86c033331a19d9bbea8dde343ca6f8320be
@@ -92,6 +92,9 @@ fi
 
 if [[ "$BUILD_ENVIRONMENT" == *slow-gradcheck* ]]; then
   export PYTORCH_TEST_WITH_SLOW_GRADCHECK=1
+  # TODO: slow gradcheck tests run out of memory a lot recently, so setting this
+  # to run them sequentially with only one process to mitigate the issue
+  export PYTORCH_TEST_CUDA_MEM_LEAK_CHECK=1
 fi
 
 if [[ "$BUILD_ENVIRONMENT" == *cuda* || "$BUILD_ENVIRONMENT" == *rocm* ]]; then

@@ -43,8 +43,9 @@ jobs:
       docker-image-name: pytorch-linux-bionic-cuda11.7-cudnn8-py3-gcc7
       test-matrix: |
         { include: [
-          { config: "default", shard: 1, num_shards: 2, runner: "linux.4xlarge.nvidia.gpu" },
-          { config: "default", shard: 2, num_shards: 2, runner: "linux.4xlarge.nvidia.gpu" },
+          { config: "default", shard: 1, num_shards: 3, runner: "linux.4xlarge.nvidia.gpu" },
+          { config: "default", shard: 2, num_shards: 3, runner: "linux.4xlarge.nvidia.gpu" },
+          { config: "default", shard: 3, num_shards: 3, runner: "linux.4xlarge.nvidia.gpu" },
         ]}
 
   linux-bionic-cuda11_7-py3-gcc7-slow-gradcheck-test:

diff --git a/WORKSPACE b/WORKSPACE
@@ -260,6 +260,11 @@ local_repository(
     path = "third_party/fmt/support/bazel",
 )
 
+local_repository(
+    name = "unused_kineto_fmt_bazel",
+    path = "third_party/kineto/libkineto/third_party/fmt/support/bazel",
+)
+
 local_repository(
     name = "unused_kineto_dynolog_googletest",
     path = "third_party/kineto/libkineto/third_party/dynolog/third_party/googletest",

@@ -1,5 +1,7 @@
 //  Copyright © 2022 Apple Inc.
 
+#pragma once
+
 #include <ATen/ATen.h>
 #include <ATen/Tensor.h>
 #include <ATen/Utils.h>

@@ -0,0 +1,199 @@
+#include <ATen/native/mps/OperationUtils.h>
+#include <ATen/native/BinaryOps.h>
+
+namespace at::native {
+namespace mps {
+
+static const char* METAL_BINARY = R"BINARY_METAL(
+
+#include <metal_stdlib>
+using namespace metal;
+
+template<typename T>
+kernel void fmax(constant void     * input_        [[buffer(0)]],
+                  constant void     * other_        [[buffer(1)]],
+                  device   void     * out_          [[buffer(2)]],
+                  constant uint3    * offsets       [[buffer(3)]],
+                  uint tid [[thread_position_in_grid]]) {
+  device   T* out   = (device   T*)((device uint8_t*)out_ + offsets[tid].x);
+  constant T* input = (constant T*)((constant uint8_t*)input_ + offsets[tid].y);
+  constant T* other = (constant T*
B41A
)((constant uint8_t*)other_ + offsets[tid].z);
+
+  *out = fmax(*input, *other);
+}
+
+template<typename T>
+kernel void fmin(constant void     * input_        [[buffer(0)]],
+                  constant void     * other_        [[buffer(1)]],
+                  device   void     * out_          [[buffer(2)]],
+                  constant uint3    * offsets       [[buffer(3)]],
+                  uint tid [[thread_position_in_grid]]) {
+  device   T* out   = (device   T*)((device uint8_t*)out_ + offsets[tid].x);
+  constant T* input = (constant T*)((constant uint8_t*)input_ + offsets[tid].y);
+  constant T* other = (constant T*)((constant uint8_t*)other_ + offsets[tid].z);
+
+  *out = fmin(*input, *other);
+}
+
+#define REGISTER_FMAX_OP(DTYPE)                       \
+template                                               \
+[[host_name("fmax_" #DTYPE)]]                         \
+kernel void fmax<DTYPE>(                  \
+  constant void     * input_        [[buffer(0)]],     \
+  constant void     * other_        [[buffer(1)]],     \
+  device   void     * out_          [[buffer(2)]],     \
+  constant uint3    * offsets       [[buffer(3)]],     \
+  uint tid [[thread_position_in_grid]]);
+
+#define REGISTER_FMIN_OP(DTYPE)                       \
+template                                               \
+[[host_name("fmin_" #DTYPE)]]                         \
+kernel void fmin<DTYPE>(                  \
+  constant void     * input_        [[buffer(0)]],     \
+  constant void     * other_        [[buffer(1)]],     \
+  device   void     * out_          [[buffer(2)]],     \
+  constant uint3    * offsets       [[buffer(3)]],     \
+  uint tid [[thread_position_in_grid]]);
+
+REGISTER_FMAX_OP(float);
+REGISTER_FMAX_OP(half);
+REGISTER_FMIN_OP(float);
+REGISTER_FMIN_OP(half);
+
+)BINARY_METAL";
+
+using namespace mps;
+
+static id<MTLLibrary> compileBinaryOpsLibrary(id<MTLDevice> device) {
+  static id<MTLLibrary> binaryLibrary = nil;
+  if (binaryLibrary) {
+    return binaryLibrary;
+  }
+
+  NSError *error = nil;
+  MTLCompileOptions *options = [[MTLCompileOptions new] autorelease];
+  [options setLanguageVersion: MTLLanguageVersion2_3];
+  binaryLibrary  = [device newLibraryWithSource:[NSString stringWithCString: METAL_BINARY encoding:NSASCIIStringEncoding]
+                                       options:options
+                                         error:&error];
+  TORCH_CHECK(binaryLibrary, "Failed to create metal binary library, error: ", [[error description] UTF8String]);
+  return binaryLibrary;
+}
+
+static id<MTLComputePipelineState> binaryPipelineState(id<MTLDevice> device, const std::string& kernel) {
+  static std::unordered_map<std::string, id<MTLComputePipelineState>> psoCache;
+  id<MTLComputePipelineState> pso = psoCache[kernel];
+  if (pso) {
+    return pso;
+  }
+
+  NSError* error = nil;
+  id<MTLLibrary> binaryLib = compileBinaryOpsLibrary(device);
+  id<MTLFunction> binaryFunc = [binaryLib newFunctionWithName:[NSString stringWithUTF8String:kernel.c_str()]];
+  TORCH_CHECK(binaryFunc, "Failed to create function state object for: ", kernel);
+  pso = [device newComputePipelineStateWithFunction:binaryFunc error:&error];
+  TORCH_CHECK(pso, "Failed to created pipeline state object, error: ", [[error description] UTF8String]);
+
+  psoCache[kernel] = pso;
+  return pso;
+}
+
+void fmax_fmin_mps_impl(TensorIteratorBase& iter, const std::string max_min) {
+  TORCH_CHECK(iter.common_dtype() != at::kDouble, "float64 is not supported on MPS");
+
+  Tensor input = iter.input(0);
+  Tensor other = iter.input(1);
+  Tensor out = iter.output(0);
+  id<MTLBuffer> inputBuffer  = getMTLBufferStorage(input);
+  id<MTLBuffer> otherBuffer  = getMTLBufferStorage(other);
+  id<MTLBuffer> outputBuffer = getMTLBufferStorage(out);
+  id<MTLDevice> device = MPSDevice::getInstance()->device();
+  MPSStream* mpsStream = getCurrentMPSStream();
+  const uint32_t nDim = iter.ndim();
+  constexpr uint32_t nOffsets = 3;
+  const uint32_t numThreads = iter.numel();
+  dispatch_sync(mpsStream->queue(), ^(){
+    @autoreleasepool {
+      NSError* error = nil;
+      id<MTLCommandBuffer> commandBuffer = mpsStream->commandBuffer();
+      id<MTLComputeCommandEncoder> computeEncoder = [commandBuffer computeCommandEncoder];
+      MTLSize gridSize = MTLSizeMake(numThreads, 1, 1);
+      const IntArrayRef& iterShape = iter.shape();
+      std::vector<uint32_t> iterShapeData(iterShape.size());
+      std::vector<std::array<uint32_t, nOffsets>> strides(nDim);
+
+      for (const auto i: c10::irange(iterShape.size())) {
+        TORCH_CHECK(i <= UINT32_MAX);
+        iterShapeData[i] = (uint32_t)(iterShape[i]);
+      }
+
+      for (const auto i: c10::irange(nDim)) {
+        for (const auto offset: c10::irange(nOffsets)) {
+            strides[i][offset] = iter.strides(offset)[i];
+        }
+      }
+
+      id<MTLFunction> kernelDataOffsetsFunction = MPSDevice::getInstance()->metalIndexingFunction("kernel_index_offsets", nil);
+      id<MTLComputePipelineState> kernelDataOffsetsPSO = [[device newComputePipelineStateWithFunction: kernelDataOffsetsFunction
+                                                                                                error: &error] autorelease];
+      id<MTLBuffer> kernelDataOffsets = [[device newBufferWithLength: numThreads * sizeof(simd_uint3)
+                                                             options: 0] autorelease];
+      TORCH_CHECK(kernelDataOffsetsPSO, "Failed to created pipeline state object, error: ", [[error description] UTF8String]);
+      [computeEncoder setComputePipelineState:kernelDataOffsetsPSO];
+      [computeEncoder setBytes:strides.data() length:sizeof(uint32_t) * nDim * nOffsets atIndex:0];
+      [computeEncoder setBuffer:kernelDataOffsets offset:0 atIndex:1];
+      [computeEncoder setBytes:iterShapeData.data() length:sizeof(uint32_t) * iterShape.size() atIndex:2];
+      [computeEncoder setBytes:&nDim length:sizeof(uint32_t) atIndex:3];
+      [computeEncoder setBytes:&nOffsets length:sizeof(uint32_t) atIndex:4];
+
+      NSUInteger kernelOffsetsTGSize = kernelDataOffsetsPSO.maxTotalThreadsPerThreadgroup;
+      if (kernelOffsetsTGSize > numThreads)
+          kernelOffsetsTGSize = numThreads;
+
+      MTLSize kernelOffsetsThreadGroupSize = MTLSizeMake(kernelOffsetsTGSize, 1, 1);
+      [computeEncoder dispatchThreads: gridSize
+                threadsPerThreadgroup: kernelOffsetsThreadGroupSize];
+
+      const std::string kernel = "f" + max_min + "_" + scalarToMetalTypeString(out.scalar_type());
+      id<MTLComputePipelineState> fmaxfminPSO = binaryPipelineState(device, kernel);
+      [computeEncoder setComputePipelineState:fmaxfminPSO];
+      [computeEncoder setBuffer:inputBuffer  offset:input.storage_offset() * input.element_size() atIndex:0];
+      [computeEncoder setBuffer:otherBuffer  offset:other.storage_offset() * other.element_size() atIndex:1];
+      [computeEncoder setBuffer:outputBuffer offset:out.storage_offset() * out.element_size() atIndex:2];
+      [computeEncoder setBuffer:kernelDataOffsets offset:0 atIndex:3];
+
+      NSUInteger tgSize = fmaxfminPSO.maxTotalThreadsPerThreadgroup;
+      if (tgSize > numThreads) {
+          tgSize = numThreads;
+      }
+
+      MTLSize threadGroupSize = MTLSizeMake(tgSize, 1, 1);
+      [computeEncoder dispatchThreads: gridSize
+                threadsPerThreadgroup: threadGroupSize];
+
+      [computeEncoder endEncoding];
+      mpsStream->commit(true);
+    }
+  });
+}
+} // namespace mps
+
+void fmax_mps_kernel(TensorIteratorBase& iter) {
+    if (isFloatingType(iter.common_dtype())) {
+        mps::fmax_fmin_mps_impl(iter, "max");
+    } else {
+        at::maximum_out(const_cast<Tensor&>(iter.output()), iter.input(0), iter.input(1));
+    }
+}
+void fmin_mps_kernel(TensorIteratorBase& iter) {
+    if (isFloatingType(iter.common_dtype())) {
+        mps::fmax_fmin_mps_impl(iter, "min");
+    } else {
+        at::minimum_out(const_cast<Tensor&>(iter.output()), iter.input(0), iter.input(1));
+    }
+}
+
+REGISTER_DISPATCH(fmax_stub, &fmax_mps_kernel);
+REGISTER_DISPATCH(fmin_stub, &fmin_mps_kernel);
+
+} // namespace at::native
@@ -140,7 +140,9 @@ void reduction_out_mps(
   const std::string& func_name) {
 
   // issue 103641234, reduction ops does not have int64 support
-  TORCH_WARN_ONCE(input_t.scalar_type() != ScalarType::Long, "MPS: no support for int64 reduction ops, casting it to int32");
+  if (input_t.scalar_type() == ScalarType::Long) {
+    TORCH_WARN_ONCE("MPS: no support for int64 reduction ops, casting it to int32");
+  }
   IntArrayRef input_shape = input_t.sizes();
 
   if (opt_dim.has_value()) {
@@ -1266,7 +1268,9 @@ Tensor std_mps(
   (const Tensor& input_t,
    MPSReductionType reduction_type,
    const std::string& func_name) {
-  TORCH_WARN_ONCE(input_t.scalar_type() != ScalarType::Long, "MPS: no support for int64 min/max ops, casting it to int32");
+  if (input_t.scalar_type() == ScalarType::Long) {
+    TORCH_WARN_ONCE("MPS: no support for int64 min/max ops, casting it to int32");
+  }
 
   using CachedGraph = MPSUnaryCachedGraph;
 

@@ -233,7 +233,7 @@ Tensor repeat_interleave_mps(const Tensor& repeat_, c10::optional<int64_t> outpu
   if (repeat.scalar_type() == kLong) {
     // #103810551: `repeat_interleave_common` uses cumsum to calculate the final shape of output,
     // which currently doesn't support int64_t as input. Casting internally the indices to int32_t.
-    TORCH_WARN_ONCE(false, "MPS: no support for int64 repeats mask, casting it to int32");
+    TORCH_WARN_ONCE("MPS: no support for int64 repeats mask, casting it to int32");
     repeat = repeat.to(kInt);
   }
   AT_DISPATCH_INDEX_TYPES(repeat.scalar_type(), "repeat_interleave_mps", [&]() {
@@ -243,4 +243,4 @@ Tensor repeat_interleave_mps(const Tensor& repeat_, c10::optional<int64_t> outpu
   return output;
 }
 
-}  // namespace at::native
+}  // namespace at::native
@@ -35,7 +35,9 @@
     indices.copy_(cpu_indices);
     return;
   }
-  TORCH_WARN_ONCE(self.scalar_type() != ScalarType::Long, "MPS: no support for int64 min/max ops, casting it to int32");
+  if (self.scalar_type() == ScalarType::Long) {
+    TORCH_WARN_ONCE("MPS: no support for int64 min/max ops, casting it to int32");
+  }
 
   MPSStream* stream = getCurrentMPSStream();
   struct CachedGraph : public MPSCachedGraph {

diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -9301,7 +9301,7 @@
   structured_inherits: TensorIteratorBase
   device_check: NoCheck   # TensorIterator
   dispatch:
-    CPU, CUDA: fmin_out
+    CPU, CUDA, MPS: fmin_out
   tags: pointwise
 
 - func: max(Tensor self) -> Tensor
@@ -9323,7 +9323,7 @@
   structured_inherits: TensorIteratorBase
   device_check: NoCheck   # TensorIterator
   dispatch:
-    CPU, CUDA: fmax_out
+    CPU, CUDA, MPS: fmax_out
   tags: pointwise
 
 - func: maximum(Tensor self, Tensor other) -> Tensor

diff --git a/benchmarks/dynamo/common.py b/benchmarks/dynamo/common.py
@@ -1947,6 +1947,8 @@ def run(runner, args, original_dir=None):
             # TODO - Using train mode for timm_models. Move to train mode for HF and Torchbench as well.
             args.use_eval_mode = True
         inductor_config.fallback_random = True
+        torch.backends.cudnn.allow_tf32 = False
+        torch.backends.cudnn.benchmark = False
         torch.backends.cudnn.deterministic = True
 
         # Remove randomeness when torch manual seed is called

diff --git a/c10/core/impl/PyInterpreter.cpp b/c10/core/impl/PyInterpreter.cpp
@@ -97,9 +97,16 @@ struct NoopPyInterpreterVTable final : public PyInterpreterVTable {
   };
 };
 
+// Construct this in Global scope instead of within `disarm`
+// where it will be only initialized first time `disarm` is called.
+// This increases the likelihood `noop_vtable` lives longer than
+// any object that refers to it.
+
+// If `noop_vtable` goes out of scope first, other objects will have dangling
+// reference to it.
+static NoopPyInterpreterVTable noop_vtable;
+
 void PyInterpreter::disarm() noexcept {
-  // Intentionally leaked
-  static NoopPyInterpreterVTable noop_vtable;
   vtable_ = &noop_vtable;
 }