// Copyright 2025 The IREE Authors
//
// Licensed 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

//===- Transforms.h - Transformations for the IREE PCF dialect ------------===//
//
// Defines transformations that apply to IREE PCF ops for use in multiple
// places.
//
//===----------------------------------------------------------------------===//
#ifndef IREE_COMPILER_CODEGEN_DIALECT_PCF_TRANSFORMS_TRANSFORMS_H_
#define IREE_COMPILER_CODEGEN_DIALECT_PCF_TRANSFORMS_TRANSFORMS_H_

#include "iree/compiler/Codegen/Dialect/PCF/IR/PCFOps.h"
#include "mlir/IR/PatternMatch.h"

// Forward declares.
namespace mlir {
class TilingInterface;
} // namespace mlir
namespace mlir::scf {
class ForallOp;
} // namespace mlir::scf
namespace mlir::tensor {
class CollapseShapeOp;
class ExtractSliceOp;
} // namespace mlir::tensor

namespace mlir::iree_compiler::IREE::PCF {

/// Converts scf.forall ops to pcf.loop by linearizing/delinearizing ids beyond
/// |numIds| into the slowest varying id. Uses DeviceMappingAttrInterface to
/// infer the order of ids from slowest to fastest varying. If |numIds| <= 0,
/// then no linearization/delinearization is done. Returns the newly created
/// pcf.loop; callers are responsible for replacing or erasing the original
/// scf.forall op as needed.
FailureOr<PCF::LoopOp> convertForallToPCFLoop(RewriterBase &rewriter,
                                              scf::ForallOp forallOp,
                                              PCF::ScopeAttrInterface scope,
                                              int64_t numIds = -1);

/// Converts an scf.forall operation to a nest of pcf.generic operations with
/// an inner scf.forall handling spillover iterations.
///
/// Each scope in |scopes| generates one pcf.generic op (outermost first).
/// Worker IDs are linearized, chunk bounds computed for work distribution,
/// and delinearized back to multi-dimensional forall bounds. Uneven chunks
/// are distributed greedily, with only the last worker lexicographically
/// getting fewer loop iterations. Returns the newly created outermost
/// pcf.generic; callers are responsible for replacing or erasing the original
/// scf.forall op as needed.
///
FailureOr<PCF::GenericOp>
convertForallToGenericNest(RewriterBase &rewriter, scf::ForallOp forallOp,
                           ArrayRef<PCF::ScopeAttrInterface> scopes);

struct ConsumerFusionParams {
  // List of operands in the consumer that are fused along.
  SmallVector<unsigned> operands;
  // List of results of the producer that are fused along.
  SetVector<unsigned> results;
  // List of slices that produce the results. This has two possible
  // interpretations. If |results| > 1, then |slices| == |operands| must hold
  // and each slice corresponds to the sole write to the result consumed by the
  // corresponding operand. If |results| == 1, then each slice is a different
  // write to the same sole result.
  //
  // In the first case, slices[0] is the most dominant slice (and thus the
  // insertion point for the fused op).
  SmallVector<PCF::WriteSliceOp> slices;
};

// Helpers to match a consumer as fusable into a producer. There are two
// supported cases:
//   1. The tilable |target| only consumes a single result of the producer but
//      the produced operand may be constructed out of multiple writes within
//      the producer.
//   1. The tilable |target| consumes multiple results of the producer but only
//      a single writing op constructs each consumed result.
// Populates |params| with the matched information needed to perform a fusion
// upon success. On failure |params| may be partially populated. It is the
// caller's responsibility to pass an empty struct on subsequent matches.
//
// Note that currently multiple producers split across block/region boundaries
// is unsupported. We need to guarantee the existence of a point in
// the control flow of the IR where the fused op is guaranteed to
// produce its original results in their entirety.
//
// For example:
//
//   pcf.write_slice
//   scf.if {
//     %0 = ...
//     pcf.write_slice ...[%0]
//
// This IR is problematic because the most dominated `write_slice` is
// the one inside the `if`, however if we put the fused op there then
// there is no guarantee we actually produce the full original result
// in the aggregate since some of the writes will be masked off.
LogicalResult matchTilableConsumer(RewriterBase &rewriter,
                                   PCF::GenericOp genericOp,
                                   TilingInterface target,
                                   ConsumerFusionParams &params);
LogicalResult matchTilableConsumer(RewriterBase &rewriter, PCF::LoopOp loopOp,
                                   TilingInterface target,
                                   ConsumerFusionParams &params);

void fuseTilableConsumer(RewriterBase &rewriter, PCF::GenericOp genericOp,
                         TilingInterface target, ConsumerFusionParams &params);
void fuseTilableConsumer(RewriterBase &rewriter, PCF::LoopOp loopOp,
                         TilingInterface target, ConsumerFusionParams &params);

struct ProducerFusionParams {
  // Which result index of the scoped op has a fusable producer init.
  unsigned resultIndex;
  // The producer op to fuse (implements TilingInterface + DPS). Stored as
  // Operation* to avoid requiring the full TilingInterface definition here.
  Operation *producer;
  // Read sites on the corresponding sref arg that consume the init value.
  SmallVector<PCF::ReadSliceOp> readSlices;
};

// Helpers to match a DPS producer as fusable into a pcf.generic/loop op
// through its tied init argument. The producer must:
//   1. Implement TilingInterface and DestinationStyleOpInterface.
//   2. Have a single result.
//   3. Have all operands dominating the generic/loop op.
//   4. Feed into a tied init of the generic/loop op whose corresponding sref
//      arg has only read_slice and write_slice users.
//
// On success, |params| is populated with the result index, producer op, and
// list of read_slice ops to replace with tiled producer computations.
LogicalResult matchTilableProducer(RewriterBase &rewriter,
                                   PCF::GenericOp genericOp,
                                   ProducerFusionParams &params);
LogicalResult matchTilableProducer(RewriterBase &rewriter, PCF::LoopOp loopOp,
                                   ProducerFusionParams &params);

// Fuses the matched producer into the pcf.generic/loop op by:
//   1. Replacing the tied init with the producer's DPS init.
//   2. Generating tiled producer computations at each read_slice site via
//      TilingInterface::generateResultTileValue.
//   3. Erasing the original producer if it has no remaining uses.
void fuseTilableProducer(RewriterBase &rewriter, PCF::GenericOp genericOp,
                         const ProducerFusionParams &params);
void fuseTilableProducer(RewriterBase &rewriter, PCF::LoopOp loopOp,
                         const ProducerFusionParams &params);

// Pattern set for dropping unused results from scoped ops. Due to memory
// effects this requires cascading operation erasure and is unsuitable for
// a canonicalization pattern.
void populatePCFDropUnusedResultPatterns(RewritePatternSet &patterns);

// Fuse a tensor.extract_slice consumer into a pcf.loop producer. This shrinks
// the result to the extracted size and clamps all write_slice ops accordingly.
LogicalResult
fuseExtractSliceIntoProducerLoop(RewriterBase &rewriter, PCF::LoopOp loopOp,
                                 tensor::ExtractSliceOp extractSliceOp);

// Fuse a tensor.extract_slice consumer into a pcf.generic producer. This
// shrinks the result to the extracted size and clamps all write_slice ops
// accordingly.
LogicalResult
fuseExtractSliceIntoProducerGeneric(RewriterBase &rewriter,
                                    PCF::GenericOp genericOp,
                                    tensor::ExtractSliceOp extractSliceOp);

// Fuse a tensor.collapse_shape consumer into a pcf.loop producer. This
// changes the result to the collapsed shape and linearizes all write_slice
// offsets/sizes accordingly. Supports both static and dynamic shapes.
//
// For each reassociation group, the innermost dimension is always kept as a
// contiguous write. In the static case, consecutive inner dimensions that
// fully cover the producer shape are folded into the contiguous chunk. Any
// remaining outer dimensions are emitted as unmapped scf.forall loops.
LogicalResult
fuseCollapseShapeIntoProducerLoop(RewriterBase &rewriter, PCF::LoopOp loopOp,
                                  tensor::CollapseShapeOp collapseOp);

// Fuse a tensor.collapse_shape consumer into a pcf.generic producer. This
// changes the result to the collapsed shape and linearizes all write_slice
// offsets/sizes accordingly. Supports both static and dynamic shapes.
//
// For each reassociation group, the innermost dimension is always kept as a
// contiguous write. In the static case, consecutive inner dimensions that
// fully cover the producer shape are folded into the contiguous chunk. Any
// remaining outer dimensions are emitted as unmapped scf.forall loops.
LogicalResult
fuseCollapseShapeIntoProducerGeneric(RewriterBase &rewriter,
                                     PCF::GenericOp genericOp,
                                     tensor::CollapseShapeOp collapseOp);

// Composes two nested slice parameter sets so the resulting slice addresses the
// outer base directly. The outer slice describes how an intermediate value is
// embedded in the final destination, and the inner slice describes how a value
// is embedded in that intermediate.
//
// For each dimension present in both slices:
//   offsets = outerOffset + innerOffset * outerStride
//   sizes = innerSizes
//   strides = outerStride * innerStride
//
// Any remaining outer dimensions are forwarded unchanged.
void composeNestedSliceParameters(
    RewriterBase &rewriter, Location loc, ArrayRef<OpFoldResult> outerOffsets,
    ArrayRef<OpFoldResult> outerSizes, ArrayRef<OpFoldResult> outerStrides,
    ArrayRef<OpFoldResult> innerOffsets, ArrayRef<OpFoldResult> innerSizes,
    ArrayRef<OpFoldResult> innerStrides,
    SmallVectorImpl<OpFoldResult> &composedOffsets,
    SmallVectorImpl<OpFoldResult> &composedSizes,
    SmallVectorImpl<OpFoldResult> &composedStrides);

// Composes a pcf.write_slice with a tensor.parallel_insert_slice from an
// scf.forall terminator. The write_slice's destination must be produced by the
// forall op, and the parallel_insert_slice must be inserting into that result.
// Returns the newly created write_slice op on success.
FailureOr<PCF::WriteSliceOp>
composeWriteSliceWithParallelInsert(RewriterBase &rewriter,
                                    PCF::WriteSliceOp writeSliceOp);

/// Folds an scf.forall containing a pcf.loop into a single pcf.generic.
///
/// Validates structural requirements:
/// - All ops in scf.forall.in_parallel are tensor.parallel_insert_slice.
/// - All insert sources come from the same pcf.loop result.
/// - All insert destinations are scf.forall shared_outs.
/// - The pcf.loop is the last op before terminator in forall body.
/// - The pcf.loop has single count argument (linearized).
/// - All pcf.loop region ref arg users are pcf.write_slice ops.
/// - Ref args have SyncOnReturnAttr sync scope.
///
/// Does NOT validate mapping or scope attributes (caller's responsibility).
/// Does NOT create workgroup_count_hint (caller's responsibility).
///
/// On success, the forall and inner loop are replaced with a pcf.generic.
FailureOr<PCF::GenericOp> foldForallIntoPCFLoop(RewriterBase &rewriter,
                                                scf::ForallOp forallOp);

} // namespace mlir::iree_compiler::IREE::PCF

#endif // IREE_COMPILER_CODEGEN_DIALECT_PCF_TRANSFORMS_TRANSFORMS_H_
