Add a flag for the upstream global reassociation algorithm change

This PR (#6598) pulls the upstream global reassociation algorithm change in DXC and can reduce redundant calculations obviously.

However, from the testing result of a large offline suite of shaders, some shaders got worse compilation results and couldn't benefit from this upstream change.

This PR adds a flag for the upstream global reassociation change. It would be easier to roll back if a shader get worse compilation result due to this upstream change.

This is part 2 of the fix for #6593.

Author: lizhengxing

include/dxc/Support/DxcOptToggles.h

@@ -38,6 +38,8 @@ enum {
 static const Toggle TOGGLE_GVN = {"gvn", DEFAULT_ON};
 static const Toggle TOGGLE_LICM = {"licm", DEFAULT_ON};
 static const Toggle TOGGLE_SINK = {"sink", DEFAULT_ON};
+static const Toggle TOGGLE_ENABLE_AGGRESSIVE_REASSOCIATION = {
+    "enable-aggressive-reassociation", DEFAULT_ON};
 static const Toggle TOGGLE_LIFETIME_MARKERS = {"lifetime-markers", DEFAULT_ON};
 static const Toggle TOGGLE_PARTIAL_LIFETIME_MARKERS = {
     "partial-lifetime-markers", DEFAULT_OFF};

include/llvm/Transforms/IPO/PassManagerBuilder.h

@@ -134,6 +134,7 @@ class PassManagerBuilder {
   unsigned ScanLimit = 0; // HLSL Change
   bool EnableGVN = true; // HLSL Change
   bool StructurizeLoopExitsForUnroll = false; // HLSL Change
+  bool HLSLEnableAggressiveReassociation = true; // HLSL Change
   bool HLSLEnableLifetimeMarkers = false; // HLSL Change
   bool HLSLEnablePartialLifetimeMarkers = false; // HLSL Change
   bool HLSLEnableDebugNops = false; // HLSL Change

include/llvm/Transforms/Scalar.h

@@ -331,6 +331,8 @@ extern char &DemoteRegisterToMemoryHlslID;
 // For example:  4 + (x + 5)  ->  x + (4 + 5)
 //
 FunctionPass *createReassociatePass();
+FunctionPass *
+createReassociatePass(bool HLSLEnableAggressiveReassociation); // HLSL Change
 
 //===----------------------------------------------------------------------===//
 //

lib/Transforms/IPO/PassManagerBuilder.cpp

@@ -469,7 +469,8 @@ void PassManagerBuilder::populateModulePassManager(
   //MPM.add(createTailCallEliminationPass()); // Eliminate tail calls
   // HLSL Change Ends.
   MPM.add(createCFGSimplificationPass());     // Merge & remove BBs
-  MPM.add(createReassociatePass());           // Reassociate expressions
+  MPM.add(createReassociatePass(
+      HLSLEnableAggressiveReassociation)); // Reassociate expressions
   // Rotate Loop - disable header duplication at -Oz
   MPM.add(createLoopRotatePass(SizeLevel == 2 ? 0 : -1));
   // HLSL Change - disable LICM in frontend for not consider register pressure.

lib/Transforms/Scalar/Reassociate.cpp

@@ -176,6 +176,16 @@ namespace {
       initializeReassociatePass(*PassRegistry::getPassRegistry());
     }
 
+    // HLSL Change - begin
+    // Enable global reassociation when HLSLEnableAggressiveReassociation is
+    // set
+    bool m_HLSLEnableAggressiveReassociation = true;
+    Reassociate(bool HLSLEnableAggressiveReassociation) : FunctionPass(ID) {
+      initializeReassociatePass(*PassRegistry::getPassRegistry());
+      m_HLSLEnableAggressiveReassociation = HLSLEnableAggressiveReassociation;
+    }
+    // HLSL Change - end
+
     bool runOnFunction(Function &F) override;
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
@@ -242,6 +252,13 @@ INITIALIZE_PASS(Reassociate, "reassociate",
 // Public interface to the Reassociate pass
 FunctionPass *llvm::createReassociatePass() { return new Reassociate(); }
 
+// HLSL Change - begin
+FunctionPass *
+llvm::createReassociatePass(bool HLSLEnableAggressiveReassociation) {
+  return new Reassociate(HLSLEnableAggressiveReassociation);
+}
+// HLSL Change - end
+
 /// Return true if V is an instruction of the specified opcode and if it
 /// only has one use.
 static BinaryOperator *isReassociableOp(Value *V, unsigned Opcode) {
@@ -2243,7 +2260,8 @@ void Reassociate::ReassociateExpression(BinaryOperator *I) {
     return;
   }
 
-  if (Ops.size() > 2 && Ops.size() <= GlobalReassociateLimit) {
+  if (m_HLSLEnableAggressiveReassociation && // HLSL Change
+      (Ops.size() > 2 && Ops.size() <= GlobalReassociateLimit)) {
     // Find the pair with the highest count in the pairmap and move it to the
     // back of the list so that it can later be CSE'd.
     // example:

tools/clang/lib/CodeGen/BackendUtil.cpp

@@ -357,6 +357,8 @@ void EmitAssemblyHelper::CreatePasses() {
       OptToggles.IsEnabled(hlsl::options::TOGGLE_LIFETIME_MARKERS);
   PMBuilder.HLSLEnablePartialLifetimeMarkers =
       OptToggles.IsEnabled(hlsl::options::TOGGLE_PARTIAL_LIFETIME_MARKERS);
+  PMBuilder.HLSLEnableAggressiveReassociation =
+      OptToggles.IsEnabled(hlsl::options::TOGGLE_ENABLE_AGGRESSIVE_REASSOCIATION);
   // HLSL Change - end
 
   PMBuilder.DisableUnitAtATime = !CodeGenOpts.UnitAtATime;

tools/clang/test/HLSLFileCheck/passes/transforms/reassociate/reassociation-flag.hlsl

@@ -0,0 +1,31 @@
+// RUN: %dxc -T cs_6_3 -E cs_main %s -opt-enable enable-aggressive-reassociation  | FileCheck %s -check-prefixes=CHECK,COMMON_FACTOR
+// RUN: %dxc -T cs_6_3 -E cs_main %s -opt-disable enable-aggressive-reassociation | FileCheck %s -check-prefixes=CHECK,NO_COMMON_FACTOR
+
+// Make sure DXC recognize the common factor and generate optimized dxils if the enable-aggressive-reassociation is true.
+
+// CHECK: [[FACTOR_SRC1:%.*]] = call i32 @dx.op.flattenedThreadIdInGroup.i32(i32 96)
+// CHECK: [[FACTOR_SRC0:%.*]] = call i32 @dx.op.threadIdInGroup.i32(i32 95, i32 0)
+
+// COMMON_FACTOR: [[FACTOR:%.*]] = mul i32 [[FACTOR_SRC0]], [[FACTOR_SRC1]]
+// COMMON_FACTOR: mul i32 [[FACTOR]],
+// COMMON_FACTOR: mul i32 [[FACTOR]],
+
+// NO_COMMON_FACTOR: [[EXPRESSION_0:%.*]] = mul i32 [[FACTOR_SRC1]],
+// NO_COMMON_FACTOR:                        mul i32 [[EXPRESSION_0]], [[FACTOR_SRC0]]
+// NO_COMMON_FACTOR: [[EXPRESSION_1:%.*]] = mul i32 [[FACTOR_SRC0]], [[FACTOR_SRC1]]
+// NO_COMMON_FACTOR:                        mul i32 [[EXPRESSION_1]],  
+
+
+RWTexture1D < float2 > outColorBuffer : register ( u0 ) ;
+
+[ numthreads ( 8 , 8 , 1 ) ]
+void cs_main ( uint3 GroupID : SV_GroupID , uint GroupIndex : SV_GroupIndex , uint3 GTID : SV_GroupThreadID , uint3 DispatchThreadID : SV_DispatchThreadID )
+{
+      // DXC should recognize (GroupIndex * GTID.x) is a common factor
+      uint a = GroupIndex * GroupID.x;
+      uint b = GroupIndex * DispatchThreadID.x;
+      uint c = a * GTID.x;
+      uint d = b * GTID.x;
+
+      outColorBuffer [ DispatchThreadID.y ] = float2(c, d);
+}