openxla · rupprecht · Jun 5, 2023 · Jun 12, 2023 · Jun 14, 2023 · Jun 27, 2023
diff --git a/BUILD b/BUILD
diff --git a/include/triton/Analysis/Alias.h b/include/triton/Analysis/Alias.h
@@ -63,11 +63,11 @@ class AliasInfo {
 // Shared Memory Alias Analysis
 //===----------------------------------------------------------------------===//
 class SharedMemoryAliasAnalysis
-    : public dataflow::SparseDataFlowAnalysis<dataflow::Lattice<AliasInfo>> {
+    : public dataflow::SparseForwardDataFlowAnalysis<dataflow::Lattice<AliasInfo>> {
 public:
-  using dataflow::SparseDataFlowAnalysis<
-      dataflow::Lattice<AliasInfo>>::SparseDataFlowAnalysis;
-  using dataflow::SparseDataFlowAnalysis<
+  using dataflow::SparseForwardDataFlowAnalysis<
+      dataflow::Lattice<AliasInfo>>::SparseForwardDataFlowAnalysis;
+  using dataflow::SparseForwardDataFlowAnalysis<
       dataflow::Lattice<AliasInfo>>::getLatticeElement;
 
   /// XXX(Keren): Compatible interface with MLIR AliasAnalysis for future use.

diff --git a/include/triton/Analysis/AxisInfo.h b/include/triton/Analysis/AxisInfo.h
@@ -272,7 +272,7 @@ class AxisInfoVisitorList {
 };
 
 class AxisInfoAnalysis
-    : public dataflow::SparseDataFlowAnalysis<dataflow::Lattice<AxisInfo>> {
+    : public dataflow::SparseForwardDataFlowAnalysis<dataflow::Lattice<AxisInfo>> {
 private:
   AxisInfoVisitorList visitors;
 
@@ -284,7 +284,7 @@ class AxisInfoAnalysis
 
 public:
   AxisInfoAnalysis(DataFlowSolver &solver);
-  using dataflow::SparseDataFlowAnalysis<
+  using dataflow::SparseForwardDataFlowAnalysis<
       dataflow::Lattice<AxisInfo>>::getLatticeElement;
   using FuncAxisInfoMapT = DenseMap<FunctionOpInterface, AxisInfo>;
 

diff --git a/include/triton/Analysis/Utility.h b/include/triton/Analysis/Utility.h
@@ -179,9 +179,9 @@ SetVector<Operation *>
 multiRootTopologicalSort(const SetVector<Operation *> &toSort);
 
 /// This uses the toplogicalSort above
-SetVector<Operation *>
-multiRootGetSlice(Operation *op, TransitiveFilter backwardFilter = nullptr,
-                  TransitiveFilter forwardFilter = nullptr);
+SetVector<Operation *> multiRootGetSlice(
+    Operation *op, SliceOptions::TransitiveFilter backwardFilter = nullptr,
+    SliceOptions::TransitiveFilter forwardFilter = nullptr);
 
 /// Create a basic DataFlowSolver with constant and dead code analysis included.
 std::unique_ptr<DataFlowSolver> createDataFlowSolver();

diff --git a/include/triton/Dialect/Triton/IR/TritonOps.td b/include/triton/Dialect/Triton/IR/TritonOps.td
@@ -633,6 +633,9 @@ def CallOp : TT_Op<"call", [CallOpInterface, /*MemRefsNormalizable, */DeclareOpI
     operand_range getArgOperands() {
       return {arg_operand_begin(), arg_operand_end()};
     }
+    MutableOperandRange getArgOperandsMutable() {
+      return getOperandsMutable();
+    }
 
     operand_iterator arg_operand_begin() { return operand_begin(); }
     operand_iterator arg_operand_end() { return operand_end(); }

diff --git a/lib/Analysis/AxisInfo.cpp b/lib/Analysis/AxisInfo.cpp
@@ -857,7 +857,7 @@ class MaxMinOpAxisInfoVisitor final : public AxisInfoVisitorImpl<OpTy> {
 //===----------------------------------------------------------------------===//
 
 AxisInfoAnalysis::AxisInfoAnalysis(DataFlowSolver &solver)
-    : dataflow::SparseDataFlowAnalysis<dataflow::Lattice<AxisInfo>>(solver) {
+    : dataflow::SparseForwardDataFlowAnalysis<dataflow::Lattice<AxisInfo>>(solver) {
   // UnrealizedConversionCast:
   // This is needed by TritonGPUToLLVM, to get AxisInfo when the graph is
   // in the process of a PartialConversion, where UnrealizedConversionCast

diff --git a/lib/Analysis/Utility.cpp b/lib/Analysis/Utility.cpp
@@ -125,7 +125,7 @@ unsigned ReduceOpHelper::getScratchSizeInBytes() {
 
   unsigned bytesPerElem = 0;
   for (const auto &ty : srcElementTypes) {
-    bytesPerElem += ty.getIntOrFloatBitWidth() / 8;
+    bytesPerElem += (ty.getIntOrFloatBitWidth() + 7) / 8;
   }
   return bytesPerElem * elems;
 }
@@ -473,9 +473,9 @@ multiRootTopologicalSort(const SetVector<Operation *> &toSort) {
   return res;
 }
 
-SetVector<Operation *> multiRootGetSlice(Operation *op,
-                                         TransitiveFilter backwardFilter,
-                                         TransitiveFilter forwardFilter) {
+SetVector<Operation *> multiRootGetSlice(
+    Operation *op, SliceOptions::TransitiveFilter backwardFilter,
+    SliceOptions::TransitiveFilter forwardFilter) {
   SetVector<Operation *> slice;
   slice.insert(op);
 
@@ -486,12 +486,12 @@ SetVector<Operation *> multiRootGetSlice(Operation *op,
     auto *currentOp = (slice)[currentIndex];
     // Compute and insert the backwardSlice starting from currentOp.
     backwardSlice.clear();
-    getBackwardSlice(currentOp, &backwardSlice, backwardFilter);
+    getBackwardSlice(currentOp, &backwardSlice, {{backwardFilter}});
     slice.insert(backwardSlice.begin(), backwardSlice.end());
 
     // Compute and insert the forwardSlice starting from currentOp.
     forwardSlice.clear();
-    getForwardSlice(currentOp, &forwardSlice, forwardFilter);
+    getForwardSlice(currentOp, &forwardSlice, {{forwardFilter}});
     slice.insert(forwardSlice.begin(), forwardSlice.end());
     ++currentIndex;
   }

diff --git a/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt b/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt
@@ -35,7 +35,7 @@ add_mlir_conversion_library(TritonGPUToLLVM
     LINK_LIBS PUBLIC
     MLIRIR
     MLIRPass
-    MLIRGPUOps
+    MLIRGPUDialect
     MLIRGPUToNVVMTransforms
     MLIRGPUToROCDLTransforms
     MLIRGPUTransforms

diff --git a/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cpp
@@ -533,7 +533,7 @@ struct ExtractSliceOpConversion
     // newShape = rank_reduce(shape)
     // Triton only supports static tensor sizes
     SmallVector<Value, 4> strideVals;
-    for (auto i = 0; i < op.static_sizes().size(); ++i) {
+    for (auto i = 0; i < op.getStaticSizes().size(); ++i) {
       if (op.getStaticSize(i) == 1) {
         offsetVals.erase(offsetVals.begin() + i);
       } else {

diff --git a/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVMBase.h b/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVMBase.h
@@ -130,7 +130,7 @@ struct FuncOpConversionBase : public ConvertOpToLLVMPattern<triton::FuncOp> {
     }
     auto newFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
         funcOp.getLoc(), funcOp.getName(), llvmType, linkage,
-        /*dsoLocal*/ false, LLVM::CConv::C, attributes);
+        /*dsoLocal*/ false, LLVM::CConv::C, /*comdat=*/nullptr, attributes);
     rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
                                 newFuncOp.end());
     if (failed(rewriter.convertRegionTypes(&newFuncOp.getBody(), *typeConverter,

diff --git a/lib/Dialect/TritonGPU/IR/CMakeLists.txt b/lib/Dialect/TritonGPU/IR/CMakeLists.txt
@@ -7,6 +7,6 @@ add_mlir_dialect_library(TritonGPUIR
   TritonGPUAttrDefsIncGen
 
   LINK_LIBS PUBLIC
-  MLIRGPUOps
+  MLIRGPUDialect
   TritonIR
 )
diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp
@@ -1218,10 +1218,10 @@ LogicalResult ConvertLayoutOp::canonicalize(ConvertLayoutOp op,
     auto newArg = rewriter.create<triton::gpu::ConvertLayoutOp>(
         op->getLoc(), newType, extract_slice.getSource());
     rewriter.replaceOpWithNewOp<triton::gpu::ExtractSliceOp>(
-        op, resType, newArg.getResult(), extract_slice.offsets(),
-        extract_slice.sizes(), extract_slice.strides(),
-        extract_slice.static_offsets(), extract_slice.static_sizes(),
-        extract_slice.static_strides());
+        op, resType, newArg.getResult(), extract_slice.getOffsets(),
+        extract_slice.getSizes(), extract_slice.getStrides(),
+        extract_slice.getStaticOffsets(), extract_slice.getStaticSizes(),
+        extract_slice.getStaticStrides());
     return mlir::success();
   }
 

diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
@@ -45,7 +45,7 @@ SmallVector<int64_t, 2> mmaVersionToShapePerWarp(int version) {
 SmallVector<unsigned, 2> warpsPerTileV2(triton::DotOp dotOp,
                                         const ArrayRef<int64_t> shape,
                                         int numWarps) {
-  auto filter = [&dotOp](Operation *op) {
+  mlir::TransitiveFilter filter = [&dotOp](Operation *op) {
     return op->getParentRegion() == dotOp->getParentRegion();
   };
   auto slices = mlir::getSlice(dotOp, {filter});
@@ -91,7 +91,7 @@ class BlockedToMMA : public mlir::RewritePattern {
     int finalBitWidth = getElementTypeOrSelf(x).getIntOrFloatBitWidth();
     int origBitWidth = finalBitWidth;
     SetVector<Operation *> slice;
-    mlir::getBackwardSlice(x, &slice, bwdFilter);
+    mlir::getBackwardSlice(x, &slice, {{bwdFilter}});
     Operation *firstOp = slice.empty() ? nullptr : *slice.begin();
     if (firstOp)
       if (Value arg = firstOp->getOperand(0))
@@ -137,7 +137,9 @@ class BlockedToMMA : public mlir::RewritePattern {
     triton::gpu::MmaEncodingAttr mmaEnc;
     if (versionMajor == 1) {
       SetVector<Operation *> aBwdSlices, bBwdSlices;
-      auto isCvt = [](Operation *op) { return isa<ConvertLayoutOp>(op); };
+      mlir::TransitiveFilter isCvt = [](Operation *op) {
+        return isa<ConvertLayoutOp>(op);
+      };
       getBackwardSlice(a, &aBwdSlices, {isCvt});
       getBackwardSlice(b, &bBwdSlices, {isCvt});
       // get the source of the first conversion found in slices

diff --git a/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp b/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp
@@ -354,7 +354,7 @@ class RematerializeForward : public mlir::RewritePattern {
       return failure();
 
     SetVector<Operation *> cvtSlices;
-    auto filter = [&](Operation *op) {
+    mlir::TransitiveFilter filter = [&](Operation *op) {
       return op->getBlock() == cvt->getBlock() &&
              !isa<triton::gpu::ConvertLayoutOp, scf::YieldOp>(op) &&
              !(isa<triton::ReduceOp>(op) &&

diff --git a/lib/Dialect/TritonGPU/Transforms/Utility.cpp b/lib/Dialect/TritonGPU/Transforms/Utility.cpp
@@ -93,6 +93,11 @@ bool isExpensiveLoadOrStore(Operation *op, Attribute &targetEncoding) {
   // same
   if (isSingleValue(op->getOperand(0)))
     return false;
+  // TODO(manany): Investigate with Openai why the change here
+  // https://github.com/openai/triton/commit/640f3c392184cd14291c1bca6a4795eb0f32a61a
+  // which introduces Case 2 causes breakage to this test
+  // //third_party/py/jax_triton/tests:pallas_test_sm80 --test_filter=test_fused_attention_bwd
+  return true;
   // Case 2: Tensor of pointers has more threads than elements
   // we can presume a high hit-rate that makes it cheap to load
   auto ptrType = op->getOperand(0).getType().cast<RankedTensorType>();

diff --git a/lib/Target/PTX/PTXTranslation.cpp b/lib/Target/PTX/PTXTranslation.cpp
@@ -49,7 +49,7 @@ std::string translateLLVMIRToPTX(llvm::Module &module, int cc, int version) {
   auto *shortPtr =
       static_cast<llvm::cl::opt<bool> *>(options["nvptx-short-ptr"]);
   assert(shortPtr);
-  shortPtr->setValue(true);
+  shortPtr->setValue(false);
   std::string sm = cc == 90 ? "sm_90a" : "sm_" + std::to_string(cc);
   // max PTX version
   int ptxMajor = maxPTX / 10;

diff --git a/llvm-hash.txt b/llvm-hash.txt
@@ -0,0 +1 @@
+50665511c79f62d97c0d6602e1577b0abe1b982f
diff --git a/python/MANIFEST.in b/python/MANIFEST.in
diff --git a/python/examples/copy_strided.py b/python/examples/copy_strided.py
diff --git a/python/examples/empty.py b/python/examples/empty.py
diff --git a/python/pyproject.toml b/python/pyproject.toml