[Pass] RISCVDotprodSplitter Passed

Author: chenqian

llvm/lib/Target/RISCV/CMakeLists.txt

@@ -46,6 +46,7 @@ add_llvm_target(RISCVCodeGen
   RISCVEsp32P4MemIntrin.cpp
   RISCVIndirectBranchTracking.cpp
   RISCVIntLoopUnrollAndRemainder.cpp
+  RISCVDotprodSplitter.cpp
   RISCVInsertReadWriteCSR.cpp
   RISCVInsertVSETVLI.cpp
   RISCVInsertWriteVXRM.cpp

llvm/lib/Target/RISCV/RISCVDotprodSplitter.cpp

@@ -0,0 +1,82 @@
+//===-- RISCVDotprodSplitter.h - RISC-V Dotprod Splitter Pass -*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the RISCVDotprodSplitterPass class.
+// This pass identifies a specific pattern often associated with calls to inner
+// dot product computation functions, where the result is passed via a pointer
+// argument (typically an alloca in the caller). The pattern involves a
+// sequence of lifetime start, the call instruction, a load from the result
+// pointer, and lifetime end, all within the same basic block.
+//
+// If this unique pattern is found, the pass restructures the control flow
+// graph (CFG) to create specialized paths for common constant "step" or
+// "stride" values (specifically 1, 2, and 3) passed as arguments to the
+// inner call.
+//
+// It introduces conditional branches based on the runtime values of the step
+// arguments. If the steps match one of the specialized constant pairs (e.g.,
+// image_step=1 and filter_step=1), control flows to a duplicated version of
+// the call sequence where the step arguments are replaced with constants.
+// Otherwise, control flows to the original call sequence (generic path).
+//
+// A PHI node merges the results from the specialized and generic paths.
+// This transformation aims to enable further optimizations like constant
+// propagation and function specialization for the common step values within
+// the called function, potentially improving performance on targets like
+// RISC-V with specific dot product acceleration capabilities.
+// The pass is controlled by the `-riscv-dotprod-splitter` command-line option.
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_RISCV_RISCVDOTPRODSPLITTER_H
+#define LLVM_LIB_TARGET_RISCV_RISCVDOTPRODSPLITTER_H
+
+#include "llvm/IR/PassManager.h"
+#include "llvm/Support/CommandLine.h"
+
+namespace llvm {
+
+class Function;
+class Module;
+
+extern cl::opt<bool> EnableRISCVDotprodSplitter;
+
+/// Pass that specializes dot product calls for common step values.
+///
+/// This pass identifies patterns where inner dot product functions are called
+/// with runtime step parameters and creates specialized versions for common
+/// constant step values (1, 2, 3) to enable better optimization.
+struct RISCVDotprodSplitterPass
+    : public PassInfoMixin<RISCVDotprodSplitterPass> {
+
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+
+  static bool isRequired() { return true; }
+
+  /// Check if the function contains patterns that can be processed by this
+  /// pass.
+  static bool hasProcessablePattern(Function &F);
+};
+
+/// Conditional LoopExtractor Pass that only runs when dotprod patterns exist.
+///
+/// This pass runs LoopExtractor only on modules that contain processable
+/// dot product patterns, avoiding unnecessary loop extraction on modules
+/// that won't benefit from the dotprod splitter optimization.
+struct RISCVConditionalLoopExtractorPass
+    : public PassInfoMixin<RISCVConditionalLoopExtractorPass> {
+
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+  static bool isRequired() { return true; }
+
+private:
+  bool moduleHasProcessablePatterns(Module &M);
+};
+
+} // namespace llvm
+
+#endif // LLVM_LIB_TARGET_RISCV_RISCVDOTPRODSPLITTER_H

llvm/lib/Target/RISCV/RISCVDotprodSplitter.h

@@ -16,6 +16,7 @@
 #include "RISCVCustomLICM.h"
 #include "RISCVEsp32P4MemIntrin.h"
 #include "RISCVIntLoopUnrollAndRemainder.h"
+#include "RISCVDotprodSplitter.h"
 #include "RISCVLoopUnrollAndRemainder.h"
 #include "RISCVMachineFunctionInfo.h"
 #include "RISCVTargetObjectFile.h"
@@ -681,6 +682,10 @@ void RISCVTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
           FPM.addPass(RISCVIntLoopUnrollAndRemainderPass());
           return true;
         }
+        if (Name == "riscv-dotprod-splitter") {
+          FPM.addPass(RISCVDotprodSplitterPass());
+          return true;
+        }
         return false;
       });
 

llvm/lib/Target/RISCV/RISCVTargetMachine.cpp

@@ -1,10 +1,10 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -S -mtriple=riscv32-esp-unknown-elf  < %s | FileCheck %s
+; RUN: opt -S -mtriple=riscv32-esp-unknown-elf -passes=riscv-dotprod-splitter -riscv-dotprod-splitter=true < %s | FileCheck %s
 
 ; Function Attrs: nofree norecurse nosync nounwind memory(read, argmem: readwrite, inaccessiblemem: none)
 define dso_local range(i32 0, 458756) i32 @dspi_dotprod_off_s16_ansi(ptr nocapture noundef readonly %in_image, ptr nocapture noundef readonly %filter, ptr nocapture noundef writeonly %out_value, i32 noundef %count_x, i32 noundef %count_y, i32 noundef %shift, i16 noundef signext %offset) local_unnamed_addr {
 ; CHECK-LABEL: define dso_local range(i32 0, 458756) i32 @dspi_dotprod_off_s16_ansi(
-; CHECK-SAME: ptr nocapture noundef readonly [[IN_IMAGE:%.*]], ptr nocapture noundef readonly [[FILTER:%.*]], ptr nocapture noundef writeonly [[OUT_VALUE:%.*]], i32 noundef [[COUNT_X:%.*]], i32 noundef [[COUNT_Y:%.*]], i32 noundef [[SHIFT:%.*]], i16 noundef signext [[OFFSET:%.*]]) local_unnamed_addr {
+; CHECK-SAME: ptr noundef readonly captures(none) [[IN_IMAGE:%.*]], ptr noundef readonly captures(none) [[FILTER:%.*]], ptr noundef writeonly captures(none) [[OUT_VALUE:%.*]], i32 noundef [[COUNT_X:%.*]], i32 noundef [[COUNT_Y:%.*]], i32 noundef [[SHIFT:%.*]], i16 noundef signext [[OFFSET:%.*]]) local_unnamed_addr {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
 ; CHECK-NEXT:    [[ADD38_US_LOC:%.*]] = alloca i64, align 8
 ; CHECK-NEXT:    [[STEP_X:%.*]] = getelementptr inbounds i8, ptr [[IN_IMAGE]], i32 4
@@ -53,15 +53,51 @@ define dso_local range(i32 0, 458756) i32 @dspi_dotprod_off_s16_ansi(ptr nocaptu
 ; CHECK-NEXT:    br label %[[FOR_COND25_PREHEADER_US:.*]]
 ; CHECK:       [[FOR_COND25_PREHEADER_US]]:
 ; CHECK-NEXT:    [[Y_082_US:%.*]] = phi i32 [ [[INC41_US:%.*]], %[[FOR_COND25_FOR_COND_CLEANUP27_CRIT_EDGE_US:.*]] ], [ 0, %[[FOR_COND25_PREHEADER_US_PREHEADER]] ]
-; CHECK-NEXT:    [[ACC_081_US:%.*]] = phi i64 [ [[ADD38_US_RELOAD_GENERIC:%.*]], %[[FOR_COND25_FOR_COND_CLEANUP27_CRIT_EDGE_US]] ], [ 0, %[[FOR_COND25_PREHEADER_US_PREHEADER]] ]
+; CHECK-NEXT:    [[ACC_081_US:%.*]] = phi i64 [ [[ADD_US_RELOAD:%.*]], %[[FOR_COND25_FOR_COND_CLEANUP27_CRIT_EDGE_US]] ], [ 0, %[[FOR_COND25_PREHEADER_US_PREHEADER]] ]
 ; CHECK-NEXT:    [[I_DATA_080_US:%.*]] = phi ptr [ [[ADD_PTR_US:%.*]], %[[FOR_COND25_FOR_COND_CLEANUP27_CRIT_EDGE_US]] ], [ [[TMP8]], %[[FOR_COND25_PREHEADER_US_PREHEADER]] ]
 ; CHECK-NEXT:    [[F_DATA_079_US:%.*]] = phi ptr [ [[ADD_PTR39_US:%.*]], %[[FOR_COND25_FOR_COND_CLEANUP27_CRIT_EDGE_US]] ], [ [[TMP9]], %[[FOR_COND25_PREHEADER_US_PREHEADER]] ]
-; CHECK-NEXT:    br label %[[CODEREPL:.*]]
-; CHECK:       [[CODEREPL]]:
+; CHECK-NEXT:    br label %[[CODEREPL_ENTRY:.*]]
+; CHECK:       [[CODEREPL_ENTRY]]:
+; CHECK-NEXT:    [[CMP_IMG_STEP1:%.*]] = icmp eq i32 [[TMP0]], 1
+; CHECK-NEXT:    [[CMP_FILT_STEP1:%.*]] = icmp eq i32 [[TMP4]], 1
+; CHECK-NEXT:    [[COND_STEP1:%.*]] = and i1 [[CMP_IMG_STEP1]], [[CMP_FILT_STEP1]]
+; CHECK-NEXT:    br i1 [[COND_STEP1]], label %[[CALL_STEP1:.*]], label %[[CHECK_STEP2:.*]]
+; CHECK:       [[CHECK_STEP2]]:
+; CHECK-NEXT:    [[CMP_IMG_STEP2:%.*]] = icmp eq i32 [[TMP0]], 2
+; CHECK-NEXT:    [[CMP_FILT_STEP2:%.*]] = icmp eq i32 [[TMP4]], 2
+; CHECK-NEXT:    [[COND_STEP2:%.*]] = and i1 [[CMP_IMG_STEP2]], [[CMP_FILT_STEP2]]
+; CHECK-NEXT:    br i1 [[COND_STEP2]], label %[[CALL_STEP2:.*]], label %[[CHECK_STEP3:.*]]
+; CHECK:       [[CHECK_STEP3]]:
+; CHECK-NEXT:    [[CMP_IMG_STEP3:%.*]] = icmp eq i32 [[TMP0]], 3
+; CHECK-NEXT:    [[CMP_FILT_STEP3:%.*]] = icmp eq i32 [[TMP4]], 3
+; CHECK-NEXT:    [[COND_STEP3:%.*]] = and i1 [[CMP_IMG_STEP3]], [[CMP_FILT_STEP3]]
+; CHECK-NEXT:    br i1 [[COND_STEP3]], label %[[CALL_STEP3:.*]], label %[[CALL_GENERIC:.*]]
+; CHECK:       [[CALL_STEP1]]:
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 -1, ptr [[ADD38_US_LOC]])
+; CHECK-NEXT:    call void @dspi_dotprod_off_s16_ansi.for.body28.us(i64 [[ACC_081_US]], i32 1, ptr [[I_DATA_080_US]], i32 1, ptr [[F_DATA_079_US]], i32 [[CONV35]], i32 [[COUNT_X]], ptr [[ADD38_US_LOC]])
+; CHECK-NEXT:    [[TMP10:%.*]] = load i64, ptr [[ADD38_US_LOC]], align 8
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 -1, ptr [[ADD38_US_LOC]])
+; CHECK-NEXT:    br label %[[CALL_MERGE:.*]]
+; CHECK:       [[CALL_STEP2]]:
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 -1, ptr [[ADD38_US_LOC]])
+; CHECK-NEXT:    call void @dspi_dotprod_off_s16_ansi.for.body28.us(i64 [[ACC_081_US]], i32 2, ptr [[I_DATA_080_US]], i32 2, ptr [[F_DATA_079_US]], i32 [[CONV35]], i32 [[COUNT_X]], ptr [[ADD38_US_LOC]])
+; CHECK-NEXT:    [[TMP11:%.*]] = load i64, ptr [[ADD38_US_LOC]], align 8
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 -1, ptr [[ADD38_US_LOC]])
+; CHECK-NEXT:    br label %[[CALL_MERGE]]
+; CHECK:       [[CALL_STEP3]]:
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 -1, ptr [[ADD38_US_LOC]])
+; CHECK-NEXT:    call void @dspi_dotprod_off_s16_ansi.for.body28.us(i64 [[ACC_081_US]], i32 3, ptr [[I_DATA_080_US]], i32 3, ptr [[F_DATA_079_US]], i32 [[CONV35]], i32 [[COUNT_X]], ptr [[ADD38_US_LOC]])
+; CHECK-NEXT:    [[TMP12:%.*]] = load i64, ptr [[ADD38_US_LOC]], align 8
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 -1, ptr [[ADD38_US_LOC]])
+; CHECK-NEXT:    br label %[[CALL_MERGE]]
+; CHECK:       [[CALL_GENERIC]]:
 ; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 -1, ptr [[ADD38_US_LOC]])
 ; CHECK-NEXT:    call void @dspi_dotprod_off_s16_ansi.for.body28.us(i64 [[ACC_081_US]], i32 [[TMP0]], ptr [[I_DATA_080_US]], i32 [[TMP4]], ptr [[F_DATA_079_US]], i32 [[CONV35]], i32 [[COUNT_X]], ptr [[ADD38_US_LOC]])
-; CHECK-NEXT:    [[ADD38_US_RELOAD_GENERIC]] = load i64, ptr [[ADD38_US_LOC]], align 8
+; CHECK-NEXT:    [[ADD38_US_RELOAD_GENERIC:%.*]] = load i64, ptr [[ADD38_US_LOC]], align 8
 ; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 -1, ptr [[ADD38_US_LOC]])
+; CHECK-NEXT:    br label %[[CALL_MERGE]]
+; CHECK:       [[CALL_MERGE]]:
+; CHECK-NEXT:    [[ADD_US_RELOAD]] = phi i64 [ [[TMP10]], %[[CALL_STEP1]] ], [ [[TMP11]], %[[CALL_STEP2]] ], [ [[TMP12]], %[[CALL_STEP3]] ], [ [[ADD38_US_RELOAD_GENERIC]], %[[CALL_GENERIC]] ]
 ; CHECK-NEXT:    br label %[[FOR_COND25_FOR_COND_CLEANUP27_CRIT_EDGE_US]]
 ; CHECK:       [[FOR_COND25_FOR_COND_CLEANUP27_CRIT_EDGE_US]]:
 ; CHECK-NEXT:    [[ADD_PTR_US]] = getelementptr inbounds i16, ptr [[I_DATA_080_US]], i32 [[MUL20]]
@@ -70,7 +106,7 @@ define dso_local range(i32 0, 458756) i32 @dspi_dotprod_off_s16_ansi(ptr nocaptu
 ; CHECK-NEXT:    [[EXITCOND85_NOT:%.*]] = icmp eq i32 [[INC41_US]], [[COUNT_Y]]
 ; CHECK-NEXT:    br i1 [[EXITCOND85_NOT]], label %[[FOR_COND_CLEANUP]], label %[[FOR_COND25_PREHEADER_US]]
 ; CHECK:       [[FOR_COND_CLEANUP]]:
-; CHECK-NEXT:    [[ACC_0_LCSSA:%.*]] = phi i64 [ 0, %[[IF_END16]] ], [ 0, %[[FOR_COND25_PREHEADER_LR_PH]] ], [ [[ADD38_US_RELOAD_GENERIC]], %[[FOR_COND25_FOR_COND_CLEANUP27_CRIT_EDGE_US]] ]
+; CHECK-NEXT:    [[ACC_0_LCSSA:%.*]] = phi i64 [ 0, %[[IF_END16]] ], [ 0, %[[FOR_COND25_PREHEADER_LR_PH]] ], [ [[ADD_US_RELOAD]], %[[FOR_COND25_FOR_COND_CLEANUP27_CRIT_EDGE_US]] ]
 ; CHECK-NEXT:    [[SUB:%.*]] = add nsw i32 [[SHIFT]], -1
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i32 1, [[SUB]]
 ; CHECK-NEXT:    [[CONV43:%.*]] = sext i32 [[SHL]] to i64
@@ -176,19 +212,19 @@ return:                                           ; preds = %for.cond.cleanup, %
 }
 
 ; Function Attrs: nofree norecurse nounwind
-define internal void @dspi_dotprod_off_s16_ansi.for.body28.us(i64 %acc.081.us, i32 %0, ptr %i_data.080.us, i32 %1, ptr %f_data.079.us, i32 %conv35, i32 %count_x, ptr %add38.us.out) {
+define internal void @dspi_dotprod_off_s16_ansi.for.body28.us(i64 %acc.081.us, i32 %img_step, ptr %i_data.080.us, i32 %filt_step, ptr %f_data.079.us, i32 %conv35, i32 %count_x, ptr %add38.us.out) {
 ; CHECK-LABEL: define internal void @dspi_dotprod_off_s16_ansi.for.body28.us(
-; CHECK-SAME: i64 [[ACC_081_US:%.*]], i32 [[TMP0:%.*]], ptr [[I_DATA_080_US:%.*]], i32 [[TMP1:%.*]], ptr [[F_DATA_079_US:%.*]], i32 [[CONV35:%.*]], i32 [[COUNT_X:%.*]], ptr [[ADD38_US_OUT:%.*]]) {
+; CHECK-SAME: i64 [[ACC_081_US:%.*]], i32 [[IMG_STEP:%.*]], ptr [[I_DATA_080_US:%.*]], i32 [[FILT_STEP:%.*]], ptr [[F_DATA_079_US:%.*]], i32 [[CONV35:%.*]], i32 [[COUNT_X:%.*]], ptr [[ADD38_US_OUT:%.*]]) {
 ; CHECK-NEXT:  [[NEWFUNCROOT:.*]]:
 ; CHECK-NEXT:    br label %[[FOR_BODY28_US:.*]]
 ; CHECK:       [[FOR_BODY28_US]]:
 ; CHECK-NEXT:    [[X_077_US:%.*]] = phi i32 [ 0, %[[NEWFUNCROOT]] ], [ [[INC_US:%.*]], %[[FOR_BODY28_US]] ]
 ; CHECK-NEXT:    [[ACC_176_US:%.*]] = phi i64 [ [[ACC_081_US]], %[[NEWFUNCROOT]] ], [ [[ADD38_US:%.*]], %[[FOR_BODY28_US]] ]
-; CHECK-NEXT:    [[MUL30_US:%.*]] = mul nsw i32 [[X_077_US]], [[TMP0]]
+; CHECK-NEXT:    [[MUL30_US:%.*]] = mul nsw i32 [[X_077_US]], [[IMG_STEP]]
 ; CHECK-NEXT:    [[ARRAYIDX_US:%.*]] = getelementptr inbounds i16, ptr [[I_DATA_080_US]], i32 [[MUL30_US]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = load i16, ptr [[ARRAYIDX_US]], align 2
 ; CHECK-NEXT:    [[CONV_US:%.*]] = sext i16 [[TMP2]] to i32
-; CHECK-NEXT:    [[MUL32_US:%.*]] = mul nsw i32 [[X_077_US]], [[TMP1]]
+; CHECK-NEXT:    [[MUL32_US:%.*]] = mul nsw i32 [[X_077_US]], [[FILT_STEP]]
 ; CHECK-NEXT:    [[ARRAYIDX33_US:%.*]] = getelementptr inbounds i16, ptr [[F_DATA_079_US]], i32 [[MUL32_US]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = load i16, ptr [[ARRAYIDX33_US]], align 2
 ; CHECK-NEXT:    [[CONV34_US:%.*]] = sext i16 [[TMP3]] to i32
@@ -209,11 +245,11 @@ newFuncRoot:
 for.body28.us:                                    ; preds = %newFuncRoot, %for.body28.us
   %x.077.us = phi i32 [ 0, %newFuncRoot ], [ %inc.us, %for.body28.us ]
   %acc.176.us = phi i64 [ %acc.081.us, %newFuncRoot ], [ %add38.us, %for.body28.us ]
-  %mul30.us = mul nsw i32 %x.077.us, %0
+  %mul30.us = mul nsw i32 %x.077.us, %img_step
   %arrayidx.us = getelementptr inbounds i16, ptr %i_data.080.us, i32 %mul30.us
   %2 = load i16, ptr %arrayidx.us, align 2
   %conv.us = sext i16 %2 to i32
-  %mul32.us = mul nsw i32 %x.077.us, %1
+  %mul32.us = mul nsw i32 %x.077.us, %filt_step
   %arrayidx33.us = getelementptr inbounds i16, ptr %f_data.079.us, i32 %mul32.us
   %3 = load i16, ptr %arrayidx33.us, align 2
   %conv34.us = sext i16 %3 to i32