[GlobalISel] An optimizing legal MIR builder

llvm/include/llvm/CodeGen/GlobalISel/OptMIRBuilder.h

@@ -0,0 +1,83 @@
+//===-- llvm/CodeGen/GlobalISel/OptMIRBuilder.h  --*- C++ -*-==---------------//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+/// \file This file implements a legal version of MachineIRBuilder which
+/// optimizes insts while building.
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_CODEGEN_GLOBALISEL_OPTMIRBUILDER_H
+#define LLVM_CODEGEN_GLOBALISEL_OPTMIRBUILDER_H
+
+#include "llvm/CodeGen/GlobalISel/CSEMIRBuilder.h"
+
+namespace llvm {
+
+class LegalizerInfo;
+struct LegalityQuery;
+
+/// OptMIRBuilder optimizes instructions while building them. It
+/// checks its operands whether they are constant or
+/// G_IMPLICIT_DEF. It never checks whether an operand is defined by
+/// G_FSINCOS. It checks operands and registers for G_IMPLICIT_DEF,
+/// G_CONSTANT, G_BUILD_VECTOR, and G_SPLAT_VECTOR and nothing else.
+/// Based on G_IMPLICIT_DEF, the constants and their values, it folds
+/// instructions into constants, G_IMPLICIT_DEF, or other
+/// instructions. For optmizations and constant folding it relies on
+/// GIConstant. It can fold G_MUL into G_ADD and G_SUB. Before folding
+/// it always queries the legalizer. When it fails to fold, it
+/// delegates the building to the CSEMIRBuilder. It is the user's
+/// responsibility to only attempt to build legal instructions pass
+/// the legalizer. OptMIRBuilder can safely be used in optimization
+/// passes pass the legalizer.
+class OptMIRBuilder : public CSEMIRBuilder {
+  const LegalizerInfo *LI;
+  const bool IsPrelegalize;
+
+  /// Legality tests.
+  bool isPrelegalize() const;
+  bool isLegal(const LegalityQuery &Query) const;
+  bool isConstantLegal(LLT);
+  bool isLegalOrBeforeLegalizer(const LegalityQuery &Query) const;
+  bool isConstantLegalOrBeforeLegalizer(LLT);
+
+  /// Returns true if the register \p R is defined by G_IMPLICIT_DEF.
+  bool isUndef(Register R) const;
+
+  /// Builds 0 - X.
+  MachineInstrBuilder buildNegation(const DstOp &, const SrcOp &);
+
+  // Constants.
+  MachineInstrBuilder buildGIConstant(const DstOp &DstOp, const GIConstant &);
+
+  /// Integer.
+  MachineInstrBuilder optimizeAdd(unsigned Opc, ArrayRef<DstOp> DstOps,
+                                  ArrayRef<SrcOp> SrcOps,
+                                  std::optional<unsigned> Flag = std::nullopt);
+  MachineInstrBuilder optimizeSub(unsigned Opc, ArrayRef<DstOp> DstOps,
+                                  ArrayRef<SrcOp> SrcOps,
+                                  std::optional<unsigned> Flag = std::nullopt);
+  MachineInstrBuilder optimizeMul(unsigned Opc, ArrayRef<DstOp> DstOps,
+                                  ArrayRef<SrcOp> SrcOps,
+                                  std::optional<unsigned> Flag = std::nullopt);
+
+public:
+  OptMIRBuilder(MachineFunction &MF, GISelCSEInfo *CSEInfo,
+                GISelChangeObserver &Observer, const LegalizerInfo *LI,
+                bool IsPrelegalize)
+      : LI(LI), IsPrelegalize(IsPrelegalize) {
+    setMF(MF);
+    setCSEInfo(CSEInfo);
+    setChangeObserver(Observer);
+  };
+
+  MachineInstrBuilder
+  buildInstr(unsigned Opc, ArrayRef<DstOp> DstOps, ArrayRef<SrcOp> SrcOps,
+             std::optional<unsigned> Flag = std::nullopt) override;
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_GLOBALISEL_OPTMIRBUILDER_H

llvm/include/llvm/CodeGen/GlobalISel/Utils.h

@@ -635,8 +635,28 @@ class GIConstant {
   /// Returns the value, if this constant is a scalar.
   APInt getScalarValue() const;
 
+  /// Returns the value, if this constant is a scalable vector.
+  APInt getSplatValue() const;
+
+  /// Returns the values, if this constant is a fixed vector.
+  ArrayRef<APInt> getAsArrayRef() const;
+
   static std::optional<GIConstant> getConstant(Register Const,
                                                const MachineRegisterInfo &MRI);
+
+  /// Returns a new constant where add(this, x) was applied.
+  GIConstant add(const GIConstant &) const;
+
+  /// Returns a new constant where sub(this, x) was applied.
+  GIConstant sub(const GIConstant &) const;
+
+  /// Returns a new constant where mul(this, x) was applied.
+  GIConstant mul(const GIConstant &) const;
+
+  bool isZero() const;
+  bool isOne() const;
+  bool isTwo() const;
+  bool isAllOnes() const;
 };
 
 /// An floating-point-like constant.

llvm/lib/CodeGen/GlobalISel/CMakeLists.txt

@@ -26,6 +26,7 @@ add_llvm_component_library(LLVMGlobalISel
   Localizer.cpp
   LostDebugLocObserver.cpp
   MachineIRBuilder.cpp
+  OptMIRBuilder.cpp
   RegBankSelect.cpp
   Utils.cpp
 

llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp

@@ -321,8 +321,12 @@ MachineInstrBuilder MachineIRBuilder::buildConstant(const DstOp &Res,
   assert(EltTy.getScalarSizeInBits() == Val.getBitWidth() &&
          "creating constant with the wrong size");
 
-  assert(!Ty.isScalableVector() &&
-         "unexpected scalable vector in buildConstant");
+  if (Ty.isScalableVector()) {
+    auto Const = buildInstr(TargetOpcode::G_CONSTANT)
+                     .addDef(getMRI()->createGenericVirtualRegister(EltTy))
+                     .addCImm(&Val);
+    return buildSplatVector(Res, Const);
+  }
 
   if (Ty.isFixedVector()) {
     auto Const = buildInstr(TargetOpcode::G_CONSTANT)

llvm/lib/CodeGen/GlobalISel/OptMIRBuilder.cpp

@@ -0,0 +1,206 @@
+//===-- llvm/CodeGen/GlobalISel/OptMIRBuilder.cpp -----------------*- C++-*-==//
+// 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
+//
+//===----------------------------------------------------------------------===//
+/// \file
+/// This file implements the OptMIRBuilder class which optimizes as it builds
+/// instructions.
+//===----------------------------------------------------------------------===//
+//
+
+#include "llvm/CodeGen/GlobalISel/OptMIRBuilder.h"
+#include "llvm/CodeGen/GlobalISel/CSEMIRBuilder.h"
+#include "llvm/CodeGen/GlobalISel/LegalizerInfo.h"
+#include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
+#include "llvm/CodeGen/GlobalISel/Utils.h"
+#include "llvm/CodeGen/TargetOpcodes.h"
+
+using namespace llvm;
+
+bool OptMIRBuilder::isPrelegalize() const { return IsPrelegalize; }
+
+bool OptMIRBuilder::isLegal(const LegalityQuery &Query) const {
+  assert(LI != nullptr && "legalizer info is not available");
+  return LI->isLegal(Query);
+}
+
+bool OptMIRBuilder::isConstantLegal(LLT Ty) {
+  if (Ty.isScalar())
+    return isLegal({TargetOpcode::G_CONSTANT, {Ty}});
+
+  LLT EltTy = Ty.getElementType();
+  if (Ty.isFixedVector())
+    return isLegal({TargetOpcode::G_BUILD_VECTOR, {Ty, EltTy}}) &&
+           isLegal({TargetOpcode::G_CONSTANT, {EltTy}});
+
+  // scalable vector
+  assert(Ty.isScalableVector() && "Unexpected LLT");
+  return isLegal({TargetOpcode::G_SPLAT_VECTOR, {Ty, EltTy}}) &&
+         isLegal({TargetOpcode::G_CONSTANT, {EltTy}});
+}
+
+bool OptMIRBuilder::isLegalOrBeforeLegalizer(const LegalityQuery &Query) const {
+  return isPrelegalize() || isLegal(Query);
+}
+
+bool OptMIRBuilder::isConstantLegalOrBeforeLegalizer(LLT Ty) {
+  return isPrelegalize() || isConstantLegal(Ty);
+}
+
+bool OptMIRBuilder::isUndef(Register Reg) const {
+  const MachineInstr *MI = getMRI()->getVRegDef(Reg);
+  return MI->getOpcode() == TargetOpcode::G_IMPLICIT_DEF;
+}
+
+MachineInstrBuilder OptMIRBuilder::buildNegation(const DstOp &DstOp,
+                                                 const SrcOp &SrcOp) {
+  LLT DstTy = DstOp.getLLTTy(*getMRI());
+
+  auto Zero = buildConstant(DstTy, 0);
+  return buildSub(DstOp, Zero, SrcOp);
+}
+
+MachineInstrBuilder OptMIRBuilder::buildGIConstant(const DstOp &DstOp,
+                                                   const GIConstant &Const) {
+  LLT DstTy = DstOp.getLLTTy(*getMRI());
+
+  switch (Const.getKind()) {
+  case GIConstant::GIConstantKind::Scalar:
+    return buildConstant(DstOp, Const.getScalarValue());
+  case GIConstant::GIConstantKind::FixedVector:
+    return buildBuildVectorConstant(DstOp, Const.getAsArrayRef());
+  case GIConstant::GIConstantKind::ScalableVector: {
+    auto Constant =
+        buildConstant(DstTy.getElementType(), Const.getSplatValue());
+    return buildSplatVector(DstOp, Constant);
+  }
+  }
+}
+
+MachineInstrBuilder OptMIRBuilder::optimizeAdd(unsigned Opc,
+                                               ArrayRef<DstOp> DstOps,
+                                               ArrayRef<SrcOp> SrcOps,
+                                               std::optional<unsigned> Flag) {
+  assert(SrcOps.size() == 2 && "Invalid sources");
+  assert(DstOps.size() == 1 && "Invalid dsts");
+
+  LLT DstTy = DstOps[0].getLLTTy(*getMRI());
+
+  if (isUndef(SrcOps[1].getReg()) || isUndef(SrcOps[0].getReg()))
+    return buildUndef(DstTy);
+
+  std::optional<GIConstant> RHS =
+      GIConstant::getConstant(SrcOps[1].getReg(), *getMRI());
+  if (!RHS)
+    return CSEMIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
+
+  if (RHS->isZero())
+    return buildCopy(DstOps[0], SrcOps[0]);
+
+  if (!isConstantLegalOrBeforeLegalizer(DstTy))
+    return CSEMIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
+
+  std::optional<GIConstant> LHS =
+      GIConstant::getConstant(SrcOps[0].getReg(), *getMRI());
+  if (!LHS)
+    return CSEMIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
+
+  GIConstant Add = LHS->add(*RHS);
+
+  return buildGIConstant(DstOps[0], Add);
+}
+
+MachineInstrBuilder OptMIRBuilder::optimizeSub(unsigned Opc,
+                                               ArrayRef<DstOp> DstOps,
+                                               ArrayRef<SrcOp> SrcOps,
+                                               std::optional<unsigned> Flag) {
+  assert(SrcOps.size() == 2 && "Invalid sources");
+  assert(DstOps.size() == 1 && "Invalid dsts");
+
+  LLT DstTy = DstOps[0].getLLTTy(*getMRI());
+
+  if (isUndef(SrcOps[1].getReg()) || isUndef(SrcOps[0].getReg()))
+    return buildUndef(DstTy);
+
+  std::optional<GIConstant> RHS =
+      GIConstant::getConstant(SrcOps[1].getReg(), *getMRI());
+  if (!RHS)
+    return CSEMIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
+
+  if (RHS->isZero())
+    return buildCopy(DstOps[0], SrcOps[0]);
+
+  if (!isConstantLegalOrBeforeLegalizer(DstTy))
+    return CSEMIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
+
+  std::optional<GIConstant> LHS =
+      GIConstant::getConstant(SrcOps[0].getReg(), *getMRI());
+  if (!LHS)
+    return CSEMIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
+
+  GIConstant Sub = LHS->sub(*RHS);
+
+  return buildGIConstant(DstOps[0], Sub);
+}
+
+MachineInstrBuilder OptMIRBuilder::optimizeMul(unsigned Opc,
+                                               ArrayRef<DstOp> DstOps,
+                                               ArrayRef<SrcOp> SrcOps,
+                                               std::optional<unsigned> Flag) {
+  assert(SrcOps.size() == 2 && "Invalid sources");
+  assert(DstOps.size() == 1 && "Invalid dsts");
+
+  LLT DstTy = DstOps[0].getLLTTy(*getMRI());
+
+  if ((isUndef(SrcOps[1].getReg()) || isUndef(SrcOps[0].getReg())) &&
+      isConstantLegalOrBeforeLegalizer(DstTy))
+    return buildConstant(DstTy, 0);
+
+  std::optional<GIConstant> RHS =
+      GIConstant::getConstant(SrcOps[1].getReg(), *getMRI());
+  if (!RHS)
+    return CSEMIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
+
+  if (RHS->isZero() && isConstantLegalOrBeforeLegalizer(DstTy))
+    return buildConstant(DstTy, 0);
+
+  if (RHS->isOne())
+    return buildCopy(DstOps[0], SrcOps[0]);
+
+  if (RHS->isAllOnes() && isConstantLegalOrBeforeLegalizer(DstTy) &&
+      isLegalOrBeforeLegalizer({TargetOpcode::G_SUB, {DstTy}}))
+    return buildNegation(DstOps[0], SrcOps[0]);
+
+  if (RHS->isTwo() && isLegalOrBeforeLegalizer({TargetOpcode::G_ADD, {DstTy}}))
+    return buildAdd(DstOps[0], SrcOps[0], SrcOps[0]);
+
+  if (!isConstantLegalOrBeforeLegalizer(DstTy))
+    return CSEMIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
+
+  std::optional<GIConstant> LHS =
+      GIConstant::getConstant(SrcOps[0].getReg(), *getMRI());
+  if (!LHS)
+    return CSEMIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
+
+  GIConstant Mul = LHS->mul(*RHS);
+
+  return buildGIConstant(DstOps[0], Mul);
+}
+
+MachineInstrBuilder OptMIRBuilder::buildInstr(unsigned Opc,
+                                              ArrayRef<DstOp> DstOps,
+                                              ArrayRef<SrcOp> SrcOps,
+                                              std::optional<unsigned> Flag) {
+  switch (Opc) {
+  case TargetOpcode::G_ADD:
+    return optimizeAdd(Opc, DstOps, SrcOps, Flag);
+  case TargetOpcode::G_SUB:
+    return optimizeSub(Opc, DstOps, SrcOps, Flag);
+  case TargetOpcode::G_MUL:
+    return optimizeMul(Opc, DstOps, SrcOps, Flag);
+  default:
+    return CSEMIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
+  }
+}