[LAYOUTS] [NFC] Just accept DistributedEncodings in SliceLayout (triton-lang#6004)

lezcano · loislo · commit 2daef6c2de24 · 2025-03-04T16:28:38.000+01:00
diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
@@ -1267,8 +1267,7 @@ def SliceEncodingAttr : DistributedEncoding<"SliceEncoding", "slice_encoding"> {
   let parameters = (
     ins
     "unsigned":$dim,
-    // TODO: constraint here to only take distributed encodings
-    "Attribute":$parent
+    "DistributedEncodingTrait":$parent
   );
 
   let extraClassDeclaration = extraDistributedDeclaration # [{
diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp
@@ -1592,7 +1592,12 @@ Attribute SliceEncodingAttr::parse(AsmParser &parser, Type type) {
   if (parser.parseGreater().failed())
     return {};
   unsigned dim = mlir::cast<IntegerAttr>(attrs.get("dim")).getInt();
-  Attribute parent = attrs.get("parent");
+  auto parent = mlir::dyn_cast<DistributedEncodingTrait>(attrs.get("parent"));
+  if (!parent) {
+    parser.emitError(parser.getNameLoc(),
+                     "expected a distributed encoding trait");
+    return {};
+  }
   return parser.getChecked<SliceEncodingAttr>(parser.getContext(), dim, parent);
 }
 
@@ -2285,8 +2290,9 @@ struct TritonGPUInferLayoutInterface
   LogicalResult
   inferReduceOpEncoding(Attribute operandEncoding, unsigned axis,
                         Attribute &resultEncoding) const override {
-    resultEncoding = SliceEncodingAttr::get(getDialect()->getContext(), axis,
-                                            operandEncoding);
+    resultEncoding =
+        SliceEncodingAttr::get(getDialect()->getContext(), axis,
+                               cast<DistributedEncodingTrait>(operandEncoding));
     return success();
   }
 
diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
@@ -428,8 +428,8 @@ static bool canUseTwoCTAs(triton::DotOp dotOp) {
   return true;
 }
 
-static Attribute
-replaceCTALayout(Attribute layout,
+static DistributedEncodingTrait
+replaceCTALayout(DistributedEncodingTrait layout,
                  const triton::gpu::CTALayoutAttr &newCTALayout) {
   if (auto blockedLayout = mlir::dyn_cast<BlockedEncodingAttr>(layout)) {
     return BlockedEncodingAttr::get(
@@ -454,7 +454,7 @@ static Value splitBOperand(Value b, mlir::PatternRewriter &rewriter) {
   auto loadOp = b.getDefiningOp<triton::LoadOp>();
   assert(loadOp && "expected LoadOp");
   RankedTensorType bType = cast<RankedTensorType>(b.getType());
-  Attribute currentLayout = bType.getEncoding();
+  auto currentLayout = cast<DistributedEncodingTrait>(bType.getEncoding());
   auto newCTALayout =
       CTALayoutAttr::get(ctx, {1, 2}, {1, 2}, getCTAOrder(currentLayout));
   Attribute newLayout = replaceCTALayout(currentLayout, newCTALayout);
diff --git a/lib/Dialect/TritonGPU/Transforms/OptimizeThreadLocality.cpp b/lib/Dialect/TritonGPU/Transforms/OptimizeThreadLocality.cpp
@@ -544,7 +544,8 @@ class TritonGPUOptimizeThreadLocalityPass
     return viewOpTensorShape;
   }
 
-  Attribute getThreadLocalityOptimizedEncoding(triton::ReduceOp reduce) const {
+  BlockedEncodingAttr
+  getThreadLocalityOptimizedEncoding(triton::ReduceOp reduce) const {
     auto srcType = cast<RankedTensorType>(reduce.getOperands()[0].getType());
     auto rank = srcType.getShape().size();
     auto srcEncoding = srcType.getEncoding();
diff --git a/lib/Dialect/TritonGPU/Transforms/Utility.cpp b/lib/Dialect/TritonGPU/Transforms/Utility.cpp
@@ -302,8 +302,9 @@ std::string GraphLayoutMarker::getColor(const Type &type) const {
 // -------------------------------------------------------------------------- //
 
 static Attribute inferDstEncoding(triton::ReduceOp op, Attribute encoding) {
-  return triton::gpu::SliceEncodingAttr::get(op->getContext(), op.getAxis(),
-                                             encoding);
+  return triton::gpu::SliceEncodingAttr::get(
+      op->getContext(), op.getAxis(),
+      cast<ttg::DistributedEncodingTrait>(encoding));
 }
 
 static Attribute inferDstEncoding(triton::ExpandDimsOp op, Attribute encoding) {
@@ -351,8 +352,9 @@ static Attribute inferSrcEncoding(triton::ReduceOp op, Attribute encoding) {
 }
 
 static Attribute inferSrcEncoding(triton::ExpandDimsOp op, Attribute encoding) {
-  return triton::gpu::SliceEncodingAttr::get(op->getContext(), op.getAxis(),
-                                             encoding);
+  return triton::gpu::SliceEncodingAttr::get(
+      op->getContext(), op.getAxis(),
+      cast<ttg::DistributedEncodingTrait>(encoding));
 }
 
 static Attribute inferSrcEncoding(JoinOp op, Attribute dstEnc) {
diff --git a/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp b/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp
@@ -57,8 +57,10 @@ Type replaceLayout(const Type &type, const Attribute &newLayout) {
   return curType;
 }
 
-Attribute replaceCTALayout(Attribute layout, llvm::ArrayRef<int64_t> shape,
-                           const ttg::CTALayoutAttr &newCTALayout) {
+ttg::DistributedEncodingTrait
+replaceCTALayout(ttg::DistributedEncodingTrait layout,
+                 llvm::ArrayRef<int64_t> shape,
+                 const ttg::CTALayoutAttr &newCTALayout) {
   if (auto blockedLayout = mlir::dyn_cast<ttg::BlockedEncodingAttr>(layout)) {
     return ttg::BlockedEncodingAttr::get(
         layout.getContext(), shape, blockedLayout.getSizePerThread(),
@@ -120,9 +122,9 @@ class CTAPlanner {
 
   bool processBroadcast(triton::BroadcastOp broadcast, Attribute layout);
   bool processExpandDimsBackward(triton::ExpandDimsOp expandDims,
-                                 Attribute newResultLayout);
+                                 ttg::DistributedEncodingTrait newResultLayout);
   bool processExpandDimsForward(triton::ExpandDimsOp expandDims,
-                                Attribute newSrcLayout);
+                                ttg::DistributedEncodingTrait newSrcLayout);
 
   bool processConvertLayoutBackward(ttg::ConvertLayoutOp convertLayout,
                                     CastOp cast);
@@ -361,7 +363,8 @@ bool CTAPlanner::processReduce(triton::FuncOp &funcOp) {
         ttg::CTALayoutAttr::get(context, CTAsPerCGA, CTASplitNum, CTAOrder);
     if (!tiled)
       setTiling(CTALayout.getCTAsPerCGA());
-    auto newSrcLayout = replaceCTALayout(srcLayout, srcShape, CTALayout);
+    auto newSrcLayout = replaceCTALayout(
+        cast<ttg::DistributedEncodingTrait>(srcLayout), srcShape, CTALayout);
     auto newResultLayout =
         ttg::SliceEncodingAttr::get(context, axis, newSrcLayout);
     unsigned numOperands = reduce.getNumOperands();
@@ -393,8 +396,9 @@ void CTAPlanner::processStoreLikeOps(triton::FuncOp &funcOp) {
         CTALayout = ttg::getCTALayout(tensorTy.getEncoding());
         setTiling(CTALayout.getCTAsPerCGA());
       }
-      auto newLayout = replaceCTALayout(tensorTy.getEncoding(),
-                                        tensorTy.getShape(), CTALayout);
+      auto newLayout = replaceCTALayout(
+          cast<ttg::DistributedEncodingTrait>(tensorTy.getEncoding()),
+          tensorTy.getShape(), CTALayout);
       processElementwise(store, newLayout);
     }
   }
@@ -421,7 +425,8 @@ bool CTAPlanner::propagateBackward(CastOp cast) {
     Type outTy = output.getType();
     if (auto ptrTy = dyn_cast<triton::PointerType>(outTy))
       outTy = ptrTy.getPointeeType();
-    Attribute layout = mlir::cast<RankedTensorType>(outTy).getEncoding();
+    auto layout = mlir::cast<ttg::DistributedEncodingTrait>(
+        mlir::cast<RankedTensorType>(outTy).getEncoding());
     Operation *op = input.getDefiningOp();
     if (op == nullptr) {
       assert(isa<BlockArgument>(input) &&
@@ -626,8 +631,10 @@ bool CTAPlanner::processLoadStore(Operation *op, Attribute layout) {
     if (auto ptrTy = dyn_cast<triton::PointerType>(type))
       type = ptrTy.getPointeeType();
     auto tensorTy = cast<RankedTensorType>(type);
-    auto newLayout = replaceCTALayout(tensorTy.getEncoding(),
-                                      tensorTy.getShape(), CTALayout);
+    auto oldLayout =
+        cast<ttg::DistributedEncodingTrait>(tensorTy.getEncoding());
+    auto newLayout =
+        replaceCTALayout(oldLayout, tensorTy.getShape(), CTALayout);
     newOperandLayouts.push_back(newLayout);
   }
 
@@ -637,8 +644,10 @@ bool CTAPlanner::processLoadStore(Operation *op, Attribute layout) {
     if (auto ptrTy = dyn_cast<triton::PointerType>(type))
       type = ptrTy.getPointeeType();
     auto tensorTy = cast<RankedTensorType>(type);
-    auto newLayout = replaceCTALayout(tensorTy.getEncoding(),
-                                      tensorTy.getShape(), CTALayout);
+    auto oldLayout =
+        cast<ttg::DistributedEncodingTrait>(tensorTy.getEncoding());
+    auto newLayout =
+        replaceCTALayout(oldLayout, tensorTy.getShape(), CTALayout);
     newResultLayouts.push_back(newLayout);
   }
 
@@ -725,16 +734,18 @@ bool CTAPlanner::processBroadcast(triton::BroadcastOp broadcast,
   return true;
 }
 
-bool CTAPlanner::processExpandDimsBackward(triton::ExpandDimsOp expandDims,
-                                           Attribute newResultLayout) {
+bool CTAPlanner::processExpandDimsBackward(
+    triton::ExpandDimsOp expandDims,
+    ttg::DistributedEncodingTrait newResultLayout) {
   auto newSrcLayout = ttg::SliceEncodingAttr::get(
       newResultLayout.getContext(), expandDims.getAxis(), newResultLayout);
   insertCasts(expandDims.getOperation(), {newSrcLayout}, {newResultLayout});
   return true;
 }
 
-bool CTAPlanner::processExpandDimsForward(triton::ExpandDimsOp expandDims,
-                                          Attribute newSrcLayout) {
+bool CTAPlanner::processExpandDimsForward(
+    triton::ExpandDimsOp expandDims,
+    ttg::DistributedEncodingTrait newSrcLayout) {
   llvm::report_fatal_error("processExpandDimsForward not implemented yet");
   return true;
 }
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/DecomposeUnsupportedConversions.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/DecomposeUnsupportedConversions.cpp
@@ -67,7 +67,8 @@ struct DecomposeUnsupportedAMDConversions
       auto srcType = cvtOp.getSrc().getType();
       auto dstType = cvtOp.getType();
 
-      auto srcEnc = srcType.getEncoding();
+      auto srcEnc =
+          cast<triton::gpu::DistributedEncodingTrait>(srcType.getEncoding());
       auto dstBlocked =
           dyn_cast<triton::gpu::BlockedEncodingAttr>(dstType.getEncoding());
 
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUsage.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUsage.cpp
@@ -92,8 +92,10 @@ class OptimizeAMDLDSUsage
     auto srcType = cvtOp.getSrc().getType();
     auto dstType = cvtOp.getType();
 
-    auto srcEnc = srcType.getEncoding();
-    auto dstEnc = dstType.getEncoding();
+    auto srcEnc =
+        cast<triton::gpu::DistributedEncodingTrait>(srcType.getEncoding());
+    auto dstEnc =
+        cast<triton::gpu::DistributedEncodingTrait>(dstType.getEncoding());
 
     auto ctx = srcEnc.getContext();
     auto rank = srcType.getRank();
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUtility.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUtility.cpp
@@ -49,7 +49,9 @@ std::vector<SmallVector<unsigned>> factorizePowerOf2(int n, int rank) {
   return factors;
 }
 
-Attribute createTmpLayout(Attribute layout, ArrayRef<unsigned> warpsPerCTA) {
+triton::gpu::DistributedEncodingTrait
+createTmpLayout(triton::gpu::DistributedEncodingTrait layout,
+                ArrayRef<unsigned> warpsPerCTA) {
   auto ctx = layout.getContext();
   if (auto src = dyn_cast<triton::gpu::AMDMfmaEncodingAttr>(layout))
     return triton::gpu::AMDMfmaEncodingAttr::get(
@@ -65,8 +67,9 @@ Attribute createTmpLayout(Attribute layout, ArrayRef<unsigned> warpsPerCTA) {
         ctx, src.getSizePerThread(), src.getThreadsPerWarp(), warpsPerCTA,
         src.getOrder(), src.getCTALayout());
   if (auto src = dyn_cast<triton::gpu::DotOperandEncodingAttr>(layout)) {
+    auto parent = cast<triton::gpu::DistributedEncodingTrait>(src.getParent());
     return triton::gpu::DotOperandEncodingAttr::get(
-        ctx, src.getOpIdx(), createTmpLayout(src.getParent(), warpsPerCTA),
+        ctx, src.getOpIdx(), createTmpLayout(parent, warpsPerCTA),
         src.getKWidth());
   }
   if (auto src = dyn_cast<triton::gpu::SliceEncodingAttr>(layout)) {
@@ -77,7 +80,7 @@ Attribute createTmpLayout(Attribute layout, ArrayRef<unsigned> warpsPerCTA) {
         ctx, src.getDim(), createTmpLayout(src.getParent(), parentWarpsPerCTA));
   }
   assert("Encountered unsupported layout");
-  return Attribute();
+  return {};
 }
 
 std::pair<triton::gpu::ConvertLayoutOp, triton::gpu::ConvertLayoutOp>
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUtility.h b/third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUtility.h
@@ -16,7 +16,9 @@ std::vector<SmallVector<unsigned>> factorizePowerOf2(int n, int rank);
 /// \param layout original layout
 /// \param warpsPerCTA new warpsPerCTA
 /// \returns create layout
-Attribute createTmpLayout(Attribute layout, ArrayRef<unsigned> warpsPerCTA);
+triton::gpu::DistributedEncodingTrait
+createTmpLayout(triton::gpu::DistributedEncodingTrait layout,
+                ArrayRef<unsigned> warpsPerCTA);
 
 /// Creates two chained convert layout operations
 ///
diff --git a/unittest/Dialect/TritonGPU/LinearLayoutConversionsTest.cpp b/unittest/Dialect/TritonGPU/LinearLayoutConversionsTest.cpp
@@ -85,7 +85,7 @@ class LinearLayoutConversionsTest : public ::testing::Test {
     return DotOperandEncodingAttr::get(&ctx, opIdx, wmma, kWidth);
   }
 
-  SliceEncodingAttr slice(Attribute parent, int dim) {
+  SliceEncodingAttr slice(DistributedEncodingTrait parent, int dim) {
     return SliceEncodingAttr::get(&ctx, dim, parent);
   }
 

Original file line number	Diff line number	Diff line change
`@@ -544,7 +544,8 @@ class TritonGPUOptimizeThreadLocalityPass`
`544`	`544`	`return viewOpTensorShape;`
`545`	`545`	`}`
`546`	`546`
`547`		`- Attribute getThreadLocalityOptimizedEncoding(triton::ReduceOp reduce) const {`
	`547`	`+ BlockedEncodingAttr`
	`548`	`+ getThreadLocalityOptimizedEncoding(triton::ReduceOp reduce) const {`
`548`	`549`	`auto srcType = cast<RankedTensorType>(reduce.getOperands()[0].getType());`
`549`	`550`	`auto rank = srcType.getShape().size();`
`550`	`551`	`auto srcEncoding = srcType.getEncoding();`
Original file line number	Diff line number	Diff line change
`@@ -302,8 +302,9 @@ std::string GraphLayoutMarker::getColor(const Type &type) const {`
`302`	`302`	`// -------------------------------------------------------------------------- //`
`303`	`303`
`304`	`304`	`static Attribute inferDstEncoding(triton::ReduceOp op, Attribute encoding) {`
`305`		`- return triton::gpu::SliceEncodingAttr::get(op->getContext(), op.getAxis(),`
`306`		`- encoding);`
	`305`	`+ return triton::gpu::SliceEncodingAttr::get(`
	`306`	`+ op->getContext(), op.getAxis(),`
	`307`	`+ cast<ttg::DistributedEncodingTrait>(encoding));`
`307`	`308`	`}`
`308`	`309`
`309`	`310`	`static Attribute inferDstEncoding(triton::ExpandDimsOp op, Attribute encoding) {`
`@@ -351,8 +352,9 @@ static Attribute inferSrcEncoding(triton::ReduceOp op, Attribute encoding) {`
`351`	`352`	`}`
`352`	`353`
`353`	`354`	`static Attribute inferSrcEncoding(triton::ExpandDimsOp op, Attribute encoding) {`
`354`		`- return triton::gpu::SliceEncodingAttr::get(op->getContext(), op.getAxis(),`
`355`		`- encoding);`
	`355`	`+ return triton::gpu::SliceEncodingAttr::get(`
	`356`	`+ op->getContext(), op.getAxis(),`
	`357`	`+ cast<ttg::DistributedEncodingTrait>(encoding));`
`356`	`358`	`}`
`357`	`359`
`358`	`360`	`static Attribute inferSrcEncoding(JoinOp op, Attribute dstEnc) {`