Fix tt.dot to support integers

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

@@ -150,6 +150,34 @@ static std::optional<unsigned> getBitWidth(Type a) {
   return std::nullopt;
 }
 
+// true means tensor elements are zeros
+// false means not zero or it cannot be determined
+static bool isZeroTensor(Value &v, bool integers) {
+    if (auto splatOp = v.getDefiningOp<triton::SplatOp>()) {
+      if (auto constOp = splatOp.getSrc().getDefiningOp<arith::ConstantOp>()) {
+        if (auto val = dyn_cast<FloatAttr>(constOp.getValue())) {
+          return val.getValueAsDouble() == 0.;
+        }
+        if (auto val = dyn_cast<IntegerAttr>(constOp.getValue())) {
+          return val.getValue() == 0;
+        }
+      }
+      return false;
+    }
+
+    if (auto constOp = v.getDefiningOp<arith::ConstantOp>()) {
+      if (auto denseAttr = dyn_cast<DenseElementsAttr>(constOp.getValue())) {
+        if (denseAttr.isSplat()) {
+          if (integers)
+            return denseAttr.getSplatValue<APInt>().isZero();
+          return denseAttr.getSplatValue<APFloat>().isZero();
+        }
+      }
+    }
+
+    return false;
+}
+
 //===----------------------------------------------------------------------===//
 // Op Lowering Patterns
 //===----------------------------------------------------------------------===//
@@ -970,36 +998,23 @@ struct MatmulConverter : public OpConversionPattern<triton::DotOp> {
   LogicalResult
   matchAndRewrite(triton::DotOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
     auto opa = adaptor.getA();
     auto opb = adaptor.getB();
     auto opc = adaptor.getC();
-    auto opcOrig = op.getC();
-
-    bool skipC = false;
-    if (auto splatOp = opcOrig.getDefiningOp<triton::SplatOp>()) {
-      if (auto val = splatOp.getSrc().getDefiningOp<arith::ConstantOp>()) {
-        if (cast<FloatAttr>(val.getValue()).getValueAsDouble() == 0.) {
-          skipC = true;
-        }
-      }
-    } else if (auto constOp = opcOrig.getDefiningOp<arith::ConstantOp>()) {
-      if (auto denseAttr = dyn_cast<DenseElementsAttr>(constOp.getValue())) {
-        if (denseAttr.isSplat() &&
-            denseAttr.getSplatValue<FloatAttr>().getValueAsDouble() == 0.) {
-          skipC = true;
-        }
-      }
-    }
 
     auto dstType = cast<RankedTensorType>(op.getType());
-    auto elemType = dstType.getElementType();
-    auto loc = op.getLoc();
-
+    auto elementType = dstType.getElementType();
+    bool integers = elementType.isInteger();
+    bool skipC = isZeroTensor(opc, integers);
     auto init =
-        rewriter.create<tensor::EmptyOp>(loc, dstType.getShape(), elemType);
+        rewriter.create<tensor::EmptyOp>(loc, dstType.getShape(), elementType);
+    TypedAttr constantAttr = integers ?
+      static_cast<TypedAttr>(rewriter.getIntegerAttr(elementType, 0)) :
+      static_cast<TypedAttr>(rewriter.getFloatAttr(elementType, 0));
 
     auto zero = rewriter.create<mlir::arith::ConstantOp>(
-        op.getLoc(), elemType, rewriter.getFloatAttr(elemType, 0));
+        op.getLoc(), elementType, constantAttr);
 
     auto zeroes =
         rewriter.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{init})
@@ -1011,7 +1026,11 @@ struct MatmulConverter : public OpConversionPattern<triton::DotOp> {
                    .getResult(0);
 
     if (!skipC) {
-      res = rewriter.create<arith::AddFOp>(loc, res, opc);
+      if (integers) {
+        res = rewriter.create<arith::AddIOp>(loc, res, opc);
+      } else {
+        res = rewriter.create<arith::AddFOp>(loc, res, opc);
+      }
     }
 
     rewriter.replaceOp(op, res);

python/examples/conftest.py

@@ -64,7 +64,6 @@ def device(request):
     "test_chain_reduce",
     "test_generic_reduction",
     "test_trans_4d",
-    "test_dot",
     "test_dot3d",
     "test_constexpr",
     "test_arange",
@@ -90,6 +89,8 @@ def pytest_collection_modifyitems(config, items):
     skip_marker = pytest.mark.skip(reason="CPU backend does not support it yet")
     # There is a dependency issue on build machine which breaks bfloat16
     skip_marker_bfloat = pytest.mark.skip(reason="bfloat16 linking issue")
+    skip_marker_tf32 = pytest.mark.skip(reason="tf32 is not supported on CPU")
+    skip_marker_float8 = pytest.mark.skip(reason="float8 is not supported on CPU")
 
     for item in items:
         test_func_name = item.originalname if item.originalname else item.name
@@ -101,4 +102,8 @@ def pytest_collection_modifyitems(config, items):
         if "parametrize" in item.keywords:
             for param_name, param_value in item.callspec.params.items():
                 if param_name.startswith('dtype') and param_value == 'bfloat16':
-                    item.add_marker(skip_marker_bfloat)
+                    item.add_marker(skip_marker_bfloat)
+                if param_name.startswith('input_precision') and param_value.startswith('tf32'):
+                    item.add_marker(skip_marker_tf32)
+                if param_name.endswith('dtype') and param_value.startswith('float8'):
+                    item.add_marker(skip_marker_float8)