[AMD] disable pointer-canonicalization for large-tensor (#8359)

This commit disables pointer-canonicalization for pointer pointing to large
tensors. The large tensors refers to JIT specialization for those tensor
argument over 2GB.

It is disabled on the ground that is has some tricky bugs. We are trying
to come up a better approach that address several conflicting
performance aspects.

---------

Co-authored-by: Shuxin Yang <[email protected]>

Author: yangshuxin

test/TritonGPU/amd/amd-canonicalize-pointers-dont-run-mlir-canonicalizer.mlir

@@ -1,6 +1,6 @@
 // NOTE: Assertions have been autogenerated by mlir/utils/generate-test-checks.py
 
-// RUN: triton-opt %s -allow-unregistered-dialect -split-input-file -tritonamdgpu-canonicalize-pointers -verify-diagnostics | FileCheck %s
+// RUN: triton-opt %s -allow-unregistered-dialect -split-input-file -tritonamdgpu-canonicalize-pointers="enable-large-tensor-ptr-canon=true" -verify-diagnostics | FileCheck %s
 
 module attributes {"ttg.num-warps" = 4 : i32} {
   tt.func @ifOpTwoYields(%arg0: !tt.ptr<f32>, %arg1: tensor<1024xf32>, %arg2: i1) -> (tensor<1024xf32>, tensor<1024xf32>) {

test/TritonGPU/amd/amd-canonicalize-pointers-no-large-tensor.mlir

@@ -0,0 +1,20 @@
+// RUN: triton-opt %s -allow-unregistered-dialect -split-input-file -tritonamdgpu-canonicalize-pointers="enable-large-tensor-ptr-canon=false" -canonicalize -verify-diagnostics | FileCheck %s
+
+// this case is copied from amd-canonicalize-pointers-no-large-tensor.mlir. With
+// enable-large-tensor-ptr-canon=false, the input is not changed at all.
+module attributes {"ttg.num-warps" = 4 : i32} {
+  tt.func @conversion1(%arg0: !tt.ptr<f32>) -> tensor<1024xf32> {
+    %c1024_i32 = arith.constant 1024 : i32
+    %0 = tt.get_program_id x : i32
+    %1 = arith.muli %0, %c1024_i32 : i32
+    %2 = tt.splat %1 : i32 -> tensor<1024xi32>
+    %3 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>>
+    %4 = tt.addptr %3, %2 : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32>
+    %5 = tt.load %4 : tensor<1024x!tt.ptr<f32>>
+    tt.return %5 : tensor<1024xf32>
+  }
+}
+
+// CHECK-LABEL:   tt.func @conversion1
+// CHECK: %[[ADDPTR:.*]] = tt.addptr
+// CHECK:                = tt.load %[[ADDPTR]]

test/TritonGPU/amd/amd-canonicalize-pointers.mlir

@@ -1,6 +1,6 @@
 // NOTE: Assertions have been autogenerated by mlir/utils/generate-test-checks.py
 
-// RUN: triton-opt %s -allow-unregistered-dialect -split-input-file -tritonamdgpu-canonicalize-pointers -canonicalize -verify-diagnostics | FileCheck %s
+// RUN: triton-opt %s -allow-unregistered-dialect -split-input-file -tritonamdgpu-canonicalize-pointers="enable-large-tensor-ptr-canon=true" -canonicalize -verify-diagnostics | FileCheck %s
 
 module attributes {"ttg.num-warps" = 4 : i32} {
   tt.func @conversion1(%arg0: !tt.ptr<f32>) -> tensor<1024xf32> {

third_party/amd/include/TritonAMDGPUTransforms/Passes.td

@@ -123,6 +123,11 @@ def TritonAMDGPUCanonicalizePointers : Pass<"tritonamdgpu-canonicalize-pointers"
 
   let dependentDialects = [];
 
+  let options = [
+    Option<"enableLargeTensorPtrCanon", "enable-large-tensor-ptr-canon",
+           "bool", /*default=*/"false",
+           "Whether to enable canonicalization for pointers pointing to large-tensors (a specialization for tensors over 2GB)">
+  ];
 }
 
 def TritonAMDGPUReorderInstructions: Pass<"tritonamdgpu-reorder-instructions", "mlir::ModuleOp"> {

third_party/amd/lib/TritonAMDGPUTransforms/CanonicalizePointers.cpp

@@ -1654,8 +1654,10 @@ static const std::string kInitFuncArgsRewritten =
 /// (ConvertUnimplementedOpUnrealizedCasts) if it wasn't DCEd (via a user
 /// extracting the tt.ptr and c0 operands).
 struct InitFuncPtrArgs : OpRewritePattern<tt::FuncOp> {
-  InitFuncPtrArgs(MLIRContext *context, FatPointers &fatPtrs)
-      : OpRewritePattern(context, 0), fatPtrs(fatPtrs) {}
+  InitFuncPtrArgs(MLIRContext *context, FatPointers &fatPtrs,
+                  bool enableLargeTensorPtrCanon_)
+      : OpRewritePattern(context, 0), fatPtrs(fatPtrs),
+        enableLargeTensorPtrCanon(enableLargeTensorPtrCanon_) {}
 
   LogicalResult matchAndRewrite(tt::FuncOp newOp,
                                 PatternRewriter &rewriter) const override {
@@ -1673,7 +1675,11 @@ struct InitFuncPtrArgs : OpRewritePattern<tt::FuncOp> {
               newOp.getArgAttrOfType<IntegerAttr>(idx, "tt.pointer_range"))
         bitness = pointerRangeAttr.getInt();
 
-      LDBG(idx << "-th argument: " << arg << ", bitness: " << bitness << "\n");
+      LDBG(idx << "-th argument: " << arg << ", bitness: " << bitness);
+      if (!enableLargeTensorPtrCanon && (bitness == 64)) {
+        LDBG("Do not init argument of large-tensor pointer: " << arg);
+        continue;
+      }
 
       Value zeroOffset =
           rewriter.create<arith::ConstantIntOp>(newOp.getLoc(), 0, bitness);
@@ -1690,6 +1696,7 @@ struct InitFuncPtrArgs : OpRewritePattern<tt::FuncOp> {
   }
 
   FatPointers &fatPtrs;
+  bool enableLargeTensorPtrCanon;
 };
 
 /// No-op to make conversion framework happy.
@@ -1816,6 +1823,8 @@ class ConvertUnimplementedOpUnrealizedCasts
 class TritonAMDGPUCanonicalizePointersPass
     : public impl::TritonAMDGPUCanonicalizePointersBase<
           TritonAMDGPUCanonicalizePointersPass> {
+  using Base::Base;
+
 public:
   void runOnOperation() override;
 };
@@ -1905,18 +1914,29 @@ void TritonAMDGPUCanonicalizePointersPass::runOnOperation() {
   FatPointers fatPrs;
   PatternRewriter rewriter(&getContext());
   // Convert tt.func; %1 = unrealize_cast(%arg0: tt.ptr, c0: i32) -> tt.ptr
-  InitFuncPtrArgs pat(&getContext(), fatPrs);
+  InitFuncPtrArgs pat(&getContext(), fatPrs, enableLargeTensorPtrCanon);
   if (failed(pat.matchAndRewrite(func, rewriter)))
     return signalPassFailure();
 
   llvm::SetVector<Operation *> opsToRewrite;
-  for (auto arg : func.getArguments()) {
-    if (llvm::isa<tt::PointerType>(arg.getType())) {
-      // NB: reusing the same SetVector invalidates the topo order implied by
-      // getForwardSlice
-      for (auto &use : arg.getUses())
-        getForwardSliceImpl(&use, use.getOwner(), &opsToRewrite);
+  for (auto [idx, arg] : llvm::enumerate(func.getArguments())) {
+    if (!llvm::isa<tt::PointerType>(arg.getType()))
+      continue;
+
+    int64_t bitness = 64;
+    if (auto pointerRangeAttr =
+            func.getArgAttrOfType<IntegerAttr>(idx, "tt.pointer_range"))
+      bitness = pointerRangeAttr.getInt();
+
+    if (!enableLargeTensorPtrCanon && (bitness == 64)) {
+      LDBG("ignore " << idx << "-th argument of large-tensor ptr: " << arg);
+      continue;
     }
+
+    // NB: reusing the same SetVector invalidates the topo order implied by
+    // getForwardSlice
+    for (auto &use : arg.getUses())
+      getForwardSliceImpl(&use, use.getOwner(), &opsToRewrite);
   }
 
   ConversionConfig config;