[MPS] Add native im2col #135706

malfet · 2024-09-11T16:36:01Z

It's called from torch.unfold and one of the few remaining vestiges in MPSFallback.mm

Strongly inspired by CUDA implementation from

pytorch/aten/src/ATen/native/cuda/im2col.cuh

Lines 40 to 61 in 09519eb

    
             int64_t idx = index / width_col; 
        
             int64_t h_out = idx % height_col; 
        
             int64_t channel_in = idx / height_col; 
        
             int64_t channel_out = channel_in * kernel_height * kernel_width; 
        
             int64_t h_in = h_out * stride_height - pad_height; 
        
             int64_t w_in = w_out * stride_width - pad_width; 
        
             dt* col = data_col + (channel_out * height_col + h_out) * width_col + w_out; 
        
             const dt* im = data_im + (channel_in * height + h_in) * width + w_in; 
        
             for (int64_t i = 0; i < kernel_height; ++i) { 
        
               for (int64_t j = 0; j < kernel_width; ++j) { 
        
                 int64_t h = h_in + i * dilation_height; 
        
                 int64_t w = w_in + j * dilation_width; 
        
                 *col = (h >= 0 && w >= 0 && h < height && w < width) 
        
                     ? im[i * dilation_height * width + j * dilation_width] 
        
                     : static_cast<dt>(0); 
        
                 col += height_col * width_col; 
        
               } 
        
             } 
        
           }

It's called from `torch.unfold` and one of the few remaining vestiges in MPSFallback

pytorch-bot · 2024-09-11T16:36:04Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/135706

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (2 Unrelated Failures)

As of commit 478c8e9 with merge base fc88ba2 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 5, 5, lf.linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
'test/profiler/test_cpp_thread.py::CppThreadTest::test_with_enable_profiler_in_child_thread'
pull / linux-focal-py3.9-clang10 / test (dynamo, 3, 3, lf.linux.2xlarge) (gh) (disabled by #132861, #132862, #134600)
test_transformers.py::TestSDPAPrivateUse1Only::test_fused_sdp_choice_privateuseone

This comment was automatically generated by Dr. CI and updates every 15 minutes.

github-actions · 2024-09-11T16:40:48Z

Attention! native_functions.yaml was changed

If you are adding a new function or defaulted argument to native_functions.yaml, you cannot use it from pre-existing Python frontend code until our FC window passes (two weeks). Split your PR into two PRs, one which adds the new C++ functionality, and one that makes use of it from Python, and land them two weeks apart. See https://github.com/pytorch/pytorch/wiki/PyTorch's-Python-Frontend-Backward-and-Forward-Compatibility-Policy#forwards-compatibility-fc for more info.

Caused by:

aten/src/ATen/native/native_functions.yaml

This reverts commit 5cf8956.

albanD · 2024-09-13T17:30:50Z

aten/src/ATen/native/mps/operations/Im2Col.mm

+    }
+  });
+  if (!batched_input) {
+    output.resize_({n_output_plane, output_length});


Suggested change

output.resize_({n_output_plane, output_length});

output.squeeze_(0);

same on the cuda code + moving it out of the AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND3 block

aten/src/ATen/native/mps/operations/Im2Col.mm

albanD · 2024-09-13T17:37:40Z

aten/src/ATen/native/mps/operations/Im2Col.mm

+INSTANTIATE_IM2COL(float2);
+INSTANTIATE_IM2COL(half);
+INSTANTIATE_IM2COL(half2);
+#if __METAL_VERSION__ >= 310


What's going to happen if I try to call this with bfloat on older metal? Do I get a nice error?

Yes, see

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

Lines 59 to 60 in 835e7bb

TORCH_CHECK_TYPE(is_macos_13_or_newer(MacOSVersion::MACOS_VER_14_0_PLUS),

"MPS bfloat16 type is supported on MacOS 14.0 or newer.");

Which will be triggered from here

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

Lines 215 to 217 in 835e7bb

case ScalarType::BFloat16:

checkSupportsBFloat16();

return "bfloat";

albanD

Sounds good !

malfet · 2024-09-14T06:07:46Z

@pytorchbot merge -f "Lint and MPS are green"

pytorchmergebot · 2024-09-14T06:09:21Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

After #135706 `getGatherScatterScalarType` returns exactly the same results as `scalarToMetalTypeString` , so delete the function and call `scalarToMetalTypeString`

After #135706 `getGatherScatterScalarType` returns exactly the same results as `scalarToMetalTypeString` , so delete the function and call `scalarToMetalTypeString` Pull Request resolved: #136295 Approved by: https://github.com/kit1980

It's called from `torch.unfold` and one of the few remaining vestiges in `MPSFallback.mm` Strongly inspired by CUDA implementation from https://github.com/pytorch/pytorch/blob/09519eb1959978af7f2f909acf768485772596ab/aten/src/ATen/native/cuda/im2col.cuh#L40-L61 Pull Request resolved: pytorch#135706 Approved by: https://github.com/albanD

After pytorch#135706 `getGatherScatterScalarType` returns exactly the same results as `scalarToMetalTypeString` , so delete the function and call `scalarToMetalTypeString` Pull Request resolved: pytorch#136295 Approved by: https://github.com/kit1980

malfet and others added 3 commits September 11, 2024 06:43

[MPS] Add native im2col

7983d6b

It's called from `torch.unfold` and one of the few remaining vestiges in MPSFallback

Will this work?

6016333

Lot more fixes

d4ea51a

malfet requested a review from kulinseth as a code owner September 11, 2024 16:36

pytorch-bot bot added ciflow/mps Run MPS tests (subset of trunk) release notes: mps Release notes category labels Sep 11, 2024

malfet added 2 commits September 11, 2024 09:49

And extend it to complex

577f1f1

Fix lint

acaec91

malfet added the topic: improvements topic category label Sep 11, 2024

malfet requested a review from albanD September 11, 2024 18:10

malfet and others added 2 commits September 11, 2024 11:12

And guard tracing

5cf8956

Revert "And guard tracing"

e032d3a

This reverts commit 5cf8956.

albanD reviewed Sep 13, 2024

View reviewed changes

Capture output

dfa4b8d

albanD approved these changes Sep 13, 2024

View reviewed changes

Use unsqueeze and squeeze

478c8e9

pytorchmergebot added the merging label Sep 14, 2024

pytorchmergebot added the Merged label Sep 14, 2024

pytorchmergebot closed this in 1a67e2b Sep 14, 2024

pytorchmergebot removed the merging label Sep 14, 2024

malfet added a commit that referenced this pull request Sep 18, 2024

[BE][MPS] Delete duplicated code in View.mm

4004267

After #135706 `getGatherScatterScalarType` returns exactly the same results as `scalarToMetalTypeString` , so delete the function and call `scalarToMetalTypeString`

malfet mentioned this pull request Sep 18, 2024

[BE][MPS] Delete duplicated code in View.mm #136295

Closed

github-actions bot deleted the malfet/mps-add-im2col branch October 14, 2024 06:25

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MPS] Add native im2col #135706

[MPS] Add native im2col #135706

Uh oh!

malfet commented Sep 11, 2024 •

edited

Loading

Uh oh!

pytorch-bot bot commented Sep 11, 2024 •

edited

Loading

Uh oh!

github-actions bot commented Sep 11, 2024

Uh oh!

albanD Sep 13, 2024

Uh oh!

albanD Sep 13, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

albanD Sep 13, 2024

Uh oh!

malfet Sep 13, 2024

Uh oh!

albanD left a comment

Uh oh!

malfet commented Sep 14, 2024

Uh oh!

pytorchmergebot commented Sep 14, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	int64_t idx = index / width_col;

	int64_t h_out = idx % height_col;
	int64_t channel_in = idx / height_col;
	int64_t channel_out = channel_in * kernel_height * kernel_width;
	int64_t h_in = h_out * stride_height - pad_height;
	int64_t w_in = w_out * stride_width - pad_width;

	dt* col = data_col + (channel_out * height_col + h_out) * width_col + w_out;
	const dt* im = data_im + (channel_in * height + h_in) * width + w_in;

	for (int64_t i = 0; i < kernel_height; ++i) {
	for (int64_t j = 0; j < kernel_width; ++j) {
	int64_t h = h_in + i * dilation_height;
	int64_t w = w_in + j * dilation_width;
	*col = (h >= 0 && w >= 0 && h < height && w < width)
	? im[i * dilation_height * width + j * dilation_width]
	: static_cast<dt>(0);
	col += height_col * width_col;
	}
	}
	}

	output.resize_({n_output_plane, output_length});
	output.squeeze_(0);

	TORCH_CHECK_TYPE(is_macos_13_or_newer(MacOSVersion::MACOS_VER_14_0_PLUS),
	"MPS bfloat16 type is supported on MacOS 14.0 or newer.");

	case ScalarType::BFloat16:
	checkSupportsBFloat16();
	return "bfloat";

[MPS] Add native im2col #135706

[MPS] Add native im2col #135706

Uh oh!

Conversation

malfet commented Sep 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/135706

✅ You can merge normally! (2 Unrelated Failures)

Uh oh!

github-actions bot commented Sep 11, 2024

Attention! native_functions.yaml was changed

Uh oh!

albanD Sep 13, 2024

Choose a reason for hiding this comment

Uh oh!

albanD Sep 13, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

albanD Sep 13, 2024

Choose a reason for hiding this comment

Uh oh!

malfet Sep 13, 2024

Choose a reason for hiding this comment

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

malfet commented Sep 14, 2024

Uh oh!

pytorchmergebot commented Sep 14, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

malfet commented Sep 11, 2024 •

edited

Loading

pytorch-bot bot commented Sep 11, 2024 •

edited

Loading