[TritonArithToLinalg] Lower batched matmul from tt.dot to linalg.batch_matmul (#35)

**Problem:**
Triton uses the same tt.dot() operation for both matmul and batched
matmul: for 3D operands, the first dimension is a batch dimension and
tt.dot performs a batched matrix product. The current lowering did not
account for this and always emitted linalg.matmul, which only supports
2D operands.

Running triton-shared-opt --triton-arith-to-linalg on a 3D tt.dot
therefore failed with:
```
  error: 'linalg.matmul' op expected operand #0 rank (3) to match the
  result rank of indexing_map (2)
```

**Solution:**
`MatmulConverter` now inspects the rank of the result tensor and emits
`linalg.batch_matmul` for rank-3 operands (and linalg.matmul for rank-2,
as before). Rank > 3 is not supported and now fails the pattern instead
of producing an invalid op.

**Testing:**
Adds a lit test (batched_matmul.mlir) covering the 3D case.

Author: mitekoth

include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp

@@ -1099,9 +1099,19 @@ struct MatmulConverter : public OpConversionPattern<triton::DotOp> {
                                          ValueRange{init})
                       .result();
 
-    auto res = linalg::MatmulOp::create(rewriter, loc, ValueRange{opa, opb},
-                                        ValueRange{zeroes})
-                   .getResult(0);
+    auto rank = dstType.getRank();
+    Value res;
+    if (rank > 3) {
+      return failure();
+    } else if (rank == 3) {
+      res = linalg::BatchMatmulOp::create(rewriter, loc, ValueRange{opa, opb},
+                                          ValueRange{zeroes})
+                .getResult(0);
+    } else {
+      res = linalg::MatmulOp::create(rewriter, loc, ValueRange{opa, opb},
+                                     ValueRange{zeroes})
+                .getResult(0);
+    }
 
     if (!skipC) {
       if (integers) {

test/Conversion/TritonArithToLinalg/batched_matmul.mlir

@@ -0,0 +1,28 @@
+// RUN: triton-shared-opt --triton-arith-to-linalg %s | FileCheck %s
+
+module {
+  tt.func @kernel(
+    %arg0: tensor<32x64x128x!tt.ptr<f16>>, %arg1: tensor<32x128x256x!tt.ptr<f16>>, %arg2: tensor<32x64x256x!tt.ptr<f32>>, %arg3: tensor<32x64x256x!tt.ptr<f32>>
+  )
+  {
+    %0 = tt.load %arg0: tensor<32x64x128x!tt.ptr<f16>>
+    %1 = tt.load %arg1: tensor<32x128x256x!tt.ptr<f16>>
+    %2 = tt.load %arg2: tensor<32x64x256x!tt.ptr<f32>>
+    %3 = tt.dot %0, %1, %2 : tensor<32x64x128xf16> * tensor<32x128x256xf16> -> tensor<32x64x256xf32>
+    tt.store %arg3, %3 : tensor<32x64x256x!tt.ptr<f32>>
+    tt.return
+  }
+}
+
+// CHECK-LABEL:   func.func @kernel(
+// CHECK-SAME:      %[[ARG0:.*]]: tensor<32x64x128x!tt.ptr<f16>>, %[[ARG1:.*]]: tensor<32x128x256x!tt.ptr<f16>>, %[[ARG2:.*]]: tensor<32x64x256x!tt.ptr<f32>>, %[[ARG3:.*]]: tensor<32x64x256x!tt.ptr<f32>>,
+// CHECK:           %[[LOAD_A:.*]] = tt.load %[[ARG0]] : tensor<32x64x128x!tt.ptr<f16>>
+// CHECK:           %[[LOAD_B:.*]] = tt.load %[[ARG1]] : tensor<32x128x256x!tt.ptr<f16>>
+// CHECK:           %[[LOAD_C:.*]] = tt.load %[[ARG2]] : tensor<32x64x256x!tt.ptr<f32>>
+// CHECK:           %[[EMPTY:.*]] = tensor.empty() : tensor<32x64x256xf32>
+// CHECK:           %[[FILL:.*]] = linalg.fill {{.*}} outs(%[[EMPTY]] : tensor<32x64x256xf32>) -> tensor<32x64x256xf32>
+// CHECK:           %[[BATCH_MATMUL:.*]] = linalg.batch_matmul ins(%[[LOAD_A]], %[[LOAD_B]] : tensor<32x64x128xf16>, tensor<32x128x256xf16>) outs(%[[FILL]] : tensor<32x64x256xf32>) -> tensor<32x64x256xf32>
+// CHECK:           %[[RESULT:.*]] = linalg.generic
+// CHECK-SAME:      ins(%[[LOAD_C]], %[[BATCH_MATMUL]] : tensor<32x64x256xf32>, tensor<32x64x256xf32>)
+// CHECK:           tt.store %[[ARG3]], %[[RESULT]] : tensor<32x64x256x!tt.ptr<f32>>
+// CHECK:           return