Merge commit '91302ea36ddc0b61629da975732b2753ef270764'

Author: whitneywhtsang

README.md

@@ -253,5 +253,5 @@ Supported Platforms:
 Supported Hardware:
 
 - NVIDIA GPUs (Compute Capability 8.0+)
-- AMD GPUs (ROCm 5.2+)
+- AMD GPUs (ROCm 6.2+)
 - Under development: CPUs

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -53,6 +53,7 @@ using namespace mlir::triton;
 #define fmin(...) rewriter.create<LLVM::MinNumOp>(loc, __VA_ARGS__)
 #define shl(...) rewriter.create<LLVM::ShlOp>(loc, __VA_ARGS__)
 #define lshr(...) rewriter.create<LLVM::LShrOp>(loc, __VA_ARGS__)
+#define ashr(...) rewriter.create<LLVM::AShrOp>(loc, __VA_ARGS__)
 #define and_(...) rewriter.create<LLVM::AndOp>(loc, __VA_ARGS__)
 #define xor_(...) rewriter.create<LLVM::XOrOp>(loc, __VA_ARGS__)
 #define or_(...) rewriter.create<LLVM::OrOp>(loc, __VA_ARGS__)

include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td

@@ -762,8 +762,6 @@ for
   ];
 
   let extraClassDeclaration = extraDistributedDeclaration # [{
-    SliceEncodingAttr squeeze(int axis);
-
     SmallVector<unsigned> getContigPerThread() {
       // Block encoding is dense stride layout. The elements per thread are contiguous.
       return getSizePerThread();

include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h

@@ -52,11 +52,6 @@ mlir::LogicalResult createTMADesc(mlir::Value tmaPtr,
       loc, builder.getI64Type(), builder.getI64IntegerAttr(elemSize));
   Value globalStride = builder.template create<arith::MulIOp>(
       loc, op.getStrides()[0], elemSizeVal);
-  // TODO: Workaround for ptxas bug, remove when we update ptxas
-  Value four = builder.template create<arith::ConstantOp>(
-      loc, builder.getI64Type(), builder.getI64IntegerAttr(4));
-  globalStride =
-      builder.template create<arith::ShRSIOp>(loc, globalStride, four);
 
   int elemTypeEnum;
   switch (elemSize) {

lib/Dialect/Triton/IR/Ops.cpp

@@ -221,7 +221,7 @@ LogicalResult TransOp::inferReturnTypes(
   Attribute retEncoding;
   if (argEncoding) {
     Dialect &dialect = argEncoding.getDialect();
-    auto inferLayoutInterface = dyn_cast<DialectInferLayoutInterface>(&dialect);
+    auto inferLayoutInterface = cast<DialectInferLayoutInterface>(&dialect);
     if (inferLayoutInterface
             ->inferTransOpEncoding(argEncoding, order, retEncoding)
             .failed()) {
@@ -250,7 +250,7 @@ DotOp::inferReturnTypes(MLIRContext *context, std::optional<Location> location,
   if (aEnc) {
     assert(bEnc && retEnc);
     Dialect &dialect = retEnc.getDialect();
-    auto interface = dyn_cast<DialectInferLayoutInterface>(&dialect);
+    auto interface = cast<DialectInferLayoutInterface>(&dialect);
     if (interface->inferDotOpEncoding(aEnc, 0, retEnc, location).failed())
       return failure();
     if (interface->inferDotOpEncoding(bEnc, 1, retEnc, location).failed())
@@ -331,8 +331,7 @@ inferReduceReturnShape(RankedTensorType argTy, Type retEltTy, int axis,
     Attribute retEncoding;
     if (argEncoding) {
       Dialect &dialect = argEncoding.getDialect();
-      auto inferLayoutInterface =
-          dyn_cast<DialectInferLayoutInterface>(&dialect);
+      auto inferLayoutInterface = cast<DialectInferLayoutInterface>(&dialect);
       if (inferLayoutInterface
               ->inferReduceOpEncoding(argEncoding, axis, retEncoding)
               .failed()) {
@@ -565,7 +564,7 @@ LogicalResult ExpandDimsOp::inferReturnTypes(
   Attribute retEncoding;
   if (argEncoding) {
     Dialect &dialect = argEncoding.getDialect();
-    auto inferLayoutInterface = dyn_cast<DialectInferLayoutInterface>(&dialect);
+    auto inferLayoutInterface = cast<DialectInferLayoutInterface>(&dialect);
     if (inferLayoutInterface
             ->inferExpandDimsOpEncoding(argEncoding, axis, retEncoding, loc)
             .failed())
@@ -604,7 +603,7 @@ LogicalResult ExpandDimsOp::canonicalize(ExpandDimsOp op,
     // Infer the encoding of the new expand op, if encodings are present.
     Attribute newExpandEnc;
     if (auto srcEnc = srcTy.getEncoding()) {
-      if (dyn_cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
+      if (cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
               ->inferExpandDimsOpEncoding(srcEnc, op.getAxis(), newExpandEnc,
                                           op.getLoc())
               .failed()) {
@@ -975,7 +974,6 @@ JoinOp::inferReturnTypes(MLIRContext *context, std::optional<Location> location,
   assert(isa<RankedTensorType>(operands[1].getType()));
 
   Value lhs = operands[0];
-  Value rhs = operands[1];
   auto srcTy = cast<RankedTensorType>(lhs.getType());
 
   SmallVector<int64_t> retShape(srcTy.getShape());
@@ -984,7 +982,7 @@ JoinOp::inferReturnTypes(MLIRContext *context, std::optional<Location> location,
   Attribute srcEnc = srcTy.getEncoding();
   Attribute retEnc;
   if (srcEnc) {
-    if (dyn_cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
+    if (cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
             ->inferJoinOpEncoding(srcEnc, retEnc, location)
             .failed()) {
       return failure();
@@ -1017,7 +1015,7 @@ LogicalResult SplitOp::inferReturnTypes(
   Attribute srcEnc = srcTy.getEncoding();
   Attribute retEnc;
   if (srcEnc) {
-    if (dyn_cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
+    if (cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
             ->inferSplitOpEncoding(srcEnc, retEnc, location)
             .failed()) {
       return failure();