[Codegen] Cleanup VectorLayoutAnalysis details (iree-org#22418)

This patch moves the implementation details of VectorLayoutAnalysis to
it's implementation file. The class needed to be exposed in an initial
implementation, but exposing these details isn't required anymore, and a
simple function call is enough.

Author: Groverkss

compiler/src/iree/compiler/Codegen/Common/BUILD.bazel

@@ -184,7 +184,6 @@ iree_compiler_cc_library(
         "TileInferenceUtils.h",
         "Transforms.h",
         "UserConfig.h",
-        "VectorLayoutAnalysis.h",
     ],
     deps = [
         ":PassHeaders",

compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt

@@ -65,7 +65,6 @@ iree_cc_library(
     "TileInferenceUtils.h"
     "Transforms.h"
     "UserConfig.h"
-    "VectorLayoutAnalysis.h"
   SRCS
     "AddFastMathFlags.cpp"
     "BlockDynamicDimensions.cpp"

compiler/src/iree/compiler/Codegen/Common/GPU/GPUNestedLayoutDistributionPatterns.cpp

@@ -7,7 +7,6 @@
 #include <cstdint>
 #include "iree/compiler/Codegen/Common/GPU/GPUPatterns.h"
 #include "iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h"
-#include "iree/compiler/Codegen/Common/VectorLayoutAnalysis.h"
 #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"

compiler/src/iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.cpp

@@ -5,7 +5,7 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include "iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h"
-#include "iree/compiler/Codegen/Common/VectorLayoutAnalysis.h"
+#include "iree/compiler/Codegen/Common/Transforms.h"
 #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
 #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
@@ -17,6 +17,8 @@
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
+#include <deque>
+
 #define DEBUG_TYPE "iree-codegen-gpu-vector-distribution"
 
 using namespace mlir::iree_compiler::IREE::VectorExt;
@@ -34,15 +36,16 @@ constexpr StringLiteral kVectorLayoutRedistributeAttrName =
 /// Set signature for the operation based on the analysis. Returns failure if
 /// an operation contains vectors that cannot be distributed i.e. they have no
 /// layout.
-LogicalResult setOpSignature(Operation *op, VectorLayoutAnalysis &analysis,
-                             const VectorLayoutOptions &options) {
+LogicalResult
+setOpSignature(Operation *op,
+               const llvm::MapVector<Value, VectorLayoutInterface> &layouts,
+               const VectorLayoutOptions &options) {
   SmallVector<Attribute> operands;
   SmallVector<Attribute> results;
 
   for (Value operand : op->getOperands()) {
     if (auto vectorOperand = dyn_cast<VectorValue>(operand)) {
-      if (auto layout =
-              analysis.getLayout<VectorLayoutInterface>(vectorOperand)) {
+      if (auto layout = layouts.lookup(vectorOperand)) {
         operands.push_back(layout);
         continue;
       }
@@ -57,8 +60,7 @@ LogicalResult setOpSignature(Operation *op, VectorLayoutAnalysis &analysis,
 
   for (Value result : op->getResults()) {
     if (auto vectorResult = dyn_cast<VectorValue>(result)) {
-      if (auto layout =
-              analysis.getLayout<VectorLayoutInterface>(vectorResult)) {
+      if (auto layout = layouts.lookup(vectorResult)) {
         results.push_back(layout);
         continue;
       }
@@ -356,17 +358,19 @@ LogicalResult distributeVectorOps(Operation *root,
                                   VectorLayoutOptions &options) {
   // Run the analysis and determine the layouts.
   LLVM_DEBUG(llvm::dbgs() << "Running Layout Analysis\n");
-  VectorLayoutAnalysis analysis(root);
-  if (failed(analysis.run()))
+  llvm::MapVector<Value, VectorLayoutInterface> layouts;
+  if (failed(propagateVectorLayoutInfo(root, layouts))) {
+    LLVM_DEBUG(llvm::dbgs() << "Layout Analysis Failed\n");
     return failure();
+  }
   LLVM_DEBUG(llvm::dbgs() << "Layout Analysis Succeded\n");
   LLVM_DEBUG(llvm::dbgs() << "\n\n");
 
   // Go to each operation, and set its distribution signature.
   LLVM_DEBUG(
       llvm::dbgs() << "Setting distribution signatures for operations\n");
   root->walk([&](Operation *op) {
-    if (failed(setOpSignature(op, analysis, options))) {
+    if (failed(setOpSignature(op, layouts, options))) {
       LLVM_DEBUG({
         llvm::dbgs() << "Skipping operation because not all vector "
                         "operands/results have a layout:\n";

compiler/src/iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h

@@ -7,7 +7,6 @@
 #ifndef IREE_COMPILER_CODEGEN_COMMON_GPU_VECTOR_DISTRIBUTION_H_
 #define IREE_COMPILER_CODEGEN_COMMON_GPU_VECTOR_DISTRIBUTION_H_
 
-#include "iree/compiler/Codegen/Common/VectorLayoutAnalysis.h"
 #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.h"
 #include "llvm/Support/Debug.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
@@ -16,6 +15,8 @@
 
 namespace mlir::iree_compiler {
 
+using IREE::VectorExt::VectorLayoutInterface;
+
 /// A signature describing the layout for each value of vector type which is
 /// an operand or result of this operation.
 ///

compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp

@@ -12,7 +12,6 @@
 #include "iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h"
 #include "iree/compiler/Codegen/Common/GPU/Passes.h"
 #include "iree/compiler/Codegen/Common/Transforms.h"
-#include "iree/compiler/Codegen/Common/VectorLayoutAnalysis.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
 #include "iree/compiler/Codegen/Dialect/GPU/Transforms/Transforms.h"

compiler/src/iree/compiler/Codegen/Common/Transforms.h

@@ -77,6 +77,76 @@ FailureOr<IREETilingResult>
 tileDispatchUsingSCFForOp(RewriterBase &rewriter, TilingInterface op,
                           linalg::LinalgTilingOptions options);
 
+namespace IREE::VectorExt {
+class VectorLayoutInterface;
+} // namespace IREE::VectorExt
+
+/// Analyzes the root op and it's nested ops to propagate vector layouts
+/// originating from to_vector operations. Example:
+///
+///    %root = vector.transfer_read
+///      |
+///      --> anchored to layout L (using a to_layout op)
+///    %root2 = vector.transfer_read
+///    %c = arith.mulf %root, %b
+///          |
+///          --> %root, %b and %c must have the same layout
+///    %e = arith.divf %b, %root2
+///          |
+///          --> %root2, %b and %e must have the same layout
+///
+/// Here, the user provided an anchor point for %root, fixing it's layout to L.
+/// The layout then uses it's inference rules to find the layout of other
+/// values:
+///
+///    %root = vector.transfer_read
+///     |
+///     --> infered to layout L
+///    %root2 = vector.transfer_read
+///     |
+///     --> infered to layout L
+///    %c = arith.mulf %root, %b
+///     |
+///     --> infered to layout L
+///    %e = arith.divf %b, %root2
+///     |
+///     --> infered to layout L
+///
+/// If at any point, a value has a layout, but the user of that value requires
+/// a different layout, the analysis inserts a resolution operation. This
+/// resolution operation is `iree_vector_ext.to_layout`.
+/// For Example:
+///
+/// %0 = vector.transfer_read
+///  |
+///  --> anchored to layout L
+/// %1 = vector.transfer_read
+///  |
+///  --> anchored to layout L'
+///  arith.addf %0, %1
+///     |
+///     --> %0 and %1 must have the same layout
+///
+/// To resolve the conflict, the analysis chooses one of the layouts, say
+/// L, and inserts a resolution operation to convert the other layout to L.
+///
+/// %0 = vector.transfer_read
+///  |
+///  --> anchored to layout L
+/// %1 = vector.transfer_read
+///  |
+///  --> anchored to layout L'
+/// %resolved = iree_vector_ext.to_layout %1
+///  |
+///  --> infered to layout L
+/// arith.addf %0, %resolved
+///
+/// The analysis itself will not try to resolve the conflict, but instead
+/// will leave it as a to_layout op, which can be rewritten by the caller.
+LogicalResult propagateVectorLayoutInfo(
+    Operation *root,
+    llvm::MapVector<Value, IREE::VectorExt::VectorLayoutInterface> &layouts);
+
 /// Transform a `scf.for` loop with a strictly positive step
 ///   for %i = %lb to %ub step %s
 /// into a 0-based loop with step 1

compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp

@@ -4,8 +4,9 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "iree/compiler/Codegen/Common/VectorLayoutAnalysis.h"
 #include "iree/compiler/Codegen/Common/Passes.h"
+#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
+#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.h"
 
 #include <cassert>
 
@@ -249,7 +250,7 @@ ChangeResult DistributionLayout::resolveWithPossibleConflict(
   // Create a new value for the resolved value and subscribe it to propagation
   // and enforcement.
   // We possibly don't need to subscribe this since this value has already
-  // reached the top of the lattice and shouldn't do anything else. But it's
+  // reached the top of the lattice and shouldn't do anything else. But its
   // nicer to do it to have consistency.
   DistributionLayout *resolvedLayout =
       propagation->getLatticeElement(resolvedValue);
@@ -296,7 +297,7 @@ ChangeResult DistributionLayout::resolve(const VectorLayoutInterface &rhs,
   }
   }
 
-  // This return will never be reached, but it's here to make the compiler
+  // This return will never be reached, but its here to make the compiler
   // happy.
   return ChangeResult::NoChange;
 }
@@ -698,7 +699,7 @@ static void enforceLayoutToMultiReductionOp(
     return;
   }
   // Reductions should always propagate value layout to result. Result can
-  // enforce it's layout on init.
+  // enforce its layout on init.
   const DistributionLayout *result = resultLattices[0];
   DistributionLayout *init = operandLattices[1];
 
@@ -1352,78 +1353,55 @@ DistributionLayout *EnforceLayout::getLatticeElement(Value val) {
   return layout;
 }
 
-/// ==========================================================================
-///        VectorLayoutAnalysis
-/// ==========================================================================
+namespace mlir::iree_compiler {
 
-LogicalResult VectorLayoutAnalysis::run() {
+LogicalResult propagateVectorLayoutInfo(
+    Operation *root, llvm::MapVector<Value, VectorLayoutInterface> &layouts) {
+  DataFlowSolver solver;
   // The order of loading matters here, because propagateLayout does anchoring
   // initialization which needs the lattice to know both enforcement and
   // propagation.
   solver.load<PropagateLayout>(root->getContext());
   solver.load<EnforceLayout>(root->getContext());
-  return solver.initializeAndRun(root);
-}
-
-VectorLayoutInterface VectorLayoutAnalysis::getLayout(Value val) {
-  const DistributionLayout *layout =
-      solver.lookupState<DistributionLayout>(val);
-  if (!layout) {
-    return VectorLayoutInterface();
+  if (failed(solver.initializeAndRun(root))) {
+    return failure();
   }
-  return layout->getLayout();
-}
-
-void VectorLayoutAnalysis::debugAnnotateLayouts() {
-  // Annotate each operation with the layout of it's result.
+  // Iterate over all values and extract their layouts.
   root->walk([&](Operation *op) {
-    if (op->getNumResults() == 0) {
-      return;
+    for (Value result : op->getResults()) {
+      const DistributionLayout *layout =
+          solver.lookupState<DistributionLayout>(result);
+      if (layout && layout->hasLayout()) {
+        layouts[result] = layout->getLayout();
+      }
     }
 
-    for (auto [index, result] : llvm::enumerate(op->getResults())) {
-      if (!isa<VectorType>(result.getType())) {
+    for (Value operand : op->getOperands()) {
+      // Some operands may not have been visited as results (e.g., block
+      // arguments).
+      if (layouts.contains(operand)) {
         continue;
       }
-
-      // Do not annotate to_layout operations since they already have
-      // this information in their attributes.
-      if (isa<IREE::VectorExt::ToLayoutOp>(op)) {
-        continue;
+      const DistributionLayout *layout =
+          solver.lookupState<DistributionLayout>(operand);
+      if (layout && layout->hasLayout()) {
+        layouts[operand] = layout->getLayout();
       }
-
-      Attribute layout = getLayout<Attribute>(result);
-      if (!layout) {
-        continue;
-      }
-
-      op->setAttr("layout_result_" + std::to_string(index), layout);
     }
   });
+  return success();
 }
 
-void VectorLayoutAnalysis::print(raw_ostream &os) {
-  debugAnnotateLayouts();
-  root->print(os);
-}
-
-void VectorLayoutAnalysis::dump() {
-  print(llvm::dbgs());
-  llvm::dbgs() << "\n";
-}
-
-namespace mlir::iree_compiler {
-
 #define GEN_PASS_DEF_TESTVECTORLAYOUTANALYSISPASS
 #include "iree/compiler/Codegen/Common/Passes.h.inc"
 
 struct TestVectorLayoutAnalysisPass final
     : impl::TestVectorLayoutAnalysisPassBase<TestVectorLayoutAnalysisPass> {
   void runOnOperation() override {
     Operation *root = getOperation();
-    VectorLayoutAnalysis analysis(getOperation());
-    if (failed(analysis.run())) {
-      root->emitError("layout analysis failed");
+    llvm::MapVector<Value, VectorLayoutInterface> layouts;
+    if (failed(propagateVectorLayoutInfo(root, layouts))) {
+      root->emitError("Layout Analysis Failed");
       return signalPassFailure();
     }
 
@@ -1433,14 +1411,9 @@ struct TestVectorLayoutAnalysisPass final
       }
 
       for (OpResult result : op->getOpResults()) {
-        if (auto layout = analysis.getLayout<Attribute>(result)) {
-          // Print layout attr to a string.
-          std::string layoutStr;
-          llvm::raw_string_ostream s(layoutStr);
-          s << layout;
-          // Emit remark.
-          op->emitRemark("layout of result #" +
-                         Twine(result.getResultNumber()) + " is " + s.str());
+        if (layouts.contains(result)) {
+          op->emitRemark("layout of result #")
+              << result.getResultNumber() << " is " << layouts[result];
         }
       }
     });