OpenXLA-specific changes.

PiperOrigin-RevId: 741163570

Author: Copybara

BUILD

@@ -587,15 +587,17 @@ We call each individual tile "rep".
                     "unsigned",
                     "getTotalElemsPerThread",
                      (ins "ArrayRef<int64_t>":$shape),
+                     /*methodBody=*/[{}],
                      /*defaultImplementation=*/[{
-                         return toLinearEncoding($_self, shape).getTotalElemsPerThread(shape);
+                         return toLinearEncoding($_attr, shape).getTotalElemsPerThread(shape);
                      }]>,
     InterfaceMethod<"Return element size per thread in each dimension.",
                     "SmallVector<unsigned>",
                     "getElemsPerThread",
                      (ins "ArrayRef<int64_t>":$shape),
+                     /*methodBody=*/[{}],
                      /*defaultImplementation=*/[{
-                         return toLinearEncoding($_self, shape).getElemsPerThread(shape);
+                         return toLinearEncoding($_attr, shape).getElemsPerThread(shape);
                      }]>,
     InterfaceMethod<"Convert to LinearLayout.",
                     "LinearLayout",

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

@@ -57,12 +57,14 @@ LinearEncodingAttr toLinearEncoding(Attribute layout, ArrayRef<int64_t> shape) {
 }
 
 unsigned getTotalElemsPerThread(Attribute layout, ArrayRef<int64_t> shape) {
-  return toLinearEncoding(layout, shape).getTotalElemsPerThread(shape);
+  auto distributedEncoding = mlir::cast<DistributedEncodingTrait>(layout);
+  return distributedEncoding.getTotalElemsPerThread(shape);
 }
 
 SmallVector<unsigned> getElemsPerThread(Attribute layout,
                                         ArrayRef<int64_t> shape) {
-  return toLinearEncoding(layout, shape).getElemsPerThread(shape);
+  auto distributedEncoding = mlir::cast<DistributedEncodingTrait>(layout);
+  return distributedEncoding.getElemsPerThread(shape);
 }
 
 SmallVector<unsigned> getElemsPerThread(Type type) {

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -160,6 +160,11 @@ struct CanonicalizeConvertFromAlloc
     auto convert = op.getSrc().getDefiningOp<ConvertLayoutOp>();
     if (!convert)
       return failure();
+    // LocalAllocOp lowering doesn't support going from DotOperandEncoding
+    // to SharedEncoding, so we want to keep this layout conversion.
+    if (mlir::isa<triton::gpu::DotOperandEncodingAttr>(
+            convert.getSrc().getType().getEncoding()))
+      return failure();
     rewriter.replaceOpWithNewOp<triton::gpu::LocalAllocOp>(
         op, op->getResult(0).getType(), convert.getSrc());
     return mlir::success();

lib/Dialect/TritonGPU/IR/Ops.cpp

@@ -185,6 +185,21 @@ getSharedMemoryMMAOperand(Value v, mlir::PatternRewriter &rewriter, int opIdx,
   auto newType = MemDescType::get(argType.getShape(), argType.getElementType(),
                                   newLayout, SharedMemorySpace);
   rewriter.setInsertionPointAfterValue(arg);
+
+  // LocalAllocOp lowering doesn't support going from DotOperandEncoding
+  // to SharedEncoding.
+  if (auto dotOpEnc = mlir::dyn_cast<DotOperandEncodingAttr>(
+          argType.getEncoding())) {
+    // Create a layout conversion from DotOperandEncoding to BlockedEncoding
+    // then pass it to the LocalAllocOp.
+    auto newArgType = RankedTensorType::get(
+        argType.getShape(), argType.getElementType(), dotOpEnc.getParent());
+    auto dotOperandToBlockedCvt =
+        rewriter.create<ConvertLayoutOp>(arg.getLoc(), newArgType, arg);
+    return rewriter.create<LocalAllocOp>(arg.getLoc(), newType,
+                                              dotOperandToBlockedCvt);
+  }
+
   return rewriter.create<LocalAllocOp>(arg.getLoc(), newType, arg);
 }
 
@@ -222,9 +237,22 @@ getWarpsPerTile(DotOp dotOp, const ArrayRef<int64_t> shape, int version,
 }
 
 static bool bwdFilter(Operation *op) {
+  // Dot operand layout assignment to Predicates are not currently supported
+  // during lowering from TritonGPU to LLVM in Triton for MMA cases. This
+  // condition limits visibility of the original bit-width so that predicate
+  // are not considered, hence, kwidth can never be = 32.
+  if (isa<arith::UIToFPOp>(op)) {
+    Type srcType = getElementTypeOrSelf(op->getOperand(0));
+    if (srcType.isInteger(1))
+      return false;
+  }
+
+  // b/405045790: We don't want to propagate through the BroadcastOp because we
+  // probably don't care about the load before a broadcast as it would likely be
+  // small. This is just a heuristic to avoid a regression.
   return (op->hasTrait<OpTrait::Elementwise>() && isMemoryEffectFree(op)) ||
          isView(op) ||
-         isa<Fp4ToFpOp, LoadOp, DescriptorLoadOp, BroadcastOp, ConvertLayoutOp>(
+         isa<Fp4ToFpOp, LoadOp, DescriptorLoadOp, /*BroadcastOp,*/ ConvertLayoutOp>(
              op);
 }
 

lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp

@@ -147,8 +147,14 @@ Value Prefetcher::generatePrefetch(Value v, unsigned opIdx, bool isPrologue,
           type.getMutableMemory(), type.getAllocShape()),
       v, offsetsVal);
 
+  // We need to assign kwidth to zero in the case where the parent layout is
+  // Blocked, otherwise the verifier emits a failure. The parent layout is
+  // Blocked only when Tensor Cores are disabled.
+  int kwidth = dyn_cast<triton::gpu::BlockedEncodingAttr>(dotEncoding)
+                   ? 0
+                   : prefetchWidth / 8;
   auto dotOperandEnc = triton::gpu::DotOperandEncodingAttr::get(
-      builder.getContext(), opIdx, dotEncoding, prefetchWidth / 8);
+      builder.getContext(), opIdx, dotEncoding, kwidth);
   Value prefetchSlice = builder.create<triton::gpu::LocalLoadOp>(
       v.getLoc(), RankedTensorType::get(shape, elementType, dotOperandEnc),
       newSmem);
@@ -198,6 +204,22 @@ LogicalResult Prefetcher::initialize() {
         break;
       if (!op->getResult(0).hasOneUse())
         break;
+      // Similar to issues faced in HoistLayoutConversion pattern in
+      // OptimizeDotOperands.cpp, we can't propagate through type casts from
+      // predicates as they aren't supported in Triton when encoded with dot_op
+      // layout.
+      if (isa<arith::UIToFPOp>(op)) {
+        Type srcType = getElementTypeOrSelf(op->getOperand(0));
+        if (srcType.isInteger(1))
+          break;
+      }
+      // Propagation through ExpandDims is currently not supported. This blindly
+      // replaces the encoding with dot encoding & but ExpandDims requires a
+      // SliceEncoding. This could be rewritten to support it somehow, but I
+      // don't think it's trivial & it's currently crashing.
+      if (isa<ExpandDimsOp>(op)) {
+        break;
+      }
       rets.push_back(op->getOperand(0));
       if (auto cvt = dyn_cast<triton::gpu::LocalLoadOp>(op)) {
         // NYI for other encodings, for example if we have transpose

lib/Dialect/TritonGPU/Transforms/Prefetch.cpp

@@ -1121,8 +1121,7 @@ void LayoutRematerialization::hoistConvertDotOperand(
   // The pass is targeted to Nvidia mma/wgmma dot operands
 
   auto canBePipelined = [&](ConvertLayoutOp convertOp) {
-    // FIXME: Check that the parent is a for loop
-    auto parent = convertOp->getParentOp();
+    auto parent = dyn_cast<scf::ForOp>(convertOp->getParentOp());
     if (!parent)
       return false;
 

lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp

@@ -41,6 +41,7 @@ struct FenceInsertionPass
     if (::triton::tools::getBoolEnv("DISABLE_MMA_V3"))
       return;
     ModuleOp mod = getOperation();
+    DenseSet<std::pair<Operation *, unsigned>> trace;
     mod.walk([&](Operation *op) {
       if (!isa<ttng::WarpGroupDotOp>(op))
         return WalkResult::advance();
@@ -51,8 +52,8 @@ struct FenceInsertionPass
           cast<RankedTensorType>(op->getResult(0).getType()).getEncoding());
       if (!mmaEncoding || !mmaEncoding.isHopper())
         return WalkResult::advance();
-      bool aDependsOnShared = dependOnSharedEncOperand(a);
-      bool bDependsOnShared = dependOnSharedEncOperand(b);
+      bool aDependsOnShared = dependOnSharedEncOperand(a, trace);
+      bool bDependsOnShared = dependOnSharedEncOperand(b, trace);
       if (!aDependsOnShared && !bDependsOnShared)
         return WalkResult::advance();
       Operation *fence = builder.create<ttng::FenceAsyncSharedOp>(
@@ -73,8 +74,7 @@ struct FenceInsertionPass
   }
 
 private:
-  bool dependOnSharedEncOperand(Value operand) {
-    static DenseSet<std::pair<Operation *, unsigned>> trace;
+  bool dependOnSharedEncOperand(Value operand, DenseSet<std::pair<Operation *, unsigned>> &trace) {
     auto op = operand.getDefiningOp();
     // avoid redundant insertion
     if (op && isa<mlir::triton::DotOpInterface>(op))
@@ -89,7 +89,7 @@ struct FenceInsertionPass
     // op and not BlockArgument
     if (op && !isa<BlockArgument>(operand)) {
       for (auto v : op->getOperands()) {
-        if (dependOnSharedEncOperand(v))
+        if (dependOnSharedEncOperand(v, trace))
           return true;
       }
     }
@@ -104,7 +104,7 @@ struct FenceInsertionPass
         auto iterOperands = forOp.getInitArgs();
         if (argNum == 0)
           return false;
-        if (dependOnSharedEncOperand(iterOperands[argNum - 1]))
+        if (dependOnSharedEncOperand(iterOperands[argNum - 1], trace))
           return true;
         // yield
         auto yieldOp = forOp.getBody()->getTerminator();
@@ -117,7 +117,7 @@ struct FenceInsertionPass
         else
           trace.insert(entry);
 
-        if (dependOnSharedEncOperand(v))
+        if (dependOnSharedEncOperand(v, trace))
           return true;
       } else if (auto whileOp = dyn_cast<scf::WhileOp>(argOwner)) {
         assert(false && "FenceInsertionPass does not supported WhileOp");

lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp

@@ -220,6 +220,7 @@ EncodingInfo combineEncodings(const EncodingInfo &lhs, const EncodingInfo &rhs,
     break;
   case 1:
     result.ctaLayout = ctaLayouts[0];
+    break;
   default:
     break;
   }
@@ -237,6 +238,7 @@ EncodingInfo combineEncodings(const EncodingInfo &lhs, const EncodingInfo &rhs,
     break;
   case 1:
     result.desiredEncoding = desiredEncodings[0];
+    break;
   default:
     break;
   }

lib/Dialect/TritonNvidiaGPU/Transforms/OptimizeDescriptorEncoding.cpp

@@ -0,0 +1,79 @@
+# NOTE: Do not depend on any targets from this directory,
+# but use //third_party/py/triton instead.
+
+load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
+
+package(
+    default_applicable_licenses = ["@triton//:license"],
+    default_visibility = [
+        "//third_party/py/triton:__pkg__",
+        "@triton//python:__subpackages__",
+    ],
+)
+
+cc_library(
+    name = "passes",
+    hdrs = ["src/passes.h"],
+    includes = ["src"],
+    visibility = ["@triton//third_party:__subpackages__"],
+)
+
+pybind_extension(
+    name = "libtriton",
+    srcs = [
+        "src/interpreter.cc",
+        "src/ir.cc",
+        "src/llvm.cc",
+        "src/main.cc",
+        "src/passes.cc",
+    ],
+    copts = ["-DTRITON_BACKENDS_TUPLE=(nvidia)"],
+    deps = [
+        ":passes",
+        "@llvm-project//llvm:Core",
+        "@llvm-project//llvm:IPO",
+        "@llvm-project//llvm:IRReader",
+        "@llvm-project//llvm:InstCombine",
+        "@llvm-project//llvm:Instrumentation",
+        "@llvm-project//llvm:Linker",
+        "@llvm-project//llvm:MC",
+        "@llvm-project//llvm:Passes",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:Target",
+        "@llvm-project//mlir:BuiltinToLLVMIRTranslation",
+        "@llvm-project//mlir:BytecodeWriter",
+        "@llvm-project//mlir:ControlFlowDialect",
+        "@llvm-project//mlir:ConversionPasses",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:IndexDialect",
+        "@llvm-project//mlir:LLVMDialect",
+        "@llvm-project//mlir:LLVMIRTransforms",
+        "@llvm-project//mlir:LLVMToLLVMIRTranslation",
+        "@llvm-project//mlir:NVVMToLLVMIRTranslation",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:ToLLVMIRTranslation",
+        "@llvm-project//mlir:Transforms",
+        "@llvm-project//mlir:UBDialect",
+        "@triton//:TritonDialects",
+        "@triton//:TritonGPUToLLVM",
+        "@triton//:TritonGPUTransforms",
+        "@triton//:TritonHSACO",
+        "@triton//:TritonLLVMIR",
+        "@triton//:TritonNvidiaGPUTransforms",
+        "@triton//:TritonPTX",
+        "@triton//:TritonToTritonGPU",
+        "@triton//:TritonTools",
+        "@triton//:TritonTransforms",
+        "@triton//third_party/nvidia:triton_nvidia",
+        "@triton//third_party/proton:ProtonIRDialect",
+    ],
+)
+
+filegroup(
+    name = "files",
+    srcs = glob(
+        include = ["triton/**/*.py"],
+    ),
+)