LLVM IR lowering: support 'dim' operation.
Add support for translating 'dim' opreation on MemRefs to LLVM IR. For a
static size, this operation merely defines an LLVM IR constant value that may
not appear in the output IR if not used (and had not been removed before by
DCE). For a dynamic size, this operation is translated into an access to the
MemRef descriptor that contains the dynamic size.
PiperOrigin-RevId: 223160774
diff --git a/lib/Target/LLVMIR/ConvertToLLVMIR.cpp b/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
index 9ade711..2e3ea53 100644
--- a/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
+++ b/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
@@ -506,6 +506,35 @@
element);
return false;
}
+ if (auto dimOp = inst.dyn_cast<DimOp>()) {
+ const SSAValue *container = dimOp->getOperand();
+ MemRefType type = container->getType().dyn_cast<MemRefType>();
+ if (!type)
+ return dimOp->emitError("only memref types are supported"), true;
+
+ auto shape = type.getShape();
+ auto index = dimOp->getIndex();
+ assert(index < shape.size() && "out-of-bounds 'dim' operation");
+
+ // If the size is a constant, just define that constant.
+ if (shape[index] != -1) {
+ valueMapping[dimOp->getResult()] = getIndexConstant(shape[index]);
+ return false;
+ }
+
+ // Otherwise, compute the position of the requested index in the list of
+ // dynamic sizes stored in the MemRef descriptor and extract it from there.
+ unsigned numLeadingDynamicSizes = 0;
+ for (unsigned i = 0; i < index; ++i) {
+ if (shape[i] == -1)
+ ++numLeadingDynamicSizes;
+ }
+ llvm::Value *memRefDescriptor = valueMapping.lookup(container);
+ llvm::Value *dynamicSize = builder.CreateExtractValue(
+ memRefDescriptor, 1 + numLeadingDynamicSizes);
+ valueMapping[dimOp->getResult()] = dynamicSize;
+ return false;
+ }
if (auto callOp = inst.dyn_cast<CallOp>()) {
auto operands = functional::map(
diff --git a/test/Target/llvmir.mlir b/test/Target/llvmir.mlir
index 645a266..3c476ab 100644
--- a/test/Target/llvmir.mlir
+++ b/test/Target/llvmir.mlir
@@ -536,3 +536,26 @@
return %3 : memref<10x?xf32>
}
+
+// CHECK-LABEL: define i64 @memref_dim({ float*, i64, i64 })
+cfgfunc @memref_dim(memref<42x?x10x?xf32>) -> index {
+bb0(%arg0: memref<42x?x10x?xf32>):
+// Expecting this to create an LLVM constant.
+ %d0 = dim %arg0, 0 : memref<42x?x10x?xf32>
+// CHECK-NEXT: %2 = extractvalue { float*, i64, i64 } %0, 1
+ %d1 = dim %arg0, 1 : memref<42x?x10x?xf32>
+// Expecting this to create an LLVM constant.
+ %d2 = dim %arg0, 2 : memref<42x?x10x?xf32>
+// CHECK-NEXT: %3 = extractvalue { float*, i64, i64 } %0, 2
+ %d3 = dim %arg0, 3 : memref<42x?x10x?xf32>
+// Checking that the constant for d0 has been created.
+// CHECK-NEXT: %4 = add i64 42, %2
+ %d01 = addi %d0, %d1 : index
+// Checking that the constant for d2 has been created.
+// CHECK-NEXT: %5 = add i64 10, %3
+ %d23 = addi %d2, %d3 : index
+// CHECK-NEXT: %6 = add i64 %4, %5
+ %d0123 = addi %d01, %d23 : index
+// CHECK-NEXT: ret i64 %6
+ return %d0123 : index
+}
