[snippy] Support per-operand imm-hist

Adds the ability to separately specify values for each immediate operand separately in imm‑hist.

Author: ruslan-sc

llvm/test/tools/llvm-snippy/Inputs/immediate-histogram-per-operand.yaml

@@ -0,0 +1,19 @@
+include:
+  - sections.yaml
+
+options:
+  mtriple: riscv64-unknown-elf
+  num-instrs: 100
+  dump-mf: on
+  mattr: +Zicsr,+f
+
+histogram:
+    - [CSRRWI, 1.0]
+
+imm-hist:
+  opcodes:
+    - 'CSRRWI':
+        operands:
+          - - [0x2, 1.0] # First operand  (CSR)
+          - - [0x0, 1.0] # Second operand (Value to write)
+            - [0x1, 1.0]

llvm/test/tools/llvm-snippy/immediate-histogram-operands-uniform.yaml

@@ -0,0 +1,24 @@
+# RUN: llvm-snippy %s -model-plugin=None |& FileCheck %s
+
+include:
+  - Inputs/sections.yaml
+
+options:
+  mtriple: riscv64-unknown-elf
+  num-instrs: 100
+  dump-mf: on
+  mattr: +Zicsr
+
+histogram:
+    - [CSRRWI, 1.0]
+
+imm-hist:
+  opcodes:
+    - 'CSRRWI':
+        operands:
+          - - [0x2, 1.0] # First operand  (CSR)
+          - uniform      # Second operand (Value to write)
+
+# NOTE: We expect that the second imm-operand will be 5-bit immediate
+# CHECK-COUNT-100: CSRRWI 2, {{[0-9]|[1-2][0-9]|3[0-1]}}
+

llvm/test/tools/llvm-snippy/immediate-histogram-per-operand.yaml

@@ -0,0 +1,3 @@
+# RUN: llvm-snippy %S/Inputs/immediate-histogram-per-operand.yaml -model-plugin=None |& FileCheck %s
+
+# CHECK-COUNT-100: CSRRWI 2, {{[0|1]}}

llvm/test/tools/llvm-snippy/immediate-histogram-regex-wrong-opcode-setting-error.yaml

@@ -15,6 +15,4 @@ imm-hist:
     - "ADDI":
         nonexistent: mapping
 
-# CHECK: error: Immediate histogram opcode setting should be either sequence or scalar.
-# CHECK-SAME: But map was encountered
-# CHECK: nonexistent: mapping
+# CHECK: error: missing required key 'operands'

llvm/test/tools/llvm-snippy/immediate-histogram-too-few-operands.yaml

@@ -0,0 +1,25 @@
+# RUN: not llvm-snippy %s -model-plugin=None |& FileCheck %s
+
+include:
+  - Inputs/sections.yaml
+
+options:
+  mtriple: riscv64-unknown-elf
+  num-instrs: 100
+  dump-mf: on
+  mattr: +Zicsr
+
+histogram:
+    - [CSRRWI, 1.0]
+
+imm-hist:
+  opcodes:
+    - 'CSRRWI':
+        operands:
+          - - [0x2, 1.0] # First operand  (CSR)
+
+# CHECK: error: Immediate histogram: The number of operands
+# CHECK-SAME: entries is not equal to the number of immediate
+# CHECK-SAME: operands for the "CSRRWI" opcode. Expected 2
+# CHECK-SAME: entries but 1 were specified
+

llvm/test/tools/llvm-snippy/immediate-histogram-too-many-operands.yaml

@@ -0,0 +1,29 @@
+# RUN: not llvm-snippy %s -model-plugin=None |& FileCheck %s
+
+include:
+  - Inputs/sections.yaml
+
+options:
+  mtriple: riscv64-unknown-elf
+  num-instrs: 100
+  dump-mf: on
+  mattr: +Zicsr
+
+histogram:
+    - [CSRRWI, 1.0]
+
+imm-hist:
+  opcodes:
+    - 'CSRRWI':
+        operands:
+          - - [0x2, 1.0] # First operand  (CSR)
+          - - [0x0, 1.0] # Second operand (Value to write)
+            - [0x1, 1.0]
+            - [0x2, 1.0]
+          - - [0x0, 1.0] # Third operand  (nonexistent)
+
+# CHECK: error: Immediate histogram: The number of operands
+# CHECK-SAME: entries is not equal to the number of immediate
+# CHECK-SAME: operands for the "CSRRWI" opcode. Expected 2
+# CHECK-SAME: entries but 3 were specified
+

llvm/test/tools/llvm-snippy/immediate-histogram-wrong-operand-setting.yaml

@@ -0,0 +1,24 @@
+# RUN: not llvm-snippy %s |& FileCheck %s
+
+include:
+  - Inputs/sections.yaml
+
+options:
+  mtriple: riscv64-unknown-elf
+  num-instrs: 100
+  dump-mf: on
+  mattr: +Zicsr,+f
+
+histogram:
+    - [CSRRWI, 1.0]
+
+imm-hist:
+  opcodes:
+    - 'CSRRWI':
+        operands:
+          - - [0x2, 1.0] # First operand  (CSR)
+          - nonexistent: mapping
+
+# CHECK: error: Immediate histogram operands entry
+# CHECK-SAME: should be either sequence or scalar
+# CHECK: nonexistent: mapping

llvm/tools/llvm-snippy/include/snippy/Config/ImmediateHistogram.h

@@ -44,30 +44,99 @@ struct ImmediateHistogramSequence final {
   bool empty() const { return Values.empty(); }
 };
 
+class ImmHistOperandsEntry final {
+public:
+  enum class Kind { Uniform, Sequence };
+
+  ImmHistOperandsEntry() = default;
+
+  ImmHistOperandsEntry(const ImmediateHistogramSequence &Seq)
+      : Underlying(Seq), EntryKind(Kind::Sequence) {}
+
+  bool isSequence() const { return EntryKind == Kind::Sequence; }
+
+  bool isUniform() const { return EntryKind == Kind::Uniform; }
+
+  Kind getKind() const { return EntryKind; }
+
+  const ImmediateHistogramSequence &getSequence() const & {
+    assert(isSequence() &&
+           std::holds_alternative<ImmediateHistogramSequence>(Underlying));
+    return std::get<ImmediateHistogramSequence>(Underlying);
+  }
+
+private:
+  std::variant<std::monostate, ImmediateHistogramSequence> Underlying;
+  Kind EntryKind = Kind::Uniform;
+};
+
+struct ImmHistOperandsEntryHolder final {
+  ImmHistOperandsEntry Underlying;
+};
+
+class ImmHistOperands final {
+  std::vector<ImmHistOperandsEntryHolder> Operands;
+  friend yaml::MappingTraits<ImmHistOperands>;
+
+public:
+  bool isUniformForOperand(unsigned OperandIdx) const {
+    assert(Operands.size() > OperandIdx);
+    return Operands[OperandIdx].Underlying.isUniform();
+  }
+
+  bool hasIHForOperand(unsigned OperandIdx) const {
+    assert(Operands.size() > OperandIdx);
+    return Operands[OperandIdx].Underlying.isSequence();
+  }
+
+  const ImmediateHistogramSequence &
+  getIHForOperand(unsigned OperandIdx) const & {
+    assert(Operands.size() > OperandIdx);
+    assert(hasIHForOperand(OperandIdx));
+    const auto &Entry = Operands[OperandIdx].Underlying;
+    return Entry.getSequence();
+  }
+
+  auto size() const { return Operands.size(); }
+};
+
 class ImmHistOpcodeSettings final {
 public:
-  enum class Kind { Uniform, Custom };
+  enum class Kind { Uniform, Custom, Operands };
 
 private:
-  std::optional<ImmediateHistogramSequence> Seq;
+  std::variant<std::monostate, ImmediateHistogramSequence, ImmHistOperands>
+      Underlying;
   Kind SettingsKind = Kind::Uniform;
 
 public:
   ImmHistOpcodeSettings() = default;
 
   ImmHistOpcodeSettings(const ImmediateHistogramSequence &Sequence)
-      : Seq(Sequence), SettingsKind(Kind::Custom) {}
+      : Underlying(Sequence), SettingsKind(Kind::Custom) {}
 
   ImmHistOpcodeSettings(ImmediateHistogramSequence &&Sequence)
-      : Seq(std::move(Sequence)), SettingsKind(Kind::Custom) {}
+      : Underlying(std::move(Sequence)), SettingsKind(Kind::Custom) {}
+
+  ImmHistOpcodeSettings(const ImmHistOperands &Map)
+      : Underlying(Map), SettingsKind(Kind::Operands) {}
 
   bool isUniform() const { return SettingsKind == Kind::Uniform; }
 
   bool isSequence() const { return SettingsKind == Kind::Custom; }
 
-  const ImmediateHistogramSequence &getSequence() const {
-    assert(isSequence() && Seq.has_value());
-    return *Seq;
+  bool isPerOperand() const { return SettingsKind == Kind::Operands; }
+
+  const ImmediateHistogramSequence &getSequence() const & {
+    assert(isSequence() &&
+           std::holds_alternative<ImmediateHistogramSequence>(Underlying));
+    return std::get<ImmediateHistogramSequence>(Underlying);
+  }
+
+  const ImmHistOperands &getOperandsMap() const & {
+    assert(isPerOperand() &&
+           std::holds_alternative<ImmHistOperands>(Underlying));
+    return std::get<ImmHistOperands>(Underlying);
   }
 
   Kind getKind() const { return SettingsKind; }
@@ -395,13 +464,25 @@ class OpcodeToImmHistSequenceMap final {
 
   OpcodeToImmHistSequenceMap() = default;
 
-  const ImmHistOpcodeSettings &
-  getConfigForOpcode(unsigned Opc, const OpcodeCache &OpCC) const {
+  const ImmHistOpcodeSettings &getConfigForOpcode(unsigned Opc) const & {
     assert(Data.count(Opc) && "Opcode was not found in immediate histogram");
     return Data.at(Opc);
   }
+
+  bool empty() const { return Data.empty(); }
 };
 
+struct ImmHistOperandsEntryMapMapper final {};
+
 } // namespace snippy
+
+namespace yaml {
+void yamlize(IO &IO, const snippy::ImmHistOperandsEntryMapMapper &, bool,
+             EmptyContext &);
+} // namespace yaml
+
+LLVM_SNIPPY_YAML_DECLARE_MAPPING_TRAITS(snippy::ImmHistOperands);
+LLVM_SNIPPY_YAML_DECLARE_MAPPING_TRAITS(snippy::ImmHistOperandsEntryHolder);
+LLVM_SNIPPY_YAML_IS_SEQUENCE_ELEMENT(snippy::ImmHistOperandsEntryHolder, false);
 LLVM_SNIPPY_YAML_DECLARE_MAPPING_TRAITS(snippy::ImmediateHistogramSequence);
 } // namespace llvm

llvm/tools/llvm-snippy/include/snippy/Target/Target.h

@@ -355,8 +355,8 @@ class SnippyTarget {
   virtual MachineOperand
   generateTargetOperand(SnippyProgramContext &ProgCtx,
                         const CommonPolicyConfig &Cfg, unsigned OpCode,
-                        unsigned OpType,
-                        const StridedImmediate &StridedImm) const = 0;
+                        unsigned OpType, const StridedImmediate &StridedImm,
+                        unsigned OperandIdx) const = 0;
 
   virtual bool canProduceNaN(const MCInstrDesc &InstrDesc) const = 0;
 
@@ -376,6 +376,8 @@ class SnippyTarget {
     return AccessMaskBit::None;
   }
 
+  virtual size_t getNumImmOperands(const MCInstrDesc &InstrDesc) const = 0;
+
   virtual unsigned getTransformSequenceLength(InstructionGenerationContext &IGC,
                                               APInt OldValue, APInt NewValue,
                                               MCRegister Register) const = 0;