[MLIR] Fix memref region compute for 0-d memref accesses

[bondhugula]

Summary

Fix memref region compute for 0-d memref accesses in certain cases (when
there are loops surrounding such 0-d accesses).

[ftynse]

mlir/lib/Analysis/Utils.cpp

214

Nit: this function is a misnomer, it returns loops not induction variables

[bondhugula]

mlir/lib/Analysis/Utils.cpp

214

Thanks for noticing - will fix this in an NFC commit soon.

mlir/lib/Analysis/Utils.cpp

LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
   memref = access.memref;
   write = access.isStore();
 
   unsigned rank = access.getRank();
 
   LLVM_DEBUG(llvm::dbgs() << "MemRefRegion::compute: " << *op
                           << "depth: " << loopDepth << "\n";);
 
+  // 0-d memrefs.
   if (rank == 0) {
     SmallVector<AffineForOp, 4> ivs;
     getLoopIVs(*op, &ivs);
-    SmallVector<Value, 8> regionSymbols;
+    assert(loopDepth <= ivs.size() && "invalid 'loopDepth'");
+    // The first 'loopDepth' IVs are symbols for this region.
+    ivs.resize(loopDepth);
+    SmallVector<Value, 4> regionSymbols;
     extractForInductionVars(ivs, &regionSymbols);
-    // A rank 0 memref has a 0-d region.
-    cst.reset(rank, loopDepth, 0, regionSymbols);
+    // A 0-d memref has a 0-d region.
+    cst.reset(rank, loopDepth, /*numLocals=*/0, regionSymbols);
     return success();
   }
 
   // Build the constraints for this region.
   AffineValueMap accessValueMap;
   access.getAccessMap(&accessValueMap);
   AffineMap accessMap = accessValueMap.getAffineMap();
 

mlir/test/Transforms/memref-bound-check.mlir

func @non_composed_bound_operand(%arg0: memref<1024xf32>) {
   affine.for %i0 = 4 to 1028 step 4 {
     %i1 = affine.apply affine_map<(d0) -> (d0 - 4)> (%i0)
     affine.for %i2 = #map_lb(%i1) to #map_ub(%i1) {
         %0 = affine.load %arg0[%i2] : memref<1024xf32>
     }
   }
   return
 }
+
+// CHECK-LABEL: func @zero_d_memref
+func @zero_d_memref() {
+  %Z = alloc() : memref<f32>
+  affine.for %i = 0 to 100 {
+    affine.load %Z[] : memref<f32>
+  }
+  return
+}