[IR] Add type checks for atomic_cas (#7578)

Author: Jokeren

include/triton/Dialect/Triton/IR/TritonOps.td

@@ -383,8 +383,14 @@ def TT_AtomicRMWOp : TT_Op<"atomic_rmw", [
     }];
 }
 
-def TT_AtomicCASOp : TT_Op<"atomic_cas", [SameOperandsAndResultShape,
-                                          SameOperandsAndResultEncoding]> {
+def TT_AtomicCASOp : TT_Op<"atomic_cas", [
+  SameOperandsAndResultShape,
+  SameOperandsAndResultEncoding,
+  TypesMatchWith<"ptr type matches cmp type", "cmp", "ptr",
+                  "getPointerTypeSameShape($_self)">,
+  TypesMatchWith<"ptr type matches value type", "val", "ptr",
+                  "getPointerTypeSameShape($_self)">
+]> {
     let summary = "atomic cas";
 
     let description = [{

test/Triton/invalid.mlir

@@ -546,3 +546,21 @@ tt.func @unsplat_invalid(%arg0: tensor<128xf32>) {
   %0 = tt.unsplat %arg0 : tensor<128xf32>
   tt.return
 }
+
+// -----
+
+tt.func @atomic_cas_different_elem_types(%arg0: tensor<128x!tt.ptr<f32>>, %arg1: tensor<128xi32>) {
+  %cmp = arith.constant dense<0> : tensor<128xi32>
+  // expected-error @below {{'tt.atomic_cas' op failed to verify that ptr type matches cmp type}}
+  %0 = tt.atomic_cas relaxed, gpu, %arg0, %cmp, %arg1 : (tensor<128x!tt.ptr<f32>>, tensor<128xi32>, tensor<128xi32>) -> tensor<128xi32>
+  tt.return
+}
+
+// -----
+
+tt.func @atomic_cas_different_elem_types(%arg0: tensor<128x!tt.ptr<f32>>, %arg1: tensor<128xi32>) {
+  %cmp = arith.constant dense<0.0> : tensor<128xf32>
+  // expected-error @below {{'tt.atomic_cas' op failed to verify that ptr type matches value type}}
+  %0 = tt.atomic_cas relaxed, gpu, %arg0, %cmp, %arg1 : (tensor<128x!tt.ptr<f32>>, tensor<128xf32>, tensor<128xi32>) -> tensor<128xi32>
+  tt.return
+}