[UPCP][release/2.4] TunableOp fix for batched MM with views.

aten/src/ATen/cuda/tunable/GemmCommon.h

@@ -177,19 +177,19 @@ struct GemmStridedBatchedParams : OpParams {
   }
 
   size_t GetSizeA() const {
-    size_t size_stride = std::min(lda, stride_a) * ((transa == 'n' || transa == 'N') ? k : m) * batch;
+    size_t size_stride = stride_a * batch;
     size_t size_dense = m * k * batch;
     return sizeof(T) * (size_stride > size_dense ? size_stride : size_dense);
   }
 
   size_t GetSizeB() const {
-    size_t size_stride = std::min(ldb, stride_b) * ((transb == 'n' || transb == 'N') ? n : k) * batch;
+    size_t size_stride = stride_b * batch;
     size_t size_dense = k * n * batch;
     return sizeof(T) * (size_stride > size_dense ? size_stride : size_dense);
   }
 
   size_t GetSizeC() const {
-    size_t size_stride = std::min(ldc, stride_c) * n * batch;
+    size_t size_stride = stride_c * batch;
     size_t size_dense = m * n * batch;
     return sizeof(T) * (size_stride > size_dense ? size_stride : size_dense);
   }

test/test_linalg.py

@@ -4577,6 +4577,26 @@ def test_bmm_tunableop_rocm(self, device, dtype):
         i2 = torch.randn((M, B, K), device=device, dtype=dtype)
         i2 = torch.permute(i2, (1, 2, 0))
         out = torch.bmm(i1, i2)
+        # case 4
+        input_tensor = torch.rand((1920, 1, 100), device=device, dtype=dtype)
+        input_tensor = torch.as_strided(
+            input_tensor, size=(1920, 1, 100), stride=(100, 100, 1)
+        )
+        batch1_tensor = torch.rand((1920, 256, 512), device=device, dtype=dtype)
+        batch1_tensor = torch.as_strided(
+            batch1_tensor, size=(1920, 256, 512), stride=(512, 983040, 1)
+        )
+        batch2_tensor = torch.rand((1920, 512, 100), device=device, dtype=dtype)
+        batch2_tensor = torch.as_strided(
+            batch2_tensor, size=(1920, 512, 100), stride=(51200, 100, 1)
+        )
+        out = torch.baddbmm(input_tensor, batch1_tensor, batch2_tensor)
+        # case 5
+        q = torch.randn([16, 16, 1024, 64], device=device, dtype=dtype)
+        k = torch.randn([16, 16, 1024, 64], device=device, dtype=dtype)
+        q_chunks = q.split(512, dim=-2)
+        k_chunks = k.split(64, dim=-2)
+        C = torch.matmul(q_chunks[0], k_chunks[0])
         # clean up, remove any file that was generated
         try:
             import os