//===- TestXeGPUTransforms.cpp -- Test Vector transforms and lowerings ----===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Value.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/Support/raw_ostream.h"
#include <optional>

using namespace mlir;
using namespace mlir::xegpu;

namespace {

#define DEBUG_TYPE "test-xegpu-unroll"

struct TestXeGPUUnrollingPatterns
    : public PassWrapper<TestXeGPUUnrollingPatterns,
                         OperationPass<gpu::GPUModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPUUnrollingPatterns)

  StringRef getArgument() const final {
    return "test-xegpu-unrolling-patterns";
  }

  StringRef getDescription() const final {
    return "Test lowering patterns to unroll ops in the xegpu dialect";
  }

  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
    registry.insert<memref::MemRefDialect>();
    registry.insert<xegpu::XeGPUDialect>();
    registry.insert<vector::VectorDialect>();
  }

  TestXeGPUUnrollingPatterns() = default;
  TestXeGPUUnrollingPatterns(const TestXeGPUUnrollingPatterns &pass) = default;

  void runOnOperation() override {
    MLIRContext *ctx = &getContext();
    xegpu::UnrollOptions options;
    options.setNativeShapeFn([&](Operation *op)
                                 -> std::optional<SmallVector<int64_t>> {
      if (isa<xegpu::CreateNdDescOp, xegpu::PrefetchNdOp, xegpu::LoadNdOp,
              xegpu::StoreNdOp, xegpu::PrefetchOp, xegpu::LoadGatherOp,
              xegpu::StoreScatterOp>(op)) {
        xegpu::TensorDescType tdescTy;
        if (auto createNdOp = dyn_cast<xegpu::CreateNdDescOp>(op)) {
          tdescTy = createNdOp.getType();
        } else if (auto prefetchNdOp = dyn_cast<xegpu::PrefetchNdOp>(op)) {
          tdescTy = prefetchNdOp.getTensorDescType();
        } else if (auto loadNdOp = dyn_cast<xegpu::LoadNdOp>(op)) {
          tdescTy = loadNdOp.getTensorDescType();
        } else if (auto storeNdOp = dyn_cast<xegpu::StoreNdOp>(op)) {
          tdescTy = storeNdOp.getTensorDescType();
        } else if (isa<xegpu::PrefetchOp, xegpu::LoadGatherOp,
                       xegpu::StoreScatterOp>(op)) {
          auto anchorOp = cast<xegpu::AnchorLayoutInterface>(op);
          auto layout =
              dyn_cast_or_null<xegpu::LayoutAttr>(anchorOp.getAnchorLayout());
          if (layout && layout.isForSubgroup()) {
            auto inst_data = layout.getEffectiveInstDataAsInt();
            if (!inst_data.empty())
              return SmallVector<int64_t>(inst_data.begin(), inst_data.end());
          }
          return std::nullopt;
        }

        if (auto layout = tdescTy.getLayoutAttr()) {
          auto inst_data = layout.getEffectiveInstDataAsInt();
          if (!inst_data.empty() && layout.isForSubgroup())
            return SmallVector<int64_t>(inst_data.begin(), inst_data.end());
        }
      }

      if (isa<xegpu::DpasOp>(op))
        return SmallVector<int64_t>{8, 16, 16};

      // For vector.multi_reduction, read tile shape from the layout attribute
      // on the source operand (layout_operand_0).
      if (isa<vector::MultiDimReductionOp>(op)) {
        xegpu::DistributeLayoutAttr layout =
            xegpu::getDistributeLayoutAttr(op->getOpOperand(0));
        if (layout) {
          auto instData = layout.getEffectiveInstDataAsInt();
          if (!instData.empty())
            return instData;
        }
        return std::nullopt;
      }

      return std::nullopt;
    });

    options.setUnrolledTypesFn(
        [&](ShapedType type, ArrayRef<int64_t> tileShape,
            bool returnSingleType = false) -> SmallVector<Type> {
          Type elemTy = type.getElementType();
          Type newTy;

          // TensorDescType needs to drop the inst_data field in the layout
          // attribute
          if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(type)) {
            Attribute encoding = tdescTy.getEncoding();
            auto layout = tdescTy.getLayoutAttr();

            if (layout) {
              if (layout.getEffectiveLaneLayoutAsInt().empty())
                layout = xegpu::LayoutAttr();
              else
                layout = layout.dropInstData();
            }

            newTy = xegpu::TensorDescType::get(ctx, tileShape, elemTy, encoding,
                                               layout);

          } else {
            newTy = type.clone(tileShape, elemTy);
          }

          if (returnSingleType)
            return SmallVector<Type>{newTy};
          std::optional<SmallVector<int64_t>> ratio =
              computeShapeRatio(type.getShape(), tileShape);
          assert(ratio && "Expecting the ratio to be valid.");
          return SmallVector<Type>(computeProduct(*ratio), newTy);
        });

    RewritePatternSet patterns(ctx);

    populateXeGPUUnrollPatterns(patterns, options);
    (void)applyPatternsGreedily(getOperation(), std::move(patterns));
  }
};

#undef DEBUG_TYPE
#define DEBUG_TYPE "test-xegpu-layout-interface"

// Test pattern for distributing vector::StepOp from workgroup to subgroup.
// Validates DistributeLayoutAttr interfaces for offset computation
// abstraction between LayoutAttr and SliceAttr.
class TestStepOpPattern : public OpConversionPattern<vector::StepOp> {
  using OpConversionPattern<vector::StepOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::StepOp op, OneToNOpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {

    auto layoutName = xegpu::getTemporaryLayoutName(op->getResult(0));
    auto sliceAttr = op->getAttrOfType<xegpu::SliceAttr>(layoutName);
    if (!sliceAttr || sliceAttr.getRank() != 1)
      return failure();

    std::optional<SmallVector<int64_t>> sgShape =
        sliceAttr.getEffectiveSgDataAsInt();
    if (!sgShape)
      return failure();

    Location loc = op.getLoc();
    VectorType type = op.getResult().getType();
    auto wgShape = type.getShape();

    Value sgId =
        gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);
    auto maybeOffsets =
        sliceAttr.computeDistributedCoords(rewriter, loc, sgId, wgShape);
    if (failed(maybeOffsets))
      return failure();

    VectorType newTy = type.cloneWith(*sgShape, type.getElementType());
    Value base = vector::StepOp::create(rewriter, loc, newTy);
    SmallVector<Value> newOps;
    for (auto offsets : *maybeOffsets) {
      Value bcast =
          vector::BroadcastOp::create(rewriter, loc, newTy, offsets[0]);
      Value add = arith::AddIOp::create(rewriter, loc, base, bcast);
      newOps.push_back(add);
    }
    rewriter.replaceOpWithMultiple(op, {newOps});
    return success();
  }
};

struct TestXeGPURecoverTemporaryLayouts
    : public PassWrapper<TestXeGPURecoverTemporaryLayouts,
                         OperationPass<gpu::GPUModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPURecoverTemporaryLayouts)

  StringRef getArgument() const final {
    return "test-xegpu-recover-temporary-layouts";
  }

  StringRef getDescription() const final {
    return "Test the implementation of XeGPU temporary layout recovery";
  }

  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
    registry.insert<arith::ArithDialect>();
    registry.insert<memref::MemRefDialect>();
    registry.insert<xegpu::XeGPUDialect>();
    registry.insert<vector::VectorDialect>();
    registry.insert<gpu::GPUDialect>();
  }

  TestXeGPURecoverTemporaryLayouts() = default;
  TestXeGPURecoverTemporaryLayouts(const TestXeGPURecoverTemporaryLayouts &pass)
      : PassWrapper(pass) {}

  void runOnOperation() override {
    Operation *op = getOperation();
    if (!xegpu::recoverTemporaryLayouts(op))
      signalPassFailure();
  }
};

/// This test pass is intended to test the subgroup to lane distribution of
/// xegpu/vector/arith operations in isolation, it does not handle any
/// structural ops like scf.for etc.
struct TestXeGPUSgToLaneDistribute
    : public PassWrapper<TestXeGPUSgToLaneDistribute,
                         OperationPass<gpu::GPUModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPUSgToLaneDistribute)

  StringRef getArgument() const final {
    return "test-xegpu-sg-to-lane-distribute";
  }

  StringRef getDescription() const final {
    return "Test the implementation of XeGPU Subgroup to Lane Distribution";
  }

  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
    registry.insert<arith::ArithDialect>();
    registry.insert<memref::MemRefDialect>();
    registry.insert<xegpu::XeGPUDialect>();
    registry.insert<vector::VectorDialect>();
    registry.insert<index::IndexDialect>();
    registry.insert<gpu::GPUDialect>();
  }

  TestXeGPUSgToLaneDistribute() = default;
  TestXeGPUSgToLaneDistribute(const TestXeGPUSgToLaneDistribute &pass)
      : PassWrapper(pass) {}

  void runOnOperation() override {
    Operation *op = getOperation();
    if (!xegpu::recoverTemporaryLayouts(op)) {
      signalPassFailure();
      return;
    }

    MLIRContext *ctx = &getContext();
    TypeConverter typeConverter;
    // Define type materializations using UnrealizedConversionCastOp.
    auto materializeCast = [&](mlir::OpBuilder &builder, mlir::Type type,
                               mlir::ValueRange inputs,
                               mlir::Location loc) -> mlir::Value {
      return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
          .getResult(0);
    };
    typeConverter.addSourceMaterialization(materializeCast);
    typeConverter.addTargetMaterialization(materializeCast);

    ConversionTarget target(*ctx);
    RewritePatternSet patterns(ctx);
    xegpu::populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
        typeConverter, patterns, target);
    (void)applyPartialConversion(op, target, std::move(patterns));
  }
};

struct TestXeGPUPropagateLayouts
    : public PassWrapper<TestXeGPUPropagateLayouts,
                         OperationPass<gpu::GPUModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPUPropagateLayouts)

  StringRef getArgument() const final { return "test-xegpu-propagate-layouts"; }

  StringRef getDescription() const final {
    return "Test the implementation of XeGPU propagate layouts.";
  }

  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
    registry.insert<xegpu::XeGPUDialect>();
    registry.insert<gpu::GPUDialect>();
  }

  TestXeGPUPropagateLayouts() = default;
  TestXeGPUPropagateLayouts(const TestXeGPUPropagateLayouts &pass)
      : PassWrapper(pass) {}

  Option<std::string> layoutKind{
      *this, "layout-kind",
      llvm::cl::desc("Propagate `subgroup` / `inst` / `lane` level of xegpu "
                     "layouts."),
      llvm::cl::init("lane")};

  void runOnOperation() override {
    OpBuilder builder(getOperation());
    LayoutKind kind;
    if (layoutKind == "subgroup")
      kind = LayoutKind::Subgroup;
    else if (layoutKind == "inst")
      kind = LayoutKind::InstData;
    else if (layoutKind == "lane")
      kind = LayoutKind::Lane;
    else {
      signalPassFailure();
      return;
    }
    if (failed(xegpu::propagateLayouts(builder, getOperation(), kind, 32))) {
      signalPassFailure();
    }
  }
};

struct TestXeGPUResolveLayoutConflicts
    : public PassWrapper<TestXeGPUResolveLayoutConflicts,
                         OperationPass<gpu::GPUModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPUResolveLayoutConflicts)

  StringRef getArgument() const final {
    return "test-xegpu-resolve-layout-conflicts";
  }

  StringRef getDescription() const final {
    return "Test the implementation of XeGPU layout conflict resolution.";
  }

  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
    registry.insert<xegpu::XeGPUDialect>();
    registry.insert<gpu::GPUDialect>();
  }

  TestXeGPUResolveLayoutConflicts() = default;
  TestXeGPUResolveLayoutConflicts(const TestXeGPUResolveLayoutConflicts &pass) =
      default;

  void runOnOperation() override {
    if (failed(xegpu::resolveLayoutConflicts(getOperation())))
      signalPassFailure();
  }
};

struct TestXeGPUArrayLengthOptimization
    : public PassWrapper<TestXeGPUArrayLengthOptimization,
                         OperationPass<gpu::GPUModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPUArrayLengthOptimization)

  StringRef getArgument() const final {
    return "test-xegpu-array-length-optimization";
  }

  StringRef getDescription() const final {
    return "Test XeGPU 2D block array load optimization patterns in isolation";
  }

  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
    registry.insert<xegpu::XeGPUDialect>();
    registry.insert<vector::VectorDialect>();
  }

  TestXeGPUArrayLengthOptimization() = default;
  TestXeGPUArrayLengthOptimization(const TestXeGPUArrayLengthOptimization &pass)
      : PassWrapper(pass) {}

  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    xegpu::populateXeGPUArrayLengthOptimizationPatterns(patterns);
    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
      signalPassFailure();
  }
};

struct TestXeGPULayoutInterface
    : public PassWrapper<TestXeGPULayoutInterface,
                         OperationPass<gpu::GPUModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPULayoutInterface)

  StringRef getArgument() const final { return "test-xegpu-layout-interface"; }

  StringRef getDescription() const final {
    return "Test the implementation of XeGPU Layout interfaces";
  }

  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
    registry.insert<arith::ArithDialect>();
    registry.insert<memref::MemRefDialect>();
    registry.insert<xegpu::XeGPUDialect>();
    registry.insert<vector::VectorDialect>();
    registry.insert<index::IndexDialect>();
  }

  TestXeGPULayoutInterface() = default;
  TestXeGPULayoutInterface(const TestXeGPULayoutInterface &pass) = default;

  void runOnOperation() override {
    MLIRContext *ctx = &getContext();

    TypeConverter typeConverter;
    auto materializeCast = [&](mlir::OpBuilder &builder, mlir::Type type,
                               mlir::ValueRange inputs,
                               mlir::Location loc) -> mlir::Value {
      return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
          .getResult(0);
    };
    typeConverter.addSourceMaterialization(materializeCast);
    typeConverter.addTargetMaterialization(materializeCast);

    RewritePatternSet patterns(ctx);
    patterns.add<TestStepOpPattern>(typeConverter, ctx);

    ConversionTarget target(*ctx);
    auto isLegal = [&](xegpu::SliceAttr layout) -> bool {
      return !layout || !layout.isForWorkgroup();
    };

    target.addDynamicallyLegalOp<vector::StepOp>(
        [&](vector::StepOp op) -> bool {
          auto layoutName = xegpu::getTemporaryLayoutName(op->getResult(0));
          auto sliceAttr = op->getAttrOfType<xegpu::SliceAttr>(layoutName);
          return isLegal(sliceAttr);
        });

    target.markUnknownOpDynamicallyLegal([](Operation *op) { return true; });

    (void)applyPartialConversion(getOperation(), target, std::move(patterns));
  }
};

} // namespace

namespace mlir {
namespace test {
void registerTestXeGPULowerings() {
  PassRegistration<TestXeGPUUnrollingPatterns>();
  PassRegistration<TestXeGPULayoutInterface>();
  PassRegistration<TestXeGPURecoverTemporaryLayouts>();
  PassRegistration<TestXeGPUSgToLaneDistribute>();
  PassRegistration<TestXeGPUPropagateLayouts>();
  PassRegistration<TestXeGPUResolveLayoutConflicts>();
  PassRegistration<TestXeGPUArrayLengthOptimization>();
}
} // namespace test
} // namespace mlir