Fix GlobalLoad 4b lowering (#1764)

In some cases GlobalLoad would have been replaced, so we can't access it

Author: dhernandez0

mlir/lib/Dialect/Rock/Transforms/SugarToLoops.cpp

@@ -1133,9 +1133,9 @@ std::tuple<SmallVector<Value>, Type> getCoordsAndType(PatternRewriter &b,
 
 // A helper to select the right i4 element if it was supposed to
 // be a scalar i4 load.
-Value selectDataIf4b(PatternRewriter &b, GlobalLoadOp op, Value loadedVec) {
-  MemRefType srcType = op.getSource().getType();
-  Type originalLoadedType = op.getResult().getType();
+Value selectDataIf4b(Location loc, PatternRewriter &b,
+                     SmallVector<Value> &coords, MemRefType srcType,
+                     Type originalLoadedType, Value loadedVec) {
   if (srcType.getElementType().getIntOrFloatBitWidth() >= 8) {
     return loadedVec;
   }
@@ -1150,8 +1150,6 @@ Value selectDataIf4b(PatternRewriter &b, GlobalLoadOp op, Value loadedVec) {
   assert(srcType.getElementType().getIntOrFloatBitWidth() == 4 &&
          "we only support 4bits in narrow types");
   assert(isa<VectorType>(loadedVec.getType()));
-  Location loc = op.getLoc();
-  SmallVector<Value, 5> coords(op.getSourceCoord());
   ArrayRef<int64_t> shape = srcType.getShape();
   Value flatAddress = flattenCoords(b, loc, coords, shape);
   Type coordType = flatAddress.getType();
@@ -1198,6 +1196,12 @@ struct GlobalLoadRewritePattern : public OpRewritePattern<GlobalLoadOp> {
       source = zeroDMemrefAsOneD(b, source);
       coords.push_back(b.createOrFold<ConstantIndexOp>(loc, 0));
     }
+    // We need to copy these params here, because the next if might replace
+    // "op". So, we can't safely access it after that.
+    // TODO: refactor this code
+    MemRefType srcType = op.getSource().getType();
+    Type originalLoadedType = op.getResult().getType();
+    SmallVector<Value> sourceCoords(op.getSourceCoord());
 
     PatternRewriter::InsertionGuard insertGuard(b);
     if (emitOobChecks && !useBufferOps) {
@@ -1207,11 +1211,12 @@ struct GlobalLoadRewritePattern : public OpRewritePattern<GlobalLoadOp> {
             loc, arith::CmpIPredicate::uge, coords[0], numElems);
         cond = b.create<arith::AndIOp>(loc, fallsOffEnd, cond);
       }
-      auto guard = b.create<scf::IfOp>(loc, loadedType, cond, true, true);
+      auto guard =
+          b.create<scf::IfOp>(loc, originalLoadedType, cond, true, true);
       b.replaceOp(op, guard);
 
       b.setInsertionPointToEnd(guard.getBody(1));
-      Value zeroes = createZeroConstantOp(b, loc, loadedType);
+      Value zeroes = createZeroConstantOp(b, loc, originalLoadedType);
       b.create<scf::YieldOp>(loc, zeroes);
       b.setInsertionPointToEnd(guard.getBody(0));
     }
@@ -1249,15 +1254,18 @@ struct GlobalLoadRewritePattern : public OpRewritePattern<GlobalLoadOp> {
         else
           loaded = thisLoad;
       });
-      loaded = selectDataIf4b(b, op, loaded);
+      loaded = selectDataIf4b(loc, b, sourceCoords, srcType, originalLoadedType,
+                              loaded);
       b.replaceOp(op, loaded);
     } else {
       Value loaded;
       if (isa<VectorType>(loadedType))
         loaded = b.create<vector::LoadOp>(loc, loadedType, source, coords);
       else
         loaded = b.create<memref::LoadOp>(loc, loadedType, source, coords);
-      loaded = selectDataIf4b(b, op, loaded);
+
+      loaded = selectDataIf4b(loc, b, sourceCoords, srcType, originalLoadedType,
+                              loaded);
       if (emitOobChecks)
         b.create<scf::YieldOp>(loc, loaded);
       else

mlir/test/Dialect/Rock/lowering_global_load_store.mlir

@@ -111,6 +111,26 @@ func.func @load_scalar_oob_large(%mem: memref<1073741825xf32>, %valid: i1) -> f3
     return %ret : f32
 }
 
+// CHECK-LABEL: func.func @load_scalar_oob_large_i4
+// CHECK-SAME: (%[[mem:.*]]: memref<1073741825xi4>, %[[valid:.*]]: i1)
+func.func @load_scalar_oob_large_i4(%mem: memref<1073741825xi4>, %valid: i1) -> i4 {
+    %c0 = arith.constant 0 : index
+    // CHECK: %[[zero:.*]] = arith.constant 0 : i4
+    // CHECK: %[[cast:.*]] = memref.memory_space_cast %[[mem]]
+    // CHECK-SAME: #gpu.address_space<global>
+    // CHECK: %[[ret:.*]] = scf.if %[[valid]] -> (i4)
+    // CHECK: %[[load:.*]] = vector.load %[[cast]]
+    // CHECK-SAME: memref<1073741825xi4, #gpu.address_space<global>>, vector<2xi4>
+    // CHECK: %[[element:.*]] = vector.extractelement %[[load]]
+    // CHECK: scf.yield %[[element]]
+    // CHECK: } else {
+    // CHECK: scf.yield %[[zero]] : i4
+    %ret = rock.global_load %mem[%c0] if %valid {needs64BitIdx}
+        : memref<1073741825xi4> -> i4
+    // CHECK: return %[[ret]]
+    return %ret : i4
+}
+
 // CHECK-LABEL: func.func @store_scalar_in_bounds
 // CHECK-SAME: (%[[source:.*]]: memref<5xf32, #gpu.address_space<private>>, %[[mem:.*]]: memref<1x2x3x4x8xf32>)
 func.func @store_scalar_in_bounds(%source: memref<5xf32, #gpu.address_space<private>>, %mem: memref<1x2x3x4x8xf32>) {