[MLIR] Enhance folding capability of spirv::CompositeExtractOp::fold

[nbpatel]

Summary

This PR improves the spirv::CompositeExtractOp::fold function by adding a backtracking mechanism.
The updated function can now traverse a chain of CompositeInsertOps to find a match.

[kuhar]

This looks fine but could you add a LIT test demonstrating how it works?

[nbpatel]

Added a test case

[kuhar]

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

146–147

nit: since insertOp is not used outside of the loop, you should be able to reduce the scope like so:

while (auto insertOp = compositeOp.getDefiningOp<spirv::CompositeInsertOp>()) {
   if (getIndices() == insertOp.getIndices())
     return insertOp.getObject();
   compositeOp = insertOp.getComposite();
}

153–154

Would we look through insertions and then look through composite constructs? I think we should be able to resume from compositeOp updated in the loop above.

153–154

*Should

mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir

199

nit: I think CHECK-NEXT: return .... would work here and we would not have to use CHECK-NOT.

[nbpatel]

Address Feedback

[nbpatel]

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

153–154

not exactly sure what you mean here

[kuhar]

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

151

We shouldn't need this since insertOp gets reassigned in the loop condition

153–154

I mean supporting sequences like this:

%a = spirv.CompositeConstruct %0, %1, %2
%b = spirv.CompositeInsert %3, %a[1]
%c = spirv.CompositeInsert %4, %b[2]
%val = spirv.CompositeExtract %c[0]
return %val

==>

return %0

[nbpatel]

clean up and enchance pattern

[nbpatel]

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

151

ah sorry, my bad, missed to remove that line

153–154

added the changes and a test case

[nbpatel]

@kuhar not sure why it shows patch application failed. Should I create a revision again?

[kuhar]

Awesome, thanks for the fixes!

In D151536#4377248, @nbpatel wrote:

@kuhar not sure why it shows patch application failed. Should I create a revision again?

Can you share some more context? What did you do and where did you see this message? In general, phabricator doesn't require creating new revisions when things go out of sync -- you should be able to rebase and upload an updated diff based on that.

[nbpatel]

in the diff details it shows this
"Diff Detail
Build Status
Buildable 234970
Build 363444: pre-merge checks patch application failed"

Do I have to rebase on latest master? I used the gui for creating the PR and revisions, it worked fine until the last revision
Also, can you help merge the patch for me since I dont think I have permission nor do I know how to merge it in upstream llvm repo

[kuhar]

In D151536#4377351, @nbpatel wrote:

in the diff details it shows this
"Diff Detail
Build Status
Buildable 234970
Build 363444: pre-merge checks patch application failed"

Do I have to rebase on latest master? I used the gui for creating the PR and revisions, it worked fine until the last revision
Also, can you help merge the patch for me since I dont think I have permission nor do I know how to merge it in upstream llvm repo

I think it could be that the patch is too far behind for the CI to pick up, and then the fix would be to rebase and upload a new diff.

Sure, I can commit it for you. What name and email would you like to use?

[nbpatel]

Name : Nishant Patel
Email: [e-mail]

Thanks @kuhar

[kuhar]

Rebase.

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

OpFoldResult spirv::BitcastOp::fold(FoldAdaptor /*adaptor*/) {
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.CompositeExtractOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult spirv::CompositeExtractOp::fold(FoldAdaptor adaptor) {
-  if (auto insertOp =
-          getComposite().getDefiningOp<spirv::CompositeInsertOp>()) {
+  Value compositeOp = getComposite();
+
+  while (auto insertOp =
+             compositeOp.getDefiningOp<spirv::CompositeInsertOp>()) {
     if (getIndices() == insertOp.getIndices())
       return insertOp.getObject();
+    compositeOp = insertOp.getComposite();
   }
 
   if (auto constructOp =
-          getComposite().getDefiningOp<spirv::CompositeConstructOp>()) {
+          compositeOp.getDefiningOp<spirv::CompositeConstructOp>()) {
     auto type = llvm::cast<spirv::CompositeType>(constructOp.getType());
     if (getIndices().size() == 1 &&
         constructOp.getConstituents().size() == type.getNumElements()) {
-      auto i = llvm::cast<IntegerAttr>(*getIndices().begin());
+      auto i = getIndices().begin()->cast<IntegerAttr>();
+      if (static_cast<size_t>(i.getValue().getSExtValue()) <
+          constructOp.getConstituents().size())
         return constructOp.getConstituents()[i.getValue().getSExtValue()];
     }
   }
 
   auto indexVector =
       llvm::to_vector<8>(llvm::map_range(getIndices(), [](Attribute attr) {
         return static_cast<unsigned>(llvm::cast<IntegerAttr>(attr).getInt());
       }));
   return extractCompositeElement(adaptor.getComposite(), indexVector);

mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir

func.func @extract_construct(%val1: vector<2xf32>, %val2: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) {
   %1 = spirv.CompositeExtract %construct[0 : i32] : !spirv.array<2xvector<2xf32>>
   %2 = spirv.CompositeExtract %construct[1 : i32] : !spirv.array<2xvector<2xf32>>
   // CHECK: return %[[VAL1]], %[[VAL2]]
   return %1, %2 : vector<2xf32>, vector<2xf32>
 }
 
 // -----
 
+ // CHECK-LABEL: fold_composite_op
+ //  CHECK-SAME: (%[[COMP:.+]]: !spirv.struct<(f32, f32)>, %[[VAL1:.+]]: f32, %[[VAL2:.+]]: f32)
+  func.func @fold_composite_op(%composite: !spirv.struct<(f32, f32)>, %val1: f32, %val2: f32) -> f32 {
+    %insert = spirv.CompositeInsert %val1, %composite[0 : i32] : f32 into !spirv.struct<(f32, f32)>
+    %1 = spirv.CompositeInsert %val2, %insert[1 : i32] : f32 into !spirv.struct<(f32, f32)>
+    %2 = spirv.CompositeExtract %1[0 : i32] : !spirv.struct<(f32, f32)>
+    // CHECK-NEXT: return  %[[VAL1]]
+    return %2 : f32
+  }
+
+// -----
+
+ // CHECK-LABEL: fold_composite_op
+ //  CHECK-SAME: (%[[VAL1:.+]]: f32, %[[VAL2:.+]]: f32, %[[VAL3:.+]]: f32)
+  func.func @fold_composite_op(%val1: f32, %val2: f32, %val3: f32) -> f32 {
+    %composite = spirv.CompositeConstruct %val1, %val1, %val1 : (f32, f32, f32) -> !spirv.struct<(f32, f32, f32)>
+    %insert = spirv.CompositeInsert %val2, %composite[1 : i32] : f32 into !spirv.struct<(f32, f32, f32)>
+    %1 = spirv.CompositeInsert %val3, %insert[2 : i32] : f32 into !spirv.struct<(f32, f32, f32)>
+    %2 = spirv.CompositeExtract %1[0 : i32] : !spirv.struct<(f32, f32, f32)>
+    // CHECK-NEXT: return  %[[VAL1]]
+    return %2 : f32
+  }
+
+// -----
+
 // Not yet implemented case
 
 // CHECK-LABEL: extract_construct
 func.func @extract_construct(%val1: vector<3xf32>, %val2: f32) -> (f32, f32) {
   // CHECK: spirv.CompositeConstruct
   %construct = spirv.CompositeConstruct %val1, %val2 : (vector<3xf32>, f32) -> vector<4xf32>
   // CHECK: spirv.CompositeExtract
   %1 = spirv.CompositeExtract %construct[0 : i32] : vector<4xf32>