[AMD] Extended AxisAnalysis to account ExtractSliceOps

Author: ravil-mobile

lib/Analysis/AxisInfo.cpp

@@ -3,6 +3,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include "third_party/amd/include/Dialect/TritonAMDGPU/IR/Dialect.h"
 #include "triton/Analysis/AxisInfo.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
@@ -1019,6 +1020,73 @@ class MaxMinOpAxisInfoVisitor final : public AxisInfoVisitorImpl<OpTy> {
   }
 };
 
+class ExtractSliceOpInfoVisitor final
+    : public AxisInfoVisitorImpl<triton::amdgpu::ExtractSliceOp> {
+public:
+  using AxisInfoVisitorImpl<
+      triton::amdgpu::ExtractSliceOp>::AxisInfoVisitorImpl;
+
+  AxisInfo
+  getAxisInfo(triton::amdgpu::ExtractSliceOp op,
+              ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override {
+
+    auto srcType = cast<RankedTensorType>(op.getOperand().getType());
+    auto srcShape = srcType.getShape();
+
+    auto dstType = cast<RankedTensorType>(op.getResult().getType());
+    auto dstShape = dstType.getShape();
+
+    auto offsets = op.getStaticOffsets();
+
+    AxisInfo opInfo = operands[0]->getValue();
+    auto origContiguity = opInfo.getContiguity();
+    auto origDivisibility = opInfo.getDivisibility();
+    auto origConstancy = opInfo.getConstancy();
+
+    auto recompute = [](ArrayRef<int64_t> vec, int64_t c) {
+      auto result = std::numeric_limits<int64_t>::max();
+      for (auto &v : vec) {
+        // compute the upper bound of `v` based on `contiguity`
+        auto newC = ((v + c - 1) / c) * c - v;
+        // make sure that the new value is not broken because
+        // of the sliced boundaries
+        newC = newC == 0 ? c : newC;
+
+        // conside the minumal value along each dimension
+        result = result > newC ? newC : result;
+      }
+      assert(vec.size() == 2);
+      const auto dimSize = vec[1] - vec[0];
+
+      // make sure that the value doesn't exceed the dimension size
+      return result > dimSize ? dimSize : result;
+    };
+
+    SmallVector<int64_t> contiguity(origContiguity.size());
+    SmallVector<int64_t> divisibility(opInfo.getDivisibility().size());
+    SmallVector<int64_t> constancy(opInfo.getConstancy().size());
+
+    for (size_t dim = 0; dim < opInfo.getRank(); ++dim) {
+      auto start = offsets[dim];
+      auto end = start + dstShape[dim];
+      contiguity[dim] = recompute({start, end}, origContiguity[dim]);
+      // note: contiguity cannot increase while slicing a tensor
+      assert(contiguity[dim] <= origContiguity[dim]);
+
+      constancy[dim] = recompute({start, end}, origConstancy[dim]);
+
+      // note: preserve divisibility if contiguity stays the same.
+      // otherwise, set it to 1 b/c we are not able to infer this property from
+      // the given data
+      divisibility[dim] =
+          contiguity[dim] == origContiguity[dim] ? origDivisibility[dim] : 1;
+    }
+
+    return AxisInfo(contiguity, divisibility, constancy,
+                    opInfo.getConstantValue());
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // AxisInfoAnalysis
 //===----------------------------------------------------------------------===//
@@ -1062,6 +1130,7 @@ AxisInfoAnalysis::AxisInfoAnalysis(DataFlowSolver &solver)
                   MaxMinOpAxisInfoVisitor<arith::MinSIOp>,
                   MaxMinOpAxisInfoVisitor<arith::MinUIOp>>();
   visitors.append<LoadOpAxisInfoVisitor>();
+  visitors.append<ExtractSliceOpInfoVisitor>();
 }
 
 LogicalResult AxisInfoAnalysis::visitOperation(

lib/Analysis/CMakeLists.txt

@@ -17,4 +17,5 @@ add_triton_library(TritonAnalysis
   TritonIR
   TritonGPUIR
   TritonNvidiaGPUIR
+  TritonAMDGPUIR
 )