[BACKEND][AMD] Disable linear layout due to perf regression (#4126)

We have identified a 20% perf regression in our downstream flash
attention perf kernel after switching to linear layout. Initial analysis
shows register pressure is increased to cause spills. Further analysis
is still ongoing.

So this commit introduces a minimal way to selectively disable linear
layout only on AMD backend to avoid affecting NVIDIA backend while
continuing bring it up on AMD side.

include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h

@@ -56,6 +56,11 @@ class TargetInfoBase {
                           StringRef message, StringRef file, StringRef func,
                           int line) const = 0;
 
+  // Whether to enable linear layout. This is a per-backend temporary escape
+  // hatch to disable linear layout while figuring out issues. Eventually we
+  // want to enable linear layout everywhere and delete this control.
+  virtual bool enableLinearLayout() const { return true; }
+
   virtual ~TargetInfoBase() {}
 };
 } // namespace mlir::triton

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -1219,7 +1219,7 @@ emitIndices(Location loc, RewriterBase &rewriter, const TargetInfoBase &target,
             bool allowLL = true) {
   // Eventually the LinearLayout path will be the only one.  For now we allow
   // both paths so we can test that they produce the same results.
-  if (allowLL) {
+  if (allowLL && target.enableLinearLayout()) {
     std::optional<SmallVector<SmallVector<Value>>> llOffsets =
         emitIndicesUsingLinearLayouts(loc, rewriter, target, layout, type,
                                       withCTAOffset);

third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.h

@@ -58,6 +58,8 @@ class TargetInfo : public mlir::triton::TargetInfoBase {
                   StringRef message, StringRef file, StringRef func,
                   int line) const override;
 
+  bool enableLinearLayout() const override { return false; }
+
 private:
   void printfImpl(Value formatStrStart, int formatStrByteCount, ValueRange args,
                   ConversionPatternRewriter &rewriter, bool useStdErr) const;