[FLANG] Change loop versioning to use shift instead of divide

[Leporacanthicus]

Summary

Despite me being convinced that the use of divide didn't produce any
divide instructions, it does in fact add more instructions than using
a plain shift operation.

This patch simply changes the divide to a shift right, with an
assert to check that the "divisor" is a power of two.

[kiranchandramohan]

flang/lib/Optimizer/Transforms/LoopVersioning.cpp

269

Would it be better to use an if-else to handle the powerOf2 and the other case?

For the types handled now, I think REAL(KIND=10) 80-bit extended precision with explicit normalization bit (15e64m) will not be a power of 2.

[tblah]

Thanks!

[Leporacanthicus]

flang/lib/Optimizer/Transforms/LoopVersioning.cpp

269

It will work for RELA(KIND=10) as is, as those are aligned to 16 bytes, although the actual data is only 10 bytes.

So as far as I can tell, the only reason we'd need non-power of two is if we allow data-types that are not the simple types currently supported.

[kiranchandramohan]

LG.

flang/lib/Optimizer/Transforms/LoopVersioning.cpp

269

OK. Thanks.

flang/lib/Optimizer/Transforms/LoopVersioning.cpp

for (auto &arg : op.argsAndDims) {
             mlir::Value scale =
                 builder.createConvert(loc, idxTy, arg.dims[i].getResult(2));
             curIndex =
                 builder.create<mlir::arith::MulIOp>(loc, scale, curIndex);
             totalIndex = (totalIndex) ? builder.create<mlir::arith::AddIOp>(
                                             loc, curIndex, totalIndex)
                                       : curIndex;
           }
-          mlir::Value elemSize =
-              builder.createIntegerConstant(loc, idxTy, arg.size);
           // This is the lowest dimension - which doesn't need scaling
           mlir::Value finalIndex =
               builder.createConvert(loc, idxTy, coop->getOperand(1));
           if (totalIndex) {
+            assert(llvm::isPowerOf2_32(arg.size) &&
+                   "Expected power of two here");
+            unsigned bits = llvm::Log2_32(arg.size);
+            mlir::Value elemShift =
+                builder.createIntegerConstant(loc, idxTy, bits);
             totalIndex = builder.create<mlir::arith::AddIOp>(
                 loc,
-                builder.create<mlir::arith::DivSIOp>(loc, totalIndex, elemSize),
+                builder.create<mlir::arith::ShRSIOp>(loc, totalIndex,
+                                                     elemShift),
                 finalIndex);
           } else {
             totalIndex = finalIndex;
           }
           auto newOp = builder.create<fir::CoordinateOp>(
               loc, builder.getRefType(elementType), caddr,
               mlir::ValueRange{totalIndex});
           LLVM_DEBUG(newOp->dump());

flang/test/Transforms/loop-versioning.fir

 // CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[DIMS0]]#2, %[[SIZE]]
 // CHECK: %[[IF_RES:.*]]:2 = fir.if %[[CMP]] -> {{.*}}
 // CHECK: %[[NEWARR:.*]] = fir.convert %[[ARG0]]
 // CHECK: %[[BOXADDR:.*]] = fir.box_addr %[[NEWARR]] : {{.*}} -> !fir.ref<!fir.array<?xf64>>
 // CHECK: %[[LOOP_RES:.*]]:2 = fir.do_loop {{.*}}
 // Check the 2D -> 1D coordinate conversion, should have a multiply and a final add.
 // Some other operations are checked to synch the different parts.
 // CHECK: %[[OUTER_IDX:.*]] = arith.muli %[[DIMS1]]#2, {{.*}}
-// CHECK: %[[ITEMSIZE:.*]] = arith.constant 8 : index
 // CHECK: %[[INNER_IDX:.*]] = fir.convert {{.*}}
-// CHECK: %[[OUTER_DIV:.*]] = arith.divsi %[[OUTER_IDX]], %[[ITEMSIZE]]
+// CHECK: %[[ITEMSHIFT:.*]] = arith.constant 3 : index
+// CHECK: %[[OUTER_DIV:.*]] = arith.shrsi %[[OUTER_IDX]], %[[ITEMSHIFT]]
 // CHECK: %[[C2D:.*]] = arith.addi %[[OUTER_DIV]], %[[INNER_IDX]]
 // CHECK: %[[COORD:.*]] = fir.coordinate_of %[[BOXADDR]], %[[C2D]] : (!fir.ref<!fir.array<?xf64>>, index) -> !fir.ref<f64>
 // CHECK: %{{.*}} = fir.load %[[COORD]] : !fir.ref<f64>
 // CHECK: fir.result %{{.*}}, %{{.*}}
 // CHECK: }
 // CHECK  fir.result %[[LOOP_RES]]#0, %[[LOOP_RES]]#1
 // CHECK: } else {
 // CHECK: %[[LOOP_RES2:.*]]:2 = fir.do_loop {{.*}}

 // CHECK: %[[NEWARR:.*]] = fir.convert %[[ARG0]]
 // CHECK: %[[BOXADDR:.*]] = fir.box_addr %[[NEWARR]] : {{.*}} -> !fir.ref<!fir.array<?xf64>>
 // CHECK: %[[LOOP_RES:.*]]:2 = fir.do_loop {{.*}}
 // Check the 3D -> 1D coordinate conversion, should have a multiply and a final add.
 // Some other operations are checked to synch the different parts.
 // CHECK: %[[OUTER_IDX:.*]] = arith.muli %[[DIMS2]]#2, {{.*}}
 // CHECK: %[[MIDDLE_IDX:.*]] = arith.muli %[[DIMS1]]#2, {{.*}}
 // CHECK: %[[MIDDLE_SUM:.*]] = arith.addi %[[MIDDLE_IDX]], %[[OUTER_IDX]]
-// CHECK: %[[ITEMSIZE:.*]] = arith.constant 8 : index
 // CHECK: %[[INNER_IDX:.*]] = fir.convert {{.*}}
-// CHECK: %[[MIDDLE_DIV:.*]] = arith.divsi %[[MIDDLE_SUM]], %[[ITEMSIZE]]
+// CHECK: %[[ITEMSHIFT:.*]] = arith.constant 3 : index
+// CHECK: %[[MIDDLE_DIV:.*]] = arith.shrsi %[[MIDDLE_SUM]], %[[ITEMSHIFT]]
 // CHECK: %[[C3D:.*]] = arith.addi %[[MIDDLE_DIV]], %[[INNER_IDX]]
 // CHECK: %[[COORD:.*]] = fir.coordinate_of %[[BOXADDR]], %[[C3D]] : (!fir.ref<!fir.array<?xf64>>, index) -> !fir.ref<f64>
 // CHECK: %{{.*}} = fir.load %[[COORD]] : !fir.ref<f64>
 // CHECK: fir.result %{{.*}}, %{{.*}}
 // CHECK: }
 // CHECK  fir.result %[[LOOP_RES]]#0, %[[LOOP_RES]]#1
 // CHECK: } else {
 // CHECK: %[[LOOP_RES2:.*]]:2 = fir.do_loop {{.*}}