Revert "[BACKEND] Disable vectorization for atomic_cas on all backends (#7711)"

This reverts commit a372a80.

Author: whitneywhtsang

python/test/unit/language/test_core.py

@@ -1969,33 +1969,24 @@ def serialized_add(data, Lock, SEM: tl.constexpr):
 
 
 @pytest.mark.interpreter
-@pytest.mark.parametrize("sem", [None, "acquire", "release", "acq_rel", "relaxed"])
+@pytest.mark.parametrize("sem", [None, 'acquire', 'release', 'acq_rel', 'relaxed'])
 @pytest.mark.parametrize("num_ctas", num_ctas_list)
-@pytest.mark.parametrize("size", [4, 128, 512])
-@pytest.mark.parametrize("dtype_str", ['bfloat16', 'float16', 'float32', 'uint64', 'int64', 'float64'])
-def test_tensor_atomic_cas(sem, size, dtype_str, num_ctas, device):
-    check_type_supported(dtype_str, device)
-    if "float" in dtype_str and is_hip():
-        pytest.skip("HIP does not support atomic cas with float types")
+def test_tensor_atomic_cas(sem, num_ctas, device):
 
     @triton.jit
-    def change_value(X, BLOCK_SIZE: tl.constexpr, sem: tl.constexpr, dtype: tl.constexpr):
+    def change_value(X, BLOCK_SIZE: tl.constexpr, sem: tl.constexpr):
         pid = tl.program_id(axis=0)
         block_start = pid * BLOCK_SIZE
         offsets = block_start + tl.arange(0, BLOCK_SIZE)
-        t1 = tl.full((BLOCK_SIZE, ), 0, dtype=dtype)
-        t2 = tl.full((BLOCK_SIZE, ), 2, dtype=dtype)
+        t1 = tl.full((BLOCK_SIZE, ), 0, dtype=tl.int64)
+        t2 = tl.full((BLOCK_SIZE, ), 2, dtype=tl.int64)
         tl.atomic_cas(X + offsets, t1, t2, sem=sem)
 
-    torch_dtype = getattr(torch, dtype_str)
-    X = torch.zeros((size, ), device=device, dtype=torch_dtype)
-    X[1::2] = 1
-    Y = X.clone()
-    Y[0::2] = 2
+    X = torch.tensor([0, 1, 0, 1, 0, 1, 0, 1], device=device, dtype=torch.int64)
+    Y = torch.tensor([2, 1, 2, 1, 2, 1, 2, 1], device=device, dtype=torch.int64)
 
-    tl_dtype = getattr(tl, dtype_str)
-    change_value[(2, )](X, BLOCK_SIZE=size // 2, sem=sem, dtype=tl_dtype)
-    assert torch.equal(X, Y)
+    change_value[(2, )](X, 4, sem)
+    assert (torch.equal(X, Y))
 
 
 @pytest.mark.interpreter

third_party/amd/lib/TritonAMDGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -1361,16 +1361,33 @@ struct AtomicCASOpConversion
                  : valueTy;
     auto valueElemNBits = valueElemTy.getIntOrFloatBitWidth();
     auto elemsPerThread = getTotalElemsPerThread(op.getVal().getType());
+    // vec = 1 for scalar
+    auto vec = getVectorSize(op.getPtr(), axisAnalysisPass);
+    // tensor
+    if (tensorTy) {
+      auto valTy = cast<RankedTensorType>(op.getVal().getType());
+      vec = std::min<unsigned>(vec, valTy.getElementType().isF16() ? 2 : 1);
+    }
+
+    auto vecTy = vec_ty(valueElemTy, vec);
     SmallVector<Value> resultVals(elemsPerThread);
 
     // atomic ops
-    for (size_t i = 0; i < elemsPerThread; i += 1) {
-      Value casVal = valElements[i];
-      Value casCmp = cmpElements[i];
+    for (size_t i = 0; i < elemsPerThread; i += vec) {
+      Value casVal = b.undef(vecTy);
+      for (int ii = 0; ii < vec; ++ii) {
+        Value iiVal = createIndexAttrConstant(
+            rewriter, loc, getTypeConverter()->getIndexType(), ii);
+        casVal = b.insert_element(vecTy, casVal, valElements[i + ii], iiVal);
+      }
+
       Value casPtr = ptrElements[i];
+      Value casCmp = cmpElements[i];
+      casVal = valElements[i];
+
       // use op
       if (tensorTy) { // for tensor
-        auto retType = valueElemTy;
+        auto retType = vec == 1 ? valueElemTy : vecTy;
         // TODO: USE ATOMIC CAS OP on Tensor
         auto successOrdering = *atomicMemOrdering;
         auto failureOrdering = LLVM::AtomicOrdering::monotonic;
@@ -1380,7 +1397,12 @@ struct AtomicCASOpConversion
 
         // Extract the new_loaded value from the pair.
         Value ret = b.extract_val(valueElemTy, cmpxchg, i);
-        resultVals[i] = ret;
+
+        for (int ii = 0; ii < vec; ++ii) {
+          resultVals[i + ii] =
+              vec == 1 ? ret
+                       : b.extract_element(valueElemTy, ret, b.i32_val(ii));
+        }
       } else { // for scalar
         // Build blocks to bypass the atomic instruction for ~rmwMask.
         auto *curBlock = rewriter.getInsertionBlock();

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -611,27 +611,52 @@ struct AtomicCASOpConversion
                  : valueTy;
     auto valueElemNBits = valueElemTy.getIntOrFloatBitWidth();
     auto elemsPerThread = getTotalElemsPerThread(op.getVal().getType());
+    // vec = 1 for scalar
+    auto vec = getVectorSize(op.getPtr());
+    auto vecOrig = vec;
+    // tensor
+    if (tensorTy) {
+      auto valTy = cast<RankedTensorType>(op.getVal().getType());
+      vec = std::min<unsigned>(vec, valTy.getElementType().isF16() ? 2 : 1);
+    }
+
+    if (vec == 1 && elemsPerThread > 1)
+      op->emitRemark() << "Warning: vectorization fails vec = " << vec
+                       << " origin vec = " << vecOrig
+                       << " elemsPerThread = " << elemsPerThread << "\n";
+
     auto freeVarMasks = getFreeVariableMasks(op.getPtr().getType());
     Value threadPred =
         emitRedundantThreadPredicate(freeVarMasks, rewriter, loc, targetInfo);
     uint32_t regMask = freeVarMasks[str_attr("reg")];
 
+    auto vecTy = vec_ty(valueElemTy, vec);
     SmallVector<Value> resultVals(elemsPerThread);
 
-    for (size_t i = 0; i < elemsPerThread; i += 1) {
-      if (auto canonicalStart = getCanonicalIndex(i, regMask);
-          canonicalStart != i) {
+    for (size_t i = 0; i < elemsPerThread; i += vec) {
+      if (auto canonicalVecStart = getCanonicalIndex(i, regMask);
+          canonicalVecStart != i) {
         // For redundant registers, refer back to the canonical result
-        resultVals[i] = resultVals[canonicalStart];
+        for (auto iVec = 0; iVec < vec; ++iVec) {
+          resultVals[i + iVec] = resultVals[canonicalVecStart + iVec];
+        }
         continue;
       }
 
-      Value casVal = valElements[i];
-      Value casCmp = cmpElements[i];
+      Value casVal = b.undef(vecTy);
+      for (int ii = 0; ii < vec; ++ii) {
+        Value iiVal = createIndexAttrConstant(
+            rewriter, loc, getTypeConverter()->getIndexType(), ii);
+        casVal = b.insert_element(vecTy, casVal, valElements[i + ii], iiVal);
+      }
+
       Value casPtr = ptrElements[i];
+      Value casCmp = cmpElements[i];
+      casVal = valElements[i];
       PTXBuilder ptxBuilderAtomicCAS;
-      std::string tyId =
-          valueElemNBits == 64 ? "l" : (valueElemNBits == 32 ? "r" : "h");
+      std::string tyId = valueElemNBits * vec == 64
+                             ? "l"
+                             : (valueElemNBits * vec == 32 ? "r" : "h");
       auto *dstOpr = ptxBuilderAtomicCAS.newOperand("=" + tyId, /*init=*/true);
       auto *ptrOpr = ptxBuilderAtomicCAS.newAddrOperand(casPtr, "l");
       auto *cmpOpr = ptxBuilderAtomicCAS.newOperand(casCmp, tyId);
@@ -646,9 +671,13 @@ struct AtomicCASOpConversion
       atom(dstOpr, ptrOpr, cmpOpr, valOpr).maybePredicate(threadPred);
 
       if (tensorTy) {
-        auto retType = valueElemTy;
+        auto retType = vec == 1 ? valueElemTy : vecTy;
         auto ret = ptxBuilderAtomicCAS.launch(rewriter, loc, retType);
-        resultVals[i] = ret;
+        for (int ii = 0; ii < vec; ++ii) {
+          resultVals[i + ii] =
+              vec == 1 ? ret
+                       : b.extract_element(valueElemTy, ret, b.i32_val(ii));
+        }
       } else {
         auto old = ptxBuilderAtomicCAS.launch(rewriter, loc, valueElemTy);
         if (op.getResult().use_empty()) {