AffineHelper.h

Go to the documentation of this file.

//===-- AffineHelper.h ------------------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "llzk/Util/Compare.h"
 
#include <mlir/IR/Builders.h>
#include <mlir/IR/OpImplementation.h>
#include <mlir/IR/Operation.h>
#include <mlir/IR/ValueRange.h>
 
#include <llvm/ADT/SmallVector.h>

namespace llzk::affineMapHelpers {

mlir::ParseResult parseDimAndSymbolList(
    mlir::OpAsmParser &parser,
    mlir::SmallVectorImpl<mlir::OpAsmParser::UnresolvedOperand> &mapOperands,
    mlir::IntegerAttr &numDims
);

void printDimAndSymbolList(
    mlir::OpAsmPrinter &printer, mlir::Operation *op, mlir::OperandRange mapOperands,
    mlir::IntegerAttr numDims
);

mlir::ParseResult parseMultiDimAndSymbolList(
    mlir::OpAsmParser &parser,
    mlir::SmallVectorImpl<mlir::SmallVector<mlir::OpAsmParser::UnresolvedOperand>>
        &multiMapOperands,
    mlir::DenseI32ArrayAttr &numDimsPerMap
);

void printMultiDimAndSymbolList(
    mlir::OpAsmPrinter &printer, mlir::Operation *op, mlir::OperandRangeRange multiMapOperands,
    mlir::DenseI32ArrayAttr numDimsPerMap
);

mlir::ParseResult parseAttrDictWithWarnings(
    mlir::OpAsmParser &parser, mlir::NamedAttrList &extraAttrs, mlir::OperationState &state
);
 
template <typename ConcreteOp>
inline void printAttrDictWithWarnings(
    mlir::OpAsmPrinter &printer, ConcreteOp /*op*/, mlir::DictionaryAttr extraAttrs,
    typename ConcreteOp::Properties /*state*/
) {
  printer.printOptionalAttrDict(extraAttrs.getValue(), ConcreteOp::getAttributeNames());
}

mlir::LogicalResult verifySizesForMultiAffineOps(
    mlir::Operation *op, int32_t segmentSize, mlir::ArrayRef<int32_t> mapOpGroupSizes,
    mlir::OperandRangeRange mapOperands, mlir::ArrayRef<int32_t> numDimsPerMap
);

mlir::LogicalResult verifyAffineMapInstantiations(
    mlir::OperandRangeRange mapOps, mlir::ArrayRef<int32_t> numDimsPerMap,
    mlir::ArrayRef<mlir::AffineMapAttr> mapAttrs, mlir::Operation *origin
);

template <typename OpClass>
inline typename OpClass::Properties &buildInstantiationAttrs(
    mlir::OpBuilder &odsBuilder, mlir::OperationState &odsState,
    mlir::ArrayRef<mlir::ValueRange> mapOperands, mlir::DenseI32ArrayAttr numDimsPerMap,
    int32_t firstSegmentSize = 0
) {
  int32_t mapOpsSegmentSize = 0;
  mlir::SmallVector<int32_t> rangeSegments;
  for (mlir::ValueRange r : mapOperands) {
    odsState.addOperands(r);
    int32_t s = llzk::checkedCast<int32_t>(r.size());
    rangeSegments.push_back(s);
    mapOpsSegmentSize += s;
  }
  typename OpClass::Properties &props = odsState.getOrAddProperties<typename OpClass::Properties>();
  props.setMapOpGroupSizes(odsBuilder.getDenseI32ArrayAttr(rangeSegments));
  props.setOperandSegmentSizes({firstSegmentSize, mapOpsSegmentSize});
  if (numDimsPerMap) {
    props.setNumDimsPerMap(numDimsPerMap);
  }
  return props;
}

template <typename OpClass>
inline void buildInstantiationAttrsNoSegments(
    mlir::OpBuilder &odsBuilder, mlir::OperationState &odsState,
    mlir::ArrayRef<mlir::ValueRange> mapOperands, mlir::DenseI32ArrayAttr numDimsPerMap
) {
  mlir::SmallVector<int32_t> rangeSegments;
  for (mlir::ValueRange r : mapOperands) {
    odsState.addOperands(r);
    int32_t s = llzk::checkedCast<int32_t>(r.size());
    rangeSegments.push_back(s);
  }
  typename OpClass::Properties &props = odsState.getOrAddProperties<typename OpClass::Properties>();
  props.setMapOpGroupSizes(odsBuilder.getDenseI32ArrayAttr(rangeSegments));
  if (numDimsPerMap) {
    props.setNumDimsPerMap(numDimsPerMap);
  }
}

template <typename OpClass>
inline typename OpClass::Properties &buildInstantiationAttrsEmpty(
    mlir::OpBuilder &odsBuilder, mlir::OperationState &odsState, int32_t firstSegmentSize = 0
) {
  typename OpClass::Properties &props = odsState.getOrAddProperties<typename OpClass::Properties>();
  // `operandSegmentSizes` = [ firstSegmentSize, mapOperands.size ]
  props.setOperandSegmentSizes({firstSegmentSize, 0});
  // There are no affine map operands so initialize the related properties as empty arrays.
  props.setMapOpGroupSizes(odsBuilder.getDenseI32ArrayAttr({}));
  props.setNumDimsPerMap(odsBuilder.getDenseI32ArrayAttr({}));
  return props;
}

template <typename OpClass>
inline typename OpClass::Properties &buildInstantiationAttrsEmptyNoSegments(
    mlir::OpBuilder &odsBuilder, mlir::OperationState &odsState
) {
  typename OpClass::Properties &props = odsState.getOrAddProperties<typename OpClass::Properties>();
  // There are no affine map operands so initialize the related properties as empty arrays.
  props.setMapOpGroupSizes(odsBuilder.getDenseI32ArrayAttr({}));
  props.setNumDimsPerMap(odsBuilder.getDenseI32ArrayAttr({}));
  return props;
}
 
} // namespace llzk::affineMapHelpers