[MPS][BE] Use conveinence methods to set args (#145736)

malfet · pytorchmergebot · commit 30dea8429d97 · 2025-01-27T17:42:01.000Z
It's better to call `mtl_setArgs` rather than set arguments one by one with the risk of making a typo Also, all interactions with MTLCommandBuffer must be serialized, which is commonly done using dispatch queues Pull Request resolved: #145736 Approved by: https://github.com/Skylion007
diff --git a/aten/src/ATen/native/mps/operations/LinearAlgebra.mm b/aten/src/ATen/native/mps/operations/LinearAlgebra.mm
@@ -815,21 +815,16 @@ static void linalg_lu_factor_out_mps_impl(const Tensor& A, bool pivot, Tensor& L
   int64_t numBlocks = (N + NB - 1) / NB;
 
   Tensor success = at::empty({B}, input.options().dtype(kInt)).fill_(1);
-  id<MTLBuffer> successBuffer = getMTLBufferStorage(success);
 
   MTLSize threadGroupSize = MTLSizeMake(256, 1, 1);
-  id<MTLBuffer> outBuffer = getMTLBufferStorage(out);
-  id<MTLComputeCommandEncoder> computeEncoder = stream->commandEncoder();
-  [computeEncoder setBuffer:outBuffer offset:0 atIndex:0];
-  [computeEncoder setBytes:&N length:sizeof(int64_t) atIndex:2];
-  [computeEncoder setBytes:&NB length:sizeof(int64_t) atIndex:3];
 
   @autoreleasepool {
     dispatch_sync_with_rethrow(stream->queue(), ^() {
+      auto computeEncoder = stream->commandEncoder();
+      mtl_setArgs(computeEncoder, out, success, N, NB);
       for (int64_t k = 0; k < numBlocks; k++) {
         [computeEncoder setComputePipelineState:factorDiagonalPSO];
-        [computeEncoder setBuffer:successBuffer offset:0 atIndex:1];
-        [computeEncoder setBytes:&k length:sizeof(int64_t) atIndex:4];
+        mtl_setBytes(computeEncoder, k, 4);
         MTLSize gridSize = MTLSizeMake(B, 1, 1);
         [computeEncoder dispatchThreadgroups:gridSize threadsPerThreadgroup:threadGroupSize];
 
